AWS发布MoE模型内核优化：提升Amazon SageMaker AI推理性能

摘要：AWS工程师团队在Amazon SageMaker AI平台上实现了针对融合型混合专家（MoE）模型的内核优化，包括Swap-AB矩阵转置优化以提升WGMMA利用率、基于Triton后端的FP8量化路由优化，以及消除Triton JIT重新编译。这些优化使H200 GPU上的MoE模型吞吐量提升高达10%，首次Token生成时间缩短9%。这项技术突破展示了AWS在AI推理优化领域的持续创新，为大模型的高效部署提供了重要支持。

---

### 混合专家（MoE）模型的内核优化 混合专家（MoE）是一种稀疏激活的架构，旨在扩展大语言模型（LLM）的容量，同时保持每个Token的计算成本较低。由于其操作占据了总计算时间的大部分，因此任何效率上的提升都能直接降低请求延迟和成本。因此，通过内核工程优化MoE操作是提升性能的最大机会。 #### Swap-AB矩阵转置优化 为了最大化张量处理吞吐量，NVIDIA Hopper架构依赖于Warp Group Matrix Multiply Accumulate（WGMMA）指令协议，该协议在128个线程的协作warpgroup上执行集体操作。在MoE网络的解码阶段，标准矩阵布局会导致严重的性能瓶颈，因为在这种布局下，矩阵A的维度太小，无法充分利用硬件的固定64行布局。通过将矩阵乘法转置，我们重新调整了软件以匹配Hopper芯片的确切拓扑结构。这种转置在GPU上触发了两个架构优势： 1. **物理布局饱和**：扩展的专家维度（N=8192）被提升到第一个操作数位置（B^T），由于N很大，它轻松地超过了硬件的强制64行阈值，从而使Tensor Cores以最大局部密度运行，没有浪费的执行周期。 2. **异步预取激活**：Hopper的专用硬件复制引擎可以异步地从共享内存预取内存地址到寄存器中。将连续的、大型的权重块（B^T）放在前导槽中，可以让这些引擎平滑地跟踪数据布局，并保持计算管道的饱和，而无需CPU干预。这在H200上的大规模MoE模型中产生了2-4%的端到端输出吞吐量提升。 #### 基于Triton后端的FP8量化路由优化 除了内核内部，选择哪个后端执行MoE计算对性能也有很大影响。vLLM MoE oracle将不同的量化方案路由到Triton、FlashInfer+CUTLASS或其他后端。分析表明，对于Hopper上的FP8张量量化模型，Triton后端始终优于默认的FLASHINFER_CUTLASS路径。Triton后端通过将整个管道合并为一个硬件通道，消除了慢速全局VRAM读写往返，从而在低并发MoE工作负载中表现优异。这一优化为FP8量化的融合MoE模型带来了10%的端到端吞吐量提升。 #### 消除Triton JIT重新编译 在GPU推理中，不必要的内核重新编译是一个常见但容易被忽视的性能瓶颈。在MoE模型服务中，_fused_inv_rope_fp8_quant_per_head内核在每次新的num_tokens值出现时都会重新编译。通过向可变长度参数添加do_not_specialize注释，可以完全消除这些重新编译，从而在H200上以8路张量并行性将TTFT提高了9%。在多GPU集群中，消除整个集群的重新编译可以平滑多GPU协调开销，为用户的初始响应延迟提供9%的加速。 ### 结论 通过与vLLM社区的合作，我们实现了并开源了多个针对融合MoE模型的内核优化。这些优化应用于Amazon SageMaker AI上的LMI深度学习容器中。这些优化共同实现了高达10%的吞吐量提升和高达9%的TTFT加速。 --- > **消息来源**：[AWS Machine Learning Blog](https://aws.amazon.com/blogs/machine-learning/mixture-of-experts-kernel-engineering-on-amazon-sagemaker-ai/) (2026-10-13) —— 完 ——

主题标签： #AWS 关注 #MoE 关注 #内核优化 关注 #AI推理 关注 #vLLM 关注