DeepSeek做大→Mega MoE，Tri Dao團隊加快→SonicMoE

編輯｜Panda

「快！」

說到索尼克，不管是刺猬索尼克還是音速索尼克，大家的第一印象多半就是「快」，而「快」也是現在許多 AI 模型和應用優化的一大核心目標。

近日，由普林斯頓大學 Tri Dao（FlashAttention 的一作）和加州大學伯克利分校 Ion Stoica 領導的一個聯合研究團隊也做出了一個超快的索尼克：SonicMoE

一作 Wentao Guo 的推文，他目前正在普林斯頓大學就讀計算機科學博士

據介紹，SonicMoE 能在英偉達 Blackwell GPU 上以峰值吞吐量運行！并且運算性能超過了 DeepSeek 之前開源并引發巨大轟動的 DeepGEMM。

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

下面我們就基于官方技術博客，簡單了解下 SonicMoE。

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

MoE 與它的隱患

要理解 SonicMoE 解決的是什么問題，先得認識一種正在主導前沿 AI 的架構設計 —— 混合專家模型（Mixture of Experts，MoE）。

細粒度 MoE 架構

想象一家醫院。面對每一位患者，醫院不會讓所有科室同時出動，而是先由全科醫生判斷，再分診給最合適的專科。MoE 架構的邏輯與此相似：模型內部有大量「專家」子網絡，每一個輸入的信息片段（即 token，可以理解為文字或詞語）只會被路由到其中一小部分專家處理，而不是流經所有參數。

這樣做的好處顯而易見：用相對較少的計算量，撐起了一個參數規模龐大的模型

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

然而，代價也隨之而來。

標準 MoE 實現前向傳播的工作流程。π 是存儲路由元數據的二進制掩碼。黃色框表示內核邊界。藍色框是 HBM 中的變量。紅色標簽表示在正向 / 反向傳播過程中緩存的激活值。紫色框是最終輸出。全局內存中每個變量旁邊的橙色框表示對應 Qwen3-235B-A22B-Thinking-2507 MoE 模型在處理 32k 個 token 時的張量大小比例。

標準 MoE 實現的反向激活梯度傳遞工作流程。

當專家越來越多、越來越「細」，訓練這樣的模型會遭遇兩堵越來越高的墻：

第一堵墻是顯存。在訓練神經網絡時，前向傳播的中間結果必須被保存下來，以便反向傳播時計算梯度。對于細粒度 MoE 來說，這些中間結果（激活值）的規模與專家粒度成正比 —— 專家越細，顯存占用越大，最終會逼近 GPU 顯存的物理極限。

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

現有的開源訓練工具（如 ScatterMoE 和 MoMoE）對這兩個問題都存在明顯不足，尤其是隨著模型越來越細粒度，差距愈發顯著。而SonicMoE正是為此而生。

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

核心創新：一次算法級的重新設計

SonicMoE 的關鍵洞察，乍聽簡單，卻需要深厚的系統級思維才能想到：問題的根源在于，現有 MoE 訓練框架在中間結果的存儲上過于「慷慨」—— 它們把太多臨時數據寫入了顯存，而這些數據本可以不存。

傳統方法在執行 MoE 的前向傳播和反向傳播時，會在每個計算階段之間將中間張量（即矩陣形式的中間數據）寫入 GPU 的高帶寬內存（HBM）。這就好比一個廚師每炒完一道中間步驟，就把食材裝盤放進冰箱，下一步再取出來繼續 —— 頻繁的存取本身就是大量時間的浪費。

SonicMoE 的前向計算工作流程以及與 PyTorch 中標準 MoE 實現的比較。這里還比較了兩種方法的激活內存和 IO 成本。

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？