GPU 推論加速的下一步，是核心融合。
編譯丨鄭佳美
編輯丨馬曉寧
史丹佛 Hazy Research 團隊剛公布了一項重量級優化成果：他們將開源模型 Llama-3.2-1B 的前向推論整合成一個「Megakernel」，並將低延遲推論能力推向了極限。
在某些即時性極高的應用中，例如對話式 AI 和人類參與的互動式工作流程中，大型語言模型的響應速度不僅重要，甚至可以決定使用者體驗的成敗。
團隊認為限制 LLM 推論速度的瓶頸其實是在記憶體載入的問題上，他們經過研究發現，現有的開源推論引擎（如 vLLM、SGLang），在極低延遲的單序列生成任務下，即使在頂級 GPU（如 H100）上，也只能利用不到 50% 的記憶體頻寬。
這主要是因為每層 Transformer 模組被拆解成幾十到上百個 CUDA 核心 (kernel)，每個核心執行非常小的操作（比如 RMS norm、注意力、MLP、旋轉位置嵌入等），它們之間存在大量上下文切換與等待。
更嚴重的是，這些核心啟動與收尾的成本加起來，並不會被 CUDA 圖形 (CUDA Graph) 或程式化依賴啟動 (PDL) 等機制充分隱藏，反而會在短時任務中被放大。
換句話說，GPU 花了大量時間「等著幹活」，而不是「在幹活」。Hazy 團隊的研究也正是圍繞著這個問題展開。
1 Megakernel：從零設計的融合思路
先說實驗結果，Megakernel 在 H100 上的推論延遲壓縮至不足 1 毫秒，顯存頻寬利用率高達 78%，相較於 vLLM 提升了 2.5 倍、相較 SGLang 提升 1.5 倍；而在更先進的 B200 平台上，延遲進一步降低至 600~680 微秒，逼近理論極限。
從一次完整推論的時間分布來看，250 微秒用於儲存啟用、等待一致性與資料載入，200 微秒用於 RMSNorm 與矩陣向量乘法 (matvec，其中 matvec 占比達 95%)，權重載入僅需 30 微秒，流水機制表現穩定。warp 間同步與屏障 (barrier) 帶來 40 微秒的延遲，其餘如設定、參數傳遞與頁狀態標記等雜項開銷合計約 80 微秒。
整體來看，在精心調度下，Hazy 團隊的 Megakernel 幾乎已將當前硬體性能壓榨至極限。
而能夠得到以上效果，其實都歸功於 Hazy 團隊提出的一個激進但有效的光明設計思路：將整個前向傳播過程整合為一個單一 CUDA 核心 (kernel)，也就是所謂的 Megakernel。
實驗中，他們基於已有 ThunderMLA 架構，開發了一個 GPU 上運行的輕量「指令解釋器」系統。該系統為每個串流多處理器 (SM) 預先分配一段「執行計畫」，其中包含多條按順序排列的指令，每條指令代表 Transformer 模型中的一個結構單元。
這些指令包括：
融合 RMSNorm、QKV 投影、RoPE 的複合指令；
注意力矩陣乘法與縮減運算（支援長序列 GQA）；
O-投影與殘差相加；
MLP 的 RMSNorm、閘激活 (SiLU) 與上投影；
下投影和最終殘差；
最後一層 RMSNorm + 語言模型頭。
每個指令都基於統一的 CUDA 模板建構，實現對載入、儲存、運算的標準化封裝。指令間依賴由解釋器在運行前靜態排布，每個 SM 可以重複複用同一個排程以處理多個 token。
此外，為確保高效的資料路徑，解釋器會將這些執行計畫按模型結構靜態編排，避免調度時動態分支，提升吞吐與併發執行能力。
同時為了實現流水化運算並防止共享記憶體衝突，團隊還對 GPU 的共享記憶體進行了分頁管理，例如：
將前 213KB 的共享記憶體分為 13 個 16KiB 頁面；
剩餘部分用於儲存指令參數、頁分配資訊等；
每條指令在載入前顯示請求頁，結束後歸還給解釋器調度器；
當頁被釋放時，解釋器會立即將其分配給下一條等待中的指令。
這種機制保證了下一個運算階段可以盡早開始預載入權重，從而最大化頻寬使用率並消除「閒置間隙」。
不過 Megakernel 結構無法依賴傳統的核心間隱式同步，因此 Hazy 團隊還使用了一個計數器系統：他們在全域記憶體中維護一組整數，每條指令完成後會對應計數器 +1，若某條指令依賴先前步驟的結果，它會等待計數器達到特定值才執行。
例如：在 MLP 下投影階段，團隊將中間態拆成 4 個區塊 (chunk)，每個區塊在寫入後立即觸發後續運算，從而實現並行流。此外，團隊透過精確設定依賴圖，避免了全域屏障，大幅減少了指令之間等待的浪費，使得整個核心執行盡可能地接近理論併發。
除此之外，研究團隊還對 CUDA 非同步屏障 (asynchronous barrier) 的性能進行了測量，發現即便在屏障已「通過」的狀態下，每次仍需 60ns，同步操作成本不可忽視。而在實際執行中，尤其在矩陣向量乘法 (matrix-vector) 這類關鍵操作中，他們發現：在 Hopper 架構（如 H100）上，使用常規 CUDA 核心（非 Tensor 核心）可以更有效，不過在 Blackwell 架構上，Tensor 核心性能佔優。
這也說明在硬體不同世代中，Megakernel 的最佳實現路徑也應適配微架構差異，而非一套方案通用所有平台。
更多內容，點擊下方關注：
未經「AI科技評論」授權，嚴禁以任何方式在網頁、論壇、社群進行轉載！
公眾號轉載請先在「AI科技評論」後台留言取得授權，轉載時需標註來源並插入本公眾號名片。
// 推薦閱讀
UCL 強化學習派：汪軍與他的學生們
大型模型隱藏玩家上桌：DeepSeek 向左，面壁向右
基座大型模型「六進二」：南階躍，北智譜