Mixture-of-Experts kernel engineering on Amazon SageMaker AI

Amazon SageMaker AI 上的专家混合（MoE）内核工程 | 人工智能

Amazon SageMaker AI 上的专家混合（MoE）内核工程

专家混合（MoE）是一种稀疏激活架构，能够在保持每个 token 计算量较低的同时扩展大语言模型（LLM）的容量。由于其操作占用了大量计算时间，这些低效之处会直接增加请求延迟和成本。因此，通过内核工程优化 MoE 操作，是提升性能的最大机会。

融合 MoE 是 MoE 架构的关键优化技术。融合 MoE 系统将 token 路由、数据重排和专用网络计算合并为 GPU 上的单个批量操作。这取代了将路由和专家计算作为独立的顺序循环操作执行的方式。虽然这带来了延迟优势，但也增加了内核实现和调试的复杂性。这种架构的高 GPU 利用率需要特定于硬件的优化。

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

用于 WGMMA 利用率的 Swap-AB

为了最大化张量处理吞吐量，NVIDIA Hopper 架构依赖 Warp 组矩阵乘加（WGMMA）指令协议，该协议在 128 线程的协作 warp 组中集体执行。这种操作数不对称性本质上是指令集架构（ISA）级别的属性。在底层的 wgmma.mma_async 汇编定义中，硬件将矩阵行维度（M）固定在 64 的常量规模上，以定义其并行线程映射拓扑。列维度（N）则高度可变。操作数 A 严格映射到这个固定的 M 维度（使用寄存器或共享内存），而操作数 B 映射到可变的 N 维度（通过共享内存描述符）。因此，这些操作数插槽在结构上是不可互换的。因此，峰值硬件瓦片利用率由这些严格的形状编码决定，要求 M 维度等于或超过 64 以防止计算执行通道空闲。

这种结构偏好源于 Hopper 张量核心如何划分内存块。硬件强制第一个操作数的主维度具有固定的最小布局形状，将块高度锁定为精确的 64（M=64）。当数学矩阵在主插槽中具有大而连续的主维度时，Hopper 的内部硬件调度器可以干净地映射和分配数据到物理计算网格。

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

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

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

在解码阶段，系统每次仅处理用户序列中的一个 token。因此，批量空间或 token 行大小非常稀薄，通常下降到 M < 64 的配置，甚至经常降至单行（M = 1）。在这种标准结构布局下：

矩阵 (A) 会转换为一个细长的、碎片化的水平向量（例如 1 x 4096）。

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

由于行数 (M) 极少，矩阵 (A) 在物理上无法填满硬件固定的 (64) 行布局桶（Operand A）。当 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 利用率的 Swap-AB 方案

通过在达到 PTX 汇编层之前编程实现矩阵乘法翻转，我们将软件重新对齐到 Hopper 硅芯片的精确拓扑结构。这种转置在 GPU 上触发了两个即时的架构优势。首先，它填满了物理布局。庞大的专家维度（N = 8192）被提升到第一个操作数位置（B T），由于 (N) 很大，轻松突破了硬件强制的 (64) 行阈值。Tensor Cores 然后以最大本地密度运行，没有浪费的执行周期。其次，它激活了异步预取。Hopper 的专用硬件复制引擎可以从共享内存中直接异步预取内存地址到寄存器，而无需在数学执行块触发前进行 CPU 干预。将连续的大权重块（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寄存器和计算执行文件中处理，随后直接流式传输到最终矩阵目标。这种设计消除了CUTLASS在被迫反复处理中间张量缩放步骤时产生的缓慢全局VRAM读写往返问题。

此外，这种寄存器级融合特别配合了围绕可变令牌分桶设计的动态调度布局，解决了传统编译模板的严重结构性缺陷。CUTLASS原生优化针对密集型、大尺寸且可预测的矩阵维度。当遇到生成阶段常见的流体超薄批次配置（M < 64或M = 1）时，其默认的块分片和线程映射结构会变得极度低效。结果导致warp利用率严重不足或产生高昂的内存填充开销。相反，Triton MoE后端擅长动态地将这些小而碎片化的令牌数组跨过活跃专家权重的物理边界进行映射。这使得硬件内部预取流和Tensor Core流水线即使在处理极低批次规模的高不稳定工作负载时仍能保持完全饱和。

更新路由策略只需纯配置变更，无需修改任何内核代码。但这一改进使H200上FP8量化融合MoE模型的端到端吞吐量提升了10%。该改进已包含在AWS Large Model Inference（LMI）v28容器中。

Triton JIT重新编译消除

GPU推理中最影响性能的路径之一也是最不易察觉的：不必要的内核重新编译。作为vLLM许多自定义内核底层的编译框架，Triton默认会对整数参数进行特化处理。每次内核被调用时，若某个特化参数出现前所未有的值，Triton就会触发新的即时编译（JIT）。在序列长度和批次大小持续变化的服务工作负载中，这会持续产生编译开销，显著增加延迟，尤其是首令牌生成时间（TTFT）。

在MoE模型服务中，_fused_inv_rope_fp8_quant_per_head内核（负责融合旋转位置嵌入和FP8量化）会针对每个新的num_tokens值重新编译。向可变长度参数添加do_not_specialize注解后，完全消除了这些重新编译，在8路张量并行（PR #42849）的H200上使TTFT提升了9%。

在多 GPU 集群中，所有 8 个 GPU 必须实现微秒级的精确同步执行。如果某个 GPU 因本地 token 批次数量变化而需要暂停重新编译其本地 _fused_inv_rope_fp8_quant_per_head 内核，其余 7 个 GPU 将处于空闲状态等待其完成。通过消除整个集群范围内的重新编译，我们优化了多 GPU 协调开销，使用户首次响应延迟提升了 9%。其核心原则是将编译器作为首要级优化表面进行处理。理解 Triton 如何对语义进行特化，以及识别参数是否为无效或变量，可以在最小代码修改的情况下实现延迟优化。

我们已将此优化提交到 vLLM，该功能在 0.22.0 及以上版本中可用。

结论

通过与 vLLM 社区的合作，我们实现了融合 MoE 模型的多项内核优化并开源。这些优化已应用于 Amazon SageMaker AI 的 LMI 深度学习容器。这些优化共同带来了最高 10% 的吞吐量提升和最高 9% 的 TTFT 速度提升。要开始自定义模型托管，请参阅 Amazon SageMaker AI 托管文档。

作者简介

'"`