MoE 内核优化提升 H200 推理性能

Mixture-of-Experts（MoE）是一种稀疏激活架构，能够扩展大语言模型（LLM）的容量，同时保持较低的单 token 计算量。由于 MoE 运算占据了总计算时间的很大一部分，其中的低效会直接增加请求延迟和成本。因此，通过 kernel 工程优化 MoE 运算，是获得性能提升最有潜力的方向。

Fused MoE 是 MoE 架构的一项关键优化技术。它将 token 路由、数据重排和专家网络计算合并为 GPU 上的一次批量操作，替代原先通过独立、串行循环分别执行路由和专家计算的方式。这种方式能降低延迟，但也增加了 kernel 实现和调试的复杂度。要让这种架构实现较高的 GPU 利用率，需要针对具体硬件进行优化。

本文将介绍我们为 fused MoE 模型实现的多项 kernel 优化，并说明如何用上这些成果。我们选取了一个具有代表性的 MoE 模型作为主要示例。

通过 Swap-AB 提高 WGMMA 利用率

为了最大限度地提高张量计算吞吐量，NVIDIA Hopper 架构依赖 Warp Group Matrix Multiply Accumulate（WGMMA）指令协议。这类指令由一个包含 128 个线程、协同工作的 warpgroup 共同执行。其操作数的不对称性，本质上是指令集架构（ISA）层面的属性。在底层的 wgmma.mma_async 汇编定义中，硬件将矩阵行维度（M）固定为 64，以此定义并行线程的映射拓扑，而列维度（N）则可以灵活变化。Operand A 严格映射到这个固定的 M 维度上，使用寄存器或共享内存；Operand B 则通过共享内存描述符映射到可变的 N 维度上。因此，这两个操作数位置在结构上不能互换。硬件 tile 能否达到最高利用率，受这些固定的形状编码约束：M 维度必须大于或等于 64，才能避免计算执行通道闲置。

这种结构偏好源于 Hopper Tensor Cores 划分内存块的方式。硬件为第一个操作数的主维度强制规定了固定的最小布局形状，将块高度锁定为 64（M=64）。当放在这个首要位置的矩阵拥有较大且连续的主维度时，Hopper 内部的硬件调度器就能顺畅地将数据映射并分配到物理计算网格上。

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

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

其中，M 表示 token 行数，也就是 batch size；K 表示隐藏特征维度；N 表示较大的专家输出维度。

在 decode 阶段，系统每次只处理每条用户序列的一个 token。因此，batch 维度，也就是 token 行数，会变得非常小，通常满足 M < 64，而且经常只有一行（M = 1）。在这种标准布局下：

矩阵 A 变成一个很薄、碎片化的横向向量，例如 1 x 4096。

矩阵 B，也就是专家的核心权重，仍然是一个大型、稠密的纵向块，例如 4096 x 8192。

由于行数 M 很小，矩阵 A 根本无法填满硬件为 Operand A 固定预留的 64 行布局空间。当 WGMMA 指令在这种布局上执行时，大量物理 Tensor Core 计算资源都得不到数据。芯片把时钟周期花在处理空的空间槽位上，导致硬件利用率低下。

在标准的 dense LLM 层中，权重矩阵是静态、统一且完全连续的。但在 MoE 层中，工作负载会分散到多个独立的子网络，也就是专家中。这会带来特有的内存和布局瓶颈，使其对 Hopper GPU 的 WGMMA 约束格外敏感。在标准 dense 层中，矩阵乘法是将一个 token 向量乘以该层唯一的权重矩阵。输出维度 N 就是模型的 hidden size，例如 4096。

在 Fused MoE kernel 中，GPU 并不是只与一个专家的权重做乘法。它会将所有活跃专家的计算组织起来并同时执行，避免为每个专家单独启动 kernel。对于标准 dense 层，N 是 hidden size（4,096）；对于 fused MoE 层，N 则是选中的专家数量乘以专家的中间层维度。如果模型采用 topk=8，fused kernel 内部的有效 N 维度就会膨胀到数万甚至数十万。相较于 decode 时极小的 token 行数（M=1），N 维度大得极不成比例，布局的不对称性也因此变得极端。将矩阵转置后，这个合并后的巨大 MoE 维度会被移到首要操作数位置，从而启用标准 dense 层无法达到的规模的硬件预取。

为了修正这种硬件与布局的不匹配，同时保持模型底层数学运算不变，我们利用了矩阵转置的代数性质：

BT[N,K] x AT[K,M] —> CT[N,M]

图 1：通过 Swap-AB 提高 WGMMA 利用率

通过在矩阵乘法进入 PTX 汇编层之前，以编程方式交换操作数并转置，我们让软件布局重新对齐 Hopper 芯片的实际硬件结构。这种转置立即带来了两项 GPU 架构层面的优势。首先，它能充分填满物理布局。较大的专家维度（N = 8192）被放到第一个操作数位置（BT），由于 N 足够大，可以轻松填满硬件要求的至少 64 行。这样，Tensor Cores 就能以最高的局部计算密度工作，不再浪费执行周期。其次，它启用了异步预取。Hopper 专用的硬件拷贝引擎可以在计算执行块启动之前，将数据从共享内存异步预取到寄存器中。将连续的大型权重块 BT 放在首要位置后，这些引擎就能顺畅地沿着数据布局读取，让计算流水线保持满负荷运行，无需 CPU 干预。这项优化让 H200 上的大规模 MoE 模型端到端输出吞吐量提高了 2%–4%。我们已将其贡献到 vLLM 上游，vLLM 0.24.0 及更高版本均可使用。

为 FP8 张量量化选择 backend 路由

除了 kernel 内部实现，选择哪个 backend 执行 MoE 计算，也会显著影响性能。vLLM MoE oracle 会将不同的量化方案路由到 Triton、FlashInfer+CUTLASS 或其他 backend。性能分析发现，对于 Hopper 上采用 FP8 张量量化的模型，Triton backend 的表现始终优于默认的 FLASHINFER_CUTLASS 路径。

在低并发 MoE 工作负载中，Triton backend 优于 CUTLASS，主要是因为它用定制的算子级寄存器融合，替代了固定的框架级矩阵计算配置。CUTLASS 依赖标准 C++ 模板结构，将 FP8 张量缩放、矩阵乘法和专家聚合视为彼此独立的数学操作。Triton backend 则把整个流水线合并为一次硬件执行。数据只需从全局高带宽内存（HBM）读取一次，随后全部在 GPU 超高速的本地 SRAM 寄存器和计算执行单元中处理，最后再以流式方式写回最终矩阵的目标位置。这样就消除了缓慢的全局 VRAM 读写往返；当 CUTLASS 必须不断执行中间的张量缩放步骤时，这些往返正是拖累其性能的原因。

此外，这种寄存器级融合还专门配合了围绕可变 token 分桶设计的动态调度布局，解决了传统编译模板中的一个严重结构缺陷。CUTLASS 原生针对稠密、大型且维度可预测的矩阵进行了优化。当它遇到实时 decode 生成阶段常见的、不断变化且极小的 batch 形状（M < 64 或 M = 1）时，默认的 block tiling 和线程映射结构就会变得非常低效，造成大量 warp 利用不足，或产生昂贵的内存填充开销。相比之下，Triton MoE backend 擅长将这些小而碎片化的 token 数组，动态映射到活跃专家权重的物理边界上。即使工作负载变化剧烈、batch size 很低，也能让硬件内部的预取数据流和 Tensor Core 流水线保持充分利用。

更新路由策略只需要修改配置，不涉及任何 kernel 代码变更。即便如此，它仍让 H200 上采用 FP8 量化的 fused MoE 模型端到端吞吐量提高了 10%。这项改进已包含在 AWS Large Model Inference（LMI）v28 容器中。

消除 Triton JIT 重编译

在 GPU 推理中，有一种性能问题影响很大，却很容易被忽略：不必要的 kernel 重编译。Triton 是 vLLM 许多自定义 kernel 所使用的编译器框架，默认会对整数参数进行特化。每当 kernel 调用为某个被特化的参数传入一个此前未见过的值，Triton 就会触发一次新的即时编译（JIT）。在序列长度和 batch size 持续变化的服务工作负载中，这会不断产生编译开销，推高延迟，尤其是首 token 时间（Time to First Token，TTFT）。

在 MoE 模型服务中，负责融合旋转位置编码和 FP8 量化的 _fused_inv_rope_fp8_quant_per_head kernel，每遇到一个新的 num_tokens 值就会重编译。为可变长度参数添加 do_not_specialize 注解后，这些重编译被完全消除。在 H200 上采用 8 路 tensor parallelism 时，TTFT 改善了 9%（PR #42849）。

在多 GPU 集群中，全部 8 张 GPU 必须以微秒级精度同步执行各自的 kernel。只要其中一张 GPU 因为本地 token batch 数量发生变化，需要停下来重编译本地的 _fused_inv_rope_fp8_quant_per_head kernel，其余 7 张 GPU 就只能闲置等待。通过消除整个集群中的重编译，我们减少了多 GPU 协调过程中的波动，让用户收到初始响应的延迟降低了 9%。这里的关键是，将编译器视为一项首要优化对象。理解 Triton 的特化语义，并识别哪些参数未被实际使用、哪些参数会变化，就能以极少的代码修改降低延迟。

我们已将这项优化贡献到 vLLM 上游，vLLM 0.22.0 及更高版本均可使用。

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