Mamba作者团队SonicMoE：一个Token舍入，让MoE训练速度提升近2倍

机器之心编辑部

混合专家（MoE）模型已成为在不显著增加计算成本的情况下，实现语言模型规模化扩展的事实标准架构。

近期 MoE 模型展现出明显的高专家粒度（更小的专家中间层维度）和高稀疏性（在专家总数增加的情况下保持激活专家数不变）的趋势，这提升了单位 FLOPs 的模型质量。

这一趋势在近期的开源模型中表现尤为明显，例如 DeepSeek V3、Kimi K2 以及 Qwen3 MoE 等，它们均采用了更细粒度的专家设计（更小的中间层维度）和更高的稀疏度，在保持激活参数量不变的同时大幅增加了总参数量。

表 1：MoE 扩展趋势：在此，团队将激活率展示为每个 Token 激活的专家数 K / 专家总数 E；针对前沿开源模型，专家粒度展示为模型嵌入维度（d）/ 专家中间层大小（n）。在 MoE 稀疏度计算中未包含共享专家。趋势表明，新的开源 MoE 模型倾向于具备更高的粒度和稀疏度。

然而，这种追求极致粒度和稀疏性的设计导致了严重的硬件效率下降问题：

• 内存墙瓶颈：对于细粒度 MoE，激活内存的占用量通常随激活专家数量线性增长，导致前向和反向传播中的内存压力剧增。
• IO 瓶颈：由于专家变得更小且更分散，算术强度（Arithmetic Intensity，即计算量与数据传输量的比值）显著降低，IO 访问变得更加动态和频繁，导致模型训练进入「内存受限」区间。
• 计算浪费：在高稀疏性场景下，由于 Grouped GEMM（分组通用矩阵乘法）内核中的 Tile 量化效应，输入数据往往需要进行填充以对齐硬件 Tile 大小，这直接导致了计算资源的浪费。

针对这些问题，普林斯顿大学助理教授 Tri Dao（Mamba、FlashAttention 的核心作者）团队提出了一套名为 SonicMoE 的系统性解决方案。该方案专为 NVIDIA Hopper 和 Blackwell 架构 GPU 量身定制，其核心贡献包括：

• 内存高效算法：团队通过重新设计 MoE 的计算图，提出了一种在计算路由梯度时不缓存激活值的方法。该方法在保持与原始 MoE 公式数学等价的前提下，大幅减少了反向传播所需的激活显存。对于细粒度 7B MoE 模型，每层的激活内存占用减少了 45%，且随着专家粒度的增加，其内存占用保持恒定，效率比现有基线高出 0.20-1.59 倍。
• 计算与 IO 重叠：利用 Hopper 架构 GPU 的 WGMMA 指令与生产者 - 消费者异步范式，SonicMoE 设计了新型 GPU 内核。该内核能够将 GEMM 计算与从 HBM 加载数据的 IO 操作并行执行，有效掩盖了细粒度 MoE 带来的高昂 IO 延迟。
• Token 舍入：这是一种即插即用的创新调度策略。它将分发给每个专家的 Token 数量四舍五入为 Grouped GEMM Tile 大小（例如 128）的倍数。算法保证每个专家的偏差最多仅为一个 Tile，从而在期望意义下保持总 token 数不变。这一策略有效减少了因填充导致的算力浪费。

实验数据有力地证明了 SonicMoE 的性能优势，在针对细粒度 7B MoE 模型的测试中：前向传播相比高度优化的 DeepGEMM 基线，速度提升43%；反向传播相比最先进的 ScatterMoE 和 MoMoE 基线，速度分别提升了 83% 和 115%；端到端训练仅依靠内核优化即可将训练吞吐量提升 50%，若配合 Token 舍入路由，在扩展专家数量时可进一步获得 16% 的额外吞吐量提升。

• 论文标题：SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations
• 论文地址：https://arxiv.org/abs/2512.14080

更直观地看，团队仅使用 64 台 H100 运行 SonicMoE，便实现了每日 2130 亿 token 的训练吞吐量，这一表现已能与使用 96 台 H100 运行 ScatterMoE 的效率相媲美。此外，在高稀疏性场景下（如 1.4B 参数模型），其 Tile 感知的 Token 舍入算法在验证了不损失下游任务精度（如在 2B 规模上的推理质量）的同时，显著提升了内核执行速度。

目前，团队已将相关内核代码开源，为大模型社区加速高性能 MoE 训练提供了强有力的工具。

图 1： 即使专家粒度（d/n，其中 d 为嵌入维度，n 为专家中间维度）增加，SonicMoE 的每层激活显存占用（左图）仍保持恒定；相比其他基线，其显存效率提升了 0.20 倍至 1.59 倍。SonicMoE 的前向计算吞吐量（右图）平均达到了理论上限的 88%（最高 91%，最低 86%），该上限基于 H100 GPU 上的 「cuBLAS BMM + 激活函数 + cuBLAS BMM + 聚合操作」 计算得出。请注意，cuBLAS 上限基线未包含路由计算部分。在此，我们使用的是 30B 参数量的 MoE 配置，微批次大小为 32768 个 token，并且从左至右依次将「激活专家数 / 总专家数」设置为 2/32、4/64、8/128 和 16/256。

内存高效的 MoE 算法

团队提供了一个高效的基于 Tensor Core 的 top-K 路由，以及一个可以接受任意路由输入的接口。但需要注意的是，SonicMoE 的 MoE 计算与路由的选择无关，因此与任意路由逻辑兼容。

SonicMoE 的 MoE 计算实现具有高度模块化特性，仅由以下两部分组成：

• 经过优化的分组 GEMM 内核（带有模块化融合）
• 经过优化的专家聚合内核

主机会根据最佳 GEMM 配置和加载 / 存储策略来调度并启动上述 8 个内核。

结果显示，尽管采用了如此高度的模块化设计，SonicMoE 仍然展现出业界领先的训练吞吐量和最低的激活内存使用量。

面向 IO 的内核设计

细粒度 MoE 的表达能力来自于每个 token 在专家选择上的多样性，但这种多样性同时带来了与专家粒度线性增长的 IO 开销，为了保持高吞吐，需要尽可能做到：

• 通过融合（fusion）减少 IO 访问
• 将 IO 延迟与计算重叠

在融合这一块有两种方式，一是利用 HBM 加载进行 Gather 融合。SonicMoE 的 Grouped GEMM 既可以接受连续打包的输入，也可以接受从不同位置 gather 得到的输入。对于第二种情况，团队将输入 gather 与从全局显存（GMEM，通常是 HBM）到共享内存（SMEM）的加载过程进行融合，从而能够将这些数据批量化，利用 Tensor Core 执行 GEMM。

这一过程包括两个步骤：

• 获取每个 expert 对应的被路由 token 的索引；
• 使用这些索引，通过 Blackwell 和 Hopper 架构的 cp.async 指令，从 HBM gather 激活值。

二是 Epilogue 融合，通过以下设计充分利用 epilogue 计算，以最大化减少不必要的 IO 访问：将 SwiGLU 以及 SwiGLU 的反向（dSwiGLU），分别与前向 up-proj 内核的 epilogue、反向 down-proj 激活梯度内核的 epilogue 进行融合；在反向 down-proj 激活梯度（dH）内核的 epilogue 中计算 dH 和 dS。

结果显示，这种「重量级 epilogue 融合」使 SonicMoE 相比其他方案获得显著加速。

Token rounding 路由方法

团队在分析稀疏 MoE 训练模式下的硬件效率时发现，随着 MoE 变得更加稀疏，因填充而产生的 GEMM tile 计算浪费会累计到不可忽略的程度，这被称为「tile 量化效应」。为此，团队提出路由方法「token rounding」来消除这种效应，从而实现更高效的训练。

Token rounding 算法首先计算基础的 TC（token-choice）路由结果，并对每个 expert 对应的 token 按路由分数进行排序，之后在第二步排序中选择：要么丢弃第一步 TC top-K 选择中的部分 token，要么在第二步排序中为某些 expert 补齐额外的 token（填充）。

过程中，团队会对路由权重矩阵进行处理，使得 TC 选中的 token 始终优先于 EC token。结果就是，无论是丢弃还是填充，都只会影响每个 expert 的最后一个输入 tile。

实验表明，这种方法在实现更高训练吞吐量的同时，并不会影响模型质量。

更多内容，可查看论文获悉！