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

在混合专家（MoE）模型逐渐成为大语言模型主流架构的当下，推理性能的瓶颈往往不在模型结构本身，而在于底层内核的实现质量。AWS 工程团队近期分享了他们在 Amazon SageMaker AI 上针对融合 MoE 模型所做的一系列内核级优化，包括 Swap-AB 提升 WGMMA 利用率、Triton 后端路由实现 FP8 量化，以及消除 Triton JIT 重复编译。这些优化在 H200 GPU 上带来了最高 10% 的吞吐量提升和最高 9% 的首 token 延迟（TTFT）改善。

为什么 MoE 推理对内核工程如此敏感

MoE 架构的核心思路是用多个专家网络替代单一前馈层，每个 token 只激活其中少数专家。这种稀疏激活让模型参数量可以大幅扩展，而计算量增长相对可控。但代价是，推理过程中出现了大量小规模、不规则的计算任务：token 到专家的路由是动态的，每个专家接收到的 token 数量不固定，GEMM 操作的矩阵形状随批次变化。

这些特性使得通用内核难以充分利用 GPU 的算力。尤其是在 Hopper 架构（如 H200）上，WGMMA（Warpgroup Matrix Multiply-Accumulate）指令对矩阵布局和流水线有严格要求，如果内核的访存模式或数据排布与之不匹配，张量核心的利用率会明显下降。因此，MoE 推理的性能优化很大程度上是一场内核工程战役，而非简单的模型压缩或框架调参。

WGMMA（Warpgroup Matrix Multiply-Accumulate）是 NVIDIA Hopper 架构引入的异步矩阵乘加指令，由整个 warpgroup（通常 128 个线程）协同执行，可直接从共享内存读取操作数并写入张量核心累加器。相比上一代 Ampere 的 mma 指令，WGMMA 支持更大的矩阵分块、更高的吞吐，并可与 TMA（Tensor Memory Accelerator）等异步拷贝机制流水线重叠。但它的高效运行依赖严格的矩阵布局：操作数需按特定 swizzle 格式驻留共享内存，warpgroup 内线程的分工也必须与指令语义对齐。一旦内核的访存模式、数据排布或同步策略与之不匹配，张量核心就会出现空转，实际算力利用率可能远低于理论峰值。这正是 MoE 这类不规则负载在 Hopper 上难以跑满的根本原因之一。

三项关键内核优化

Swap-AB：让 WGMMA 跑满

在标准 GEMM 中，A 矩阵通常代表激活值，B 矩阵代表权重。但在 MoE 的融合内核中，权重矩阵的规模和布局与激活值的关系可能并不适合直接映射到 WGMMA 的最优执行模式。AWS 团队采用的 Swap-AB 策略，本质上是交换两个操作数的角色，使得内核在 Hopper 的 WGMMA 流水线中获得更好的数据复用和更少的同步开销。

这一改动的价值在于，它不需要改变模型的数学等价性，却能显著提升张量核心的占用率。对于 MoE 中大量存在的分组 GEMM 和批处理 GEMM，这种操作数层面的调整往往比在更高层面调整分块策略更直接有效。

GEMM 的数学形式为 C = A × B，其中 A 通常行主序代表激活，B 代表权重。但由于矩阵乘法满足 (A×B)ᵀ = Bᵀ×Aᵀ，交换两个操作数并转置结果并不改变计算语义，只改变数据在内存与寄存器中的排布方式。在 Hopper 上，WGMMA 对 A、B 操作数的来源（寄存器还是共享内存）、主序方向以及分块粒度有不同约束，某些形状下 A 更适合放激活、B 更适合放权重，而在 MoE 融合内核中往往恰好相反。Swap-AB 就是利用这种等价变换，把负载更重或布局更规整的一方放到 WGMMA 更擅长的位置，从而减少共享内存搬运与同步次数。它属于典型的“零数学代价、纯工程收益”优化。

Triton 后端路由：FP8 量化的落地路径

FP8 量化是当前降低推理成本的重要手段，但在 MoE 模型中，不同专家的数值分布可能存在差异，统一的量化策略容易导致精度损失。AWS 团队利用 Triton 的后端路由机制，让不同的计算路径可以走到各自适合的后端实现，从而在 FP8 量化下保持数值稳定性。

这种做法实际上是把量化的粒度控制下沉到内核调度层。相比在模型层面做全局量化配置，后端路由允许更细粒度的决策：哪些专家或哪些计算阶段适合 FP8，哪些需要保留更高精度。对于追求极致吞吐的生产环境，这种灵活性是必要的。

FP8 是 8 位浮点格式，常见变体有 E4M3（4 位指数、3 位尾数，动态范围较小但精度较高）和 E5M2（5 位指数、2 位尾数，范围大但精度低）。相比 FP16/BF16，FP8 可将显存占用和带宽需求减半，并在 Hopper 等支持 FP8 张量核心的 GPU 上获得更高算力。但其代价是尾数位极少，对数值分布敏感：若某组权重的动态范围超出所选格式的可表示区间，就会出现溢出或下溢。MoE 中不同专家训练数据与激活分布各异，统一量化很容易让个别专家精度崩塌。Triton 的后端路由允许同一内核的不同分支编译到不同的目标实现（如不同的量化格式或精度路径），从而按专家或按计算阶段选择最合适的数值策略。

消除 Triton JIT 重复编译

Triton 的 JIT 编译机制在开发阶段非常方便，但在生产推理中，如果每次遇到新的输入形状都触发重新编译，会带来不可预测的延迟抖动。MoE 的动态路由恰好容易产生这种问题：不同批次的路由结果可能导致内核参数变化，进而触发重复编译。

AWS 团队针对这一问题做了优化，减少了 JIT 重编译的发生频率。虽然原文没有披露具体实现细节，但从效果来看，这一优化对首 token 延迟的改善贡献明显——TTFT 对在线服务的用户体验至关重要，任何编译层面的抖动都会被放大。

Triton 是一种面向 GPU 的类 Python 内核 DSL，采用即时编译（JIT）：开发者用 Python 描述内核逻辑，Triton 在首次调用时根据张量形状、数据类型等“特化键”生成并编译 PTX/CUBIN。这一机制让开发迭代很快，但每次遇到新的形状组合都可能触发一次完整编译，耗时从数十毫秒到数秒不等。MoE 的动态路由使每个专家、每个批次接收的 token 数不断变化，形状组合极易发散，导致编译缓存命中率低。生产推理中这种偶发编译会直接表现为长尾延迟和首 token 抖动。常见缓解手段包括：预编译并缓存常见形状、对形状做分桶归并、固定最大形状配合掩码，以及在部署时预热所有可能的内核变体。

性能收益与工程启示

三项优化叠加后，在 H200 GPU 上实现了最高 10% 的吞吐量提升和最高 9% 的 TTFT 改善。这个数字单独看并不惊人，但考虑到 MoE 推理的内核优化空间往往被框架层的抽象所掩盖，这类底层改动的价值在于它们可以与上层优化正交叠加。

从工程角度看，这组优化传递了几个信号。其一，Hopper 及后续架构的张量核心能力需要通过精细的内核布局来释放，通用内核的"够用"和"高效"之间仍有距离。其二，量化不是简单的精度换速度，FP8 在 MoE 中的落地需要内核层提供足够的调度灵活性。其三，生产环境中的延迟稳定性与平均吞吐同样重要，JIT 编译这类看似边缘的问题，在规模化服务中会被放大。

对于在 SageMaker AI 上部署 MoE 模型的团队来说，这些优化很可能会以托管形式体现在服务性能中，但理解其背后的原理，有助于在自定义内核或选择实例类型时做出更明智的决策。内核工程正在重新成为大模型推理优化的核心战场，而 MoE 只是把这个趋势推到了台前。