SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

文章标题: SonicMoE: 通过IO和Tile感知优化加速MoE
作者: Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, and Tri Dao
机构: 普林斯顿大学, 加州大学伯克利分校, Together AI

---

A1 主要贡献

核心问题与研究目标

混合专家（MoE）模型已成为在不显著增加计算成本的情况下扩展语言模型参数规模的关键技术。然而，现代MoE模型呈现出高专家粒度（专家中间维度更小）和更高稀疏性（激活的专家数量不变，总专家数量更多）的趋势，这虽然能提升每FLOP的模型质量，但也带来了硬件效率低下的问题：
1. 激活内存占用增加：对于细粒度的MoE模型，激活内存大小通常与激活的专家数量成线性关系，导致内存占用激增。
2. 硬件效率降低：细粒度专家导致算术强度降低和IO成本增加。
3. 计算资源浪费：对于高稀疏性的MoE模型，Grouped GEMM（分组通用矩阵乘法）内核中的填充操作导致计算浪费。

现有SOTA的MoE内核（如ScatterMoE, MoMoE）并未针对这些高IO成本进行优化，导致训练吞吐量显著下降。本文旨在通过MoE架构、GPU内核和路由方法的协同设计，解决上述效率问题，提升MoE模型的训练效率。

创新点与主要贡献

本文提出了SonicMoE，一个硬件和模型架构协同设计的解决方案，以解决MoE训练的效率问题，主要贡献如下：
* 实现最小激活内存占用的MoE训练算法：通过分析MoE粒度对前向和反向传播的影响，本文发现增加粒度会导致反向传播所需的激活内存线性增加。基于此，本文重新设计了计算图，在不增加FLOPs的情况下，避免为路由器梯度计算缓存激活值。这使得SonicMoE在细粒度7B MoE模型中，每层激活内存使用量减少了高达45%。
* 通过重叠IO与计算实现SOTA训练吞吐量的高效MoE内核：本文指出，增加粒度和稀疏性会使MoE越来越受内存带宽限制。为了缓解这一瓶颈，本文利用GEMM和IO操作的异步性，通过重叠它们来最大化吞吐量。对于同一个细粒度7B MoE模型，与高度优化的DeepGEMM基准相比，本文方法在前向传播中相对速度提升了43%；在反向传播中，相比SOTA的ScatterMoE和MoMoE，速度分别提升了83%和115%。
* 提出“令牌舍入”路由算法消除稀疏MoE的计算浪费：本文引入了一种即插即用的路由算法，该算法将每个专家的令牌数量“舍入”到Grouped GEMM内核所用tile大小（如128）的倍数。这种舍入方法在尽可能保持原始令牌到专家分配的同时，减少了因填充而浪费的计算。该算法确保每个专家的令牌数与原始top-K结果的偏差最多不超过一个tile。在保持相同预期总令牌数的同时，有效消除了填充浪费，并在高稀疏MoE训练中保持了稳健的令牌选择准确性。

图1: 左图显示，即使专家粒度（d/n，d为嵌入维度，n为专家中间维度）增加，SonicMoE每层的激活内存占用保持不变，比其他基准内存效率高0.20-1.59倍。右图显示，SonicMoE的前向计算吞吐量平均达到上限（cuBLAS BMM + 激活 + cuBLAS BMM + 聚合 on H100）的88%（最高91%，最低86%）。注意，cuBLAS上限基准不包括路由器计算。这里使用了一个30B MoE配置，微批次大小为32768个令牌，激活专家数/总专家数从左到右分别为2/32, 4/64, 8/128, 和 16/256。

表1: MoE扩展趋势：此处展示了前沿开源模型的激活率（每个令牌激活的专家数K / 总专家数E）和专家粒度（模型嵌入维度d / 专家中间尺寸n）。MoE稀疏度计算不包括共享专家。趋势表明，新的开源MoE模型倾向于更细粒度和更稀疏。

---

A3 背景知识

2.1 使用Grouped GEMM的MoE

现代GPU的GEMM计算。现代GPU支持Tensor Cores，这是具有高矩阵乘法吞吐量的专用硬件单元【索引36，NVIDIA H100 Tensor Core GPU Architecture，2022】。一个GEMM（通用矩阵乘法）【索引27，Basic linear algebra subprograms for Fortran usage，Lawson et al.，1979，ACM Transactions on Mathematical Software (TOMS)】内核通常有3个阶段：prologue（开始输入加载）、mainloop（持续加载输入并计算GEMM）和epilogue（在GEMM输出上进行杂项IO/数学运算）。内核会将计算分块（将大矩阵划分为小块），并可选择性地填充维度，使计算与硬件允许的tile大小对齐。本文遵循大多数BLAS库【索引27，Basic linear algebra subprograms for Fortran usage，Lawson et al.，1979，ACM Transactions on Mathematical Software (TOMS)】中的标准GEMM表示法：对于$C = AB$，我们有$A \in R^{M \times K}$，$B \in R^{K \times N}$，$C \in R^{M \times N}$，问题形状为$(M, N, K)$。这种表示法被实现高效CUDA GEMM的CUTLASS【索引37，CUTLASS: CUDA Templates for Linear Algebra Subroutines，NVIDIA，2025】所采用。

Hopper GPU上的GEMM范式。在NVIDIA Hopper GPU上，GEMM是以生产者-消费者范式异步执行的【索引49，FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision，Shah et al.，2024】。生产者专门负责将数据tile从高带宽内存（HBM）或逻辑上的全局内存（GMEM）加载到共享内存（SMEM），而消费者warpgroup则负责GEMM计算。在prologue和mainloop阶段，生产者warpgroup获取一个数据tile并缓存到专用流水线，消费者warpgroup从该流水线的缓存tile中读取数据，执行分块矩阵乘法（MMA）并在GEMM的K维度上累加。mainloop之后，进入epilogue阶段，消费者warpgroup对最终的MMA结果进行后处理（如应用激活函数并将结果写回HBM）。

MoE块与Grouped GEMM。一个MoE块通常由一个令牌路由器和多个通常大小相等的小型子网络（称为“专家”）组成。路由器负责将令牌分派给专家，然后由特定专家用于计算。该层中所有专家的输出随后被聚合并传递到下一层。MoE计算可以使用Grouped GEMM（一个具有可能不同${M, N, K}$维度的GEMM列表）来执行。算法1展示了使用Grouped GEMM运行MoE前向传播的过程。

图2: MoE计算通常需要一个Grouped GEMM。每个专家从输入张量的不同位置收集输入（上图），或者读取分组输入数组上的连续块（下图）。该图改编自Tan等人（2024）的图2。

Varlen Grouped GEMM。如算法1所示，在前向传播（以及反向激活梯度计算）期间，每个专家会接收到可变数量的令牌。此时执行的Grouped GEMM操作具有固定的$(N, K)$维度（专家权重矩阵）但可变的$M$维度（令牌维度），我们称之为“varlen-M Grouped GEMM”。在反向权重梯度计算期间，嵌入维度（反向为$M$）和中间隐藏大小（反向为$N$）是恒定的，而我们在令牌维度（$K$）上进行归约，我们称之为“varlen-K Grouped GEMM”。对于每个Grouped GEMM，输入通常是从不同位置收集的或连续打包的，如图2所示。例如，在算法1中，up-proj的输入是收集的，而down-proj的输入已经是连续打包的。

2.2 MoE计算

算术强度与硬件瓶颈。算术强度，定义为FLOPs与传输字节数（IO）的比率，是量化一个内核是内存受限（内核运行时间由内存IO成本主导）还是计算受限（内核运行时间由计算吞吐量主导）的指标。

MoE计算分解。一个带有SwiGLU激活的专家e的标准MoE计算可以分解为以下部分：
$H_e = \text{up-projection}(X_e) = X_e W_{1,e} : \mathbb{R}^{T_e \times d} \rightarrow \mathbb{R}^{T_e \times 2n}$
$A_e = \text{SwiGLU}(H_e) : \mathbb{R}^{T_e \times 2n} \rightarrow \mathbb{R}^{T_e \times n}$
$Y_e = \text{down-projection}(A_e) = A_e W_{2,e} : \mathbb{R}^{T_e \times n} \rightarrow \mathbb{R}^{T_e \times d}$
其中$X_e \in R^{T_e \times d}$表示专家e接收到的输入。

算术强度分析。在这里，up-projection使用$2T_e \cdot 2n \cdot d$ FLOPs，HBM内存传输字节为$2T_ed + 2 \cdot 2n \cdot d + 2T_en$（这里忽略$H_e$的写入）。类似地，down-projection使用$2T_end$ FLOPs，字节数为$2T_en + 2nd + 2T_ed$。假设$\rho = K/E$是MoE激活率，$G = d/n$是粒度，并且路由均匀，即$T_e = T\rho$，则一个专家前向传播的算术强度（忽略$H_e$的写入）为：

$$\frac{2 T_e \cdot 2 n \cdot d+2 T_e n d}{4 T_e n+6 n d+4 T_e d}=\frac{3}{\frac{2}{d}+\frac{2}{n}+\frac{3}{T_e}}=\frac{3}{\frac{2+2 G}{d}+\frac{3}{T_\rho}}$$

对于特定的模型大小（$d$恒定），可以看出增加粒度（增加$G$）或增加稀疏性（减小$\rho$）会导致算术强度下降。这是由IO成本随专家粒度线性扩展引起的，如图3所示。因此，对于细粒度MoE（高$G$）的情况，通过最大限度地减少IO访问和隐藏IO延迟来解决增加的IO成本变得越来越重要。我们将在第3节研究内存高效的MoE内核设计，并在第4节讨论减少IO访问和延迟的技术。

图3: 在isoFLOPs训练下，从1.4B到120B的MoE配置中（配置见表9a），单层MoE前向传播的IO成本随专家粒度的变化。我们保持MoE激活率$\rho = K/E$和每个MoE层的参数数量$3dnE$恒定。当我们扩大专家粒度$d/n$时，我们相应缩小专家中间尺寸$n$，同时保持$nE$和$nK$不变。

现有MoE内核设计的局限性。目前有多种MoE实现可用：ScatterMoE【索引53，Scattered mixture-of-experts implementation，Tan et al.，2024】，MoMoE【索引10，MoMoE: Memory optimized Mixture of Experts，Costin et al.，2025】，MegaBlocks【索引18，Megablocks: Efficient sparse training with mixtureof-experts，Gale et al.，2023】，和Megatron【索引51，Megatronlm: Training multi-billion parameter language models using model parallelism，Shoeybi et al.，2019】。然而，它们并未专门针对细粒度MoE进行优化，这类MoE的IO成本随专家粒度线性增加，如图3所示。相比之下，我们的内核设计SonicMoE，旨在最小化IO成本对训练吞吐量的影响。在第4节和图14中，我们展示了当专家粒度$G$增加时，由于IO感知优化，SonicMoE相比现有MoE内核设计表现出更大的相对加速。我们将在附录B中详细阐述SonicMoE与先前MoE内核的技术差异，并在表2中提供概览。

2.3 MoE路由方法

路由方法概述。在MoE中，路由决定了为每个令牌激活哪些专家。令牌选择（Token Choice, TC）路由是MoE模型的默认方法【索引50，Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer，Shazeer et al.，2017】，即每个令牌独立选择要激活的专家。我们通常使用top-K TC路由，其中令牌$t$的路由决策是$TopK_{e \in [E]}(S_{t,e}, K)$，$S_{t,e}$是令牌$t$的专家分数。除了top-K，Huang等人【索引22，Harder Task Needs More Experts: Dynamic Routing in MoE Models，2024】引入了token-choice top-P路由，以便在训练中灵活分配计算资源，但这会引入激活专家数量和每令牌消耗FLOPs的不确定性。Zeng等人【索引66，AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models，2024b】也提出了类似的想法，使用“空专家”动态调整激活的专家数量。

专家选择（EC）路由及其挑战。除了TC路由，专家选择（Expert Choice, EC）路由被提出来以避免专家并行中的负载不平衡问题【索引71，Mixture-of-experts with expert choice routing，Zhou et al.，2022】，它让专家来选择令牌。然而，EC路由不能直接用于推理，因为它与自回归解码不兼容，在推理时切换回TC会导致不匹配。此外，EC会通过未来令牌信息泄露破坏因果关系【索引58，Auxiliary-loss-free load balancing strategy for mixture-of-experts，Wang et al.，2024】。为解决EC路由的推理问题，Raposo等人【索引44，Mixture-of-depths: Dynamically allocating compute in transformer-based language models，2024】引入了一个辅助损失来促进TC和EC路由结果的一致性，或者训练一个辅助路由器来显式预测EC路由器的路由结果，并在推理时使用这个辅助路由器。

本文提出的路由方法。本文提出了一种新颖的、感知Grouped GEMM tile的令牌舍入方法，该方法将每个专家接收的令牌数（“专家频率”）舍入到Grouped GEMM tile大小的附近倍数，并且每个专家最多只改变一个tile的令牌。这种方法有效减少了稀疏MoE训练期间由Grouped GEMM填充引起的FLOPs浪费，同时保持了训练后MoE模型的推理质量。有类似的工作提出丢弃和重路由令牌，包括Rectify-Router【索引65，Turn Waste into Worth: Rectifying Top-k Router of MoE，Zeng et al.，2024a】，但它们不关注Grouped GEMM的tile结构。其他工作如TMA-adaptive FP8 Grouped GEMM【索引17，TMA-Adaptive FP8 Grouped GEMM: Eliminating Padding Requirements in Low-Precision Training and Inference on Hopper，Fu et al.，2025】专注于减少与填充相关的加载流量，但未解决GEMM计算中因tile大小未对齐而浪费的FLOPs。

---

A2 方法细节

3 内存高效的MoE算法

3.1 SonicMoE的MoE内核概述

SonicMoE的计算流程。SonicMoE中的MoE计算会启动8个内核：前向传播有up-proj ($A$)、down-proj ($Y$)和专家聚合($O$)内核；反向传播有针对$dH$ (down-proj)、$d\tilde{X}$ (up-proj)、$dX$ (跨专家聚合$d\tilde{X}$)的激活梯度内核，以及权重梯度内核$dW_1$和$dW_2$。图4展示了这8个内核的计算工作流。我们提供了一个高效的TC top-K路由器，以及一个接受任意路由输入的接口。但需要注意的是，SonicMoE的MoE计算与MoE路由器选择无关，因此与任意路由器逻辑兼容。

模块化设计。SonicMoE的MoE计算实现是高度模块化的：它仅包含(1) 带有模块化融合的优化Grouped GEMM内核和(2) 优化的专家聚合内核。主机根据情况分派最佳的GEMM配置和加载/存储策略来启动上述8个内核。尽管具有如此高的模块化，SonicMoE仍然展现出SOTA的训练吞吐量和最小的激活内存使用量，我们将在下文描述。

图4: SonicMoE启动的8个内核的计算工作流，由黄色框分组。前向和反向计算分别启动3个和5个内核。指向黄色圆圈的箭头表示从HBM加载到SRAM的变量，向外的箭头表示存储到HBM的变量。我们将HBM上的所有变量框进行着色，紫色框表示前向和反向的输出，蓝色框表示中间变量或权重($W_1, W_2$)。我们将所有缓存的激活$X, H, \pi, S$用红色标注。算法2正式描述了SonicMoE的前向传播，算法3和5描述了反向传播。

3.2 激活内存效率

避免随粒度扩展的激活内存。MoE前向和反向计算的FLOPs为$(6+12)T nKd$。对于给定的$T, d$，我们需要保持$nK$恒定以维持恒定的FLOPs。因此，增加粒度需要减小$n$并按比例增加$K$。因此，任何内存大小为$O(TKd)$的激活值都不应为反向计算而缓存，以避免激活内存随粒度扩展。对于像ScatterMoE这样的现有MoE内核，激活值随专家粒度线性扩展。激活值$Y$（down-proj输出）和$X_e$（收集的$X$）的大小为$TKd$，避免缓存它们可以消除激活内存对粒度的依赖。我们避免将$dY$（$Y$的梯度）和$dO_e$（收集的$dO$）写入HBM，因为它们会增加反向计算期间的峰值激活内存。

内存优化策略。
* 对于$X$和$dO$，我们将gather操作与HBM加载融合，从而无需在HBM中物化和缓存激活。我们在图6和图20中展示，这种gather融合显著提高了细粒度MoE的吞吐量。
* 一个简单的实现计算$dS$和$dH$会需要$Y$和$dY$。相反，我们确定了一个替代的计算路径来计算$dS$和$dH$而不增加FLOPs。这是通过将$dS$和$dH$展开成一个不涉及使用$Y$和$dY$的方程来实现的，如附录C所示。SonicMoE的$dH$内核如算法3所示。

最终激活内存占用。因此，我们只缓存$X$和$H$以及路由元数据，每层总大小为$2Td + 4TKn$字节。这个激活内存使用量与具有相同数量激活参数的密集模型相同，这是在不进行GEMM激活重计算的情况下反向计算所需的最小激活内存。在图13中，我们分析了SonicMoE在7B MoE训练配置下的激活内存，并证明了SonicMoE的激活内存与专家粒度无关。更多从1.4B到120B的结果包含在图13中。

4 IO感知的内核设计

核心思想。细粒度MoE的表达能力源于每个令牌专家选择的多样性，这反过来导致IO成本随专家粒度线性扩展（图3）。为了维持高吞吐量，我们需要最大程度地(1) 通过融合减少IO访问，(2) 将IO延迟与计算重叠。我们将在4.1.1和4.1.2节分别研究令牌收集与计算的融合，以及数学和IO与epilogue的融合。然后，在4.2节描述将MMA与IO重叠的技术。最后在4.3节研究SonicMoE的top-K排序内核。附录B中，我们将SonicMoE与其他MoE内核设计进行了比较，总结在表2中。

表2: SonicMoE与先前MoE内核的比较。✓表示内核实现了该功能或语义上类似的功能，✗表示内核缺少该功能。“NA”表示该功能超出了预期范围。我们对Megatron使用GroupedMLP，对MegaBlocks使用ParallelDroplessMLP。更多讨论见附录B。

4.1 SonicMoE的Grouped GEMM

融合策略。SonicMoE建立在高效的varlen-M和varlen-K Grouped GEMM之上。在Grouped GEMM内部，我们将gather操作与激活加载融合（4.1.1），并将SwiGLU/dSwiGLU/dS与epilogue融合（4.1.2）。Gather融合帮助SonicMoE比那些需要单独gather内核的MoE内核设计（如MegaBlocks、Megatron和DeepGEMM++）更快，DeepGEMM++是基于DeepGEMM库【索引69，DeepGEMM: clean and efficient FP8 GEMM kernels with fine-grained scaling，Zhao et al.，2025b】的优化MoE前向传播实现。Epilogue融合使SonicMoE在反向传播中比ScatterMoE更快。这些融合减少了不必要的IO访问，并可以与计算MMA重叠，我们将在4.2节讨论。

4.1.1 Gather与HBM加载的融合

实现方式。SonicMoE的Grouped GEMM接受连续打包的输入或从不同位置收集的输入，如图2所示。对于后者，我们将输入gather与从全局内存（GMEM，通常是HBM）到共享内存（SMEM）的输入加载相融合，这样我们就可以批量处理它们以在Tensor Core上执行GEMM【索引10，MoMoE: Memory optimized Mixture of Experts，Costin et al.，2025；53，Scattered mixture-of-experts implementation，Tan et al.，2024】。这包括(1) 获取每个专家路由的令牌索引，然后(2) 使用这些索引通过Blackwell和Hopper的cp.async指令来收集激活。对于第二步通常没有更好的替代方案，但同步的索引获取仍然可以通过生产者warp之间的预取和协作获取来优化。我们在图18中展示了我们的策略。

在Hopper GPU上的优势。如图6所示，Gather融合为SonicMoE在H100上提供了相对于现有MoE内核设计（如DeepGEMM）的主要优势。尽管DeepGEMM的varlen-M Grouped GEMM内核经过高度优化，但DeepGEMM假设输入已经是连续打包并填充到128的倍数，这需要在Grouped GEMM之前单独启动一个内核来进行gather和pad。在图6中，即使我们提供了一个优化的gather内核，且DeepGEMM的varlen-M Grouped GEMM也高度优化，但收集$X$的大量IO（$2TKd$字节）仍然使DeepGEMM++（定义见图6）比SonicMoE慢。

反向传播中的Gather融合。在反向传播中，up-proj和down-proj的权重梯度（分别为$dW_1$和$dW_2$）需要收集$X$和$dO$，$dH$的激活梯度也需要收集$dO$。尽管反向传播有更多内核需要gather操作，但现有方法包括ScatterMoE【索引53，Scattered mixture-of-experts implementation，Tan et al.，2024】和MoMoE【索引10，MoMoE: Memory optimized Mixture of Experts，Costin et al.，2025】在前向传播时融合了gather，但在反向传播时仍然启动一个单独的gather内核。融合这个gather操作减少了$2TKd$字节的IO成本，并削减了细粒度MoE训练时间的很大一部分。例如在图6中，ScatterMoE和MoMoE在反向传播中的2次gather分别消耗了19.6%和20.6%的总反向传播时间，这甚至比它们的up-proj权重梯度$dW_1$内核时间还要长。

在Blackwell GPU上的实现。在撰写本文时，SonicMoE支持Blackwell GPU上的varlen-M Grouped GEMM及其gather融合。在Blackwell GPU上，当使用2-CTA集群（图5）进行GEMM计算时，与cp.async的gather融合遇到了一个架构挑战。cp.async指令只能在同一个CTA内发信号通知完成。然而，Blackwell的2-CTA GEMM要求leader CTA（CTA 0）中的MMA指令等待来自两个CTA的gather完成。为了解决这个限制，CTA 1需要一个专用的中继warp，它接收cp.async完成信号，并使用集群级同步原语（例如，具有集群范围的mbarrier）将其转发给CTA 0的MMA warp。这种中继机制增加了调度复杂性，但实现了跨2-CTA集群的高效gather融合，保持了Grouped GEMM的高吞吐量。

图5: 在Blackwell GPU上使用2-CTA集群的cp.async gather融合的流水线结构。

4.1.2 Epilogue融合

设计选择。我们利用epilogue计算来最大程度地减少不必要的IO访问，具体设计选择如下：
* SwiGLU和dSwiGLU融合：我们将SwiGLU和其反向传播分别与前向up-proj和反向down-proj激活梯度内核的epilogue融合【索引10，MoMoE: Memory optimized Mixture of Experts，Costin et al.，2025】。在图6中，尽管DeepGEMM++拥有高度优化的Grouped GEMM和SwiGLU内核，其up-proj（0.49ms, 629TFLOPS）和SwiGLU（0.11ms, 2.88TB/s）的总时间（0.60ms）仍然比SonicMoE的up-proj（0.55ms, 559TFLOPS）要长，尽管SonicMoE除了SwiGLU外还额外进行了Gather融合。
* 在反向down-proj激活梯度(dH)内核的epilogue中计算dH和dS：这种重度epilogue融合为SonicMoE带来了相对于替代设计的显著加速。我们的dH内核（0.47ms, 328TFLOPS）产生的输出与ScatterMoE的down-proj act（0.43ms, 364TFLOPS）、dS（0.24ms）和dSwiGLU（0.33ms）组合在一起（0.99ms）相比，总时间要少得多（见图6中的7B MoE训练）。此外，SonicMoE也比DeepGEMM++（0.57ms）快，后者启动了一个高效的Grouped GEMM（0.32ms, 480TFLOPS）和一个单独的优化内核（0.25ms, 2.43TB/s）来一起计算dSwiGLU和dS。

图6: 在H100上，7B MoE训练配置(T, d, n, E, K) = (24576, 1536, 256, 128, 8)下，不同MoE内核的运行时分解(ms↓)。我们为内存受限内核标注了内存带宽(TB/s↑)，为分组GEMM内核标注了计算吞吐量(TFLOPS↑)。此图按内核运行时语义分组，一个块可能包含多个实际内核计时结果。例如，左子图的“router related”包括路由器GEMM和路由元数据计算时间。此外，我们未考虑内核间的CUDA流气泡时间。我们对Megatron使用GroupedMLP，对MegaBlocks使用ParallelDroplessMLP。DeepGEMM++是基于DeepGEMM SM90 BF16 Grouped GEMM内核构建的最佳可能MoE实现，无需修改其源代码。

dS的计算方式。在附录C.1中，我们证明了SonicMoE的$dS = \langle dA, A' \rangle = \langle dA, \text{Broadcast}(s)A \rangle$是细粒度MoE在计算和激活内存上都高效的选择。然而，ScatterMoE和MoMoE都选择计算$dS$为$\langle dO, Y \rangle$，这需要额外的$2TKd$ HBM加载成本，并需要缓存$2TKd$字节的激活内存。在图6（右子图）中，ScatterMoE为$dS$启动了一个单独的内核（0.24ms），而MoMoE将$dS$与up-proj激活梯度融合（总共1.58ms, 196TFLOPS），这比SonicMoE的up-proj激活梯度（0.50ms, 618TFLOPS）花费的时间要长得多。

重度Epilogue融合的性能。反向down-proj激活梯度dH内核中重度epilogue融合的吞吐量得益于异步IO和MMA的重叠，我们将在4.2节详细阐述。这种重叠帮助SonicMoE即使在dH内核中有重度epilogue融合（加载H和S，计算dH、dS和作为dW2输入的A'）的情况下，也能同时维持合理的训练吞吐量（328TFLOPS）和内存带宽（2.14TB/s）。

4.2 GEMM MMA与异步IO的重叠

Hopper GPU上的Ping-Pong调度。在NVIDIA Hopper GPU中，GEMM以生产者-消费者范式异步执行【索引49，FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision，Shah et al.，2024】。假设我们有两个消费者warpgroup，我们可以让它们协作地以大tile大小发出WGMMA指令，或者以较小tile大小将一个warpgroup的IO与另一个warpgroup的GEMM重叠。完成后，我们交换warpgroup的角色（有效交错IO和GEMM）。这通常被称为Hopper GPU上的Ping-Pong调度【索引49，FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision，Shah et al.，2024；61，Deep Dive on CUTLASS Ping-Pong GEMM Kernel，Wright and Hoque，2024b】，如图7所示。

Ping-Pong调度的应用。Ping-Pong调度对于在重度epilogue下维持高Tensor Core吞吐量特别有用。例如，down-proj前向Y内核的epilogue相对于mainloop有大量的HBM存储IO（$2TKd$字节）。在down-proj激活梯度（dH）内核的epilogue中，我们需要加载H（$4TKn$字节）并执行多个激活和归约操作来计算和存储dH、dS和A'作为dW2的输入。我们注意到，将MMA与IO重叠以及Ping-Pong调度的概念在其他地方（如Flash Attention 3【索引49，FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision，Shah et al.，2024】）是已知的，但将Ping-Pong调度应用于解决细粒度MoE内核设计中日益增加的IO成本是新颖的。

图7: SonicMoE在Hopper GPU上的Ping-Pong warpgroup调度。绿色箭头表示一个消费者warpgroup发出epilogue开始的信号，另一个消费者warpgroup可以继续进行MMA。此步骤完成后，两个消费者warpgroup的角色交换。SonicMoE主要将Ping-Pong用于前向down-proj Y内核和反向down-proj激活梯度dH内核，因为它们都有重度epilogue。在dH内核中，SonicMoE在epilogue期间有异步TMA加载，生产者warp需要为收集dO和用TMA加载专家权重发出cp.async。此图改编自Wright和Hoque（2024a）关于Ping-Pong调度的博客。

与DeepGEMM的比较。DeepGEMM SM90 BF16 varlen-M Grouped GEMM内核没有实现Ping-Pong调度。这种设计选择适用于轻量级epilogue（例如，up-proj前向），但在down-proj前向的重度epilogue情况下表现不佳（413 TFLOPS和2.15 TB/s vs. SonicMoE的485 TFLOPS和2.52 TB/s），如图6所示。在图19中，SonicMoE的down-proj平均比DeepGEMM高出10.0%的TFLOPS。

异步TMA操作。除了Ping-Pong调度，SonicMoE还依赖异步TMA操作来执行GMEM到SMEM的加载和SMEM到GMEM的存储。我们将以下异步IO与MMA操作重叠：
* dH内核epilogue期间的异步TMA加载：在dH内核的epilogue中，我们需要加载H来从dA计算dH。我们为H的异步TMA加载创建了一个专用的流水线，以便与epilogue阶段的其他epilogue操作重叠。在图7中，消费者warpgroup中的透明TMA块说明了这种异步epilogue加载。
* 前向down-proj Y和反向up-proj激活梯度dX̃内核中的异步TMA存储：SonicMoE对所有6个Grouped GEMM都应用了异步TMA存储。在前向down-proj和反向up-proj激活梯度中，SonicMoE没有将scatter与HBM存储融合，而ScatterMoE和MoMoE都选择将HBM存储与scatter融合。这是因为scatter融合(1)有更多的同步索引获取和地址计算，并且(2)在Hopper GPU上需要一个同步的SMEM到GMEM的存储指令。对于细粒度MoE，同步的GMEM存储会阻塞下一个tile的MMA执行，并大幅降低TFLOPS（约20%），如图8所示。我们还注意到，Ping-Pong warpgroup调度不能完全恢复同步epilogue IO操作的吞吐量下降，因为epilogue消费者warpgroup会被阻塞，直到当前的同步GMEM存储完成才能与MMA warpgroup交换角色。

图8: 说明异步TMA存储（上）具有更高的内存带宽，并且可以自然地与TensorCore MMA重叠，而同步的st.global（下）PTX指令（在Hopper GPU上进行scatter融合所必需）会阻塞下一个Tensor Core MMA tile的执行，导致更长的内核运行时间。图22中“SonicMoE (gemm + gth w. sum)”（TMA存储）平均比“SonicMoE (gemm w. sct + sum)”（st.global）快20.1%，这证实了这一点。因此，SonicMoE不将scatter与HBM存储融合，而是在专家聚合内核中让每个令牌收集专家结果。ScatterMoE和MoMoE均未采用这种设计，因此SonicMoE在前向down-proj内核中平均可实现1.75倍和3.11倍的加速（如图6）。

图9: 存储结果和聚合每个令牌结果的可能策略。SonicMoE选择第一种策略（左），其中每个专家在GEMM epilogue中通过TMA直接存储连续打包的输出。在专家聚合内核中，每个令牌收集并对激活的专家输出求和。ScatterMoE和MoMoE（中）选择在epilogue中将HBM存储与scatter融合，然后启动一个求和内核。我们注意到，每个令牌收集（左）Grouped GEMM结果等同于每个专家分散（中）Grouped GEMM输出。在图22中，我们在SonicMoE上实现了两种策略，并观察到左侧策略比中间策略快17%。也可以在epilogue中融合原子加法以规避专家聚合内核的要求，如右子图所示。然而，这种原子加法操作会产生新问题，如非确定性【索引20，Defeating Nondeterminism in LLM Inference，He and Machines，2025】和数值精度问题（对于BF16原子加法）。此图改编自Tan等人（2024）的图2。

Blackwell GPU上的重叠。在NVIDIA Blackwell GPU上，GEMM内核在精神上使用相同的“Ping-Pong”调度，但实现与Hopper不同。Blackwell引入了Tensor Memory（TMEM），这是每个SM上专用的256KB片上内存【索引38，NVIDIA Blackwell Architecture Technical Brief，NVIDIA，2025b；45，CUTLASS Tutorial: Writing GEMM Kernels Using Tensor Memory For NVIDIA Blackwell GPUs，Colfax Research，2024】。矩阵乘法的累加器结果直接存储在TMEM中而不是寄存器中。这种架构变化允许epilogue warp与MMA warp并发地从不同的TMEM阶段读取和累加结果，从而比Hopper的ping-pong调度实现更好的epilogue和MMA操作重叠。

4.3 高效的MoE top-K排序内核

问题与解决方案。现有的MoE方法如ScatterMoE、MoMoE和MegaBlocks使用PyTorch top-K (torch.topk)来计算每个令牌的专家分配。我们发现PyTorch top-K内核可能占用路由器计算时间的约40%。为此，我们实现了一个高效的top-K内核来减少PyTorch top-K的开销。我们的top-K内核支持$E \le 4096$和$K \le 16$，并针对大量令牌$T$的情况进行了优化。我们还提供了一个在top-K内核内对top-K值进行softmax融合的可选功能。

实现细节。top-K内核接受形状为$(T, E)$的路由器输出，并在$T$上并行化。内核在每一行上使用双调排序（bitonic sort）【索引1，Sorting networks and their applications，Batcher，1968】（对$E$个值进行排序），并选择前$K$列作为排序输出。加载输入后，我们将前$K$列的列索引打包到FP32值的低$\log_2(E)$个尾数位中，并对基础排序情况（值数量$\le 64$）进行了特化，遵循了从最优低延迟排序网络【索引13，bert dobbelaere sorting networks，Dobbelaere，2025】中获得的比较策略。

图10: 排序是在我们将列索引位打包到较低尾数位之后的值上进行的。这种值格式确保了稳定的排序结果。Triton的官方top-K内核遵循类似的格式。

性能优势与稳定性。双调比较和合并发生在同一线程内或通过warp-shuffle在同一warp内。因此，每个交换和合并操作只使用线程内或warp内寄存器。这使得该内核相比其他内核设计（如PyTorch TopK、Triton官方示例和RTop-K）具有更高的内存带宽，如图23所示。由于每行值的分配列索引总是唯一的，将列索引打包到低尾数位后不会有任何相等的数字。因此，SonicMoE的top-K内核总是稳定的，因为在双调比较和合并期间不会有任何平局情况。

5 令牌舍入路由

本节分析了稀疏MoE训练机制下的硬件效率，并指出随着MoE变得更稀疏，由于所谓的“tile量化”效应，浪费在填充GEMM tile上的计算会累积到不可忽视的程度。为此，我们提出了一种新颖的路由方法“令牌舍入”来消除tile量化效应。

5.1 稀疏MoE的训练效率

算术强度与稀疏性。除了粒度，MoE的算术强度还取决于MoE激活率$\rho$，如公式4所示。当我们降低$\rho$时，每个专家接收到的预期令牌数$E_{e \in [E]}T_e = \bar{T}_e = T\rho$也会线性减少，GEMM计算会转向内存受限的状态。

Tile量化效应。现代GPU上的GEMM通常以tile为单位进行计算【索引36，NVIDIA H100 Tensor Core GPU Architecture，2022】，如果$M, N, K$的任何维度不能被tile大小整除，我们总是需要填充到下一个tile大小的倍数。一旦输入的大小（例如每个专家的令牌维度）变小，由填充浪费的TFLOPS就可能变得不可忽视。因此，我们建议使用令牌舍入来避免启动这些额外的tile，从而实现更高效的训练。我们还通过实验证明，我们的令牌舍入方法在不影响模型质量的同时，实现了更高的训练吞吐量。

图12: 稀疏MoE的tile量化效应演示。TR中的舍入子程序做出丢弃或填充令牌的二元决策，以保证每个专家接收到$M_{tile}$倍数的令牌数量。

图11: 在图16右下角两个子图所示的$T=16k, d=4k, n=1k, K=4$的MoE前向和反向传播中，由填充造成的FLOPs浪费。

5.2 令牌舍入路由

算法描述。因此，我们提出使用令牌舍入（TR）方法，这是一个两步排序算法，如算法4所示。令牌舍入算法首先计算常规的令牌选择（TC）路由结果，然后对每个专家的令牌按路由器分数进行排序，类似于EC的排序步骤。然后我们选择要么丢弃在第一步TC top-K路由中选择的令牌，要么在第二步排序中填充额外的令牌。在这两步之间，我们处理路由权重矩阵，使得TC令牌总是优先于EC令牌。这样做是为了确保丢弃或填充只影响每个专家的最后一个输入tile。

算法4：令牌舍入路由
输入：$X \in R^{T \times d}$；专家数$E$和每个令牌期望激活的专家数$K$；tile大小$M_{tile}$；路由器分数$S \in [0, 1]^{T \times E}$；round_and_sparsify决定向上或向下舍入。
输出：$M_{tile}$舍入后的路由器分数 $\lfloor S \rceil_{M_{tile}}$

1. // (1) 使用TC top-K路由计算每个令牌的专家
2. // (2) 计算每个专家接收到的令牌频率及其$M_{tile}$舍入的倍数
3. // (3) 构建Top-K偏好的S'用于专家级排序
   // 确保非top-K条目更小
   $S'e \leftarrow S_e - 1$
   并行处理 $t \in [T]$ 和 $k \in [K]$:
   $S'$}(t, k)} \leftarrow S_{topK, t, k
4. // (4) 对每个专家进行令牌舍入
   并行处理 $e \in [E]$:
   // 令牌排序和排序后的分数
   $\pi_e, s_e \leftarrow \text{sort}(S'e)$
   // 对专家e进行令牌舍入
   $\pi'_e, s'_e \leftarrow \text{round_and_sparsify}(\pi_e, s_e, f_e, \lceil f_e \rceil)$}}, \lfloor f_e \rfloor_{M_{tile}
   $\lfloor S \rceil_{M_{tile}, e} \leftarrow \text{Gather}(S, \pi_e)$

舍入子程序。令牌舍入需要一个round_and_sparsify子程序来做出丢弃或填充的二元决策。我们的默认选择是将专家频率舍入到最近的$M_{tile}$倍数：如果$\lceil f_e \rceil_{M_{tile}} - f_e$小于$f_e - \lfloor f_e \rfloor_{M_{tile}}$，我们选择填充EC选择的令牌。我们在表6中进行了消融研究，发现(1)我们的TR算法对底层的舍入子程序相当鲁棒，(2)这种简单的最近舍入策略通常足以产生出色的任务性能。关于不同舍入子程序的更详细讨论包含在附录F.2中。

训练与推理质量。这个简单的算法保证了对于每个专家，与令牌选择路由的最大偏差最多为1个tile。我们发现这个特性即使在稀疏MoE训练机制下也具有惊人的鲁棒性能，并且可以在稀疏MoE训练设置下作为令牌选择的替代品，如表3所示。我们还在表7和8中对微批次大小$T$和tile大小$M_{tile}$对使用TR训练的MoE模型质量的影响进行了消融研究，我们发现当$\bar{T}e/M \ge 2$时，令牌舍入路由通常是鲁棒的。

训练吞吐量。TR保证没有tile量化效应，在6.3.3节中，我们展示了在高稀疏MoE训练机制中，TR的训练吞吐量始终高于常规TC top-K，并且当我们扩大$E$而保持$K$不变时，内核运行时间可实现16%的TFLOPS提升。

---

A4 实验环境

• 数据集：

  • 预训练：使用去重版的FineWeb-Edu【索引2，SmolLM-Corpus，Ben Allal et al.，2024】作为预训练语料库，上下文长度为4096。
  • 评估：在11个下游任务上进行评估，包括WinoGrande【索引47，WinoGrande: An Adversarial Winograd Schema Challenge at Scale，Sakaguchi et al.，2020】、SocialIQA【索引48，SocialIQA: Commonsense Reasoning about Social Interactions，Sap et al.，2019】、PIQA【索引4，PIQA: Reasoning about Physical Commonsense in Natural Language，Bisk et al.，2020】、HellaSwag【索引63，HellaSwag: Can a Machine Really Finish Your Sentence?，Zellers et al.，2019】等。

• 模型架构：

  • 基于OLMoE【索引35，OLMoE: Open Mixture-of-Experts Language Models，Muennighoff et al.，2025】的基础架构构建MoE模型。
  • MoE层使用SwiGLU激活函数。
  • 使用辅助负载均衡损失，系数为0.01。
  • LM head的权重与token embedding矩阵的权重绑定。
  • 具体参数（如层数、头数、d, n, E, K）见附录H中的表10。

• 硬件配置：

  • 实验主要在NVIDIA H100 GPU上进行。
  • 端到端训练实验使用了64台H100，并与96台H100上的基线进行比较。

• 软件配置：

  • 代码实现：SonicMoE主要使用CuTe-DSL【索引39，NVIDIA CUTLASS Documentation，NVIDIA，2025c】编写，并提供PyTorch接口。
  • 依赖代码库：端到端吞吐量测试使用了lm-engine代码库【索引33，LM Engine: A Hyper-Optimized Library for Pretraining and Finetuning，Mishra，2024】。令牌舍入的质量评估实验使用了OLMoE代码库。
  • 分布式训练：端到端实验使用FSDP-2，单节点内（8x H100）使用ZeRO-3进行模型分片，并在节点间复制分片单元。

---

A4 实验结果

6.1 SonicMoE的激活内存

实验内容：在不同模型规模（1.4B至120B）和不同专家粒度下，测量并比较SonicMoE与基线方法（ScatterMoE, MoMoE, MegaBlocks, Megatron, DeepGEMM++）单层MoE的峰值激活内存占用。
实验结果：

• 如图13所示，SonicMoE在所有模型规模上都具有最低的激活内存占用。
• 对于7B模型（n=256），SonicMoE的内存使用比ScatterMoE减少了45%。对于30B和120B模型，差距更大，在120B规模下，每层节省超过3GiB内存。
• 如图1（左）所示，SonicMoE的激活内存占用随专家粒度的增加保持恒定，而其他基线则线性增长。

图13: 不同模型规模（1.4B–120B）下每层的峰值激活内存使用情况。MegaBlocks不支持小的n。基准配置列于表9b。我们只缓存X、收集的Xe、每个专家e的He以及路由元数据，这是在不进行GEMM重计算的情况下进行反向计算所需的最小激活内存量。

6.2 SonicMoE的训练吞吐量

6.2.1 完整的正向和反向吞吐量

实验内容：在多种MoE训练配置下，测量并比较单层MoE的前向和反向计算吞吐量（TFLOPS）。同时，进行端到端的7B MoE模型训练，比较实际的日处理令牌数。
实验结果：
* 内核级吞吐量：如图14所示，在所有模型规模下，SonicMoE始终实现最高的TFLOPS。在1.4B和7B设置中，TFLOPS比ScatterMoE和MoMoE提高了40%。对于30B和120B MoE，SonicMoE在前向和反向传播中均超过500 TFLOPS，而其他基线要么不支持某些n大小，要么性能显著下降。
* 端到端吞吐量：在使用FSDP-2进行7B MoE模型训练时，SonicMoE在64个H100上实现了每天2130亿令牌的吞吐量，与ScatterMoE在96个H100上每天2250亿令牌的吞吐量相当。
* 现代MoE配置下的吞吐量：如图15所示，在使用近期开源MoE模型的配置进行测试时，SonicMoE通常能达到超过550 TFLOPS，并持续优于所有基线。特别是在高稀疏和细粒度MoE配置下，相对基线的速度提升更大。

图14: 在H100上不同MoE内核的前向和后向TFLOPS。DeepGEMM没有提供高效的路由器实现、收集和专家聚合内核，我们为此使用了标准的PyTorch实现（“DeepGEMM-pt”）或我们高度优化的内核（“DeepGEMM++”）。在后向传播中，“DeepGEMM++”和“DeepGEMM-pt”都使用与SonicMoE相同的计算路径，只是我们启动了单独的内核来计算dS、A′和dSwiGLU。MoE配置与图13相同。

图15: 在H100上，针对从7B到685B参数的不同配置，单个MoE层的不同MoE内核的前向和后向TFLOPS。MoE配置从左到右分别采用了OLMoE-1B-7B-0125、gpt-oss-20b、Kimi-Linear-48B-A3B-Base、Qwen3-Next-80B-A3B-Thinking、Qwen3-235B-A22B-Thinking-2507和DeepSeek-V3.2-Exp的模型大小。为了公平比较，我们不考虑共享专家和专家偏置，并且始终使用带有softmax分数的TC top-K路由器。ScatterMoE、MoMoE、DeepGEMM-pt和DeepGEMM++在DeepSeek-V3.2-Exp配置下都无法运行（由于索引溢出或CUDA OOM错误）。

6.3 令牌舍入

6.3.1 令牌舍入的通用任务评估

实验内容：使用令牌舍入（TR）算法训练MoE模型，但在评估时切换回标准的令牌选择top-K（TC top-K）路由，以评估TR训练后与TC的兼容性。与TC、EC（专家选择）及其变体（带辅助路由器、微调TC路由器）以及TC（令牌丢弃）基线进行比较。
实验结果：
* 模型质量：如表3所示，在0.5B和1.4B模型的各种稀疏配置下，TR训练的模型在评估时与TC训练的模型质量相当。在极度稀疏的设置下（例如，K/E < 1/32），TR甚至取得了稍低的验证集困惑度和更高的平均准确率。
* 与EC的比较：TR显著优于EC及其变体。EC存在训练-测试不匹配问题，即使通过微调或辅助路由器来弥补，其最终性能仍劣于TR。
* 与令牌丢弃的比较：TR也优于简单的令牌丢弃基线（总是向下舍入），后者虽然也能减少填充，但会导致模型质量下降。

表3: 不同路由方法的任务评估比较。“Train”和“Val”分别指训练结束时的困惑度和验证集上的困惑度。接下来的11列是训练结束时评估的下游任务，我们报告了每个任务的准确率。“Avg”是这11个下游任务的平均准确率。在评估验证困惑度和任务性能时，我们对TR、令牌丢弃和EC基线使用TC top-K路由。T¯e表示每个微批次中每个专家接收的平均令牌数。

(a) 0.5B参数，20B令牌，8/64激活 ($T¯e$ = 4096, $M$ = 128)

6.3.3 令牌舍入的训练吞吐量

实验内容：在iso-FLOPs设置下，通过增加专家数E来提高稀疏度，比较令牌舍入（TR）与top-K令牌选择（TC）的MoE主内核（不含路由器）运行时的吞吐量（TFLOPS）。
实验结果：
* 吞吐量对比：如图16所示，随着稀疏度增加，TC路由的TFLOPS下降，这是因为tile量化效应（填充浪费）和IO增加。TR由于消除了填充浪费，其TFLOPS下降得更慢，始终高于TC。
* 性能提升：在高度稀疏的情况下，TR的优势更加明显。例如，对于一个256个专家的MoE（K/E = 1/128），TR在前向传播中TFLOPS提升25.7%，反向传播提升11.8%，端到端提升15.9%。
* 现代MoE配置下的吞吐量：如图17所示，在现代稀疏MoE模型配置上，TR也展现了显著的吞吐量优势。例如，在Qwen3-Next-80B-A3B-Thinking (K/E = 10/512)配置下，TR相比TC在前向和反向传播中分别快了19.6%和7.9%。

图16: 使用不同路由方法的SonicMoE MoE内核的前向和后向模型TFLOPS。我们将配备“通过专家频率最近舍入到$M_{tile}$倍数”子程序的TR与TC top-K路由进行比较。配置细节见附录G。

图17: 在H100上，针对从7B到685B参数的不同配置，配备不同路由方法的SonicMoE单MoE层的前向和后向TFLOPS。MoE配置与图15相同。我们将配备“通过专家频率最近舍入到$M_{tile}$倍数”子程序的TR与TC top-K路由进行比较。

---

A5 结论

本文提出了SonicMoE，一个协同优化MoE架构和GPU内核的解决方案，以应对细粒度和稀疏MoE模型带来的训练挑战。
主要贡献包括：
1. 一种内存高效的算法，随着MoE变得更细粒度，能将激活内存大小最小化。
2. 能将IO与计算重叠以提升吞吐量的GPU内核。
3. 一种感知tile的令牌舍入路由方法，能在不损失模型质量的情况下带来额外的速度提升。

未来的研究方向包括将该方法扩展到低精度和微缩放格式（如FP8, MXFP8, MXFP4）以进一步节省内存，以及在专家并行等分布式设置中将通信与计算重叠。作者设想未来的模型架构设计将优化每计算小时的质量，而不仅仅是每FLOP的质量，即同时考虑算法和硬件效率。

---

A6 附录

B SonicMoE与现有MoE内核设计的比较

核心差异。现有的高效MoE内核也将MoE计算框架为Grouped GEMM，但它们的构成与SonicMoE不同。以下是关键差异的概述：

• ScatterMoE【索引53，Scattered mixture-of-experts implementation，Tan et al.，2024】：为varlen-M Grouped GEMM实现了gather融合，但varlen-K没有。它不重叠MMA计算与内存IO，并且基于不支持TMA的旧版Triton。它计算dS的方式是dS = <dO, Y>，需要缓存Y，导致高IO和激活内存。其前向和反向传播的融合有限，因此比SonicMoE慢得多，尤其是在反向计算中。

• MoMoE【索引10，MoMoE: Memory optimized Mixture of Experts，Costin et al.，2025】：与ScatterMoE类似，为varlen-M实现了gather融合但varlen-K没有。尽管dS计算与up-proj激活梯度融合，但仍使用dS = <dO, Y>。与ScatterMoE一样，MoMoE不使用TMA进行IO。其scatter操作比SonicMoE慢得多（如图22所示）。

• MegaBlocks【索引18，Megablocks: Efficient sparse training with mixtureof-experts，Gale et al.，2023】：我们关注其ParallelDroplessMLP实现，该实现基于块稀疏矩阵乘法。它首先收集和填充令牌，然后为up/down-proj启动块稀疏GEMM，接着在聚合专家结果前启动一个scatter内核。这些稀疏矩阵乘法通常比高度优化的Grouped GEMM耗时长（如图6所示），并且gather和scatter内核的总IO成本为$8TKd$字节，这对细粒度MoE来说可能成为瓶颈。

• Megatron-LM【索引51，Megatron-lm: Training multi-billion parameter language models using model parallelism，Shoeybi et al.，2019】：我们关注其GroupedMLP实现，它使用CUTLASS库的Grouped GEMM作为后端，并带有JIT epilogue融合。与DeepGEMM类似，GroupedMLP不将gather与prologue融合（它假设输入是连续打包的）。最近的一个内存高效补丁在前向传播期间将S加权与SwiGLU计算融合，并在反向传播中允许PyTorch autograd引擎遵循与SonicMoE类似的计算路径。

• DeepGEMM【索引69，DeepGEMM: clean and efficient FP8 GEMM kernels with fine-grained scaling，Zhao et al.，2025b】：为连续打包的输入设计了Grouped GEMM内核。它没有为SM90（Hopper）BF16 Grouped GEMM实现任何其他融合。DeepGEMM更专注于专家并行的分布式训练，通常会启动一个单独的all2all内核，然后是连续的Grouped GEMM。DeepGEMM SM90 BF16内核还假设每个专家接收的令牌数是$M_{tile}$的倍数，因为它在Grouped GEMM计算期间没有实现TMA张量描述符的在线更新。其BF16 GEMM在SM90上也没有采用Ping-Pong调度。

实现语言与限制。此外，ScatterMoE和MoMoE都是用Triton实现的，这简化了开发，但牺牲了对Hopper和Blackwell GPU异步计算和内存IO的完全可编程性。例如，它们无法在GEMM的epilogue期间实现对异步加载和存储的细粒度控制，也无法使用Ping-Pong调度将MMA与重度epilogue操作重叠。当GEMM计算规模较小（如细粒度MoE）时，重叠epilogue中的IO操作对于实现高GPU利用率变得越来越重要。

C 梯度计算

公式推导。对于一个专家e，设
$X_e \in \mathbb{R}^{T_e \times d}, W_{1,e} \in \mathbb{R}^{d \times 2n}, W_{2,e} \in \mathbb{R}^{n \times d}$
前向激活计算由下式给出：

$$H_e = X_e W_{1,e} \in \mathbb{R}^{T_e \times 2n}, \quad A_e = \text{SwiGLU}(H_e) \in \mathbb{R}^{T_e \times n}, \quad Y_e = A_e W_{2,e} \in \mathbb{R}^{T_e \times d}.$$

带有分数$S = {s_{t,e}}$的令牌聚合由下式给出：
$O_t = \sum_e s_{t,e} Y_{e,t}, \quad dO_t \in \mathbb{R}^{T \times d} \text{ as the gathered results from } dO.$
我们知道
$dY_{e,t} = s_{t,e} \, dO_t \implies dY_e = \text{Broadcast}(s_e) \, dO_e.$
定义Grouped GEMM的输出为 $dA'e := dO_e W^\top$。} \in R^{T_e \times n
那么从公式8可得

分数$dS$的激活梯度为：
$dS_{t,e} = \langle dO_t, Y_{e,t} \rangle = \langle dO_t, W_{2,e}^\top, A_{e,t} \rangle = \langle dA'{e,t}, A \rangle.$
此外，我们可以从$dA_e$和$A_e$（由$H_e$重新计算）导出$dH_e$：
$dH_e = \text{dSwiGLU}(dA_e, H_e).$
使用公式8，

C.1 dS的计算选择

不同计算方式的对比。如果不实现自定义内核而仅依赖PyTorch的autograd引擎，我们可以在(1) down-proj前向传播之前或(2)之后添加专家加权(S)。两者在前向和反向传播中得到相同的结果，但dS的计算方式不同。对于(1)，我们需要计算$\langle dA'{e,t}, A \rangle$，这是(2)所要求的。} \rangle$，这是SonicMoE和Megatron采用的方式。MoMoE、ScatterMoE和MegaBlocks则计算$\langle dO_t, Y_{e,t

SonicMoE选择的优势。注意$dS$可以计算为$dS_{t,e} = \langle dA'{e,t}, A = \langle dA'} \rangle = \langle dO_t, Y_{e,t} \rangle$中的任意一种，但计算为$dS_{t,e{e,t}, A \rangle$在计算和激活内存方面是更高效的选择，原因如下：
* 额外的HBM流量（0 vs. $2TKd$字节）：$\langle dA'{e,t}, A \rangle$需要的$dA'{e,t}$和$A$在dH内核期间已经计算好了，所以我们可以避免额外的不必要加载。
* 额外的缓存激活内存（0 vs. $2TKd$字节）：ScatterMoE、MoMoE和MegaBlocks的缓存激活内存未能随专家粒度保持恒定的原因之一是需要缓存Y来计算dS。
* 并行归约轮数（$\log_2(n)$ vs. $\log_2(d)$）：$\langle dA'{e,t}, A \rangle$在d维度上归约。这种差异至少节省了$\log_2(d/n)$轮的归约。} \rangle$在n维度上归约，而$\langle dO_t, Y_{e,t

D SonicMoE算法

本节展示了主论文中引用的图表和算法。我们进一步描述了在Hopper GPU上的组GEMM内核中集成的索引预取策略，如图18所示。

图18: 在H100 GPU上，针对varlen-M Grouped GEMM的M维（左）和varlen-K Grouped GEMM的K维（右）进行收集时的索引预取策略。对于M维收集（左），我们让每个线程在主循环前独立地将索引预取到自己的寄存器中。对于K维收集（右），我们在SMEM上创建一个缓冲区，让4个生产者warp协作地将索引预取到SMEM，每个生产者线程再从这个SMEM缓冲区读取到自己的寄存器中。

算法5：SonicMoE的MoE内核上投影反向传播
输入：$X, \pi, W_1, dH$
输出：$dX, dW_1$
Up-proj act dX̃ 内核 ($dH, W_1$) → $dX̃$:
// varlen-M Grouped GEMM
并行处理每个专家 $e \in [E]$:
$dH_e, W_{1,e} \leftarrow \text{load}(dH_e, W_{1,e})$
$dX̃e \leftarrow dH_e W^\top$
$dX̃e \leftarrow \text{store}(dX̃_e)$
Up-proj weight dW1 内核 ($X, dH, \pi$) → $dW_1$:
// Gather + varlen-K Grouped GEMM
并行处理每个专家 $e \in [E]$:
$X, \pi, dH_e)$}, dH_e \leftarrow \text{load}(X, \pi_{:,e
$X_e \leftarrow \text{Gather}(X, \pi_{:,e})$
$dW_{1,e} \leftarrow X_e^\top dH_e$
$dW_{1,e} \leftarrow \text{store}(dW_{1,e})$
Expert aggregation dX 内核 ($dX̃, \pi$) → $dX$:
// Gather and sum
并行处理每个令牌 $t \in [T]$:
$dX̃, \pi_{t,e} \leftarrow \text{load}(dX̃, \pi_{t,e})$
$dXt \leftarrow \sum_{e \in [E]} \pi_{t,e} dX̃_{e,t}$
$dXt \leftarrow \text{store}(dXt)$

E SonicMoE内核级吞吐量消融研究

E.1 Hopper GPU上的Grouped GEMM

连续打包输入的Grouped GEMM。在图19中，我们在H100 GPU上比较了SonicMoE与DeepGEMM在没有任何其他融合的varlen-M Grouped GEMM上的性能。我们发现SonicMoE的up-proj比DeepGEMM快2.7% TFLOPS，而down-proj则快10.0%。SonicMoE的相对TFLOPS加速在30B down-proj配置中分别为57.4%、14.0%和5.3%。我们对n < 1024的down-proj使用Ping-Pong调度，而DeepGEMM使用协作调度。

图19: 在H100 GPU上，前向传播期间使用连续打包输入到up和down-proj的Varlen-M Grouped GEMM。配置与图13相同。“cuBLAS BMM”是一个密集BMM基线，相当于所有专家接收相同数量的令牌（完美负载均衡），其TFLOPS可被视为任何Grouped GEMM内核的上限。

带gather融合的Grouped GEMM。在图20中，我们报告了SonicMoE、ScatterMoE、MoMoE和DeepGEMM在有无gather融合（分别用不透明和透明条表示）情况下的性能。
* M维度gather：SonicMoE在有无gather融合的情况下，平均相对TFLOPS差异为6.3%。带gather融合的SonicMoE始终比ScatterMoE（平均高9.7%）、MoMoE（平均高30.9%）和DeepGEMM（平均高38.3%）实现更高的TFLOPS。
* K维度gather：SonicMoE在有无gather融合的情况下，平均相对TFLOPS差异为8.5%。即使没有gather融合，SonicMoE也已比ScatterMoE（平均高23.5%）、MoMoE（平均高4.3%）和DeepGEMM（平均高41.4%）快。当我们将带gather融合的SonicMoE与ScatterMoE和MoMoE及其各自的gather内核结合进行比较时，随着专家粒度增加，差距会扩大。平均而言，带gather融合的SonicMoE分别比ScatterMoE、MoMoE和DeepGEMM快55.1%、42.4%、71.8%的TFLOPS。

图20: 在H100 GPU上的前向传播up-proj（M维gather）和反向传播up-proj权重梯度dW1（K维gather）内核。SonicMoE支持从不同位置收集的输入（不透明条）和连续打包的输入（透明条）。ScatterMoE和MoMoE都对varlen-M有gather融合但对varlen-K没有，因此我们通过将其连续打包权重梯度内核的时间（透明条）与其自身gather内核的时间相加来测试其带gather的varlen-K Grouped GEMM时间（不透明条）。DeepGEMM对varlen-M和varlen-K Grouped GEMM都没有gather融合，因此我们在两种情况下都提供了一个优化的gather内核。我们还提供了一个“cuBLAS dense BMM”（透明条）基线和带gather的GEMM时间（不透明条），通过将一个高度优化的gather内核的时间与相同输入形状相加，这可以被视为任何没有gather融合的Grouped GEMM内核的TFLOPS上限。

E.2 Hopper GPU上的专家聚合

聚合带宽对比。我们在图21中对SonicMoE聚合内核的带宽进行了基准测试。我们将SonicMoE的gather-and-sum聚合（图9左）与ScatterMoE的torch.bmm和MoMoE的torch.sum聚合（图9中）进行了比较。我们还实现了一个高度优化的triton聚合内核（“triton sum (contig. Y)”）作为上限参考。尽管SonicMoE的聚合内核在HBM加载期间需要gather融合，其内存带宽仍然超过ScatterMoE（平均2.92倍）和MoMoE（平均1.05倍），并且仅比triton上限略慢（平均0.98倍）。

图21: MoE前向传播期间的专家聚合内核（O kernel）。配置与图13相同。“ScatterMoE (contig. Y)”是ScatterMoE采用的专家聚合策略。“MoMoE (contig. Y)”是MoMoE的torch.sum调用。我们还实现了一个优化的triton内核“triton sum (contig. Y)”。

图22: 在H100上Grouped GEMM和专家聚合内核的吞吐量。“SonicMoE (gemm + gth w. sum)”是SonicMoE的最终设计选择，如图9左侧策略所示。我们将其与在SonicMoE上实现图9中间策略的“SonicMoE (gemm w. sct + sum)”进行比较。我们还与ScatterMoE的设计和MoMoE的设计进行了比较。对于每种方法，我们在透明条中报告GEMM TFLOPS，在不透明条中报告GEMM和专家聚合总运行时间的TFLOPS。

E.3 Hopper GPU上的Top-K排序

Top-K内核性能。我们在图23中对SonicMoE的top-K内核在H100上的带宽进行了基准测试。我们将其与PyTorch、triton官方示例、tilelang官方示例和RTop-K在BF16和FP32输入上进行了比较。SonicMoE的基于排序网络和纯寄存器通信的设计，在大多数情况下都比其他实现更快，特别是对于MoE中常见的具有大量T和中等E、K的场景。

图23: 在MoE前向传播期间，使用BF16输入（第一行）和FP32输入（第二行）的Top-K内核。配置与图13相同。“torch”是直接的torch.topk调用。“triton”和“tilelang”取自其官方示例。“RTop-K”【索引62，RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs，Xie et al.，2025】只支持FP32输入。

F 更多实验

F.1 专家粒度的影响

实验验证。这里我们验证了采用细粒度MoE的有效性。我们固定了0.5B和1.4B模型的MoE激活率$\rho = K/E$，并在表5a和5b的第一行到第三行中，按比例扩大K和E，同时线性减小n。

结论。总的来说，我们观察到n=256的性能优于n=1024，这也与表1中提到的MoE扩展趋势一致。在图1右子图中，我们发现在iso-FLOPs下，从n=1024到n=256，SonicMoE和cuBLAS仍能维持吞吐量，但从n=256开始，FLOPs会随粒度线性下降。因此，我们在表3的所有实验中选择n=256。

表5: 在iso-FLOPs（nK恒定）和iso-params（nE恒定）设置下，评估MoE随粒度的变化。“PPL”指训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。“dense, iso-FLOPs”指中间大小为nK的密集模型，“dense, iso-params”指中间大小为nE的密集模型。

(a) 0.5B参数，20B令牌，8/64激活

F.2 令牌舍入不同舍入子程序的消融研究

实验内容。我们进行了消融研究，以探究不同路由子程序对TR训练的MoE的影响。我们将令牌舍入与最近舍入（“NR”）与各种其他舍入方法进行了比较，包括随机舍入（“SR”）、总是向上舍入（“UP”）和总是向下舍入（“DOWN”）。

舍入子程序描述：
* NR-f：通过专家频率最近舍入到$M_{tile}$倍数。如果$\lceil f_e \rceil_{M_{tile}} - f_e < f_e - \lfloor f_e \rfloor_{M_{tile}}$，则填充EC令牌。这是默认选择。
* SR-f：通过专家频率随机舍入。从伯努利分布$\text{Bernoulli}\left( \frac{f_e - \lfloor f_e \rfloor_{M_{tile}}}{M_{tile}} \right)$采样决定是否填充。
* NR-s：通过专家分数最近舍入。采样分布为：

$$ \text{Bernoulli} \left( \frac{\sum_t s_{e,t} - \sum_t |s_{e,t}| M_{\text{tile}}}{\sum_t [s_{e,t}] M_{\text{tile}} - \sum_t [s_{e,t}] M_{\text{tile}}} \right) $$

* Balance-f：通过专家频率进行平衡舍入，确保舍入后所有专家的总令牌数得以保留。
* UP：总是向上舍入到$\lceil f_e \rceil_{M_{tile}}$。
* DOWN：总是向下舍入到$\lfloor f_e \rfloor_{M_{tile}}$，相当于令牌丢弃。

结论。结果如表6所示，我们发现令牌舍入算法通常对具体的舍入子程序是鲁棒的。总是丢弃令牌（“DOWN”）虽然内核运行时间最短，但质量下降明显。总是填充（“UP”）虽然困惑度较低，但任务准确率不一定更高，且运行时间最长。为了平衡训练效率和模型质量，“NR-f”是一个很好的默认选择。

表6: 评估配备不同round_and_sparsify子程序的令牌舍入算法。“PPL”指训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。
(a) 0.5B参数，40B令牌，2/64激活 ($T¯e$ = 512, $M$ = 128)

F.3 微批次大小T和tile大小Mtile影响的消融研究

微批次大小T的影响。由于令牌舍入在微批次级别应用，微批次大小T的选择会影响TR的质量结果。在表7中，我们改变微批次大小，同时保持小批量大小（每个优化步骤消耗的令牌数）不变。我们发现当$\bar{T}e / M} \ge 2$时，TR能保持其训练的MoE质量，但如果$\bar{Te / M = 1$，则质量有明显下降，但仍优于用EC训练后微调TC路由器的基线。

表7: 评估当我们改变微批次大小T以改变每个专家的平均令牌数($\bar{T}e$)时，令牌舍入算法的性能。对于每次试验，我们将微批次大小从4（$\bar{T}_e$=512）变为1（$\bar{T}_e$=128），并保持小批量大小不变。$M$始终为128。“PPL”指训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。
(a) 0.5B参数，40B令牌，2/64激活 ($M_{tile}$ = 128)

Tile大小Mtile的影响。类似地，在表8中，我们可以发现当$\bar{T}_e / M_{tile} \ge 2$时，TR对$M_{tile}$通常是鲁棒的，而当$\bar{T}_e / M_{tile} = 1$时，有明显下降，但总体结果仍优于EC基线。

表8: 评估当我们改变令牌舍入的tile大小$M_{tile}$时，令牌舍入算法的性能。“PPL”指训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。
(a) 0.5B参数，40B令牌，2/64激活 ($\bar{T}_e$ = 512)

G 激活内存和训练吞吐量基准测试配置

配置详情。计算图3中IO成本的配置见表9a，图13和14的配置见表9b。图16中4个子图的配置如下，TR速度基准测试中$M_{tile}$始终为128：
* 左上2个子图: $(T, d, n, K) = (16384, 1536, 256, 8)$，E从64到512变化。
* 右上2个子图: $(T, d, n, K) = (16384, 1536, 1024, 2)$，E从16到128变化。
* 左下2个子图: $(T, d, n, K) = (16384, 4096, 512, 8)$，E从64到512变化。
* 右下2个子图: $(T, d, n, K) = (16384, 4096, 1024, 4)$，E从32到256变化。

表9: 图3、13和14的基准配置。
(a) 图3中内存IO成本的基准配置。

(b) 图13和14以及所有其他内核级消融研究使用的基准配置。

H LM训练的超参数细节

训练设置。我们使用OLMoE代码库及其默认配置的下游任务。所有模型使用4096令牌的上下文长度进行训练。
对于表3中所有EC带微调TC路由器的实验，我们额外使用4B令牌，并且只微调路由器权重（所有其他参数冻结）。我们总是使用2e-4的学习率，0.01的权重衰减和带10%预热步骤的余弦学习率调度器。微调期间每个小批量的令牌数为1M。
对于所有EC带辅助路由器的实验，我们使用一个2层MLP（每个线性层大小为$E \times E$，带SiLU激活），它以原始路由器logits为输入，并为所有专家做出E个独立的二元预测。
我们通过丢弃从TC top-K排序中选择的令牌，或总是向下舍入来实现“TC（令牌丢弃）”。

表10: MoE预训练实验的通用配置