Mamba团队SonicMoE：新增1个Token，MoE训练速度提升近2倍

首页

热心网友

转载

2025-12-19

免费影视、动漫、音乐、游戏、小说资源长期稳定更新！ 👉 点此立即查看 👈

机器之心编辑部

混合专家（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（填充）。