1. 鲸林向海首页
  2. 大模型训练

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

• 大模型训练 • 阅读 149

关键词：混合专家模型（MoE）、SonicMoE、GPU 内核优化、内存高效算法、令牌舍入路由、细粒度稀疏 MoE

SonicMoE 已开源，方案基于 CuTe-DSL 实现并提供 PyTorch 接口，采用宽松许可证：github.com/Dao-AILab/sonic-moe。未来研究将围绕两方面展开：一是扩展到 FP8、MXFP8、MXFP4 等低精度和微缩放格式以进一步节省内存；二是在专家并行等分布式场景下实现通信与计算重叠。

同时我们期望未来的模型架构设计能同时考虑算法效率和硬件效率，以“每计算小时质量”（quality per compute hour）而非仅“每 FLOP 质量”（quality per FLOP）作为优化目标。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升
  • SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations
  • https://arxiv.org/pdf/2512.14080
  • 代码库：https://github.com/open-lm-engine/lm-engine
  • 开源 Kernel：https://github.com/Dao-AILab/sonic-moe

Mixture of Experts（MoE）是大模型参数缩放的核心架构，能在不显著增加计算成本的前提下提升模型性能。

然而，当前MoE模型向高专家粒度（更小的专家中间维度）和高稀疏性（激活专家数固定、总专家数增多）发展时，面临三大挑战：激活内存占用线性增加、IO成本上升导致硬件效率下降、分组GEMM（Grouped GEMM）中填充操作造成计算浪费。

为此，本文提出 SonicMoE，通过算法-硬件协同设计解决上述问题。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 4：SonicMoE 的 8 个内核计算流程（黄色框分组）。前向传播启动 3 个内核，反向传播启动 5 个内核。指向黄色圆圈的箭头表示“从高带宽内存（HBM）加载变量到共享内存（SRAM）”，从黄色圆圈出发的箭头表示“将变量存储到高带宽内存（HBM）”。高带宽内存中的所有变量用不同色框标注：紫色框表示前向和反向传播的输出，蓝色框表示中间变量或权重，红色框表示所有缓存的激活值（X、H、S）。后文算法 2 详细描述 SonicMoE 的前向传播，算法 3 和 5 详细描述 SonicMoE 的反向传播。

SonicMoE 的核心贡献包括三方面：

  • 设计内存高效的前向/反向计算算法。通过重构计算图，避免缓存路由器梯度所需激活，在不增加计算量的前提下，使激活内存占用不随专家粒度增加而增长。对细粒度 7B MoE 模型，单层激活内存减少 45%。
  • 研发 IO 与计算重叠的 GPU 核。利用NVIDIA Hopper/Blackwell GPU的异步特性（如 Ping-Pong 调度、TMA 异步操作），将内存 IO 与矩阵计算（MMA）重叠。对 7B MoE 模型，前向吞吐量较 DeepGEMM 提升43%，反向较 ScatterMoE、MoMoE 分别提升83%、115%。
  • 提出“令牌舍入”（token rounding）路由方法。将每个专家接收的令牌数调整为分组 GEMM tile 大小的整数倍，消除填充浪费。高稀疏场景下核执行时间额外加速 1.16 倍，且不影响下游任务性能。

实验表明：

  • SonicMoE 在 Hopper GPU 上对细粒度 7B MoE 的计算吞吐量达 ScatterMoE BF16 核的 1.86 倍；
  • 64 个 H100 训练 7B MoE 时，日吞吐量达 2130 亿 token，堪比 96 个 H100 上 ScatterMoE 的 2250 亿 token。
  • 其前向计算吞吐量平均_达 H100 cuBLAS 理论上限的 88%。_

本文开源了所有核代码，为高效 MoE 训练提供支撑：https://github.com/Dao-AILab/sonic-moe

本文目录

  • 一、引言
  • 二、背景
    • 2.1 基于分组矩阵乘法的 MoE
    • 2.2 MoE 计算
    • 2.3 MoE 路由方法
  • 三、内存高效的 MoE 算法
    • 3.1 SonicMoE 的 MoE 内核概述
    • 3.2 激活内存效率
  • 四、面向 IO 感知的内核设计（IO-aware kernel design）
    • 4.1 SonicMoE 的分组矩阵乘法（Grouped GEMM）
    • 4.2 矩阵乘加（MMA）与异步 IO 重叠
    • 4.3 面向 MoE 的高效 Top-K 排序内核
  • 五、令牌舍入路由（Token rounding routing）
    • 5.1 稀疏 MoE 的训练效率
    • 5.2 令牌舍入路由
  • 六、实验
    • 6.1 SonicMoE 的激活内存
    • 6.2 SonicMoE 的训练吞吐量
    • 6.3 令牌舍入（Token Rounding）
  • 七、结论
  • 参考文献
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

关键问题

问题 1：令牌舍入路由在极端场景下的性能保真性

文中令牌舍入路由的性能保真性验证，【仅】覆盖 1.4B-120B 模型规模与固定 128 的 Tile size，当模型突破千亿参数且专家稀疏度进一步提升（如 K/E<1%），或 Tile size 为适配低算力 GPU 调整为更小值（如 64 以下）时，该路由方法是否仍能维持与原生 Top-K 路由一致的下游任务精度？

  1. 千亿参数模型与更高稀疏度场景：论文验证了模型规模从1.4B到120B（包括685B的DeepSeek-V3.2-Exp）的性能保真度。即使在K/E=10/512（约1.95%）的高稀疏度场景下，模型精度依然得以维持。虽然未直接测试K/E < 1%的极端情况，但由于TR（令牌舍入）的核心设计原则是“单专家最大偏差不超过1个Tile”，并且在稀疏度提升（例如E从128增至256，K固定）时仍能保持精度，因此推测在极端稀疏场景下，其下游性能【仍能维持】与原生Top-K路由一致的水平。
  2. 更小Tile size（64以下）场景：论文默认Tile size为128。作者指出，TR策略对Tile size的变化具有鲁棒性。当Tile size降至64时，只要每个专家处理的平均令牌数≥128（即满足2倍Tile size），性能就不会显著下降；若<128，则可能出现轻微的精度衰减，但其表现仍优于采用EC路由结合TC微调的方案。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表7：改变微批量大小T（即改变每个专家的平均令牌数）时，对令牌舍入算法的评估汇总。实验中，微批量大小从4（=512）调整至1（=128），同时保持小批量大小不变。始终设为128。“PPL”指训练结束时的验证困惑度，“Avg”是11个下游任务的平均准确率。两组参数的消融实验共同指向一个关键阈值：≥128——这是TR策略能稳定发挥MoE模型性能的核心条件。低于此阈值可能导致性能损失，但TR的整体效果仍优于传统EC路由方案，表明TR在参数选择上具有更强的“容错性”和更高的“性能上限”。

问题2：SonicMoE在旧架构GPU及大规模专家并行场景下的性能

SonicMoE的IO-计算重叠内核高度依赖Hopper/Blackwell GPU的硬件特性（如cp.async、TMEM）。其在Ampere及更早架构GPU上的性能衰减比例是多少？ 此外，在大规模专家并行场景下，其IO优化与跨节点通信开销的权衡是否会抵消吞吐量优势，导致64块H100的等效性能无法维持？

  • Ampere及更早架构的性能衰减：作者未提供直接的衰减数据。然而，SonicMoE的核心优化（如cp.async异步IO、TMEM利用、Ping-Pong调度）高度依赖Hopper/Blackwell的硬件特性。Ampere架构缺乏这些特性，IO与计算的重叠效果将大幅减弱，推测性能衰减比例约为30%-50%（参考DeepGEMM在无Ping-Pong调度时的吞吐量差距）。
  • 大规模专家并行场景的通信权衡：实验在64块H100（8节点）的专家并行配置中验证了吞吐量优势，但未涉及更大规模的跨节点通信。由于SonicMoE当前未针对跨节点通信与IO优化进行协同设计，跨节点数据传输开销可能抵消部分吞吐量提升。但论文已将“重叠通信与计算”列为未来研究方向，暗示当前64块H100媲美96块H100（ScatterMoE）的等效性能在更大规模下仍可维持，【仅】需补充相应的分布式通信优化。

一、引言

专家混合模型（MoE）（Shazeer等人，2017）已成为在不显著增加训练计算需求的前提下扩展模型参数规模的关键技术（Kimi等人，2025；Zhao等人，2025a）。现代Transformer模型的层通常由序列混合块（如多头注意力（Vaswani等人，2017））和通道混合块（如密集型多层感知机）组成，其中MoE因其浮点运算效率优势，成为密集型MLP的理想替代方案。

一个MoE块通常包含令牌路由器和多个规模较小且通常尺寸相等的子网络（即“专家”）。MoE通过为每个令牌仅激活部分专家，减少了训练过程中的浮点运算次数消耗。然而，浮点运算次数的减少并不直接等同于硬件利用率的提升——MoE计算存在更动态的IO访问模式（每个专家需从不同位置收集令牌嵌入，并将计算结果分散回原位置），且这种“硬件不友好性”会随专家粒度的细化（专家中间维度减小）和稀疏度的提升（总专家数增加而激活专家数固定）而加剧，具体趋势如表1所示。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表1：MoE缩放趋势：展示了前沿开源模型的“激活比”（每令牌激活专家数K与总专家数E的比值，即K/E）和“专家粒度”（模型嵌入维度D与专家中间维度D_expert的比值，即D/D_expert）。计算MoE稀疏度时未包含共享专家。趋势表明，新型开源MoE模型的粒度更细、稀疏度更高。

MoE缩放定律（Clark等人，2022；Krajewski等人，2024；Tian等人，2025）预测：随着专家粒度（模型嵌入维度与每个专家中间维度的比值）和稀疏度的提升，单位浮点运算次数所能获得的模型质量会更高。

近年来的MoE模型，如DeepSeek V3（DeepSeek-AI等人，2024）、Qwen3 MoE（QwenLM，2025）和gpt-oss-120b（OpenAI，2025）已证明，在大规模场景下，“细粒度”MoE的性能优于“粗粒度”MoE。除粒度外，为在保持计算需求不变的前提下提升模型质量，现代MoE的稀疏度也在不断提高。例如，Kimi K2（Kimi等人，2025）的激活参数数量与DeepSeek V3相当，但其总参数数量远大于后者。

总体而言，如表1所示，MoE的粒度（更小的专家中间维度）和稀疏度（激活专家数固定、总专家数增多）近年来呈持续上升趋势。值得注意的是，近期与MoE竞争的替代架构，如PEER（He，2024）、Memory Layers（Berges等人，2024）和Ultra-Mem（Huang等人，2025），也普遍采用了“提升粒度与稀疏度”的设计思路。

尽管更高粒度和更高稀疏度的MoE可以提升单位浮点运算次数的模型质量，但它们会因以下问题导致硬件效率下降：
1. 细粒度MoE模型的激活内存占用更大：激活内存大小通常随激活专家数量线性增长。
2. 细粒度专家导致算术强度（计算量与IO量的比值）降低、IO成本增加。
3. 高稀疏MoE中，分组矩阵乘法的“瓦片量化效应”会造成计算浪费。

高粒度与高稀疏度共同将MoE训练推向“内存受限”状态，需要设计【专门的MoE内核】以掩盖增加的IO成本。现有的主流MoE内核，如ScatterMoE（Tan等人，2024）和MoMoE（Costin等人，2025），并未针对这些高IO成本场景进行优化，因此会面临显著的训练吞吐量下降。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图1：SonicMoE的每层激活内存占用（左图）即使在专家粒度（D/D_expert）增加时仍保持恒定，且内存效率较其他基线方法高0.20-1.59倍；SonicMoE的前向计算吞吐量（右图）平均达到理论上限值（H100上的cuBLAS BMM+激活函数+cuBLAS BMM+聚合操作）的88%（最大值91%，最小值86%）。注：cuBLAS上限基线未包含路由器计算；本实验采用300亿参数MoE配置，微批次大小为32768个令牌，从左到右的“激活专家数/总专家数”分别为2/32、4/64、8/128、16/256。

我们提出一种MoE架构与GPU内核的协同设计方案，该方案针对NVIDIA Blackwell和Hopper系列GPU，并结合了新颖的路由方法，具体包括：

  1. 提出一种更高效的 MoE 反向传播算法，可显著降低激活内存占用，且该占用不会随专家粒度的增加而上升。
  2. 利用 Blackwell 和 Hopper GPU 的新硬件特性，将内存 IO 与计算操作重叠，该优化适用于所有 MoE 模型，尤其对细粒度 MoE 增益显著。
  3. 设计一种硬件感知的令牌舍入路由方法，确保路由到每个专家的令牌数量始终是分组矩阵乘法所用瓦片大小的整数倍。

大量实验表明：

  • 当专家数量在 30B MoE 模型基础上扩展 4 倍时，令牌舍入路由比基线令牌选择路由快 16%。
  • 在 20 亿参数规模下验证，该方法能够保持 MoE 模型的推理质量。
  • 结合上述（1）和（2）的优化，可在不改变 Top-K 令牌选择路由的前提下，将 7B MoE 模型的端到端训练吞吐量提升 50%。
  • 令牌舍入路由方法在扩展专家数量时，可进一步提升 16% 的训练吞吐量，且无任何精度损失。

主要贡献总结：我们提出 SonicMoE，一种硬件与模型架构协同设计的解决方案，用于解决 MoE 训练效率问题，具体贡献如下：

  • 最小激活内存的 MoE 训练（无额外浮点运算）：我们分析了 MoE 粒度对 MoE 层前向和反向传播的影响，发现“保持浮点运算次数不变但增加 MoE 粒度”会导致反向传播所需的激活内存线性增加。基于这一观察，我们重新设计了计算图，在保持与原始 MoE 公式数学等价的前提下，避免了缓存路由器梯度计算所需的激活值。实验结果表明，对于细粒度 7B MoE，SonicMoE 可将每层激活内存占用减少高达 45%。
  • IO 与计算重叠的高效 MoE 内核（实现最优训练吞吐量）：我们发现，粒度和稀疏度的提升会使 MoE 越来越受内存带宽限制。为缓解这一瓶颈，我们利用矩阵乘法（GEMM）与 IO 操作的异步性，通过重叠两者以最大化吞吐量。对于同一细粒度 7B MoE 模型，我们的方法在前向传播中比高度优化的 DeepGEMM 基线快 43%，在反向传播中比主流 MoE 基线（ScatterMoE 和 MoMoE）分别快 83% 和 115%。为评估这些技术的性能，我们通过全面的内核级剖析和 IO 感知的 MoE 计算路径探索，进行了深入的性能分析。
  • 消除稀疏 MoE 计算浪费的令牌舍入路由：我们提出一种“即插即用”的路由算法，可将每个专家的令牌数量舍入为分组矩阵乘法所用瓦片大小（如 128）的整数倍。这种舍入操作在尽可能保留原始令牌-专家分配关系的同时，减少了填充操作造成的计算浪费。该算法确保：对于每个专家，其令牌数量与原始 Top-K 令牌选择结果的最大偏差不超过一个瓦片大小。这种方法可有效消除分组矩阵乘法中的填充浪费，同时保证令牌总数的期望不变，且在高稀疏 MoE 训练场景下仍能保持稳定的令牌选择精度。我们在 14 亿参数的稀疏训练场景中验证了该令牌舍入策略的性能，结果表明其计算吞吐量持续超过传统 Top-K 令牌选择路由；在高稀疏场景下，端到端 MoE 计算的浮点运算效率（TFLOPS）提升高达 16%。

我们已发布 SonicMoE，主要基于 CuTe-DSL（NVIDIA，2025c）编写，提供 PyTorch 接口，并采用宽松许可证，以助力研究人员和从业者。

二、背景

首先，在 2.1 节中概述 MoE 架构和采用分组矩阵乘法的标准 MoE 内核；在 2.2 节中讨论粒度和 MoE 稀疏度对 MoE 训练效率的影响；最后在 2.3 节中分析 MoE 路由方法对 MoE 模型质量和训练效率的影响。

2.1 基于分组矩阵乘法的 MoE

现代 GPU 支持张量核心，这是一种专门用于高效矩阵乘法的硬件单元（NVIDIA，2022）。

矩阵乘法内核通常包含三个阶段：序幕（开始加载输入数据）、主循环（持续加载输入并执行矩阵乘法）和尾声（对矩阵乘法输出执行额外的 IO/数学操作）。内核会将计算“瓦片化”，并可选择对维度进行填充，以确保计算与硬件允许的瓦片大小对齐。

本文采用大多数 BLAS 库中的标准矩阵乘法符号：对于矩阵乘法 ，定义 、、，其中 为问题维度。该符号被 CUTLASS（NVIDIA，2025a）采用，后者是 CUDA 平台上高效的矩阵乘法实现库。

在 NVIDIA Hopper GPU 上，矩阵乘法采用“生产者-消费者”范式异步执行（Shah 等，2024）：

  • 生产者负责将数据瓦片从高带宽内存加载到共享内存。
  • 而消费者线程束组负责矩阵乘法计算（Shah 等，2024）。
  • 补充：线程束是 NVIDIA GPU 上的基本执行单元。Hopper 架构的矩阵乘法中，一个线程束组包含 4 个连续线程束（共 128 个线程）。Hopper GPU 提供高吞吐量的 WGMMA 指令用于 MMA，该指令由线程束组集体发起。
  • 在序幕和主循环阶段，生产者线程束组获取数据瓦片并缓存到专用流水线，消费者线程束组从该流水线读取缓存的瓦片，执行瓦片矩阵乘法并沿矩阵乘法的 维度累积结果。
  • 主循环结束后进入尾声阶段，消费者线程束组对最终 MMA 结果执行后处理（如激活函数、将结果写回 HBM）。

一个 MoE 块通常包含令牌路由器和多个更小且尺寸通常相等的子网络，即“专家”。路由器负责将令牌分配给专家，专家随后使用这些令牌进行计算；层中所有专家的输出会被聚合，再传递到下一层。特别说明，本文中：

  • “MoE 路由”指确定每个令牌激活的专家及相关路由元数据的过程。
  • “MoE 计算”指每个专家处理分配的令牌及专家结果聚合的过程。
    • MoE 计算可通过分组矩阵乘法实现，Grouped GEMM 是一组可能具有不同{M, N, K}维度的矩阵乘法操作。算法 1 展示了如何用分组矩阵乘法执行 MoE 前向传播。
    • 算法 2、3、5 是 SonicMoE 的 MoE 计算组件，可兼容任意路由算法。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

算法 1：基于分组矩阵乘法的 MoE 前向传播

如算法 1 所示，在前向传播及反向激活梯度计算中，每个专家分配到的令牌数量是可变的。此时需执行“M 维度（令牌维度）可变、N 和 K 维度（专家权重维度）固定”的分组矩阵乘法，我们称之为“varlen-M 分组矩阵乘法”。在反向权重梯度计算中，嵌入维度和中间隐藏层大小是固定的，而需沿令牌维度进行归约，我们称之为“varlen-K 分组矩阵乘法”。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 2：MoE 计算通常需要分组矩阵乘法。每个专家的输入要么是从输入张量的不同位置收集的（上图），要么是从分组输入数组的连续块中读取的（下图）

对于每个分组矩阵乘法，输入数据可能是从不同位置收集的，也可能是连续存储的，如图 2 所示。例如在算法 1 中，上投影（up-proj）的输入是“收集所得”，而下投影（down-proj）的输入已是“连续存储”。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

算法2 SonicMoE的MoE核前向传播。存储在HBM中的变量用蓝色标注。load（加载）和store（存储）分别表示从HBM加载和存储到HBM中

2.2 MoE 计算

算术强度 （arithmetic intensity）定义为“浮点运算次数（FLOPs）与数据传输字节数（IO）的比值”，是判断内核“内存受限”（内存 IO 成本主导）或“计算受限”（计算吞吐量主导）的关键指标。

对于采用 SwiGLU 激活函数的专家，标准 MoE 计算可分解为以下组件：

上投影下投影

其中，$X_j$ 表示专家 $j$ 接收的输入令牌矩阵，$n_j$ 为路由到专家 $j$ 的令牌数。

  • 上投影的浮点运算次数为 $2 cdot d_{text{model}} cdot d_{text{inter}} cdot n_j$，即矩阵乘法 $X_j W_1$ 的浮点运算次数为 $2 cdot d_{text{model}} cdot d_{text{inter}} cdot n_j$，高带宽内存（HBM）数据传输字节数为 $2 cdot d_{text{model}} cdot n_j$，这里忽略 $H_j$ 的写入字节数；
  • 下投影的浮点运算次数为 $2 cdot d_{text{inter}} cdot d_{text{model}} cdot n_j$，即矩阵乘法 $H_j W_2$ 的浮点运算次数为 $2 cdot d_{text{inter}} cdot d_{text{model}} cdot n_j$，高带宽内存数据传输字节数为 $2 cdot d_{text{inter}} cdot n_j$。

假设 $k$ 为 MoE 激活比（$k$ 为每令牌激活专家数），$g$ 为专家粒度，且令牌路由均匀（即 $n_j = frac{k}{g} cdot n$，$n$ 为总令牌数），则前向传播中专家 $j$ 的算术强度（忽略 $H_j$ 的写入字节数）为：

$$
text{算术强度} = frac{4 cdot d_{text{model}} cdot d_{text{inter}} cdot n_j}{2 cdot d_{text{model}} cdot n_j + 2 cdot d_{text{inter}} cdot n_j} = frac{2 cdot d_{text{model}} cdot d_{text{inter}}}{d_{text{model}} + d_{text{inter}}}
$$

对于特定模型规模（$d_{text{model}}$ 固定），上面公式表明：增加粒度（$g$ 增大）或增加稀疏度（$k$ 减小）会导致算术强度降低。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 3：在 14 亿（1.4B）到 1200 亿（120B）参数的等浮点运算次数（iso-FLOPs）训练下，不同 MoE 配置的单层前向传播 IO 成本与专家粒度的关系（配置详情见表 9a）。实验中保持 MoE 激活比 $k$ 和每层 MoE 参数数量 $P_{text{layer}}$ 不变；当提升专家粒度 $g$ 时，会减小专家中间维度 $d_{text{inter}}$，同时保持 $P_{text{layer}}$ 和 $d_{text{model}}$ 不变

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表9：图3、图13和图14的基准配置

这是因为 IO 成本随专家粒度线性增长，如图 3 所示。因此，对于细粒度 MoE（高 $g$），通过最大化减少 IO 访问和掩盖 IO 延迟来应对增加的 IO 成本，变得至关重要。在本文中，“细粒度 MoE”指专家中间维度 $d_{text{inter}}$ 小于嵌入维度 $d_{text{model}}$ 的 MoE，实验均基于“等浮点运算次数”和“等参数规模”设置。

我们将在第 3 节中介绍内存高效的 MoE 内核设计，并在第 4 节中讨论减少 IO 访问和延迟的技术。

现有 MoE 内核设计：目前已有多种 MoE 实现方案，包括 ScatterMoE（Tan 等，2024）、MoMoE（Costin 等，2025）、MegaBlocks（Gale 等，2023）和 Megatron（Shoeybi 等，2019）。然而，这些方案并未针对“IO 成本随专家粒度线性增长”的细粒度 MoE 场景设计，如图 3 所示。

相比之下，我们的内核设计 SonicMoE 可最大程度降低 IO 成本对训练吞吐量的影响。在第 4 节和图 14 中，我们将展示：随着专家粒度 $g$ 的增加，SonicMoE 因 IO 感知优化，较现有 MoE 内核的相对速度提升会更大。关于 SonicMoE 与现有 MoE 内核的技术差异，我们在附录 B 中详细阐述，并在表 2 中进行汇总。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

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

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 14：H100 GPU 上不同 MoE 内核的前向与反向传播 TFLOPS。DeepGEMM 未提供高效的路由实现、聚集和专家聚集内核，因此我们对这些模块分别采用标准 PyTorch 实现（“DeepGEMM-pt”）或高度优化的内核（“DeepGEMM++”）。在反向传播中，“DeepGEMM++”和“DeepGEMM-pt”采用与 SonicMoE 相同的计算路径，但需单独启动内核来计算 dS、和 dSwiGLU。MoE 配置与图 13 相同

2.3 MoE 路由方法

在 MoE 中，路由 （routing）决定每个令牌激活哪些专家。

令牌选择（TC）路由 （token choice routing）是 MoE 模型的默认路由方式（Shazeer 等，2017），该方式中每个令牌独立选择激活的专家。

  • 最常用的是Top-K 令牌选择路由 （top-K TC routing）：对于令牌 $i$，选择路由分数 $s_{i,j}$（令牌 $i$ 对专家 $j$ 的路由分数）最高的 $k$ 个专家，即 $text{top-k}j(s{i,j})$。
  • 除 Top-K 外，Huang 等（2024）提出“令牌选择 Top-P 路由”（token-choice top-P routing），可在训练中灵活分配计算资源，但该方法会导致“每令牌激活专家数”和“浮点运算次数消耗”的不确定性。Zeng 等（2024b）也提出类似思路，通过“空专家”（null experts）动态调整激活专家数。

除令牌选择（TC）路由外，专家选择（EC）路由 （expert choice routing）旨在避免专家并行时的负载不均衡（Zhou 等，2022），该方式由专家选择令牌而非令牌选择专家。然而，EC 路由无法直接用于推理（与自回归解码不兼容），若在推理时切换回 TC 路由会导致“训练-推理不匹配”；此外，EC 路由还会因“未来令牌信息泄露”破坏因果性（Wang 等，2024）。

为解决 EC 路由的推理问题，Raposo 等（2024）提出两种方案：
1. 引入辅助损失，促进 TC 与 EC 路由结果的一致性；
2. 训练辅助路由器，显式预测 EC 路由器的选择结果，并在推理时使用该辅助路由器。

本文提出一种新颖的“分组矩阵乘法瓦片感知令牌舍入” （tile-aware token rounding）方法：将每个专家接收的令牌数量（“专家频率”）舍入为“分组矩阵乘法瓦片大小”的邻近整数倍，且每个专家的令牌数量调整不超过一个瓦片。

这种方法可有效减少稀疏 MoE 训练中“分组矩阵乘法填充操作”造成的计算浪费，同时保留训练后 MoE 模型的推理质量。

现有相关工作包括“令牌丢弃与重路由”如 Rectify-Router（Zeng 等，2024a），但这些工作未关注分组矩阵乘法的瓦片结构；其他工作如 TMA 自适应 FP8 分组矩阵乘法（Fu 等）聚焦于减少填充相关的负载流量，但未解决“瓦片大小不对齐导致的矩阵乘法计算浪费”问题。

三、内存高效的 MoE 算法

首先，在第 3.1 节中概述 SonicMoE 的高层内核设计，该设计对应算法 2、3、5 所示的 SonicMoE MoE 计算流程；随后在第 3.2 节中聚焦 SonicMoE 的激活内存占用优化。

3.1 SonicMoE 的 MoE 内核概述

SonicMoE 的 MoE 计算共启动 8 个内核：
* 前向传播：上投影（A）内核、下投影（Y）内核、专家聚合（O）内核（共 3 个）；
* 反向传播：下投影激活梯度（dH）内核、上投影激活梯度（d）内核、专家聚合激活梯度（dX）内核（聚合所有专家的 d）、上投影权重梯度（dW₁）内核、下投影权重梯度（dW₂）内核（共 5 个）。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 4：SonicMoE 的 8 个内核计算流程（黄色框分组）。前向传播启动 3 个内核，反向传播启动 5 个内核。指向黄色圆圈的箭头表示“从高带宽内存（HBM）加载变量到共享内存（SRAM）”，从黄色圆圈出发的箭头表示“将变量存储到高带宽内存（HBM）”。高带宽内存中的所有变量用不同色框标注：紫色框表示前向和反向传播的输出，蓝色框表示中间变量或权重，红色框表示所有缓存的激活值（X、H、S）。算法 2 详细描述 SonicMoE 的前向传播，算法 3 和 5 详细描述 SonicMoE 的反向传播。

SonicMoE 的 MoE 计算实现具有高度模块化特性，仅包含两部分：
1. 带模块化融合的优化分组矩阵乘法内核；
2. 优化的专家聚合内核。

主机（host）会调度最优的矩阵乘法配置和加载/存储策略，以启动上述 8 个内核。尽管模块化程度高，SonicMoE 仍实现了最优的训练吞吐量和最小的激活内存占用。

3.2 激活内存效率

MoE 前向和反向传播的总浮点运算次数为（前向 6 倍，反向 12 倍）。对于固定的令牌数（T）和嵌入维度（d），需保持专家中间维度（h）与每令牌激活专家数（K）的乘积恒定以确保浮点运算次数不变。因此，增加粒度（h 增大）需同时减小 K 并按比例增大专家数（E）。

由此可知，若缓存内存占用为 O(ThK) 的激活值，会导致激活内存随粒度线性增长。现有 MoE 内核（如 ScatterMoE）的激活内存即存在这一问题。为避免激活内存依赖于粒度，需避免缓存以下激活值：
* 下投影输出 Y（内存占用 O(ThK)）；
* 收集所得的输入 X（内存占用 O(TdK)）。

此外，为避免反向传播中峰值激活内存增大，还需避免将“Y 的梯度 dY”和“O 的收集结果”写入高带宽内存（HBM）。具体优化策略如下：

  1. 输入收集与 HBM 加载融合：对于输入 X 和 O 的梯度 dO，将“收集操作”与“HBM 加载”融合，可避免在 HBM 中实例化和缓存这些变量。如图 6 所示，这种融合可显著提升细粒度 MoE 的吞吐量。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图6：在H100上进行7B MoE训练时，不同MoE内核的运行时间分解（毫秒↓），其中（T, d, n, E, K）=（24576, 1536, 256, 128, 8）。我们为内存受限的内核（gather、SwiGLU/dSwiGLU以及专家聚合内核）标注了模型内存带宽（TB/s ↑），为分组GEMM内核标注了计算吞吐量（TFLOPS ↑，图中缩写为TF/s）。请注意，此分布图是按内核运行时语义分组的，一个块可能包含多个实际的内核计时结果。例如，左侧子图中的“与路由相关”部分既包括路由GEMM的时间，也包括路由元数据计算的时间。此外，本图中未考虑内核间的CUDA流空泡时间。对于Megatron，我们使用GroupedMLP；对于MegaBlocks，我们使用ParallelDroplessMLP。DeepGEMM在正向传播过程中没有提供高效的路由实现、gather以及专家聚合内核，因此我们为这些部分使用了标准的PyTorch实现（“DeepGEMM-pt”）或我们高度优化的内核（“DeepGEMM++”）。在反向传播过程中，“DeepGEMM++”和“DeepGEMM-pt”都采用与SonicMoE相同的计算路径，不同之处在于我们会启动单独的内核来共同计算dS、A’和dSwiGLU。DeepGEMM++是在不修改DeepGEMM源代码的情况下，基于DeepGEMM SM90 BF16分组GEMM内核构建的最佳MoE实现。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图22：H100上分组通用矩阵乘法（GEMM）和专家聚合内核的吞吐量。“SonicMoE（gemm + gth w. sum）”是SonicMoE的最终设计选择，如图9左侧策略所示。我们将此设计与在SonicMoE上实现图9中间策略的“SonicMoE（gemm w. sct + sum）”进行比较。对于“SonicMoE（gemm + gth w. sum）”和“SonicMoE（gemm w. sct + sum）”，我们使用相同的瓦片大小和其他通用矩阵乘法（GEMM）配置。我们还与ScatterMoE的设计（融合散射与通用矩阵乘法（GEMM）+ torch.bmm，标记为“ScatterMoE（gemm w. sct + BMM）”）以及MoMoE的设计（融合散射与通用矩阵乘法（GEMM）+ torch.sum，标记为“MoMoE（gemm w. sct + sum）”）进行比较。对于每种方法，我们用透明条形图表示通用矩阵乘法的TFLOPS，用不透明条形图表示通用矩阵乘法和专家聚合总运行时间的TFLOPS。

  1. 梯度计算路径优化（无需 dY 和 Y）：传统方法计算路由器分数梯度 dS 和上投影输出梯度 dH 时，需依赖 Y 和 dY；而我们通过推导替代计算路径，在不增加浮点运算次数的前提下，避免使用 Y 和 dY。具体而言，我们将 dS 和 dH 展开为不包含 Y 和 dY 的等式（附录 C 详细说明）。SonicMoE 的 dH 内核如算法 3 所示。

通过上述优化，SonicMoE 每层仅需缓存以下数据，总内存占用为 O(Td) 字节：
* 输入激活值 X（内存 O(Td)）；
* 上投影输出 H（内存 O(ThK)）；
* 路由元数据（与激活值相比，大小可忽略不计）。

这一激活内存占用与“具有相同激活参数数量的密集型模型”相同，是“不通过矩阵乘法重计算激活值”的前提下，反向传播所需的最小激活内存。补充两点：
* SonicMoE 仅需在 dH 内核的尾声阶段，通过寄存器重计算 Y（从 H 推导），无需额外存储。
* 尽管 SonicMoE 仍会实例化临时变量 Y，但可在每层计算结束后回收该变量的内存。只要 MoE 层数（7B 以上 MoE 通常 ≥32 层）大于每令牌激活专家数（K），Y 的瞬时内存占用就可忽略。若要完全消除 Y 的实例化，需在全局内存中执行原子加法（图 9 右图），但这会带来新问题：确定性缺失、数值精度损失（16 位脑浮点 BF16 的原子加法）、与专家并行的全收集（all-gather）通信不兼容。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 9：结果存储与令牌结果聚集的可能策略。SonicMoE 选择第一种策略（左）：每个专家在矩阵乘法收尾阶段通过 TMA 直接存储连续打包的输出；在专家聚集内核中，每个令牌聚集并求和所有路由专家的输出。ScatterMoE 和 MoMoE（中）选择在收尾阶段将 HBM 存储与分散操作融合，随后启动求和内核。需注意：每个令牌聚集（左）分组矩阵乘法结果，与每个专家分散（中）分组矩阵乘法输出，在数学上是等价的。在图 22 中，我们在 SonicMoE 上实现了两种策略，发现左策略比中策略提速 17%。还可在收尾阶段融合原子加法（右），以避免启动专家聚集内核，但原子加法操作会带来新问题，如非确定性[20]和数值精度问题（对于 BF16 原子加法）。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图22：H100上分组通用矩阵乘法（GEMM）和专家聚合内核的吞吐量。“SonicMoE（gemm + gth w. sum）”是SonicMoE的最终设计选择，如图9左侧策略所示。我们将此设计与在SonicMoE上实现图9中间策略的“SonicMoE（gemm w. sct + sum）”进行比较。对于“SonicMoE（gemm + gth w. sum）”和“SonicMoE（gemm w. sct + sum）”，我们使用相同的瓦片大小和其他通用矩阵乘法（GEMM）配置。我们还与ScatterMoE的设计（融合散射与通用矩阵乘法（GEMM）+ torch.bmm，标记为“ScatterMoE（gemm w. sct + BMM）”）以及MoMoE的设计（融合散射与通用矩阵乘法（GEMM）+ torch.sum，标记为“MoMoE（gemm w. sct + sum）”）进行比较。对于每种方法，我们用透明条形图表示通用矩阵乘法（GEMM）的每秒万亿次浮点运算（TFLOPS），用不透明条形图表示通用矩阵乘法（GEMM）和专家聚合总运行时间的每秒万亿次浮点运算（TFLOPS）。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

算法 3：SonicMoE 的 MoE 内核反向传播（下投影部分）

四、面向 IO 感知的内核设计（IO-aware kernel design）

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表9：图3、图13和图14的基准配置。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 3：在 14 亿（1.4B）到 1200 亿（120B）参数的等浮点运算次数（iso-FLOPs）训练下，不同 MoE 配置的单层前向传播 IO 成本与专家粒度的关系（配置详情见表 9a）。实验中保持 MoE 激活比 和每层 MoE 参数数量 不变；当提升专家粒度 时，会减小专家中间维度 ，同时保持 和 不变。

细粒度 MoE（Fine-grained MoE）的表达能力源于每个令牌（token）选择专家（expert）的多样性，但这反过来也导致 IO 成本随专家粒度（expert granularity）呈线性增长见图 3。

为维持高吞吐量，我们需要最大限度地实现：（1）通过融合（fusion）减少 IO 访问；（2）将 IO 延迟与计算过程重叠。

  • 首先，我们在 4.1.1 节和 4.1.2 节分别分析令牌聚集（gather）与计算的融合，以及数学运算与收尾阶段（epilogue）的 IO 融合；
  • 随后在 4.2 节阐述将矩阵乘加（MMA）与 IO 重叠的技术；
  • 最后在 4.3 节分析 SonicMoE 的 Top-K 排序内核。在附录 B 中，我们将 SonicMoE 与其他 MoE 内核设计进行对比，并在表 2 中总结关键差异。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

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

4.1 SonicMoE 的分组矩阵乘法（Grouped GEMM）

SonicMoE 基于高效的变长 M 分组矩阵乘法（varlen-M Grouped GEMM）和变长 K 分组矩阵乘法（varlen-K Grouped GEMM）构建。

在分组矩阵乘法内部，我们将聚集操作与激活值加载融合（4.1.1 节），并将 SwiGLU/dSwiGLU/dS 与收尾阶段融合 （4.1.2 节）。

  • 聚集融合使 SonicMoE 比需要独立聚集内核的 MoE 内核（如 MegaBlocks、Megatron 以及基于 DeepGEMM 库[69]优化的 MoE 前向实现 DeepGEMM++）更快；
  • 收尾阶段融合则使 SonicMoE 在反向传播中比 ScatterMoE 更快。

这些融合操作减少了不必要的 IO 访问，且可与矩阵乘加（MMA）计算重叠（详见 4.2 节）。

4.1.1 与 HBM 加载融合的聚集操作

SonicMoE 的分组矩阵乘法既支持连续打包的输入（contiguously-packed inputs），也支持从不同位置聚集的输入（见图 2）。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 2：MoE 计算通常需要分组矩阵乘法。每个专家的输入要么是从输入张量的不同位置收集的（上图），要么是从分组输入数组的连续块中读取的（下图）。

对于后者，我们将输入聚集操作与从全局内存（GMEM，通常即 HBM）到共享内存（SMEM）的输入加载过程融合，以便批量处理并在张量核心（Tensor Core）上执行矩阵乘法[10,53]。这一过程包括：

  1. 获取每个专家的路由令牌索引；
  2. 利用这些索引，通过 Blackwell 和 Hopper 架构的cp.async指令（异步复制指令，用于高效从全局内存加载数据到共享内存）聚集激活值。

其中，第二步通常没有更优的替代方案[10]，但同步索引获取仍可通过预取（prefetching）和生产者线程束（producer warps）协同获取来优化，具体策略如图 18 所示。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图18：在H100 GPU上，针对变长-M分组GEMM的M维度（左图）和变长-K分组GEMM的K维度（右图）进行聚合时的索引预取策略。对于M维度上的聚合（左图），我们让每个线程在主循环之前独立地将索引预取到各自的寄存器中。对于K维度上的聚合（右图），我们在共享内存（SMEM）上创建一个缓冲区，并让4个生产者线程束协作将索引预取到共享内存中，每个生产者线程再从该共享内存缓冲区读取数据到各自的寄存器中。

Hopper GPU
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图6：在H100上进行7B MoE训练时，不同MoE内核的运行时间分解（毫秒↓），其中（T, d, n, E, K）=（24576, 1536, 256, 128, 8）。我们为内存受限的内核（gather、SwiGLU/dSwiGLU以及专家聚合内核）标注了模型内存带宽（TB/s ↑），为分组GEMM内核标注了计算吞吐量（TFLOPS ↑，图中缩写为TF/s）。请注意，此分布图是按内核运行时语义分组的，一个块可能包含多个实际的内核计时结果。例如，左侧子图中的“与路由相关”部分既包括路由GEMM的时间，也包括路由元数据计算的时间。此外，本图中未考虑内核间的CUDA流空泡时间。对于Megatron，我们使用GroupedMLP；对于MegaBlocks，我们使用ParallelDroplessMLP。DeepGEMM在正向传播过程中没有提供高效的路由实现、gather以及专家聚合内核，因此我们为这些部分使用了标准的PyTorch实现（“DeepGEMM-pt”）或我们高度优化的内核（“DeepGEMM++”）。在反向传播过程中，“DeepGEMM++”和“DeepGEMM-pt”都采用与SonicMoE相同的计算路径，不同之处在于我们会启动单独的内核来共同计算dS、A’和dSwiGLU。DeepGEMM++是在不修改DeepGEMM源代码的情况下，基于DeepGEMM SM90 BF16分组GEMM内核构建的最佳MoE实现。

如图 6 所示，聚集融合使 SonicMoE 在 H100 GPU 上相比 DeepGEMM 等现有 MoE 内核具有显著优势。尽管 DeepGEMM 的变长 M 分组矩阵乘法内核已高度优化，但它假设输入已连续打包且填充至 128 的倍数，因此在分组矩阵乘法之前，需要单独启动一个聚集和填充内核。图 6 显示，即使提供了优化的聚集内核，且 DeepGEMM 的变长 M 分组矩阵乘法也已高度优化，但聚集 X（2TKd 字节数据）所需的大量 IO 仍导致 DeepGEMM++ 比 SonicMoE 更慢。

在反向传播中，上投影和下投影的权重梯度需要聚集 X 和激活梯度，而 H 的梯度也需要聚集。尽管反向传播中有更多内核需要聚集操作，但现有方案仅在前向传播中融合聚集操作，在反向传播中仍需单独启动聚集内核。融合反向传播中的聚集操作可减少 2TKd 字节的 IO 成本，并大幅缩短细粒度 MoE 的训练时间。例如，在图 6 中，ScatterMoE 和 MoMoE 反向传播中的两次聚集操作分别消耗了总反向时间的 19.6% 和 20.6%，甚至比它们的上投影权重梯度内核耗时更长。

Blackwell GPU

在撰写本文时，SonicMoE 已支持 Blackwell GPU 的变长 M 分组矩阵乘法及其聚集融合。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 5：在 Blackwell GPU 上使用 2 线程块集群，通过 cp.async 实现聚集融合的流水线结构

在 Blackwell GPU 上，当使用 2 个线程块集群执行矩阵乘法时，cp.async 指令的聚集融合面临一个架构挑战：该指令仅能在同一线程块内信号完成。然而，Blackwell 的 2 线程块矩阵乘法要求主线程块中的 MMA 指令等待两个线程块的聚集操作完成。为解决这一限制，线程块 1 需要一个专用的中继线程束：该线程束接收 cp.async 的完成信号，并通过集群级同步原语将信号转发给主线程块的 MMA 线程束。这种中继机制增加了调度复杂度，但实现了跨 2 线程块集群的高效聚集融合，维持了分组矩阵乘法的高吞吐量。

4.1.2 收尾阶段融合

我们利用收尾阶段计算，通过以下设计选择最大限度减少不必要的 IO 访问：

  • SwiGLU 与 dSwiGLU 融合：我们将 SwiGLU 激活函数及其反向计算分别与前向传播上投影内核和反向传播下投影激活梯度内核的收尾阶段融合。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图6：在H100上进行7B MoE训练时，不同MoE内核的运行时间分解（毫秒↓），其中（T, d, n, E, K）=（24576, 1536, 256, 128, 8）。我们为内存受限的内核标注了模型内存带宽（TB/s ↑），为分组GEMM内核标注了计算吞吐量（TFLOPS ↑）。请注意，此分布图是按内核运行时语义分组的，一个块可能包含多个实际的内核计时结果。例如，左侧子图中的“与路由相关”部分既包括路由GEMM的时间，也包括路由元数据计算的时间。此外，本图中未考虑内核间的CUDA流空泡时间。对于Megatron，我们使用GroupedMLP；对于MegaBlocks，我们使用ParallelDroplessMLP。DeepGEMM在正向传播过程中没有提供高效的路由实现、gather以及专家聚合内核，因此我们为这些部分使用了标准的PyTorch实现或我们高度优化的内核。在反向传播过程中，“DeepGEMM++”和“DeepGEMM-pt”都采用与SonicMoE相同的计算路径，不同之处在于我们会启动单独的内核来共同计算dS、A’和dSwiGLU。DeepGEMM++是在不修改DeepGEMM源代码的情况下，基于DeepGEMM SM90 BF16分组GEMM内核构建的最佳MoE实现

如图 6 所示，尽管 DeepGEMM++ 具有高度优化的分组矩阵乘法和 SwiGLU 内核，但其上投影和 SwiGLU 的总耗时仍比 SonicMoE 的上投影更长——尽管 SonicMoE 除了 SwiGLU 外还额外进行了聚集融合。

  • 在反向传播下投影激活梯度内核的收尾阶段计算 dH 与 dS：这种重度收尾阶段融合为 SonicMoE 带来了远超其他设计的加速效果。在 7B MoE 训练中，SonicMoE 的 dH 内核产生的输出与 ScatterMoE 的下投影激活、dS 和 dSwiGLU 的组合输出完全相同，但总耗时远短于 ScatterMoE。此外，SonicMoE 相比 DeepGEMM++ 也实现了加速：DeepGEMM++ 需启动一个高效分组矩阵乘法和一个独立的优化内核来共同计算 dSwiGLU 和 dS。

在附录中，我们证明 SonicMoE 计算 dS 的方式对于细粒度 MoE 而言，在计算效率和激活内存效率上均为更优选择。然而，ScatterMoE 和 MoMoE 均选择通过计算 dS，这需要额外加载 2TKd 字节的 HBM 数据，且需缓存 2TKd 字节的激活值。在图 6 中，ScatterMoE 为计算 dS 单独启动一个内核，而 MoMoE 将 dS 与上投影激活梯度融合，这两种方式的耗时均远长于 SonicMoE 的上投影激活梯度。

反向传播下投影激活梯度内核的重度收尾阶段融合吞吐量，通过异步 IO 与 MMA 的重叠得到提升。这种重叠使 SonicMoE 即使在 dH 内核的重度收尾阶段融合中，仍能同时维持合理的训练吞吐量和内存带宽。

4.2 矩阵乘加（MMA）与异步 IO 重叠

Hopper GPU

在 NVIDIA Hopper GPU 中，矩阵乘法通过生产者-消费者范式异步执行：

  • 生产者线程束专门负责将数据瓦片从高带宽内存加载到共享内存。
  • 消费者线程束组负责矩阵乘加计算。
  • 在收尾阶段和主循环中，生产者线程束组获取数据瓦片并缓存到专用流水线，消费者线程束组则从该流水线读取缓存的瓦片，执行瓦片矩阵乘法并沿矩阵乘法的 K 维度累积结果。
  • 主循环结束后进入收尾阶段，消费者线程束组对最终 MMA 结果执行后处理，并将结果写回 HBM。

线程束（warp）是 NVIDIA GPU 上的基本执行单元。Hopper GPU 矩阵乘法中的线程束组（warpgroup）由 4 个连续线程束（共 128 个线程）组成。Hopper GPU 提供高吞吐量的 WGMMA 指令（瓦片矩阵乘加指令），该指令由线程束组集体发起。

假设存在 2 个消费者线程束组，可以有两种调度策略：让它们协同发起大瓦片尺寸的 WGMMA 指令；或者让一个线程束组执行 IO 操作，另一个执行小瓦片尺寸的 MMA 计算，从而实现 IO 与计算的重叠。完成一轮计算后，切换两个线程束组的角色（即交错执行 IO 与 MMA）。这种在 Hopper GPU 上实现 IO 与计算重叠的策略通常被称为“乒乓调度（Ping-Pong scheduling）”[49,60]，见图 7。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

乒乓调度对于在重度收尾阶段下维持张量核心高吞吐量尤为重要。
* 例如，在前向传播的下投影 Y 内核中，收尾阶段涉及相对于主循环较重的 HBM 存储 IO（2TKd 字节）。
* 在反向传播的下投影激活梯度（dH）内核中，收尾阶段需要加载 H（4TKn 字节），并执行多个激活函数计算和归约操作，以计算并存储 dH、dS 等。

需要指出的是，MMA 与 IO 重叠及乒乓调度的概念在其他领域（如 Flash Attention 3[49]）中已被熟知，但将其应用于细粒度 MoE 内核设计以应对 IO 成本增长的问题，是一项创新。

补充说明：选择协同调度还是乒乓调度，很大程度上取决于选择更大的瓦片尺寸还是更多的收尾阶段重叠。细粒度 MoE 两者都需要：和内核通常具有较长的主循环，因此通常需要最大的瓦片尺寸，此时协同调度几乎总是更优；而 Y 内核和 dH 内核具有重度收尾阶段，此时乒乓调度通常更有利。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

DeepGEMM 的 SM90 BF16 变长 M 分组矩阵乘法内核[11]未实现乒乓调度。这种设计选择适用于轻量级收尾阶段（如前向传播上投影），但在具有重度收尾阶段的前向传播下投影中表现较差：如图 6 所示，DeepGEMM 的吞吐量为 413 TFLOPS 和 2.15 TB/s，而 SonicMoE 则达到 485 TFLOPS 和 2.52 TB/s。在图 19 中，SonicMoE 的下投影吞吐量平均比 DeepGEMM 高 10.0%。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

除乒乓调度外，SonicMoE 还依赖异步 TMA 操作执行全局内存到共享内存的加载（GMEM-to-SMEM）和共享内存到全局内存的存储（SMEM-to-GMEM）。我们将以下异步 IO 与 MMA 操作重叠：

  • dH 内核收尾阶段的异步 TMA 加载：在 dH 内核的收尾阶段，需要加载 H 以从 dA 计算 dH。我们为 H 的异步 TMA 加载创建专用流水线，使其与收尾阶段各步骤中的其他操作重叠。在图 7 中，消费者线程束组中的透明 TMA 块即代表这种异步收尾阶段加载。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升
  • 前向传播下投影 Y 内核和反向传播上投影激活梯度内核的异步 TMA 存储：SonicMoE 为所有 6 个分组矩阵乘法均应用了异步 TMA 存储。在前向传播下投影和反向传播上投影激活梯度中，SonicMoE 不将分散操作（scatter）与 HBM 存储融合——而 ScatterMoE[17]和 MoMoE[18]均选择将 HBM 存储与分散操作融合。这是因为：（1）分散融合需要更多的同步索引获取和地址计算[19]；（2）在 Hopper GPU 上，分散融合需要同步的共享内存到全局内存存储指令[20]。对于细粒度 MoE 的前向传播下投影和反向传播上投影激活梯度内核计算，同步全局内存存储会阻塞下一个瓦片的 MMA 执行，导致吞吐量显著下降（约 20%），如图 8 所示。我们还注意到，乒乓线程束组调度无法完全恢复同步收尾阶段 IO 操作导致的吞吐量下降——因为收尾阶段的消费者线程束组会被阻塞，直到当前同步全局内存存储完成后，才能与 MMA 线程束组切换角色。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 8：异步 TMA 存储（上）与同步 st.global 存储（下）的对比示意图。异步 TMA 存储具有更高的内存带宽，且能自然与张量核心 MMA 操作重叠；而 Hopper GPU 上分散融合所需的同步 st.global PTX 指令会阻塞下一个张量核心 MMA 瓦片的执行，导致内核运行时间更长。这一结论得到图 22 中透明柱的支持：“SonicMoE（gemm + gth w. sum）”（TMA 存储）比“SonicMoE（gemm w. sct + sum）”（st.global存储）平均提速 20.1%。因此，SonicMoE 不将分散操作与 HBM 存储融合，而是在专家聚集内核中让每个令牌聚集专家结果。ScatterMoE 和 MoMoE 均未采用此设计，因此 SonicMoE 在前向传播下投影内核中，相比两者分别实现了 1.75 倍和 3.11 倍的平均加速（见图 6）。

对于瓦片尺寸为的矩阵乘法，每个输出令牌需计算次分散索引；而对于 TMA 存储，在矩阵乘法收尾阶段无需获取分散索引，而是在聚集内核中执行聚集操作，此时可在整个 d 维度上完全复用相同的聚集索引，每个令牌的同步索引获取次数仅为 K。若矩阵乘法与分散操作融合，且需在 d 维度上复用索引，则每个线程块（CTA）需在处理下一个 前持续处理 d 维度上的所有矩阵乘法瓦片， 若专家数量 E 较小，极易导致流处理器（Streaming Processors）利用率不足。

在 Hopper GPU 上，若不使用 1D TMA 存储，分散融合与 HBM 存储仅能通过同步的 st.global PTX 指令实现（与聚集操作不同，cp.async 是异步的，但无法用于共享内存到全局内存的存储）。尽管 Blackwell GPU 支持异步的 st.async.release.global PTX 指令，但重复的索引获取仍会使分散操作成为较劣选择。

Blackwell GPU

在 NVIDIA Blackwell GPU 上，矩阵乘法内核在本质上仍采用“乒乓”调度思想，但其实现与 Hopper GPU 不同。

Blackwell 引入了张量内存（TMEM，Tensor Memory）——每个流多处理器（SM）上专用的 256KB 片上内存，组织为 512 列 ×128 行的 32 位单元[38,45]。矩阵乘法的累加器结果直接存储在 TMEM 中（而非寄存器），512 列的结构天然支持两阶段累加器流水线 ： 每个阶段使用 256 列，一个阶段通过新的 UMMA（统一矩阵乘加，Unified Matrix Multiply-Accumulate）指令执行 MMA 操作，另一个阶段执行收尾阶段计算。

与 Hopper GPU 的 WGMMA 指令（需线程束组级协同且占用大量寄存器内存）不同，Blackwell 的 UMMA 是单线程异步操作，消除了累加过程中的寄存器压力。 这一架构改进使收尾阶段线程束能够从一个 TMEM 阶段读取并处理结果，同时 MMA 线程束向另一个 TMEM 阶段累 加，实现了比 Hopper GPU 乒乓调度更优的收尾阶段与 MMA 操作重叠。

4.3 面向 MoE 的高效 Top-K 排序内核

现有 MoE 方案（如 ScatterMoE[21]、MoMoE[22]和 MegaBlocks[23]）均使用 PyTorch 的 torch.topk 函数计算每个令牌的专家分配。

我们发现，PyTorch 的 Top-K 内核耗时约占路由计算时间的 40%。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 10：将列索引比特打包到低尾数位后，对值进行排序。这种值格式确保排序结果具有稳定性。Triton 的官方 Top-K 内核也采用类似格式。

因此，SonicMoE 实现了一个高效的 Top-K 内核，以减少 PyTorch torch.topk 带来的开销。该内核支持专家数量、每个令牌激活专家数量[24]，并针对大令牌数量 T 的场景进行了优化。此外，该内核还支持在 Top-K 值上可选地融合 softmax 操作。 此处我们限制 E 和 K 的取值，是为了仅使用线程级和线程束级原语 。更大的需要共享内存（SMEM）缓冲区和块级同步，且对于大多数 MoE 模型而言，E 通常不会这么大。

Top-K 内核接收形状为的路由输出，并沿 T 维度并行处理。该内核对每一行（共 E 个值）执行 bitonic 排序[1]，并选择前 K 列作为排序结果。加载输入后，我们将前 K 列的列索引（用于 argtopK）打包到寄存器中 FP32 值的低位尾数中[25]—— 例外情况是，对于基数排序场景（值数量 ），我们采用最优低延迟排序网络[13]的比较策略，该策略提供最少的并行操作步骤和所需的比较-交换（compare-and-swap）调用次数。

此设计假设 FP32 值的低尾数位（最多为位，最大为 12 位）不会显著影响排序结果，或者说，前 K 个位置的相对排序差距大于。另一种比特打包设计是使用 FP64 值，但会导致所有交换操作变慢，且每个线程能处理的值数量减少。

bitonic 比较和合并操作在同一线程或同一线程束内通过线程束洗牌（warp-shuffle）指令执行。因此，所有交换和合并操作仅使用线程内或线程束内的寄存器，这使得该内核的内存带宽高于其他设计（如 PyTorch TopK[41]、Triton[56]和 Tilelang[59]官方示例、RTop-K[62]），如图 23 所示。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图23：在MoE前向传播过程中，使用BF16输入（第一行）和FP32输入（第二行）的Top-K核。配置与图13相同。“torch”是直接的torch.topk调用。“triton”和“tilelang”取自它们的官方示例，并稍作修改以支持BF16输入。对于triton官方核，为了公平比较，在本示例中我们移除了不必要的位矩阵存储并禁用了softmax融合。“RTop-K”（Xie等人，2025）仅支持FP32输入。我们为RTop-K设置了ϵ = 0，最大迭代次数为8。

由于每行值的分配列索引始终唯一，将列索引打包到低尾数位后，不会出现相等的值。因此，SonicMoE 的 Top-K 内核在 bitonic 比较和合并过程中不会出现平局（tie-breaking）情况，排序结果始终具有稳定性。

五、令牌舍入路由（Token rounding routing）

在本节中，我们分析稀疏 MoE 训练场景下的硬件效率，发现随着 MoE 稀疏度增加，分组矩阵乘法瓦片中因填充（padding）导致的计算浪费会累积到显著水平，这种现象被称为“瓦片量化效应 （tile quantization effect）”。为此， 我们提出一种新颖的路由方法——“令牌舍入（token rounding）”，以消除瓦片量化效应。

5.1 稀疏 MoE 的训练效率

如公式（4）所示，MoE 的算术强度还依赖于 MoE 激活率（激活专家比例）。当激活率降低时，每个专家接收的预期令牌数量也会线性减少，矩阵乘法计算会逐渐转向内存受限（memory-bound）状态，内核运行时间主要由内存 IO 成本决定。

瓦片量化效应

现代 GPU 上的矩阵乘法通常按瓦片（tile）计算[36]。若矩阵维度 M、N、K 中任一维度的大小不能被瓦片尺寸整除，则需填充到下一个瓦片尺寸的倍数。当输入尺寸（如每个专家的令牌维度）较小时，填充导致的 TFLOPS 浪费会变得非常显著，如图 11 和图 12 所示。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 11：在 MoE 前向与反向传播中，填充导致的计算浪费（TFLOPS）。图 16 右下角两个子图对此进行了进一步说明。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 12：稀疏 MoE 的瓦片量化效应示意图。为确保每个专家接收的令牌数为瓦片尺寸的倍数，需对令牌进行填充，导致计算浪费；令牌舍入（TR）的舍入子程序通过丢弃或补充令牌，避免了这种浪费。

因此，我们提出使用令牌舍入方法避免启动此类额外瓦片，从而实现更高效的训练。实验证明，令牌舍入方法在不影响模型质量的前提下，能显著提升训练吞吐量。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

算法 4：令牌舍入路由

5.2 令牌舍入路由

我们提出令牌舍入（TR）方法，它是一种两阶段排序算法（如算法 4 所示）。

令牌舍入算法首先计算标准令牌选择（TC）路由结果，然后对每个专家的令牌路由得分进行排序（类似专家选择（EC）路由的排序步骤）。在第二阶段排序中，选择丢弃第一阶段 TC Top-K 路由选中的令牌，或补充额外令牌。

在这两个阶段之间，我们对路由权重矩阵进行处理，确保 TC 选中的令牌始终优先于 EC 选中的令牌——这样，无论是丢弃还是补充令牌，都仅影响每个专家的最后一个输入瓦片。

令牌舍入需要一个“舍入与稀疏化”子程序来决定是丢弃还是补充令牌。默认选择是将专家的令牌频率舍入到最接近的瓦片尺寸倍数。

为简化实验，表 3 和图 16 中我们始终使用固定的瓦片尺寸，但需注意该尺寸依赖于 GPU 型号，在某些情况下还依赖于其他维度。本文聚焦于变长 M 分组矩阵乘法（它占 MoE 训练总计算量的 12/18）；另外 6/18 的计算量来自权重梯度内核中的变长 K 分组矩阵乘法，此时填充发生在 K 维度，且通常影响较小，因此本文不对此展开讨论。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表6：对算法4中配备不同舍入和稀疏化子例程的令牌舍入算法的评估。“PPL”指的是训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。

我们在表 6 中进行了消融实验，发现：
1. 令牌舍入算法对底层舍入子程序具有较强的鲁棒性。
2. 基于专家频率的最近舍入策略尽管简单，但足以实现优异的下游任务性能。

更多关于不同舍入子程序的讨论见附录 F.2。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

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

MoE 训练与推理质量

该算法保证：对于每个专家，其令牌分布与令牌选择（TC）路由的最大偏差不超过 1 个瓦片。我们发现，即使在稀疏 MoE 训练场景下，这一特性仍能保证模型性能的鲁棒性，且可作为稀疏 MoE 训练场景下令牌选择（TC）路由的替代方案（见表 3）。

我们还针对微批量大小 T 和瓦片尺寸对令牌舍入路由训练的 MoE 模型质量的影响进行了消融实验（表 7 和表 8），发现当令牌数量与瓦片尺寸适配时，令牌舍入路由的性能通常具有鲁棒性。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表7：当我们改变微批次大小T以改变每个专家的平均令牌数（T̄e）时，对令牌舍入算法的评估。在每个试验中，我们将微批次大小从4（T̄e = 512）调整到1（T̄e = 128），同时保持小批次大小不变。Mtile始终保持为128。“PPL”指的是训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表8：当我们改变用于令牌舍入的块大小Mtile时，对令牌舍入算法的评估。“PPL”指的是训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。

MoE 训练吞吐量

令牌舍入（TR）确保无瓦片量化效应。在 6.3.3 节中，我们证明：
* 在高度稀疏的 MoE 训练场景下，令牌舍入的训练吞吐量始终高于标准 TC Top-K 路由。
* 当专家总数 E 增加而每个令牌激活专家数 K 保持不变时，内核运行时间的 TFLOPS 提升可达 16%。

六、实验

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 13：不同模型规模（1.4B–120B）下每层的峰值激活内存占用。MegaBlocks 不支持小尺寸 n（专家中间维度）。基准配置见表 9b。我们仅缓存每个专家 e 的输入 X、聚集后的激活值以及路由元数据——这是在不通过矩阵乘法重新计算激活值的前提下，反向传播所需的最小激活内存。

在图 13 中，我们通过剖析 7B MoE 训练配置的激活内存，验证了 SonicMoE 的激活内存与专家粒度无关；1.4B 到 120B 参数模型的更多验证结果也包含在图 13 中。

我们从以下方面评估 SonicMoE：
1. 与其他基线 MoE 实现相比的激活内存占用（6.1 节）。
2. 训练吞吐量（6.2 节）。
3. 令牌舍入路由策略的有效性——证明训练时使用令牌舍入，推理时可无缝替换为令牌选择路由（6.3.1 节），且在稀疏 MoE 配置下仍能维持训练吞吐量（6.3.3 节）。

6.1 SonicMoE 的激活内存

如图 13 所示，在所有模型规模下，SonicMoE 的单 MoE 层峰值激活内存占用均为最低。
* 对于 7B 模型，SonicMoE 相比 ScatterMoE 减少了 45% 的内存占用，相比 MoMoE 的减少幅度更大。
* 在 30B 和 120B 模型上，内存节省差距进一步扩大：在 120B 规模下，SonicMoE 每层比 MoMoE 节省超过 3GiB 内存。

我们还验证了 SonicMoE 的激活内存与专家粒度无关，见图 1。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 1：SonicMoE 的每层激活内存占用（左图）在专家粒度（其中 为嵌入维度， 为专家中间维度）增加时保持恒定，且内存效率较其他基线方法高 0.20-1.59 倍；其前向计算吞吐量（右图）平均达到上限值（H100 上的 cuBLAS BMM+激活函数+cuBLAS BMM+聚合操作）的 88%（最大值 91%，最小值 86%）。注：cuBLAS 上限基线未包含路由器计算；本实验采用 300 亿参数（30B）MoE 配置，微批次大小为 32768 个令牌，从左到右的“激活专家数/总专家数”分别为 2/32、4/64、8/128、16/256。

6.2 SonicMoE 的训练吞吐量

6.2.1 完整前向与反向传播吞吐量

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 14：H100 GPU 上不同 MoE 内核的前向与反向传播 TFLOPS。DeepGEMM 未提供高效的路由实现、聚集和专家聚集内核，因此我们对这些模块分别采用标准 PyTorch 实现（“DeepGEMM-pt”）或高度优化的内核（“DeepGEMM++”）。在反向传播中，“DeepGEMM++”和“DeepGEMM-pt”采用与 SonicMoE 相同的计算路径，但需单独启动内核来计算 dS、和 dSwiGLU。MoE 配置与图 13 相同。

图 14 报告了不同 MoE 训练配置下，单个 MoE 层前向与反向传播的计算吞吐量（TFLOPS）。
* 在所有模型规模下，SonicMoE 始终实现最高的 TFLOPS：在 1.4B 和 7B 配置下，相比 ScatterMoE 和 MoMoE 提升 40%的 TFLOPS；
* 在 30B 和 120B MoE 配置下，吞吐量差距进一步扩大——SonicMoE 的前向与反向传播吞吐量超过 500 TFLOPS，而其他基线要么无法支持特定 n 尺寸（如 MegaBlocks），要么性能显著下降（如 MoMoE）。

SonicMoE 在前向传播中相比 DeepGEMM++也实现了加速，这主要得益于聚集 X 内核和乒乓调度——随着 MoE 细粒度增加（图 14 中各配置从右到左），这两个特性的效果愈发显著，因此 SonicMoE 相比 DeepGEMM++的相对加速比也随之增大。

我们进一步测量了采用 FSDP-2（一种完全共享数据并行分布式训练策略）的 7B MoE 模型的实际训练吞吐量：64 块 H100 GPU 上的 SonicMoE 实现了 2130 亿令牌/天的吞吐量，与 96 块 H100 GPU 上 ScatterMoE 的 2250 亿令牌/天吞吐量相当。该吞吐量测量基于 lm-engine 代码库[27]（Mishra, 2024），实验中在单个节点（8 块 H100 GPU）内采用 ZeRO-3（一种优化内存的分布式策略）对模型进行分片，并将分片后的单元跨节点复制。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 13：不同模型规模（1.4B–120B）下每层的峰值激活内存占用。MegaBlocks 不支持小尺寸 n（专家中间维度）。基准配置见表 9b。我们仅缓存每个专家 e 的输入 X、聚集后的、激活值以及路由元数据——这是在不通过矩阵乘法重新计算激活值的前提下，反向传播所需的最小激活内存。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表9：图3、图13和图14的基准配置。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图 15：H100 GPU 上不同 MoE 内核在 7B–685B 参数配置下的单 MoE 层前向与反向传播 TFLOPS。从左到右的 MoE 配置分别对应 OLMoE-1B-7B-0125[35]、gpt-oss-20b[40]、Kimi-Linear-48B-A3B-Base[67]、Qwen3-Next-80B-A3B-Thinking[42]、Qwen3-235B-A22B-Thinking-2507[42]和 DeepSeek-V3.2-Exp[11]。为公平比较，我们未考虑共享专家和专家偏置，且始终使用带 softmax 得分的 TC Top-K 路由器。ScatterMoE、MoMoE、DeepGEMM-pt 和 DeepGEMM++均无法在 DeepSeek-V3.2-Exp 配置下运行（或因索引溢出，或因 CUDA 内存不足）。

此外，我们在图 15 中测量了单个 MoE 层在近期开源 MoE 配置下的训练吞吐量。SonicMoE 在前向和反向传播中通常能实现超过 550 TFLOPS 的吞吐量，并持续超越所有基线。
* 值得注意的是，对于 685B MoE 模型 DeepSeek-V3.2-Exp 的配置，ScatterMoE、MoMoE、DeepGEMM-pt 和 DeepGEMM++均无法运行，而 SonicMoE 在单块 H100 GPU 上（无专家并行）成功运行，前向传播吞吐量达 534.8 TFLOPS，反向传播达 480.1 TFLOPS。
* 此外，对于稀疏细粒度 MoE（如 Qwen3-Next-80B-A3B-Thinking，、，图 15 第 4 列），SonicMoE 的 IO 感知内核设计相比基线实现了更大的相对加速：前向传播相比 ScatterMoE 提速 61%、相比 MoMoE 提速 92%；反向传播相比 ScatterMoE 提速 85%、相比 MoMoE 提速 120%。

6.3 令牌舍入（Token Rounding）

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

本节评估采用令牌舍入（TR）算法训练的 MoE 模型质量：训练时使用 TR，评估时切换为令牌选择 Top-K（TC Top-K）路由——以此验证 TR 在训练后能否无缝替换为 TC[28]。我们基于 OLMoE 代码库[35]构建 MoE 模型（采用 OLMoE 基础架构），使用去重后的 FineWeb-Edu 数据集[2]训练所有模型，更多细节见附录 H。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

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

表 3 中我们始终使用，“舍入与稀疏化”子程序始终将专家频率舍入到最接近的倍数（“NR-f”，见附录 F.2），且对 TR 采用 softmax 重归一化。我们将 TR 与令牌选择（TC）Top-K 路由、专家选择（EC）路由[71]进行比较。然而，EC 路由存在未来令牌泄露问题[44,58]，导致自回归生成性能下降；为解决此问题，我们采用 MoD 的方法[44]：训练一个辅助路由器，用于在推理时预测 EC 路由器的选择[29]——该基线在表 3 各子表中记为“EC (aux router)”。

然而，对于 MoE，我们需要解决比 MoD[44]更难的 E 标签预测问题（而非二分类预测）：因为 EC 路由器可为每个令牌激活任意数量的专家，且需为所有专家独立预测标签。这种方法对于 MoE 可能不具备可扩展性，因为预测问题规模随 E 增长而增大。

我们还通过微调一个学习型 TC Top-K 路由器，将 EC 路由适配为 TC 路由，并比较其与无适配 TR 的任务性能——该基线记为“EC (ft TC router)”。最后，我们考虑一个令牌丢弃基线：将每个专家的容量设为不超过其令牌频率的最大倍数，并丢弃得分最低的令牌——该基线记为“TC (token drop)”，本质上是 TR 的始终向下舍入版本。

TR 的训练-测试差距
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

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

我们在0.5B（表3a）和1.4B（表3c）MoE模型上验证TR的性能，随后通过以下方式增加MoE稀疏度：
1. 保持专家总数E不变，降低激活专家数K（从3a到3b，从3d到3e）；
2. 保持K不变，增加E（从3a到3c）。

在这些稀疏MoE配置下，我们始终观察到TR与TC的模型质量相近。事实上，在极稀疏MoE配置下（3c和3e），TR实现了更低的验证困惑度和更高的平均准确率。EC与TC之间存在显著差距：在3c、3d和3e配置下，EC的训练与验证困惑度差距较大，而TR与TC的差距通常较小。微调TC路由器比使用辅助路由器更能缩小这一差距，但TR的任务评估性能仍始终更优。此外，与令牌丢弃基线相比，TR始终实现更低的验证困惑度，且在3a、3c、3e配置下实现更高的平均任务准确率。因此，TR可作为训练时TC路由的即插即用替代品。

6.3.2 令牌舍入路由的消融实验

有三个变量可能影响令牌舍入路由训练的MoE模型质量：
1. 舍入子程序（round and sparsify）；
2. 微批量大小T；
3. 舍入用瓦片尺寸。

我们分析它们的影响：

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表6：对算法4中配备不同舍入和稀疏化子例程的令牌舍入算法的评估。“PPL”指训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。

  • 舍入子程序的选择：表6评估了不同路由子程序对TR训练MoE的影响。结果表明，令牌舍入算法对具体舍入子程序具有较强的鲁棒性；将专家频率舍入到最接近的倍数（表6中的“NR-f”）尽管简单，但足以实现优异的下游任务性能。因此，我们选择NR-f作为默认舍入子程序。
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表7：当我们改变微批次大小T以改变每个专家的平均令牌数（T̄e）时，对令牌舍入算法的评估。在每个试验中，我们将微批次大小从4（T̄e = 512）调整到1（T̄e = 128），同时保持小批次大小不变。Mtile始终保持为128。“PPL”指训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。

  • 微批量大小T和瓦片尺寸的影响：令牌舍入在微批量级别执行，因此微批量大小T的变化会导致TR的定性结果不同（EC路由也存在此特性）。例如，对整个序列执行EC与对文本片段执行EC，会导致不同的模型质量。

表7显示：
* 当T̄e ≥ 256时，TR能保持训练后的MoE质量；
* 即使T̄e = 128（两个子表的最后一行），训练后的MoE推理质量仍优于使用EC训练并微调TC Top-K路由的模型。

类似地，表8显示，当Mtile ≥ 128时，TR对瓦片尺寸具有较强的鲁棒性；当Mtile = 64时，模型质量虽有显著下降，但仍优于EC基线。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

表8：当我们改变用于令牌舍入的块大小Mtile时，对令牌舍入算法的评估。“PPL”指训练结束时的验证困惑度。“Avg”是11个下游任务的平均准确率。

6.3.3 令牌舍入的训练吞吐量

图16中，我们基准测试了令牌舍入（TR）与Top-K令牌选择（TC）路由的MoE主内核运行时间（不含路由）。

实验采用等计算量（iso-FLOPs）设置：保持T、n、K不变，通过线性增加专家总数E（保持K不变）来提高MoE稀疏度。随着E线性增加，我们观察到TC路由的TFLOPS下降——原因包括：（1）瓦片量化效应：填充导致的计算浪费随MoE稀疏度线性增加（见图11）；（2）更多专家权重导致IO线性增加。TC和TR的TFLOPS均随E增加而下降，但TC的下降幅度更显著（见图16）。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图16：SonicMoE MoE内核在不同路由方法下的前向与反向传播模型TFLOPS。我们比较采用“基于专家频率舍入到倍数”子程序的TR与TC Top-K路由。配置细节见附录G。

  • 在图16上排第3、4列中，128个专家（E=128）、n=2048的MoE模型采用令牌舍入路由后，前向传播TFLOPS提升16.5%，反向传播提升6.1%，端到端提升9.4%。
  • 在图16下排第3、4列中，256个专家（E=256）的MoE模型采用令牌舍入路由后，前向传播TFLOPS提升25.7%，反向传播提升11.8%，端到端提升15.9%。

总体而言，随着中间维度n增大（更接近计算受限）和MoE稀疏度提高，TR与TC Top-K的吞吐量差距愈发显著。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

图17：H100 GPU上，SonicMoE在7B–685B参数配置下的单MoE层前向与反向传播TFLOPS（不同路由方法对比）。从左到右的MoE配置与图15相同（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）。我们比较采用“基于专家频率舍入到倍数”子程序的TR与TC Top-K路由。

这一趋势在近期开源MoE配置中同样成立（见图17）：当SonicMoE的MoE内核采用TR路由器而非TC Top-K路由器时，在高度稀疏MoE（如Qwen3-Next-80B-A3B-Thinking，K=2）上观察到更大的相对加速——前向传播提速19.6%，反向传播提速7.9%。

七、结论

本文提出SonicMoE——一种协同设计方案，通过联合优化MoE架构与GPU内核，解决细粒度和稀疏MoE的训练挑战。主要贡献如下：

  1. 激活内存最小化的MoE训练算法：分析MoE粒度对MoE层前向与反向传播的影响，发现保持计算量（FLOPs）不变时，MoE粒度增加会导致反向传播所需的激活内存线性增加。基于这一观察，我们重新设计计算图，在保持与原始MoE公式数学等价的前提下，避免缓存路由器梯度计算所需的激活值。结果表明，对于细粒度7B MoE，SonicMoE将每层激活内存占用降低高达45%。
  2. IO与计算重叠的高效MoE内核：证明MoE的粒度和稀疏度增加会导致其逐渐转向内存带宽受限状态。为缓解这一瓶颈，我们利用矩阵乘加（MMA）与IO操作的异步性，通过重叠两者来最大化吞吐量。对于同一细粒度7B MoE模型，该方法在前向传播中相比高度优化的DeepGEMM基线提速43%，在反向传播中相比当前最优MoE基线ScatterMoE和MoMoE分别提速83%和115%。为评估这些技术的性能，我们通过全面的内核级剖析和MoE计算路径的IO感知探索，进行了详尽的性能分析。
  3. 消除稀疏 MoE 计算浪费的瓦片感知令牌舍入：提出一种即插即用（drop-in）路由算法，将每个专家的令牌数舍入到分组矩阵乘法所用瓦片尺寸（如 128）的倍数。这种舍入在尽可能保留原始令牌-专家分配关系的同时，减少了填充导致的计算浪费。该算法确保：对于每个专家，其令牌分布与原始 Top-K 令牌选择结果的最大偏差不超过一个瓦片。此方法在保持预期令牌总数不变的前提下，有效消除了分组矩阵乘法中的填充浪费，且在高度稀疏 MoE 训练场景下仍能保持稳健的令牌选择准确性。我们在 1.4B 参数稀疏训练场景中验证了该令牌舍入策略的性能，证明其计算吞吐量持续超过标准 Top-K 令牌选择路由；在高度稀疏场景下，端到端 MoE 计算的 TFLOPS 提升高达 16%。

SonicMoE 已基于 CuTe-DSL[39]实现并开源，提供 PyTorch 接口，采用宽松许可证。GitHub 链接：https://github.com/Dao-AILab/sonic-moe。

未来研究方向包括：
* 扩展到低精度和微缩放格式（FP8、MXFP8、MXFP4）以进一步节省内存；
* 在专家并行等分布式场景下实现通信与计算重叠。

我们期望未来的模型架构设计能同时考虑算法效率和硬件效率，以“每计算小时质量”（quality per compute hour）而非仅“每 FLOP 质量”（quality per FLOP）为优化目标。

SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升
SonicMoE：开源软硬件协同加速方案，64块H100媲美96块性能，实现45%内存节省与1.86倍吞吐量提升

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：http://www.itsolotime.com/archives/14803

赞 (0)
0 0
上一篇 2025年12月22日 下午6:49
下一篇 2025年12月23日 上午7:59

相关推荐