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

低延遲LLM推理，就是專門針對(duì)這個(gè)問題的解決辦法。

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

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

結(jié)果直接炸場(chǎng)：

• H100上提速1.5倍，帶寬利用率飆到78%
• B200上僅需0.00068秒（人類眨1/3眼的時(shí)間?。?/span>
• 比vLLM快3.5倍，把SGLang也甩出尾氣

網(wǎng)友辣評(píng)：「?jìng)鹘y(tǒng)推理引擎還在騎馬，Megakernel已經(jīng)開上戰(zhàn)斗機(jī)！」

速度！使用32個(gè)token的提示詞生成128個(gè)token的結(jié)果，未使用推測(cè)機(jī)制。

Fwd/s是衡量模型推理速度的一個(gè)指標(biāo)，表示模型每秒可以執(zhí)行多少次前向傳播，數(shù)值越高，說明模型處理速度越快。

傳統(tǒng)推理引擎：GPU在「摸魚」

通常情況下，人們?cè)贕PU上運(yùn)行代碼的方式是啟動(dòng)一個(gè)「內(nèi)核（kernel）」—— 一個(gè)執(zhí)行明確定義操作的小型程序（例如RMS歸一化、MLP等）。

當(dāng)GPU運(yùn)行大模型時(shí)，通常會(huì)把任務(wù)拆成上百個(gè)小內(nèi)核，目前所有AI工作負(fù)載都是如此，表現(xiàn)為一系列相對(duì)較短小的內(nèi)核。

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

為了有一個(gè)初步認(rèn)識(shí)，來看看Llama-1B Transformer模塊中的運(yùn)算，以及它們可能被劃分的一些示例內(nèi)核邊界。

LLaMA-1B Transformer 模塊的一組示例內(nèi)核邊界。紅色框表示各個(gè)內(nèi)核執(zhí)行的工作

使用Llama-1B解碼單個(gè)序列時(shí)，是一種純粹受內(nèi)存限制的工作負(fù)載：性能取決于是否能夠持續(xù)不斷地從GPU的全局內(nèi)存中加載權(quán)重。

那么，為什么現(xiàn)有的方法距離充分利用GPU的全部帶寬還差得遠(yuǎn)呢？

關(guān)鍵問題就是，當(dāng)前基于內(nèi)核的模型運(yùn)行方法引入了停頓，阻礙了持續(xù)加載內(nèi)存：

內(nèi)核空閑

首先GPU內(nèi)核是按嚴(yán)格順序啟動(dòng)的，因此前一個(gè)內(nèi)核中的所有線程塊完全結(jié)束后，后一個(gè)內(nèi)核中的線程塊才能開始執(zhí)行。

每次啟動(dòng)一個(gè)新的內(nèi)核時(shí)，都必須等待前一個(gè)內(nèi)核中所有落后的線程塊完成。

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

148 -（512 - 148 * 3）= 80

因?yàn)榇蟛糠志€程塊已經(jīng)運(yùn)行完，只有少數(shù)幾個(gè)還在運(yùn)行，這些少數(shù)線程塊占用了全部的SM，而其他SM就只能空等，造成資源浪費(fèi)。

內(nèi)核開銷

其次，正如之前提到的，每次內(nèi)核（kernel）的啟動(dòng)和關(guān)閉都會(huì)帶來開銷。

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

各種內(nèi)核開銷的風(fēng)格化甘特圖。有時(shí)這些開銷可以忽略不計(jì)，但很多時(shí)候并非如此！

舉個(gè)例子，在H100上運(yùn)行了一個(gè)簡(jiǎn)單的「假內(nèi)核」測(cè)試（這個(gè)內(nèi)核只記錄開始時(shí)間、休眠一段時(shí)間、然后記錄結(jié)束時(shí)間）。結(jié)果發(fā)現(xiàn)：

• 如果通過普通的CUDA流（stream）來運(yùn)行，這個(gè)內(nèi)核的啟動(dòng)開銷大約是2.1微秒；
• 如果使用CUDA圖，啟動(dòng)開銷雖然下降了，但也只有降到約1.3微秒。

這段時(shí)間里，GPU并沒有做任何有用的計(jì)算，只是在準(zhǔn)備或等待，對(duì)整體性能來說是浪費(fèi)。

所有優(yōu)化目標(biāo)是：讓GPU的每一微秒都用在真正有意義的工作上。

內(nèi)核等待

最后，即便啟動(dòng)了下一個(gè)內(nèi)核，仍然需要等待權(quán)重（weights）和激活值（activations）加載完成之后，計(jì)算才能真正開始。

這些等待帶來的延遲會(huì)讓GPU空閑上成千上萬個(gè)周期！

理想情況下，希望在執(zhí)行前一個(gè)內(nèi)核的計(jì)算和數(shù)據(jù)存儲(chǔ)時(shí)，就能開始加載下一個(gè)內(nèi)核所需的權(quán)重。為此，NVIDIA 提供了一種機(jī)制，叫做 Programmatic Dependent Launch（PDL），它允許在前一個(gè)內(nèi)核還在運(yùn)行的同時(shí)，就提前為下一個(gè)內(nèi)核做準(zhǔn)備。

但是，PDL仍然會(huì)引入不必要的停頓，這是因?yàn)樗耐綑C(jī)制（cudaGridDependencySynchronize）太粗粒度了。舉個(gè)例子，它要求所有的query、key和value全部計(jì)算完畢后，attention才能開始，而不能在某個(gè)head準(zhǔn)備好時(shí)就立即開始計(jì)算。

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

綜合來看，這些形成了標(biāo)題中提到的「內(nèi)存流水線氣泡」——這也是并非始終從內(nèi)存加載數(shù)據(jù)的一個(gè)關(guān)鍵原因。

對(duì)于短時(shí)操作來說，這些暫停累積起來會(huì)浪費(fèi)大量的潛在帶寬。

部分原因在于，Llama-1B（實(shí)際為1.24B參數(shù)）在批量大小為1時(shí)實(shí)在太過……微?。喝绻總€(gè)操作本身非?？?，那么操作之間的間隔時(shí)間就開始變得至關(guān)重要。

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

但每層需要7次內(nèi)核啟動(dòng)，共16層，即使每次內(nèi)核阻塞時(shí)間樂觀估計(jì)為5微秒（包括尾部任務(wù)、內(nèi)核啟動(dòng)和內(nèi)存延遲），生成速度也將僅約為每秒770次前向傳遞。

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

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

如何設(shè)計(jì)Megakernel

如何將整個(gè)LLaMA前向傳遞過程融合到一個(gè)單一內(nèi)核中，需要解決三個(gè)關(guān)鍵問題：

1 融合數(shù)十個(gè)操作從頭開始做起來很難。需要一種機(jī)制來在Megakernel中執(zhí)行這些操作。

2 為了在相同硬件上重疊多個(gè)操作，需要防止對(duì)有限資源（例如共享內(nèi)存）的競(jìng)爭(zhēng)。

3 在傳統(tǒng)內(nèi)核模型中，GPU會(huì)在每個(gè)內(nèi)核之后進(jìn)行同步。沒有了內(nèi)核，必須自己手動(dòng)對(duì)GPU進(jìn)行同步！

問題1：融合大爆炸

傳統(tǒng)內(nèi)核融合通常只合并兩到三個(gè)操作。

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

因此，需要一個(gè)合理的抽象方式，來對(duì)Megakernel進(jìn)行編程。

一種方法是基于一個(gè)在GPU上的解釋器——本質(zhì)上是ThunderMLA底層架構(gòu)的一個(gè)更復(fù)雜版本的解釋器設(shè)計(jì)使得 GPU 中的每個(gè)流式多處理器（SM）都能接收一連串的指令（每個(gè)指令都使用相同的 CUDA 模板實(shí)現(xiàn)）并加以執(zhí)行。

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

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

• 融合的RMS歸一化、QKV和RoPE指令。
• 一個(gè)注意力計(jì)算指令。
• 一種注意力縮減指令（用ThunderGQA在長(zhǎng)序列上的處理）。
• 一個(gè)O投影加殘差指令。
• 融合的RMS歸一化、上行門控和SiLU指令。
• 一個(gè)下投影加殘差指令。
• 一個(gè)RMS歸一化和語言建模頭部指令，用于計(jì)算最終的tokenlogits。
• 使用一個(gè)通用的CUDA模板（包含加載、存儲(chǔ)、計(jì)算樣板函數(shù)）來實(shí)現(xiàn)每條這些指令，從而在的解釋器框架內(nèi)促進(jìn)互操作性。

問題2：共享內(nèi)存以消除內(nèi)存氣泡

指令與解釋器結(jié)構(gòu)使能夠清晰地組織Megakernel。

然而，尚未解決一個(gè)關(guān)鍵問題：確保模型權(quán)重始終按順序加載，以最大化內(nèi)存帶寬利用率。

Megakernel之所以能讓解決此問題，是因?yàn)榭梢栽谥噶钪g進(jìn)行內(nèi)存加載流水線操作：解釋器一旦能夠開始加載某條指令的模型權(quán)重，即使前一條指令仍在完成階段（例如將結(jié)果存儲(chǔ)到全局內(nèi)存），它也會(huì)立即開始加載。

正是這種緊密的指令間切換，最大限度地減少了因啟動(dòng)多個(gè)內(nèi)核而可能出現(xiàn)的內(nèi)存氣泡。

然而，這里有個(gè)問題：如果下一個(gè)指令沒有空間存放已加載的數(shù)據(jù)，那么從全局內(nèi)存加載權(quán)重并不會(huì)帶來太大好處！

更準(zhǔn)確地說，所有的權(quán)重矩陣都是從GPU全局內(nèi)存加載到SM的「共享內(nèi)存」中——這是NVIDIA對(duì)每個(gè)SM上快速內(nèi)存的稱呼。

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

這就需要一種方法來跟蹤哪個(gè)指令正在使用哪一部分共享內(nèi)存，并在當(dāng)前指令完成時(shí)迅速將共享內(nèi)存過渡給下一個(gè)指令使用。

通過分頁共享內(nèi)存來實(shí)現(xiàn)這一點(diǎn)。

首先將H100上的前213KB共享內(nèi)存劃分為13個(gè)16KiB的頁面，并將剩余的共享內(nèi)存用于特殊用途，例如存儲(chǔ)指令參數(shù)。

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

解釋器會(huì)自動(dòng)將已釋放的頁面?zhèn)鬟f給下一條指令，允許它們?cè)诠蚕韮?nèi)存可用后盡早開始發(fā)出內(nèi)存加載操作。

問題3：同步

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

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

然而，正是這一特性使得管理數(shù)據(jù)依賴關(guān)系變得簡(jiǎn)單。當(dāng)一個(gè)kernel啟動(dòng)時(shí)，CUDA保證該kernel所需的所有輸入張量已經(jīng)生成，并且可以安全地立即讀取。

使用Megakernel時(shí)，沒有這樣的保障：當(dāng)一個(gè)SM開始執(zhí)行新指令時(shí)，其輸入可能尚未就緒！

為了解決這個(gè)問題，在Megakernel內(nèi)部顯式地對(duì)指令進(jìn)行同步。

通過一個(gè)簡(jiǎn)單的計(jì)數(shù)器系統(tǒng)來實(shí)現(xiàn)這一點(diǎn)。

在Megakernel啟動(dòng)之前，在GPU全局內(nèi)存中初始化一個(gè)計(jì)數(shù)器數(shù)組（即整數(shù)數(shù)組），初始值為零。

每當(dāng)一條指令完成時(shí)，它會(huì)增加其中一個(gè)計(jì)數(shù)器的值。

同樣，每當(dāng)新指令開始時(shí)，它必須等待其中某些計(jì)數(shù)器達(dá)到目標(biāo)值，這表明其所有依賴項(xiàng)均已執(zhí)行完畢。

這一優(yōu)化在Llama-1B的大型多層感知機(jī)（MLP）中得以實(shí)現(xiàn)。

• 在使用PDL的樸素實(shí)現(xiàn)中，必須等待整個(gè)隱藏狀態(tài)完成之后，才能開始下投影矩陣的乘法運(yùn)算。
• 改為將中間狀態(tài)劃分為四個(gè)塊進(jìn)行處理，每個(gè)塊都有各自的計(jì)數(shù)器。這樣一來，針對(duì)下投影的指令只需等待其對(duì)應(yīng)的輸入塊完成即可。

整合所有內(nèi)容

據(jù)研究團(tuán)隊(duì)所知，H100 Megakernel代表了有人首次在GPU上實(shí)現(xiàn)以16位精度運(yùn)行參數(shù)超過10億的語言模型的前向傳播時(shí)間低于一毫秒。

而的B200實(shí)現(xiàn)更是將這一時(shí)間進(jìn)一步縮短至每次前向傳播不到680微秒！

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

• 在H100上，Megakernel運(yùn)行速度幾乎是vLLM的2.5倍，比SGLang快超過1.5倍。
• 在B200上，與vLLM的差距擴(kuò)大到3.5倍以上，仍然比SGLang快1.5倍以上。

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

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

以下是當(dāng)前B200前向計(jì)算運(yùn)行時(shí)間的分解（總運(yùn)行時(shí)間600微秒）：

• 存儲(chǔ)激活值、等待一致性以及加載這些激活值需要花費(fèi)250微秒。

這比簡(jiǎn)單模型預(yù)測(cè)的結(jié)果高出約20％：由于每條指令都依賴于前一條指令，需要支付兩次加載延遲（檢查就緒狀態(tài)，然后加載激活值）和兩次存儲(chǔ)延遲（存儲(chǔ)激活值，然后標(biāo)記為就緒）的開銷，每條指令都是如此。

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

• 實(shí)際運(yùn)行RMS歸一化和矩陣向量計(jì)算花費(fèi)了200微秒。

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

• 30微秒用于等待全局內(nèi)存中的權(quán)重（流水線工作正常?。?。

其中，40%的時(shí)間花費(fèi)在LM頭部，這是整個(gè)Megakernel中流水線效率最高的部分，因?yàn)槠渚哂懈叨纫恢滦院妄嫶蟮囊?guī)模。

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

這里的一個(gè)關(guān)鍵問題是，即使在「通過」?fàn)顟B(tài)時(shí)，CUDA的異步屏障操作速度也相對(duì)較慢，每次都需要大約60納秒的時(shí)間。

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

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

本次突破明確展示了減少內(nèi)核切換、優(yōu)化內(nèi)存流水線和精細(xì)同步的重要性，這也預(yù)示著低延遲推理技術(shù)的進(jìn)一步發(fā)展?jié)摿Α?/span>