新智元報道

編輯：定慧

【新智元導讀】斯坦福Hazy實驗室推出新一代低延遲推理引擎「Megakernel」，將Llama-1B模型前向傳播完整融合進單一GPU內核，實現推理時間低于1毫秒。在B200上每次推理僅需680微秒，比vLLM快3.5倍。

想象一下：你和AI聊天時，每句話都要等它3秒——血壓是不是瞬間飆升？

低延遲LLM推理，就是專門針對這個問題的解決辦法。

博客地址：https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles

最近斯坦福Hazy實驗室「整了個大活」：他們手搓了個叫做「Megakernel」超級玩具（推理引擎），把Llama-1B的前向傳播整個塞進單個GPU內核！

結果直接炸場：

• H100上提速1.5倍，帶寬利用率飆到78%

• B200上僅需0.00068秒（人類眨1/3眼的時間！）

• 比vLLM快3.5倍，把SGLang也甩出尾氣

網友辣評：「傳統推理引擎還在騎馬，Megakernel已經開上戰斗機！」

速度！使用32個token的提示詞生成128個token的結果，未使用推測機制。

Fwd/s是衡量模型推理速度的一個指標，表示模型每秒可以執行多少次前向傳播，數值越高，說明模型處理速度越快。

傳統推理引擎：GPU在「摸魚」

通常情況下，人們在GPU上運行代碼的方式是啟動一個「內核（kernel）」—— 一個執行明確定義操作的小型程序（例如RMS歸一化、MLP等）。

當GPU運行大模型時，通常會把任務拆成上百個小內核，目前所有AI工作負載都是如此，表現為一系列相對較短小的內核。

比如先算RMS歸一化 → 再搞注意力 → 接著MLP層...像流水線工人反復交接。

為了有一個初步認識，來看看Llama-1B Transformer模塊中的運算，以及它們可能被劃分的一些示例內核邊界。

LLaMA-1B Transformer 模塊的一組示例內核邊界。紅色框表示各個內核執行的工作

使用Llama-1B解碼單個序列時，是一種純粹受內存限制的工作負載：性能取決于是否能夠持續不斷地從GPU的全局內存中加載權重。

那么，為什么現有的方法距離充分利用GPU的全部帶寬還差得遠呢？

關鍵問題就是，當前基于內核的模型運行方法引入了停頓，阻礙了持續加載內存：

內核空閑

首先GPU內核是按嚴格順序啟動的，因此前一個內核中的所有線程塊完全結束后，后一個內核中的線程塊才能開始執行。

每次啟動一個新的內核時，都必須等待前一個內核中所有落后的線程塊完成。

例如，如果一個內核運行512個線程塊（如LLaMA-1B降維投影（down projection）），但在B200上只有148個流式多處理器（streaming multiprocessors），那么到內核執行的最后階段，就會出現80個空閑的SM——

148 -（512 - 148 * 3）= 80

因為大部分線程塊已經運行完，只有少數幾個還在運行，這些少數線程塊占用了全部的SM，而其他SM就只能空等，造成資源浪費。

內核開銷

其次，正如之前提到的，每次內核（kernel）的啟動和關閉都會帶來開銷。

理論上，NVIDIA的CUDA圖（CUDA graphs）可以在一定程度上隱藏這些開銷，但根據測量結果來看，仍然有不少資源浪費。

各種內核開銷的風格化甘特圖。有時這些開銷可以忽略不計，但很多時候并非如此！

舉個例子，在H100上運行了一個簡單的「假內核」測試（這個內核只記錄開始時間、休眠一段時間、然后記錄結束時間）。結果發現：

• 如果通過普通的CUDA流（stream）來運行，這個內核的啟動開銷大約是2.1微秒；

• 如果使用CUDA圖，啟動開銷雖然下降了，但也只有降到約1.3微秒。

這段時間里，GPU并沒有做任何有用的計算，只是在準備或等待，對整體性能來說是浪費。

所有優化目標是：讓GPU的每一微秒都用在真正有意義的工作上。

內核等待

最后，即便啟動了下一個內核，仍然需要等待權重（weights）和激活值（activations）加載完成之后，計算才能真正開始。

這些等待帶來的延遲會讓GPU空閑上成千上萬個周期！

理想情況下，希望在執行前一個內核的計算和數據存儲時，就能開始加載下一個內核所需的權重。為此，NVIDIA 提供了一種機制，叫做Programmatic Dependent Launch（PDL），它允許在前一個內核還在運行的同時，就提前為下一個內核做準備。

但是，PDL仍然會引入不必要的停頓，這是因為它的同步機制（cudaGridDependencySynchronize）太粗粒度了。舉個例子，它要求所有的query、key和value全部計算完畢后，attention才能開始，而不能在某個head準備好時就立即開始計算。

稍后會在Llama-1B的一個具體案例中展示，這種機制在哪些情況下會限制性能。

綜合來看，這些形成了標題中提到的「內存流水線氣泡」——這也是并非始終從內存加載數據的一個關鍵原因。

對于短時操作來說，這些暫停累積起來會浪費大量的潛在帶寬。

部分原因在于，Llama-1B（實際為1.24B參數）在批量大小為1時實在太過……微小：如果每個操作本身非常快，那么操作之間的間隔時間就開始變得至關重要。

為了說明問題的嚴重性：在單個H100上以16位精度進行單序列生成時，內存限制為3.35TB/s / 2.48GB = 每秒約 1350 次前向傳遞。

但每層需要7次內核啟動，共16層，即使每次內核阻塞時間樂觀估計為5微秒（包括尾部任務、內核啟動和內存延遲），生成速度也將僅約為每秒770次前向傳遞。

實際情況往往更差。在低延遲工作負載下，GPU只有很少一部分時間真正用于執行有用的工作！

雖然CUDA確實提供了一些現有功能（例如圖、流、PDL）來部分解決這些問題，但研究團隊想看看是否有一種不同的方法可以解決所有這些問題，即將整個模型的前向計算融合成一個單一的內核。

如何設計Megakernel

如何將整個LLaMA前向傳遞過程融合到一個單一內核中，需要解決三個關鍵問題：

1 融合數十個操作從頭開始做起來很難。需要一種機制來在Megakernel中執行這些操作。

2 為了在相同硬件上重疊多個操作，需要防止對有限資源（例如共享內存）的競爭。

3 在傳統內核模型中，GPU會在每個內核之后進行同步。沒有了內核，必須自己手動對GPU進行同步！

問題1：融合大爆炸

傳統內核融合通常只合并兩到三個操作。

但是這里需要融合大約一百個操作。

因此，需要一個合理的抽象方式，來對Megakernel進行編程。

一種方法是基于一個在GPU上的解釋器——本質上是ThunderMLA底層架構的一個更復雜版本的解釋器設計使得 GPU 中的每個流式多處理器（SM）都能接收一連串的指令（每個指令都使用相同的 CUDA 模板實現）并加以執行。

在Python端提前安排好每個SM的指令序列，值得注意的是，每個調度可以被重用于數百次前向傳遞！

對于端到端的Llama前向傳遞Megakernel，定義了以下指令集：

• 融合的RMS歸一化、QKV和RoPE指令。

• 一個注意力計算指令。

• 一種注意力縮減指令（用ThunderGQA在長序列上的處理）。

• 一個O投影加殘差指令。

• 融合的RMS歸一化、上行門控和SiLU指令。

• 一個下投影加殘差指令。

• 一個RMS歸一化和語言建模頭部指令，用于計算最終的tokenlogits。

• 使用一個通用的CUDA模板（包含加載、存儲、計算樣板函數）來實現每條這些指令，從而在的解釋器框架內促進互操作性。

問題2：共享內存以消除內存氣泡

指令與解釋器結構使能夠清晰地組織Megakernel。

然而，尚未解決一個關鍵問題：確保模型權重始終按順序加載，以最大化內存帶寬利用率。

Megakernel之所以能讓解決此問題，是因為可以在指令之間進行內存加載流水線操作：解釋器一旦能夠開始加載某條指令的模型權重，即使前一條指令仍在完成階段（例如將結果存儲到全局內存），它也會立即開始加載。

正是這種緊密的指令間切換，最大限度地減少了因啟動多個內核而可能出現的內存氣泡。

然而，這里有個問題：如果下一個指令沒有空間存放已加載的數據，那么從全局內存加載權重并不會帶來太大好處！

更準確地說，所有的權重矩陣都是從GPU全局內存加載到SM的「共享內存」中——這是NVIDIA對每個SM上快速內存的稱呼。

共享內存是每個SM上的稀缺資源，如果前一個指令占用了全部共享內存，就無法為新指令啟動新的加載操作。

這就需要一種方法來跟蹤哪個指令正在使用哪一部分共享內存，并在當前指令完成時迅速將共享內存過渡給下一個指令使用。

通過分頁共享內存來實現這一點。

首先將H100上的前213KB共享內存劃分為13個16KiB的頁面，并將剩余的共享內存用于特殊用途，例如存儲指令參數。

要使用這些頁面之一，指令必須顯式地從解釋器請求并釋放它們。

解釋器會自動將已釋放的頁面傳遞給下一條指令，允許它們在共享內存可用后盡早開始發出內存加載操作。

問題3：同步

雖然Megakernels能夠幫助最大限度地減少流水線氣泡，但它們也引入了一個新的問題：同步。

在常規的多kernel執行模型中，性能限制在于，直到之前所有kernel中的線程塊都完成后，下一個kernel中的線程塊才能開始執行。

然而，正是這一特性使得管理數據依賴關系變得簡單。當一個kernel啟動時，CUDA保證該kernel所需的所有輸入張量已經生成，并且可以安全地立即讀取。

使用Megakernel時，沒有這樣的保障：當一個SM開始執行新指令時，其輸入可能尚未就緒！

為了解決這個問題，在Megakernel內部顯式地對指令進行同步。

通過一個簡單的計數器系統來實現這一點。

在Megakernel啟動之前，在GPU全局內存中初始化一個計數器數組（即整數數組），初始值為零。

每當一條指令完成時，它會增加其中一個計數器的值。

同樣，每當新指令開始時，它必須等待其中某些計數器達到目標值，這表明其所有依賴項均已執行完畢。

這一優化在Llama-1B的大型多層感知機（MLP）中得以實現。

• 在使用PDL的樸素實現中，必須等待整個隱藏狀態完成之后，才能開始下投影矩陣的乘法運算。

• 改為將中間狀態劃分為四個塊進行處理，每個塊都有各自的計數器。這樣一來，針對下投影的指令只需等待其對應的輸入塊完成即可。

整合所有內容

據研究團隊所知，H100 Megakernel代表了有人首次在GPU上實現以16位精度運行參數超過10億的語言模型的前向傳播時間低于一毫秒。

而的B200實現更是將這一時間進一步縮短至每次前向傳播不到680微秒！

如文章開頭的圖片所示，Megakernel性能優于vLLM和SGLang基線（它們使用CUDA圖和Torch編譯）：

• 在H100上，Megakernel運行速度幾乎是vLLM的2.5倍，比SGLang快超過1.5倍。

• 在B200上，與vLLM的差距擴大到3.5倍以上，仍然比SGLang快1.5倍以上。

距離B200上理論極限（大約每秒3,000次前向計算）仍有相當大的差距。

部分原因在于，該理論極限純粹基于內存帶寬——但仍需等待加載激活值。盡管這些激活值體積較小（不會占用大量帶寬），但加載它們時仍然存在無法隱藏的延遲。

以下是當前B200前向計算運行時間的分解（總運行時間600微秒）：

• 存儲激活值、等待一致性以及加載這些激活值需要花費250微秒。

這比簡單模型預測的結果高出約20％：由于每條指令都依賴于前一條指令，需要支付兩次加載延遲（檢查就緒狀態，然后加載激活值）和兩次存儲延遲（存儲激活值，然后標記為就緒）的開銷，每條指令都是如此。

以每次加載/存儲大約500納秒的延遲來計算，這將帶來約200微秒的開銷。（懷疑剩余的大約50微秒中，有一部分來自在全局內存中處理原子操作所花費的時間。）

• 實際運行RMS歸一化和矩陣向量計算花費了200微秒。

這部分時間中約有95%用于矩陣向量運算。在Blackwell上，發現使用張量核心對此幫助不大；而在Hopper上，直接在CUDA核心上運行效果更好。這種差異的原因在于兩種GPU的CUDA核心性能相對接近，但Blackwell的張量核心要快得多。

• 30微秒用于等待全局內存中的權重（流水線工作正常！）。

其中，40%的時間花費在LM頭部，這是整個Megakernel中流水線效率最高的部分，因為其具有高度一致性和龐大的規模。

• 在各個線程束（warp）之間，有40微秒花費在低層次的同步開銷上。

這里的一個關鍵問題是，即使在「通過」狀態時，CUDA的異步屏障操作速度也相對較慢，每次都需要大約60納秒的時間。

• 80微秒用于設置和各種其他開銷。

例如，通過指令屏障、將頁面標記為完成等。

本次突破明確展示了減少內核切換、優化內存流水線和精細同步的重要性，這也預示著低延遲推理技術的進一步發展潛力。

參考資料：

https://news.ycombinator.com/item?id=44111673

https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles