AWS 在 Amazon SageMaker AI 上通过 Swap-AB 与 Triton 优化提升 MoE 模型推理性能

AWS 发布了一篇关于在 Amazon SageMaker AI 上针对混合专家（MoE）架构进行内核工程的教程，旨在解决稀疏激活架构中因计算效率低下导致的延迟和成本问题。文章重点介绍了三项针对 NVIDIA Hopper 架构的优化技术：Swap-AB、后端路由选择以及 Triton JIT 重编译消除。

核心事实在于 Swap-AB 技术解决了 Hopper Tensor Cores 对矩阵维度 M 必须大于等于 64 的硬性约束。在 MoE 模型的解码阶段，由于每次仅处理一个 token，导致输入矩阵 A 的行数 M 极小（如 M=1），无法填满硬件固定的 64 行布局，造成算力闲置。通过将矩阵乘法顺序交换为 B^T @ A^T，将庞大的专家维度 N 提升至第一操作数位置，从而饱和物理布局并激活异步预取功能。此外，对于 FP8 量化模型，Triton 后端因采用自定义算子级寄存器融合而非 CUTLASS 的标准模板，能更有效地处理低并发下的碎片化 token 数组，避免了全局 VRAM 读写往返。

推断层面，这些优化措施表明 AWS 正在深度参与开源推理框架 vLLM 的核心开发，并将特定于云基础设施的性能调优直接回馈给社区。例如，Swap-AB 优化已上游至 vLLM 0.24.0+ 版本，而 Triton JIT 重编译消除则已在 vLLM 0.22.0+ 中可用。这暗示了未来基于 vLLM 构建的推理服务在 Hopper 架构上将具备更高的吞吐量上限和更低的 TTFT（首字延迟）。

猜测成分在于，虽然文中提到这些优化带来了最高 10% 的吞吐量和 9% 的 TTFT 提升，但实际效果可能高度依赖于具体的模型结构（如 top-k 设置）和集群规模。对于非 Hopper 架构或不同量化精度的场景，这些特定的内核调整策略可能需要重新评估。