0
| 本文作者： 鄭佳美 | 2025-05-28 12:49 |
斯坦福 Hazy Research 團隊剛剛公布了一項重量級優化成果：他們將開源模型 Llama-3.2-1B 的前向推理整合成了一個“Megakernel”，并將低延遲推理能力推向了極限。
在某些實時性極高的應用中，例如對話式 AI 和人類參與的交互式工作流中，大語言模型的響應速度不僅重要，甚至可以決定用戶體驗的成敗。
團隊認為限制 LLM 推理速度的瓶頸其實是在內存加載的問題上，他們經過研究發現，現有的開源推理引擎（如 vLLM、SGLang），在極低延遲的單序列生成任務下，即使在頂級 GPU（如 H100）上，也只能利用不到 50% 的內存帶寬。
這主要是因為每層 Transformer 模塊被拆解成幾十到上百個 CUDA kernel，每個 kernel 執行非常小的操作（比如 RMS norm、注意力、MLP、Rotary Position Embedding 等），它們之間存在大量上下文切換與等待。
更嚴重的是，這些 kernel 啟動與收尾的成本加起來，并不會被 CUDA Graph 或 PDL（Programmatic Dependent Launch）等機制充分隱藏，反而在短時任務中被放大。
換句話說，GPU 花了大量時間“等著干活”，而不是“在干活”。Hazy 團隊的研究也正是圍繞著這個問題展開。
先說實驗結果，Megakernel在 H100 上的推理延遲壓縮至不足 1 毫秒，顯存帶寬利用率高達 78%，相較于 vLLM 提升了 2.5 倍、相較 SGLang 提升 1.5 倍；而在更先進的 B200 平臺上，延遲進一步降低至 600~680 微秒，逼近理論極限。
從一次完整推理的時間分布來看，250 微秒用于存儲激活、等待一致性與數據加載，200 微秒用于 RMSNorm 與 matvec（其中 matvec 占比達 95%），權重加載僅需 30 微秒，流水機制表現穩定。warp 間同步與 barrier 帶來 40 微秒的延遲，其余如 setup、參數傳遞與頁狀態標記等雜項開銷合計約 80 微秒。
整體來看，在精心調度下，Hazy 團隊的 Megakernel 幾乎已將當前硬件性能壓榨至極限。
而能夠得到以上效果，其實都歸功于 Hazy 團隊提出的一個激進但有效的設計思路：將整個前向傳播過程整合為一個單一 CUDA kernel，也就是所謂的 Megakernel。
實驗中，他們基于已有 ThunderMLA 架構，開發了一個 GPU 上運行的輕量“指令解釋器”系統。該系統為每個 Streaming Multiprocessor（SM）預先分配一段“執行計劃”，其中包含多條按順序排列的指令，每條指令代表 Transformer 模型中的一個結構單元。
這些指令包括：
融合 RMSNorm、QKV projection、RoPE 的復合指令；
attention 矩陣乘與縮減計算（支持長序列 GQA）；
O-projection 與 residual 相加；
MLP 的 RMSNorm、gate 激活（SiLU）與上投影；
down projection 和最終 residual；
最后一層 RMSNorm + language modeling head。
每個指令都基于統一的 CUDA 模板構建，實現對 load、store、compute 的標準化封裝。指令間依賴由解釋器在運行前靜態排布，每個 SM 可以重復復用同一個 schedule 以處理多個 token。
此外，為確保高效的數據路徑，解釋器會將這些執行計劃按模型結構靜態編排，避免調度時動態分支，提升吞吐與并發執行能力。
同時為了實現流水化計算并防止 shared memory 沖突，團隊還對 GPU 的共享內存進行了分頁管理，例如：
將前 213KB 的 shared memory 分為 13 個 16KiB 頁面；
剩余部分用于存儲指令參數、頁分配信息等；
每條指令在加載前顯示請求頁，結束后歸還給解釋器調度器；
當頁被釋放時，解釋器會立即將其分配給下一條等待中的指令。
這種機制保證了下一個計算階段可以盡早開始預加載權重，從而最大化帶寬使用率并消除“氣泡”。
不過 Megakernel 結構無法依賴傳統的 kernel 間隱式同步，因此 Hazy 團隊還使用了一個計數器系統：他們在 global memory 中維護一組整數，每條指令完成后會對對應計數器 +1，若某條指令依賴先前步驟的結果，它會等待計數器達到特定值才執行。
例如：在 MLP 下投影階段，團隊將中間態拆成 4 個 chunk，每個 chunk 在寫入后立即觸發后續計算，從而實現并行流。此外，團隊通過精確設置依賴圖，避免了全局 barrier，大幅減少了指令之間等待的浪費，使得整個內核執行盡可能地接近理論并發。
除此之外，研究團隊還對 CUDA 異步屏障（asynchronous barrier）的性能進行了測量，發現即便在 barrier 已“pass”的狀態下，每次仍需 60ns，同步操作成本不可忽視。而在實際執行中，尤其在 matrix-vector（矩陣乘向量）這類關鍵操作中，他們發現：在 Hopper 架構（如 H100）上，使用常規 CUDA 核心（非 Tensor Core）可以更有效，不過在 Blackwell 架構上，Tensor Core 性能占優。
這也說明在硬件不同世代中，Megakernel 的最佳實現路徑也應適配微架構差異，而非一套方案通用所有平臺。
在詳細展開 Megakernel 的構建之前，Hazy 團隊其實先回頭梳理了一個關鍵問題：為什么現在主流的 LLM 推理系統，在小 batch、極低延遲這種場景下，表現這么“不給力”。
他們發現，像 vLLM 和 SGLang 這樣的系統，在處理生成一個 token 這種極限情況時，GPU 的顯存帶寬利用率其實非常低。核心原因是——模型前向過程被拆成了太多太小的 CUDA kernel。也就是說，模型里的每一個小操作（比如 RMSNorm、一個 MLP 層）都是一個單獨的 kernel。這種“微核模式”，看起來很模塊化、易于維護，但其實隱藏了一個很大的性能坑。
每個 kernel 的啟動和銷毀，其實都有固定成本——你可以把它理解成“換個小任務都要重新開會安排”。在極低延遲場景下，這種“開會”的時間反而成了主開銷來源。而且 GPU 在運行這些小 kernel 的時候，還經常會卡在“尾巴”上——比如一個 kernel 要 512 個線程塊跑完，但 GPU 只有 148 個執行單元（SM），那后面的線程塊就只能排隊等前面的慢慢結束，造成很多資源空轉。
即便用上 CUDA Graphs、PDL（Programmatic Dependent Launch）等加速器，也還是得花 1.3～2.1 微秒去啟動一個 kernel。而這段時間，GPU 其實啥都沒干，就是在等待環境切換。更糟的是，由于這些 kernel 是串行排隊執行的，后面的 kernel 也沒法提前加載它要用的數據，導致 GPU 一直斷斷續續地訪問 global memory，帶寬用不上去。
這就形成了所謂的 “memory pipeline bubbles”——計算和計算之間總有空檔期，GPU 明明閑不下來，卻還是停在那等。舉個例子：H100 的帶寬是 3.35TB/s，推理 Llama-1B 每次只需要 2.48GB，理論上 1 秒鐘能跑 1350 次 forward pass。但因為每層模型得跑 7 個 kernel，一共有 16 層，哪怕每個 kernel 只帶來 5 微秒的 stall，總延遲也已經把性能拉到 770 次以內，實際可能還更低。
所以，Hazy 團隊很明確地說：這個問題不是哪個 kernel 慢的問題，而是系統性低效。一個個去優化 kernel 其實沒有用，核心是要干掉這些 kernel 邊界，別再讓 GPU 一會做這個、一會做那個地切換。這就是他們提出 Megakernel 的根本動因。
現代 LLM，動輒幾十上百層 transformer，每層又包含 RMSNorm、注意力、MLP 等等操作。主流框架為了清晰易調試，把這些都拆成一個個小 kernel，每個做一件小事，像流水線上的工人。但問題是，這流水線換手太頻繁，每次“換人”都耽誤事，還導致 GPU 的顯存訪問老是斷斷續續，帶寬效率拉垮。
更要命的是，CUDA 的一些機制雖然看起來是為優化服務的，但在這種極限場景下其實也成了“絆腳石”。比如 PDL 的 cudaGridDependencySynchronize 會強制等所有任務完成才能繼續，這就意味著，即便有些任務早就準備好了，也得一起等著。
所以歸根結底，雷峰網(公眾號：雷峰網)認為現在的推理系統架構，在“單序列、毫秒級響應”這類場景下，是低效的，而且是從系統層面低效。只有重構整個執行方式，讓 GPU 少切換、多并行，才有可能真正把它的算力榨干，這正是 Megakernel 的價值所在。
雷峰網原創文章，未經授權禁止轉載。詳情見轉載須知。
