編輯｜Panda

「快！」

說(shuō)到索尼克，不管是刺猬索尼克還是音速索尼克，大家的第一印象多半就是「快」，而「快」也是現(xiàn)在許多 AI 模型和應(yīng)用優(yōu)化的一大核心目標(biāo)。

近日，由普林斯頓大學(xué) Tri Dao（FlashAttention 的一作）和加州大學(xué)伯克利分校 Ion Stoica 領(lǐng)導(dǎo)的一個(gè)聯(lián)合研究團(tuán)隊(duì)也做出了一個(gè)超快的索尼克：SonicMoE

一作 Wentao Guo 的推文，他目前正在普林斯頓大學(xué)就讀計(jì)算機(jī)科學(xué)博士

據(jù)介紹，SonicMoE 能在英偉達(dá) Blackwell GPU 上以峰值吞吐量運(yùn)行！并且運(yùn)算性能超過(guò)了 DeepSeek 之前開(kāi)源并引發(fā)巨大轟動(dòng)的 DeepGEMM。

有趣的是，DeepSeek 前些天還在 DeepGEMM 庫(kù)中開(kāi)源了新的技術(shù) Mega MoE，即巨型 MoE—— 從名字也能看出來(lái)，這與 SonicMoE（音速 MoE）顯然是兩個(gè)不同的方向，我們也期待能看到「」與「」這兩個(gè)方向的更直接的對(duì)比。

下面我們就基于官方技術(shù)博客，簡(jiǎn)單了解下 SonicMoE。

• 博客地址：https://tridao.me/blog/2026/sonicmoe-blackwell/
• 代碼庫(kù)：https://github.com/Dao-AILab/sonic-moe
• 論文地址：https://arxiv.org/abs/2512.14080

MoE 與它的隱患

要理解 SonicMoE 解決的是什么問(wèn)題，先得認(rèn)識(shí)一種正在主導(dǎo)前沿 AI 的架構(gòu)設(shè)計(jì) —— 混合專(zhuān)家模型（Mixture of Experts，MoE）。

細(xì)粒度 MoE 架構(gòu)

想象一家醫(yī)院。面對(duì)每一位患者，醫(yī)院不會(huì)讓所有科室同時(shí)出動(dòng)，而是先由全科醫(yī)生判斷，再分診給最合適的專(zhuān)科。MoE 架構(gòu)的邏輯與此相似：模型內(nèi)部有大量「專(zhuān)家」子網(wǎng)絡(luò)，每一個(gè)輸入的信息片段（即 token，可以理解為文字或詞語(yǔ)）只會(huì)被路由到其中一小部分專(zhuān)家處理，而不是流經(jīng)所有參數(shù)。

這樣做的好處顯而易見(jiàn)：用相對(duì)較少的計(jì)算量，撐起了一個(gè)參數(shù)規(guī)模龐大的模型

2024 年發(fā)布的 Mixtral 8x22B，以及近期的 DeepSeek V3.2、Kimi K2.5、Qwen3 等明星模型，都是 MoE 架構(gòu)的忠實(shí)擁躉。按照模型縮放法則，專(zhuān)家越「細(xì)粒度」（即每個(gè)專(zhuān)家越小、數(shù)量越多），模型在同等計(jì)算量下的表現(xiàn)往往越好。于是在短短兩年間，MoE 專(zhuān)家的粒度提升了整整 9 倍，每次激活的專(zhuān)家比例則降至原來(lái)的十二分之一。

然而，代價(jià)也隨之而來(lái)。

標(biāo)準(zhǔn) MoE 實(shí)現(xiàn)前向傳播的工作流程。π 是存儲(chǔ)路由元數(shù)據(jù)的二進(jìn)制掩碼。黃色框表示內(nèi)核邊界。藍(lán)色框是 HBM 中的變量。紅色標(biāo)簽表示在正向 / 反向傳播過(guò)程中緩存的激活值。紫色框是最終輸出。全局內(nèi)存中每個(gè)變量旁邊的橙色框表示對(duì)應(yīng) Qwen3-235B-A22B-Thinking-2507 MoE 模型在處理 32k 個(gè) token 時(shí)的張量大小比例。

標(biāo)準(zhǔn) MoE 實(shí)現(xiàn)的反向激活梯度傳遞工作流程。

當(dāng)專(zhuān)家越來(lái)越多、越來(lái)越「細(xì)」，訓(xùn)練這樣的模型會(huì)遭遇兩堵越來(lái)越高的墻：

第一堵墻是顯存。在訓(xùn)練神經(jīng)網(wǎng)絡(luò)時(shí)，前向傳播的中間結(jié)果必須被保存下來(lái)，以便反向傳播時(shí)計(jì)算梯度。對(duì)于細(xì)粒度 MoE 來(lái)說(shuō)，這些中間結(jié)果（激活值）的規(guī)模與專(zhuān)家粒度成正比 —— 專(zhuān)家越細(xì)，顯存占用越大，最終會(huì)逼近 GPU 顯存的物理極限。

第二堵墻是內(nèi)存帶寬。GPU 的性能取決于兩個(gè)維度：算力（每秒能做多少次運(yùn)算）和帶寬（每秒能搬多少數(shù)據(jù)）。當(dāng)專(zhuān)家足夠細(xì)時(shí)，每個(gè)專(zhuān)家處理的數(shù)據(jù)量太少，GPU 的算力根本來(lái)不及被填滿(mǎn)，大量時(shí)間都花在了從內(nèi)存「搬運(yùn)」數(shù)據(jù)上。這正是所謂「內(nèi)存瓶頸」。對(duì)于典型的 Qwen3 細(xì)粒度 MoE，其單位計(jì)算量的內(nèi)存訪問(wèn)強(qiáng)度比等參數(shù)量的普通模型高出 12 倍。

現(xiàn)有的開(kāi)源訓(xùn)練工具（如 ScatterMoE 和 MoMoE）對(duì)這兩個(gè)問(wèn)題都存在明顯不足，尤其是隨著模型越來(lái)越細(xì)粒度，差距愈發(fā)顯著。而SonicMoE正是為此而生。

SonicMoE 的每一層激活記憶占用空間（左圖）即使在專(zhuān)家粒度（嵌入維度 / 專(zhuān)家中間維度）增加時(shí)也保持不變，并且與現(xiàn)有的 MoE 訓(xùn)練核 ScatterMoE 和 MoMoE 相比，SonicMoE 可以實(shí)現(xiàn) 1.87-4.04 倍的相對(duì)加速。

核心創(chuàng)新：一次算法級(jí)的重新設(shè)計(jì)

SonicMoE 的關(guān)鍵洞察，乍聽(tīng)簡(jiǎn)單，卻需要深厚的系統(tǒng)級(jí)思維才能想到：?jiǎn)栴}的根源在于，現(xiàn)有 MoE 訓(xùn)練框架在中間結(jié)果的存儲(chǔ)上過(guò)于「慷慨」—— 它們把太多臨時(shí)數(shù)據(jù)寫(xiě)入了顯存，而這些數(shù)據(jù)本可以不存。

傳統(tǒng)方法在執(zhí)行 MoE 的前向傳播和反向傳播時(shí)，會(huì)在每個(gè)計(jì)算階段之間將中間張量（即矩陣形式的中間數(shù)據(jù)）寫(xiě)入 GPU 的高帶寬內(nèi)存（HBM）。這就好比一個(gè)廚師每炒完一道中間步驟，就把食材裝盤(pán)放進(jìn)冰箱，下一步再取出來(lái)繼續(xù) —— 頻繁的存取本身就是大量時(shí)間的浪費(fèi)。

SonicMoE 的前向計(jì)算工作流程以及與 PyTorch 中標(biāo)準(zhǔn) MoE 實(shí)現(xiàn)的比較。這里還比較了兩種方法的激活內(nèi)存和 IO 成本。

SonicMoE 的算法重設(shè)計(jì)從根本上改變了這一流程，核心有兩點(diǎn)：

第一，激活內(nèi)存與專(zhuān)家粒度解耦

在訓(xùn)練反向傳播中，SonicMoE 通過(guò)重新設(shè)計(jì)計(jì)算順序，完全避免了緩存任何與專(zhuān)家規(guī)模成比例的中間張量。

具體來(lái)說(shuō)，它將原本需要緩存的「下投影輸出」等關(guān)鍵中間量，通過(guò)重排矩陣乘法的收縮順序來(lái)消除 —— 不再存儲(chǔ)中間結(jié)果，而是在需要時(shí)通過(guò)聰明的計(jì)算路徑直接推導(dǎo)出所需梯度。

這使得 SonicMoE 的每層激活內(nèi)存占用，在專(zhuān)家粒度大幅增加時(shí)保持恒定，相當(dāng)于一個(gè)相同激活參數(shù)量的稠密模型。

這一改進(jìn)無(wú)需任何額外的矩陣重計(jì)算代價(jià)，正面回答了此前業(yè)界一直認(rèn)為「魚(yú)和熊掌不可兼得」的問(wèn)題。

第二，IO 感知的算子融合

SonicMoE 將原本分散成多個(gè) GPU 核函數(shù)（kernel）的操作大量融合在一起。

例如，「Gather 融合」技術(shù)讓數(shù)據(jù)搬運(yùn)操作在矩陣乘法計(jì)算核的執(zhí)行過(guò)程中同步完成，而不是作為單獨(dú)步驟先把數(shù)據(jù)重排好再交給矩陣乘法 —— 這不僅省去了一次完整的內(nèi)存讀寫(xiě)，還利用了 GPU L2 緩存的局部性?xún)?yōu)勢(shì)，讓緩存命中率從約 66% 提升至約 75%，進(jìn)一步降低了訪問(wèn)慢速 HBM 的頻率。

此外，SwiGLU 激活函數(shù)的計(jì)算也被融入矩陣乘法的尾聲（epilogue）階段，在數(shù)據(jù)還駐留在寄存器時(shí)就地完成，無(wú)需額外的內(nèi)存讀寫(xiě)。

在最關(guān)鍵的反向傳播核函數(shù)（dH kernel）中，SonicMoE 還進(jìn)一步利用 GPU 的異步執(zhí)行特性，將數(shù)據(jù)搬運(yùn)的等待時(shí)間與矩陣運(yùn)算重疊起來(lái)。

SonicMoE 的 dH 工作流程圖的語(yǔ)義與標(biāo)準(zhǔn) PyTorch MoE 多核實(shí)現(xiàn)等效，同時(shí) SonicMoE 顯著降低了 IO 成本。

實(shí)測(cè)結(jié)果顯示，即便該核函數(shù)的 HBM 數(shù)據(jù)流量增加了 24%，張量核心（Tensor Core）的利用率僅下降約 10%—— 內(nèi)存開(kāi)銷(xiāo)幾乎被算力完全「吸收」。

可以利用最新的 NVIDIA 硬件特性來(lái)隱藏 SonicMoE 的 dH 內(nèi)核中的 IO 延遲，并大幅減少整體運(yùn)行時(shí)間。

軟件抽象層 QuACK：讓創(chuàng)新能跨代遷移

SonicMoE 還有一個(gè)容易被忽視的工程亮點(diǎn)：研究團(tuán)隊(duì)開(kāi)發(fā)了一套名為QuACK的統(tǒng)一軟件抽象層，將所有 MoE 矩陣乘法核函數(shù)統(tǒng)一為「主循環(huán) + 可定制尾聲」的共同結(jié)構(gòu)。

兩個(gè)使用 QuACK 實(shí)現(xiàn)的 SonicMoE 內(nèi)核。左側(cè)：內(nèi)核工作流程圖。中間：QuACK 尾聲混合類(lèi)，其中每個(gè)內(nèi)核重寫(xiě) epi_visit_subtile（dH 為 88 行代碼，上投影前向?yàn)?21 行代碼）。右側(cè)：SonicMoE 的簡(jiǎn)化內(nèi)核啟動(dòng)調(diào)用。

這樣的設(shè)計(jì)意味著，當(dāng) GPU 從上一代 Hopper 架構(gòu)（H100）升級(jí)到最新的 Blackwell 架構(gòu)（B200/B300）時(shí)，硬件特有的優(yōu)化只需要在極少數(shù)地方做局部修改，核心算法邏輯無(wú)需重寫(xiě)。

Tri Dao 與 Ion Stoica 團(tuán)隊(duì)之所以能快速將 SonicMoE 移植到英偉達(dá)最新旗艦 Blackwell GPU 并達(dá)到峰值吞吐，很大程度上正是受益于這一前瞻性的軟件架構(gòu)。

實(shí)驗(yàn)結(jié)果

研究團(tuán)隊(duì)在英偉達(dá)最新 B300 GPU 上，以六個(gè)真實(shí)開(kāi)源 MoE 模型配置為基準(zhǔn)進(jìn)行了全面測(cè)評(píng)，涵蓋從 7B 到 685B 參數(shù)的不同規(guī)模，包括 OLMoE、Qwen3-235B、DeepSeek V3.2 等當(dāng)下最受關(guān)注的 MoE 架構(gòu)。

B300 上 6 種真實(shí) 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 與基線(xiàn)模型在 B300 上針對(duì) 7B OLMoE 規(guī)模 MoE（T=32768，d=2048，n=1024，E=64，K=8）的運(yùn)行時(shí)分解情況。

結(jié)果相當(dāng)顯著：

• 與同樣針對(duì) Blackwell GPU 優(yōu)化、由 DeepSeek 開(kāi)發(fā)的 DeepGEMM 基準(zhǔn)相比，SonicMoE 在前向傳播上平均高出54%，在反向傳播上平均高出35%—— 而 DeepGEMM 本身已是業(yè)界公認(rèn)的高性能實(shí)現(xiàn)；
• 與 Triton 官方 MoE 示例相比，SonicMoE 前向傳播快21%
• 與目前學(xué)術(shù)界和工業(yè)界廣泛使用的 ScatterMoE、MoMoE 等訓(xùn)練框架相比，SonicMoE 的速度優(yōu)勢(shì)往往達(dá)到 近兩倍甚至更高。

從核函數(shù)級(jí)別的運(yùn)行時(shí)分析來(lái)看，SonicMoE 的加速主要來(lái)自?xún)蓚€(gè)方面：其一，Gather 融合消除了獨(dú)立的數(shù)據(jù)搬運(yùn)核函數(shù)，這是最主要的加速來(lái)源；其二，更快的分組矩陣乘法實(shí)現(xiàn)（得益于 Blackwell 獨(dú)有的 CLC 調(diào)度器和 2CTA MMA 技術(shù)）貢獻(xiàn)了額外約 10% 的提升。

在激活內(nèi)存方面，當(dāng)專(zhuān)家粒度從 Mixtral 時(shí)代提高到 Kimi K2.5 量級(jí)時(shí)，傳統(tǒng)方案的每層激活內(nèi)存會(huì)線(xiàn)性膨脹，而 SonicMoE 的占用則保持穩(wěn)定。這對(duì)于在有限顯存中訓(xùn)練更細(xì)粒度的未來(lái)模型，意味著更大的操作空間。

結(jié)語(yǔ)

SonicMoE 很快，同時(shí)還有更深層的意義：當(dāng)硬件的進(jìn)步受制于物理規(guī)律逐漸放緩，軟件層面的創(chuàng)新正越來(lái)越多地扮演起「平權(quán)者」的角色。

SonicMoE 的論文標(biāo)題是「硬件高效、軟件可擴(kuò)展的細(xì)粒度 MoE 藍(lán)圖」—— 這個(gè)「藍(lán)圖」二字，或許正是研究團(tuán)隊(duì)想傳遞的信號(hào)：這不只是一個(gè)工具，而是一種可以被復(fù)制和繼承的設(shè)計(jì)哲學(xué)。

SonicMoE 目前已在 GitHub 和 PyPI 開(kāi)源，支持 H100 和最新 B200/B300 GPU，未來(lái)計(jì)劃擴(kuò)展至專(zhuān)家并行、MXFP8/FP4 精度支持，以及下一代英偉達(dá) Rubin GPU。

在內(nèi)存和算力日益稀缺的今天，這種創(chuàng)新極具價(jià)值，畢竟這是在為整個(gè) AI 生態(tài)節(jié)省真金白銀的成本。

你更看好 DeepSeek 的 Mega MoE 還是今天介紹的 SonicMoE？