
「快!」
当谈到索尼克时,无论是刺猬索尼克还是音速索尼克,人们首先想到的往往是「快速」这个特点。「快」也是当前许多 AI 模型和应用优化的核心目标之一。
近日,来自普林斯顿大学的研究人员 Tri Dao 和加州大学伯克利分校的 Ion Stoica 带领的一个联合团队开发了一款名为 SonicMoE 的新工具,该工具展示了惊人的速度性能。

该项目的主要贡献者 Wentao Guo 目前在普林斯顿大学攻读计算机科学博士学位
据介绍,SonicMoE 能够在英伟达 Blackwell GPU 上实现极高的运算效率,并且其性能超过了 DeepSeek 公司之前开源的 DeepGEMM 技术。
值得一提的是,DeepSeek 最近又发布了新的技术——巨型 MoE(Mega MoE),而 SonicMoE 则是另一种快速处理方式。我们期待看到这两种方法之间的直接比较结果。
为了更好地理解 SonicMoE 的工作原理,我们将参考官方博客和技术细节进行说明。

- 关于 SonicMoE 的更多信息,请访问技术博客:https://tridao.me/blog/2026/sonicmoe-blackwell/
- SonicMoE 的代码库可在 GitHub 上找到:https://github.com/Dao-AILab/sonic-moe
- 有关该研究的详细论文,可以查阅 arXiv 网站上的相关文献:https://arxiv.org/abs/2512.14080
MoE 与它的隐患
要理解 SonicMoE 的核心问题,首先要了解混合专家模型(Mixture of Experts,简称 MoE)的架构。

细粒度 MoE 架构
想象一下医院的工作流程:面对每一位患者时,并不会同时派遣所有科室医生进行诊治,而是先由全科医生初步判断,再将病人转介到最适合治疗的专业科室。MoE 架构与此类似,模型内部包含大量的「专家」子网络,每个输入的信息片段(token)只会被分配给最合适的几个专家处理。
这种方法的优点在于能够以较小的计算成本实现大规模参数模型的功能
例如,Mixtral 8x22B、DeepSeek V3.2、Kimi K2.5 和 Qwen3 等近期发布的热门 AI 模型都采用了 MoE 架构。随着专家规模的增长和颗粒度的细化,模型在同等计算量下的性能显著提升。
然而,这种架构也带来了一些挑战。

标准 MoE 实现中前向传播的工作流程图显示了二进制掩码 π 的使用以及内核边界和缓存激活值的标记。此外,还展示了 Qwen3-235B-A22B-Thinking-2507 模型在处理 32k token 时各变量的张量大小比例。

在 MoE 架构中实现反向传播工作流程图同样重要。
当专家数量增加且变得更加细化,训练这些模型面临两个主要障碍:
第一,是显存问题。在神经网络训练过程中,前向计算的中间结果必须被保存以供后续阶段使用,对于细粒度 MoE 来说,这些数据量随着专家规模的增长而增大。
其次则是内存带宽限制。GPU 性能取决于算力和带宽两个方面,当每个专家处理的数据减少时,整体性能可能会受到限制。
SonicMoE 的开发团队利用了 NVIDIA Blackwell GPU 的优势,在前向传播和反向传播过程中实现了显著的性能提升,超过了 DeepSeek 开发的 DeepGEMM 技术和其他现有工具如 ScatterMoE 和 MoMoE。

从内核级别的运行时分析来看,SonicMoE 性能改进主要得益于 Gather 融合和更快的分组矩阵乘法实现。
在内存占用方面,当专家规模增加到 Kimi K2.5 的水平时,传统方法会导致每层激活内存显著增长,而 SonicMoE 则保持稳定。这为未来训练更细化的模型提供了更大的灵活性。
SonicMoE 不仅速度快,在软件和硬件结合优化上也具有重要意义,尤其是在硬件进步放缓的情况下。
该研究团队认为,SonicMoE 提供了一种可以被复制的设计哲学。

目前,SonicMoE 已在 GitHub 和 PyPI 上开源,并支持 H100 及以上版本的 GPU。未来计划进一步扩展到专家并行、MXFP8/FP4 等精度支持,以及下一代 NVIDIA Rubin GPU。
在资源日益稀缺的情况下,这种创新对于整个 AI 生态系统来说至关重要,因为它有助于节省成本。
您认为 DeepSeek 的 Mega MoE 和 SonicMoE 谁更有前景呢?
在训练反向传播中,SonicMoE 通过重新设计计算顺序,完全避免了缓存任何与专家规模成比例的中间张量。
具体来说,它将原本需要缓存的「下投影输出」等关键中间量,通过重排矩阵乘法的收缩顺序来消除 —— 不再存储中间结果,而是在需要时通过聪明的计算路径直接推导出所需梯度。
这使得 SonicMoE 的每层激活内存占用,在专家粒度大幅增加时保持恒定,相当于一个相同激活参数量的稠密模型。
这一改进无需任何额外的矩阵重计算代价,正面回答了此前业界一直认为「鱼和熊掌不可兼得」的问题。
第二,IO 感知的算子融合
SonicMoE 将原本分散成多个 GPU 核函数(kernel)的操作大量融合在一起。
例如,「Gather 融合」技术让数据搬运操作在矩阵乘法计算核的执行过程中同步完成,而不是作为单独步骤先把数据重排好再交给矩阵乘法 —— 这不仅省去了一次完整的内存读写,还利用了 GPU L2 缓存的局部性优势,让缓存命中率从约 66% 提升至约 75%,进一步降低了访问慢速 HBM 的频率。
此外,SwiGLU 激活函数的计算也被融入矩阵乘法的尾声(epilogue)阶段,在数据还驻留在寄存器时就地完成,无需额外的内存读写。
在最关键的反向传播核函数(dH kernel)中,SonicMoE 还进一步利用 GPU 的异步执行特性,将数据搬运的等待时间与矩阵运算重叠起来。

SonicMoE 的 dH 工作流程图的语义与标准 PyTorch MoE 多核实现等效,同时 SonicMoE 显著降低了 IO 成本。
实测结果显示,即便该核函数的 HBM 数据流量增加了 24%,张量核心(Tensor Core)的利用率仅下降约 10%—— 内存开销几乎被算力完全「吸收」。

可以利用最新的 NVIDIA 硬件特性来隐藏 SonicMoE 的 dH 内核中的 IO 延迟,并大幅减少整体运行时间。
软件抽象层 QuACK:让创新能跨代迁移
SonicMoE 还有一个容易被忽视的工程亮点:研究团队开发了一套名为QuACK的统一软件抽象层,将所有 MoE 矩阵乘法核函数统一为「主循环 + 可定制尾声」的共同结构。

两个使用 QuACK 实现的 SonicMoE 内核。左侧:内核工作流程图。中间:QuACK 尾声混合类,其中每个内核重写 epi_visit_subtile(dH 为 88 行代码,上投影前向为 21 行代码)。右侧:SonicMoE 的简化内核启动调用。
这样的设计意味着,当 GPU 从上一代 Hopper 架构(H100)升级到最新的 Blackwell 架构(B200/B300)时,硬件特有的优化只需要在极少数地方做局部修改,核心算法逻辑无需重写。
Tri Dao 与 Ion Stoica 团队之所以能快速将 SonicMoE 移植到英伟达最新旗舰 Blackwell GPU 并达到峰值吞吐,很大程度上正是受益于这一前瞻性的软件架构。
实验结果
研究团队在英伟达最新 B300 GPU 上,以六个真实开源 MoE 模型配置为基准进行了全面测评,涵盖从 7B 到 685B 参数的不同规模,包括 OLMoE、Qwen3-235B、DeepSeek V3.2 等当下最受关注的 MoE 架构。

B300 上 6 种真实 MoE 配置的前向(左)和后向(右)TFLOPS。从左到右依次为: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。Triton 官方示例不支持后向传播,Qwen3-Next-80B 的前向传播也不支持 K=10。

SonicMoE 与基线模型在 B300 上针对 7B OLMoE 规模 MoE(T=32768,d=2048,n=1024,E=64,K=8)的运行时分解情况。
结果相当显著:
- 与同样针对 Blackwell GPU 优化、由 DeepSeek 开发的 DeepGEMM 基准相比,SonicMoE 在前向传播上平均高出54%,在反向传播上平均高出35%—— 而 DeepGEMM 本身已是业界公认的高性能实现;
- 与 Triton 官方 MoE 示例相比,SonicMoE 前向传播快21%
- 与目前学术界和工业界广泛使用的 ScatterMoE、MoMoE 等训练框架相比,SonicMoE 的速度优势往往达到 近两倍甚至更高。
从核函数级别的运行时分析来看,SonicMoE 的加速主要来自两个方面:其一,Gather 融合消除了独立的数据搬运核函数,这是最主要的加速来源;其二,更快的分组矩阵乘法实现(得益于 Blackwell 独有的 CLC 调度器和 2CTA MMA 技术)贡献了额外约 10% 的提升。
在激活内存方面,当专家粒度从 Mixtral 时代提高到 Kimi K2.5 量级时,传统方案的每层激活内存会线性膨胀,而 SonicMoE 的占用则保持稳定。这对于在有限显存中训练更细粒度的未来模型,意味着更大的操作空间。
结语
SonicMoE 很快,同时还有更深层的意义:当硬件的进步受制于物理规律逐渐放缓,软件层面的创新正越来越多地扮演起「平权者」的角色。
SonicMoE 的论文标题是「硬件高效、软件可扩展的细粒度 MoE 蓝图」—— 这个「蓝图」二字,或许正是研究团队想传递的信号:这不只是一个工具,而是一种可以被复制和继承的设计哲学。
SonicMoE 目前已在 GitHub 和 PyPI 开源,支持 H100 和最新 B200/B300 GPU,未来计划扩展至专家并行、MXFP8/FP4 精度支持,以及下一代英伟达 Rubin GPU。
在内存和算力日益稀缺的今天,这种创新极具价值,毕竟这是在为整个 AI 生态节省真金白银的成本。
你更看好 DeepSeek 的 Mega MoE 还是今天介绍的 SonicMoE?
