機(jī)器之心編輯部

混合專家（MoE）模型已成為在不顯著增加計(jì)算成本的情況下，實(shí)現(xiàn)語言模型規(guī)模化擴(kuò)展的事實(shí)標(biāo)準(zhǔn)架構(gòu)。

近期 MoE 模型展現(xiàn)出明顯的高專家粒度（更小的專家中間層維度）和高稀疏性（在專家總數(shù)增加的情況下保持激活專家數(shù)不變）的趨勢，這提升了單位 FLOPs 的模型質(zhì)量。

這一趨勢在近期的開源模型中表現(xiàn)尤為明顯，例如 DeepSeek V3、Kimi K2 以及 Qwen3 MoE 等，它們均采用了更細(xì)粒度的專家設(shè)計(jì)（更小的中間層維度）和更高的稀疏度，在保持激活參數(shù)量不變的同時大幅增加了總參數(shù)量。

表 1：MoE 擴(kuò)展趨勢：在此，團(tuán)隊(duì)將激活率展示為每個 Token 激活的專家數(shù) K / 專家總數(shù) E；針對前沿開源模型，專家粒度展示為模型嵌入維度（d）/ 專家中間層大小（n）。在 MoE 稀疏度計(jì)算中未包含共享專家。趨勢表明，新的開源 MoE 模型傾向于具備更高的粒度和稀疏度。

然而，這種追求極致粒度和稀疏性的設(shè)計(jì)導(dǎo)致了嚴(yán)重的硬件效率下降問題：

內(nèi)存墻瓶頸：對于細(xì)粒度 MoE，激活內(nèi)存的占用量通常隨激活專家數(shù)量線性增長，導(dǎo)致前向和反向傳播中的內(nèi)存壓力劇增。IO 瓶頸：由于專家變得更小且更分散，算術(shù)強(qiáng)度（Arithmetic Intensity，即計(jì)算量與數(shù)據(jù)傳輸量的比值）顯著降低，IO 訪問變得更加動態(tài)和頻繁，導(dǎo)致模型訓(xùn)練進(jìn)入「內(nèi)存受限」區(qū)間。計(jì)算浪費(fèi)：在高稀疏性場景下，由于 Grouped GEMM（分組通用矩陣乘法）內(nèi)核中的 Tile 量化效應(yīng)，輸入數(shù)據(jù)往往需要進(jìn)行填充以對齊硬件 Tile 大小，這直接導(dǎo)致了計(jì)算資源的浪費(fèi)。

針對這些問題，普林斯頓大學(xué)助理教授 Tri Dao（Mamba、FlashAttention 的核心作者）團(tuán)隊(duì)提出了一套名為 SonicMoE 的系統(tǒng)性解決方案。該方案專為 NVIDIA Hopper 和 Blackwell 架構(gòu) GPU 量身定制，其核心貢獻(xiàn)包括：

內(nèi)存高效算法：團(tuán)隊(duì)通過重新設(shè)計(jì) MoE 的計(jì)算圖，提出了一種在計(jì)算路由梯度時不緩存激活值的方法。該方法在保持與原始 MoE 公式數(shù)學(xué)等價的前提下，大幅減少了反向傳播所需的激活顯存。對于細(xì)粒度 7B MoE 模型，每層的激活內(nèi)存占用減少了 45%，且隨著專家粒度的增加，其內(nèi)存占用保持恒定，效率比現(xiàn)有基線高出 0.20-1.59 倍。計(jì)算與 IO 重疊：利用 Hopper 架構(gòu) GPU 的 WGMMA 指令與生產(chǎn)者 - 消費(fèi)者異步范式，SonicMoE 設(shè)計(jì)了新型 GPU 內(nèi)核。該內(nèi)核能夠?qū)?GEMM 計(jì)算與從 HBM 加載數(shù)據(jù)的 IO 操作并行執(zhí)行，有效掩蓋了細(xì)粒度 MoE 帶來的高昂 IO 延遲。Token 舍入：這是一種即插即用的創(chuàng)新調(diào)度策略。它將分發(fā)給每個專家的 Token 數(shù)量四舍五入為 Grouped GEMM Tile 大小（例如 128）的倍數(shù)。算法保證每個專家的偏差最多僅為一個 Tile，從而在期望意義下保持總 token 數(shù)不變。這一策略有效減少了因填充導(dǎo)致的算力浪費(fèi)。

實(shí)驗(yàn)數(shù)據(jù)有力地證明了 SonicMoE 的性能優(yōu)勢，在針對細(xì)粒度 7B MoE 模型的測試中：前向傳播相比高度優(yōu)化的 DeepGEMM 基線，速度提升43%；反向傳播相比最先進(jìn)的 ScatterMoE 和 MoMoE 基線，速度分別提升了 83% 和 115%；端到端訓(xùn)練僅依靠內(nèi)核優(yōu)化即可將訓(xùn)練吞吐量提升 50%，若配合 Token 舍入路由，在擴(kuò)展專家數(shù)量時可進(jìn)一步獲得 16% 的額外吞吐量提升。

論文標(biāo)題：SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations論文地址：https://arxiv.org/abs/2512.14080

更直觀地看，團(tuán)隊(duì)僅使用 64 臺 H100 運(yùn)行 SonicMoE，便實(shí)現(xiàn)了每日 2130 億 token 的訓(xùn)練吞吐量，這一表現(xiàn)已能與使用 96 臺 H100 運(yùn)行 ScatterMoE 的效率相媲美。此外，在高稀疏性場景下（如 1.4B 參數(shù)模型），其 Tile 感知的 Token 舍入算法在驗(yàn)證了不損失下游任務(wù)精度（如在 2B 規(guī)模上的推理質(zhì)量）的同時，顯著提升了內(nèi)核執(zhí)行速度。

目前，團(tuán)隊(duì)已將相關(guān)內(nèi)核代碼開源，為大模型社區(qū)加速高性能 MoE 訓(xùn)練提供了強(qiáng)有力的工具。

圖 1： 即使專家粒度（d/n，其中 d 為嵌入維度，n 為專家中間維度）增加，SonicMoE 的每層激活顯存占用（左圖）仍保持恒定；相比其他基線，其顯存效率提升了 0.20 倍至 1.59 倍。SonicMoE 的前向計(jì)算吞吐量（右圖）平均達(dá)到了理論上限的 88%（最高 91%，最低 86%），該上限基于 H100 GPU 上的 「cuBLAS BMM + 激活函數(shù) + cuBLAS BMM + 聚合操作」 計(jì)算得出。請注意，cuBLAS 上限基線未包含路由計(jì)算部分。在此，我們使用的是 30B 參數(shù)量的 MoE 配置，微批次大小為 32768 個 token，并且從左至右依次將「激活專家數(shù) / 總專家數(shù)」設(shè)置為 2/32、4/64、8/128 和 16/256。

內(nèi)存高效的 MoE 算法

團(tuán)隊(duì)提供了一個高效的基于 Tensor Core 的 top-K 路由，以及一個可以接受任意路由輸入的接口。但需要注意的是，SonicMoE 的 MoE 計(jì)算與路由的選擇無關(guān)，因此與任意路由邏輯兼容。

SonicMoE 的 MoE 計(jì)算實(shí)現(xiàn)具有高度模塊化特性，僅由以下兩部分組成：

經(jīng)過優(yōu)化的分組 GEMM 內(nèi)核（帶有模塊化融合）經(jīng)過優(yōu)化的專家聚合內(nèi)核

主機(jī)會根據(jù)最佳 GEMM 配置和加載 / 存儲策略來調(diào)度并啟動上述 8 個內(nèi)核。

結(jié)果顯示，盡管采用了如此高度的模塊化設(shè)計(jì)，SonicMoE 仍然展現(xiàn)出業(yè)界領(lǐng)先的訓(xùn)練吞吐量和最低的激活內(nèi)存使用量。

面向 IO 的內(nèi)核設(shè)計(jì)

細(xì)粒度 MoE 的表達(dá)能力來自于每個 token 在專家選擇上的多樣性，但這種多樣性同時帶來了與專家粒度線性增長的 IO 開銷，為了保持高吞吐，需要盡可能做到：

通過融合（fusion）減少 IO 訪問將 IO 延遲與計(jì)算重疊

在融合這一塊有兩種方式，一是利用 HBM 加載進(jìn)行 Gather 融合。SonicMoE 的 Grouped GEMM 既可以接受連續(xù)打包的輸入，也可以接受從不同位置 gather 得到的輸入。對于第二種情況，團(tuán)隊(duì)將輸入 gather 與從全局顯存（GMEM，通常是 HBM）到共享內(nèi)存（SMEM）的加載過程進(jìn)行融合，從而能夠?qū)⑦@些數(shù)據(jù)批量化，利用 Tensor Core 執(zhí)行 GEMM。

這一過程包括兩個步驟：

獲取每個 expert 對應(yīng)的被路由 token 的索引；使用這些索引，通過 Blackwell 和 Hopper 架構(gòu)的 cp.async 指令，從 HBM gather 激活值。

二是 Epilogue 融合，通過以下設(shè)計(jì)充分利用 epilogue 計(jì)算，以最大化減少不必要的 IO 訪問：將 SwiGLU 以及 SwiGLU 的反向（dSwiGLU），分別與前向 up-proj 內(nèi)核的 epilogue、反向 down-proj 激活梯度內(nèi)核的 epilogue 進(jìn)行融合；在反向 down-proj 激活梯度（dH）內(nèi)核的 epilogue 中計(jì)算 dH 和 dS。

結(jié)果顯示，這種「重量級 epilogue 融合」使 SonicMoE 相比其他方案獲得顯著加速。

Token rounding 路由方法

團(tuán)隊(duì)在分析稀疏 MoE 訓(xùn)練模式下的硬件效率時發(fā)現(xiàn)，隨著 MoE 變得更加稀疏，因填充而產(chǎn)生的 GEMM tile 計(jì)算浪費(fèi)會累計(jì)到不可忽略的程度，這被稱為「tile 量化效應(yīng)」。為此，團(tuán)隊(duì)提出路由方法「token rounding」來消除這種效應(yīng)，從而實(shí)現(xiàn)更高效的訓(xùn)練。

Token rounding 算法首先計(jì)算基礎(chǔ)的 TC（token-choice）路由結(jié)果，并對每個 expert 對應(yīng)的 token 按路由分?jǐn)?shù)進(jìn)行排序，之后在第二步排序中選擇：要么丟棄第一步 TC top-K 選擇中的部分 token，要么在第二步排序中為某些 expert 補(bǔ)齊額外的 token（填充）。

過程中，團(tuán)隊(duì)會對路由權(quán)重矩陣進(jìn)行處理，使得 TC 選中的 token 始終優(yōu)先于 EC token。結(jié)果就是，無論是丟棄還是填充，都只會影響每個 expert 的最后一個輸入 tile。

實(shí)驗(yàn)表明，這種方法在實(shí)現(xiàn)更高訓(xùn)練吞吐量的同時，并不會影響模型質(zhì)量。

更多內(nèi)容，可查看論文獲悉！