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