Amazon SageMaker AI 上的 MoE 内核工程优化实践

Mixture-of-Experts (MoE) 是一种稀疏激活架构，可在降低每 token 计算量的同时扩展大语言模型 (LLM) 的容量。由于 MoE 的运算占总计算时间的很大一部分，其中的低效会直接推高请求延迟和成本。因此，通过内核工程优化 MoE 运算是获得性能提升的最大机会。

Fused MoE 是 MoE 架构的一项关键优化技术。融合的 MoE 系统将 token 路由、数据重排和专门的神经网络计算合并为 GPU 上单个批处理操作，取代了以分离的、顺序的循环操作执行这些路由和专家计算的方式。虽然这带来了延迟优势，但也给内核实现和调试能力带来了复杂性。要使这种架构实现高 GPU 利用率，需要进行特定于硬件的优化。

在本文中，我们介绍了为融合 MoE 模型实现的各种内核优化，并向您展示如何从这项工作中受益。我们使用了一个具有代表性的 MoE 模型作为主要示例。

Swap-AB 以充分利用 WGMMA

为了最大化张量处理吞吐量，NVIDIA Hopper 架构依赖于 Warp Group Matrix Multiply Accumulate (WGMMA) 指令协议，该协议由一个 128 线程的协作 warpgroup 共同执行。这种操作数不对称本质上是指令集架构 (ISA) 层面的属性。在底层的 wgmma.mma_async 汇编定义中，硬件将矩阵行维度 (M) 固定在恒定值 64，以定义其并行线程映射拓扑。列维度 (N) 则保持高度可变。操作数 A 严格映射到这个固定的 M 维度（使用寄存器或共享内存），而操作数 B 映射到可变的 N 维度（通过共享内存描述符）。因此，这些操作数槽位在结构上是不可互换的。由此，硬件的峰值 tile 利用率受这些刚性形状编码支配，要求 M 维度必须等于或大于 64，以防止计算执行通道闲置。

这种结构性偏好源于 Hopper Tensor Core 划分内存块的方式。硬件对第一个操作数的前导维度强制固定一个最小布局形状，将块高度精确锁定为 64 (M=64)。当一个数学矩阵在这个主槽位中拥有较大的、连续的前导维度时，Hopper 的内部硬件调度器就能干净地在物理计算网格上映射和分发数据。

在 MoE 网络的 decode 阶段，这种结构不对称会成为严重的性能瓶颈。在标准的 MoE 层中，控制专家计算的线性代数是用传统矩阵记法表示的：

A[M,K] x B[K,N] —> C[M,N]

其中 (M) 表示 token 行数（或批大小），(K) 表示隐藏特征维度，(N) 表示庞大的专家输出维度。

在 decode 阶段，系统对每个用户序列一次只处理一个 token。因此，批空间或 token 行数非常薄，通常降到 (M < 64) 的情形，甚至经常降到单行 (M = 1)。在这种标准结构布局下：

矩阵 (A) 变成一个薄而碎片化的水平向量（例如 1 x 4096）。

矩阵 (B)（核心专家权重）仍保持为一个大而稠密的垂直块（例如 4096 x 8192）。

由于行数 (M) 很小，矩阵 (A) 在物理上无法填满硬件为操作数 A 预设的固定 (64) 行布局桶。当 WGMMA 指令在这种布局上执行时，会让物理 Tensor Core 容量的很大一部分完全处于饥饿状态。硅片在空的空间槽位上耗费时钟周期，导致硬件利用率低下。

在标准的稠密 LLM 层中，权重矩阵是静态、统一且完全连续的。然而在 MoE 层中，工作负载被分散到一群独立的子网络（专家）之中。这造成了独特的内存和布局瓶颈，使其对 Hopper GPU 的 WGMMA 约束极度敏感。在标准稠密层中，矩阵乘法是将 token 向量与单个层权重矩阵相乘。输出维度 (N) 是模型的隐藏层大小（例如 4096）。

在融合 MoE 内核中，GPU 不仅仅是与一个专家相乘。它会对所有激活的专家同时进行分组并执行计算，以避免为每个专家单独启动内核。对于标准稠密层，N 是隐藏层大小（4,096）。对于融合 MoE 层，N 是被选中的专家数量乘以专家中间层大小。对于运行 topk=8 例程的模型，融合内核内部的有效 N 维度会爆炸到数万或数十万个参数。由于 N 维度与极小的解码 token 大小 (M=1) 相比大得不成比例，布局的不对称性变得极端。对矩阵进行转置可将这个庞大的组合 MoE 维度移到主槽位，从而解锁标准稠密层无法企及的硬件预取能力。

为了在不改变模型底层数学完整性的前提下解决这种硬件错位问题，我们利用了矩阵转置的代数性质：

B^T[N,K] x A^T[K,M] —> C^T[N,M]

图 1：为提升 WGMMA 利用率而交换 A/B（Swap-AB）

通过在矩阵乘法到达 PTX 汇编层之前以编程方式将其翻转，我们让软件与 Hopper 硅片的精确拓扑结构重新对齐。这次转置在 GPU 上触发了两个直接的结构性优势。第一，它填满了物理布局。庞大的专家维度 (N = 8192) 被提升到第一个操作数位置 (B^T)，并且由于 (N) 很大，它轻松填满了硬件强制的 (64) 行阈值。于是 Tensor Core 以最大的局部密度运行，没有浪费的执行周期。第二，它激活了异步预取。Hopper 专用的硬件拷贝引擎可以在数学执行单元启动之前，将内存地址从共享内存异步预取到寄存器中。将连续的大型权重块 (B^T) 放在主槽位，可以让这些引擎顺畅地跟踪数据布局，并在无需 CPU 干预的情况下保持计算管线饱和。这在 H200 上的大规模 MoE 模型中带来了 2–4% 的端到端输出吞吐量提升。我们将这一优化上游提交到了 vLLM，可在 0.24.0 或更高版本中使用。

FP8 张量量化的后端路由

除了内核内部机制之外，选择哪个后端来执行 MoE 计算对性能也有很大影响。vLLM 的 MoE oracle 会将不同的量化方案路由到 Triton、FlashInfer+CUTLASS 或其他后端。性能分析显示，对于 Hopper 上的 FP8 张量量化模型，Triton 后端始终优于默认的 FLASHINFER_CUTLASS 路径。

Triton 后端在低并发 MoE 工作负载中优于 CUTLASS，主要原因在于它用自定义的算子级寄存器融合取代了僵化的框架级矩阵数学配置。CUTLASS 依赖标准的 C++ 模板结构，将 FP8 张量缩放、矩阵乘法和专家聚合视为离散的数学操作。而 Triton 后端则将整个流水线合并为单次硬件处理。数据仅从全局高带宽内存（HBM）中读取一次，并完全在 GPU 超高速的本地 SRAM 寄存器和计算执行单元中处理，然后流回最终的矩阵目标位置。这消除了缓慢的全局 VRAM 读写往返——当 CUTLASS 被迫不断执行中间张量缩放步骤时，正是这些往返拖累了它的表现。

此外，这种寄存器级融合还专门与围绕可变 token 分桶设计的动态调度布局相配合，解决了传统编译模板中的一个严重结构性缺陷。CUTLASS 原生针对稠密、大型且可预测的矩阵维度进行了优化。当它遇到在线解码生成阶段典型的流动、超小批处理特征（M < 64 或 M = 1）时，其默认的块分块和线程映射结构会变得非常低效，导致严重的 warp 利用不足或昂贵的内存填充开销。相反，Triton MoE 后端擅长将这些细小、零散的 token 数组动态映射到活跃专家权重的物理边界上，即使在极低批大小下处理高度不稳定的负载，也能让硬件的内部预取流和 Tensor Core 流水线保持完全饱和。

更新路由策略是一次纯粹的配置变更，没有修改任何内核代码。即便如此，它仍为 H200 上的 FP8 量化融合 MoE 模型带来了 10% 的端到端吞吐量提升。这一改进已在 AWS Large Model Inference (LMI) v28 容器中提供。

消除 Triton JIT 重新编译

GPU 推理中最具影响的性能问题之一也是最不显眼的之一：不必要的内核重新编译。Triton 是 vLLM 许多自定义内核所依赖的编译框架，它默认会对整数参数进行特化。每当内核以一个之前未见过的特化参数值被调用时，Triton 就会触发一次全新的即时（JIT）编译。在序列长度和批大小持续变化的服务工作负载中，这会形成源源不断的编译开销，从而推高延迟，尤其是首 Token 时间（TTFT）。

在 MoE 模型服务中，_fused_inv_rope_fp8_quant_per_head 内核（负责融合旋转位置编码和 FP8 量化）会在每个新的 num_tokens 值出现时重新编译。通过为可变长度参数添加 do_not_specialize 注解，这些重新编译被完全消除，在 H200 上 8 路张量并行的情况下将 TTFT 提升了 9%（PR #42849）。

在多 GPU 集群中，所有 8 个 GPU 必须以完美、微秒级的同步执行各自的内核。如果哪怕只有一个 GPU 因其本地 token 批次数量发生变化而必须停下来重新编译其本地的 _fused_inv_rope_fp8_quant_per_head 内核，其余七个 GPU 就得闲置等待。通过消除整个集群中的重新编译，我们消除了多 GPU 协调开销，使用户的首个响应延迟获得了 9% 的加速。这里的原则是把编译器视为一等优化对象。理解 Triton 的特化语义，并识别参数何时是死参数或可变参数，就能以极小的代码改动换来延迟的改善。

我们将这项优化向上游贡献给了 vLLM，它已在 0.22.0 或更高版本中可用。

结论

通过与 vLLM 社区的合作，我们为融合 MoE 模型实现并开源了多项内核优化。我们在 Amazon SageMaker AI 的 LMI Deep Learning Container 中应用了这些优化。这些优化合计带来了最高 10% 的吞吐量提升和最高 9% 的 TTFT 缩短。要开始使用自定义模型托管，请参阅 Amazon SageMaker AI 托管文档。

关于作者