機(jī)器之心編輯部

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

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

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

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

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

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

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

• 內(nèi)存高效算法：團(tuán)隊(duì)通過(guò)重新設(shè)計(jì) MoE 的計(jì)算圖，提出了一種在計(jì)算路由梯度時(shí)不緩存激活值的方法。該方法在保持與原始 MoE 公式數(shù)學(xué)等價(jià)的前提下，大幅減少了反向傳播所需的激活顯存。對(duì)于細(xì)粒度 7B MoE 模型，每層的激活內(nèi)存占用減少了 45%，且隨著專(zhuān)家粒度的增加，其內(nèi)存占用保持恒定，效率比現(xiàn)有基線(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 帶來(lái)的高昂 IO 延遲。
• Token 舍入：這是一種即插即用的創(chuàng)新調(diào)度策略。它將分發(fā)給每個(gè)專(zhuān)家的 Token 數(shù)量四舍五入為 Grouped GEMM Tile 大小（例如 128）的倍數(shù)。算法保證每個(gè)專(zhuān)家的偏差最多僅為一個(gè) Tile，從而在期望意義下保持總 token 數(shù)不變。這一策略有效減少了因填充導(dǎo)致的算力浪費(fèi)。

實(shí)驗(yàn)數(shù)據(jù)有力地證明了 SonicMoE 的性能優(yōu)勢(shì)，在針對(duì)細(xì)粒度 7B MoE 模型的測(cè)試中：前向傳播相比高度優(yōu)化的 DeepGEMM 基線(xiàn)，速度提升43%；反向傳播相比最先進(jìn)的 ScatterMoE 和 MoMoE 基線(xiàn)，速度分別提升了 83% 和 115%；端到端訓(xùn)練僅依靠?jī)?nèi)核優(yōu)化即可將訓(xùn)練吞吐量提升 50%，若配合 Token 舍入路由，在擴(kuò)展專(zhuān)家數(shù)量時(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 臺(tái) H100 運(yùn)行 SonicMoE，便實(shí)現(xiàn)了每日 2130 億 token 的訓(xùn)練吞吐量，這一表現(xiàn)已能與使用 96 臺(tái) H100 運(yùn)行 ScatterMoE 的效率相媲美。此外，在高稀疏性場(chǎng)景下（如 1.4B 參數(shù)模型），其 Tile 感知的 Token 舍入算法在驗(yàn)證了不損失下游任務(wù)精度（如在 2B 規(guī)模上的推理質(zhì)量）的同時(shí)，顯著提升了內(nèi)核執(zhí)行速度。

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

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

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

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

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

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

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

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

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

細(xì)粒度 MoE 的表達(dá)能力來(lái)自于每個(gè) token 在專(zhuān)家選擇上的多樣性，但這種多樣性同時(shí)帶來(lái)了與專(zhuān)家粒度線(xiàn)性增長(zhǎng)的 IO 開(kāi)銷(xiāo)，為了保持高吞吐，需要盡可能做到：

• 通過(guò)融合（fusion）減少 IO 訪(fǎng)問(wèn)
• 將 IO 延遲與計(jì)算重疊

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

這一過(guò)程包括兩個(gè)步驟：

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

二是 Epilogue 融合，通過(guò)以下設(shè)計(jì)充分利用 epilogue 計(jì)算，以最大化減少不必要的 IO 訪(fǎng)問(wèn)：將 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é)果顯示，這種「重量級(jí) epilogue 融合」使 SonicMoE 相比其他方案獲得顯著加速。

Token rounding 路由方法

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

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

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

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

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