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。

實驗表明，這種方法在實現更高訓練吞吐量的同時，并不會影響模型質量。

更多內容，可查看論文獲悉！