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