機器之心報道

編輯：杜偉

在 AI 領(lǐng)域，英偉達開發(fā)的 CUDA 是驅(qū)動大語言模型（LLM）訓(xùn)練和推理的核心計算引擎。

不過，CUDA 驅(qū)動的 LLM 推理面臨著手動優(yōu)化成本高、端到端延遲高等不足，需要進一步優(yōu)化或者尋找更高效的替代方案。

近日，CMU 助理教授賈志豪（Zhihao Jia）團隊創(chuàng)新玩法，推出了一個名為「Mirage Persistent Kernel（MPK）」的編譯器，可以自動將 LLM 轉(zhuǎn)化為優(yōu)化的巨型內(nèi)核（megakernel），從而將 LLM 推理延遲降低 1.2 到 6.7 倍。

• GitHub 地址：https://github.com/mirage-project/mirage/tree/mpk
• 博客地址：https://zhihaojia.medium.com/compiling-llms-into-a-megakernel-a-path-to-low-latency-inference-cf7840913c17

MPK 將 LLM 推理延遲推近硬件極限。在單個 A100-40GB GPU 上，MPK 將 Qwen3-8B 每個 token 的延遲從 14.5 毫秒 (vLLM/SGLang) 降低到 12.5 毫秒，逼近基于內(nèi)存帶寬計算得出的 10 毫秒理論下限。

MPK 的易用性很強，你只需要幾十行 Python 代碼就能將 LLM 編譯成一個高性能巨型內(nèi)核，實現(xiàn)快速推理，整個過程無需 CUDA 編程。

評論區(qū)對 MPK 的看法也很正向，并提出了一些未來的延展方向。

引入 MPK 的必要性

降低 LLM 推理延遲最有效的方法之一，是將所有計算和通信融合進一個單一的巨型內(nèi)核，也稱為持續(xù)內(nèi)核。

在這種設(shè)計中，系統(tǒng)僅啟動一個 GPU 內(nèi)核來執(zhí)行整個模型 —— 從逐層計算到 GPU 間通信 —— 整個過程無需中斷。這種方法提供了以下幾個關(guān)鍵的性能優(yōu)勢：

• 消除內(nèi)核啟動開銷：通過避免重復(fù)的內(nèi)核調(diào)用，即使是在多 GPU 環(huán)境下，也能消除內(nèi)核啟動開銷；
• 實現(xiàn)跨層軟件 pipeline 允許內(nèi)核在計算當(dāng)前層的同時，開始為下一層加載數(shù)據(jù)；
• 重疊計算與通信：由于巨型內(nèi)核可以同時執(zhí)行計算操作和 GPU 間通信，從而隱藏通信延遲。

盡管有這些優(yōu)勢，將 LLM 編譯成巨型內(nèi)核仍然極具挑戰(zhàn)性。

現(xiàn)有的高級 ML 框架 —— 如 PyTorch、Triton 和 TVM，它們本身并不支持端到端巨型內(nèi)核生成。此外，現(xiàn)代 LLM 系統(tǒng)由各種不同的專用內(nèi)核庫構(gòu)建而成：用于通信的 NCCL 或 NVSHMEM，用于高效注意力計算的 FlashInfer 或 FlashAttention，以及用于自定義計算的 CUDA 或 Triton。

這種碎片化使得將整個推理 pipeline 整合進一個單一的、統(tǒng)一的內(nèi)核變得非常困難。

那么能否通過編譯自動化這個過程呢？受到這個問題的啟發(fā)，來自 CMU、華盛頓大學(xué)、加州大學(xué)伯克利分校、英偉達和清華大學(xué)的團隊開發(fā)出了 MPK—— 一個編譯器和運行時系統(tǒng)，它能自動將多 GPU 的 LLM 推理轉(zhuǎn)換為高性能的巨型內(nèi)核。MPK 釋放了端到端 GPU 融合的效能優(yōu)勢，同時只需要開發(fā)者付出極小的手動努力。

MPK 的優(yōu)勢

MPK 的一個關(guān)鍵優(yōu)勢在于：通過消除內(nèi)核啟動開銷，并最大程度地重疊跨層的計算、數(shù)據(jù)加載和 GPU 間通信，實現(xiàn)了極低的 LLM 推理延遲。

下圖 1 展示了 MPK 與現(xiàn)有 LLM 推理系統(tǒng)在單 GPU 和多 GPU 配置下的性能對比（具體可見上文）。

除了單 GPU 優(yōu)化，MPK 還將計算與 GPU 間通信融合進一個單一的巨型內(nèi)核。 這種設(shè)計使得 MPK 能夠最大程度地重疊計算與通信。因此，MPK 相對于當(dāng)前系統(tǒng)的性能提升隨著 GPU 數(shù)量的增加而增大，使其在多 GPU 部署場景下尤為高效。

MPK 的工作原理

MPK 的工作原理包括以下兩大部分

• Part 1：MPK 編譯器，其將 LLM 的計算圖轉(zhuǎn)化為優(yōu)化的任務(wù)圖；
• Part 2：MPK 運行時系統(tǒng)，該系統(tǒng)在單個巨型內(nèi)核內(nèi)執(zhí)行任務(wù)圖，以實現(xiàn)高吞吐量與低延遲。

編譯器 —— 將 LLM 轉(zhuǎn)化為細(xì)粒度任務(wù)圖

LLM 的計算過程通常表示為計算圖，其中每個節(jié)點對應(yīng)一個計算算子（如矩陣乘法、注意力機制）或集合通信原語（如 all-reduce），邊表示算子間的數(shù)據(jù)依賴關(guān)系。現(xiàn)有系統(tǒng)通常為每個算子啟動獨立的 GPU 內(nèi)核。

然而，這種「單算子單內(nèi)核」的執(zhí)行模型難以實現(xiàn) pipeline 優(yōu)化，因為依賴關(guān)系是在整個內(nèi)核的粗粒度層面強制執(zhí)行的，而非實際數(shù)據(jù)單元層面。

典型案例如矩陣乘法（matmul）后接 all-reduce 操作：現(xiàn)有系統(tǒng)中，all-reduce 內(nèi)核必須等待整個 matmul 內(nèi)核完成。而實際上，all-reduce 的每個數(shù)據(jù)分塊僅依賴 matmul 輸出的局部結(jié)果。這種邏輯依賴與實際依賴的錯配，嚴(yán)重限制了計算與通信的重疊潛力。

下圖 2 展示了 MPK 編譯器將 PyTorch 定義的 LLM 計算圖轉(zhuǎn)化為優(yōu)化細(xì)粒度任務(wù)圖，最大化暴露并行性。右側(cè)展示次優(yōu)方案 —— 其引入不必要的數(shù)據(jù)依賴與全局屏障，導(dǎo)致跨層流水線優(yōu)化機會受限。

為了解決此問題，MPK 引入的編譯器可將 LLM 計算圖自動轉(zhuǎn)化為細(xì)粒度任務(wù)圖。該任務(wù)圖在子內(nèi)核級別顯式捕獲依賴關(guān)系，實現(xiàn)更激進的跨層流水線優(yōu)化。

具體來講，在 MPK 任務(wù)圖中（如圖 2 所示）：

• 任務(wù)（矩形表示），代表分配給單個 GPU 流式多處理器（SM）的計算 / 通信單元。
• 事件（圓形表示），表示任務(wù)間的同步點。
• 觸發(fā)機制，每個任務(wù)發(fā)出指向觸發(fā)事件的邊，該事件在關(guān)聯(lián)任務(wù)全部完成后激活。
• 依賴機制，每個任務(wù)接收來自依賴事件的邊，表明事件激活后任務(wù)立即啟動。

任務(wù)圖使 MPK 能夠發(fā)掘計算圖中無法實現(xiàn)的 pipeline 優(yōu)化機會。例如，MPK 可以構(gòu)建優(yōu)化任務(wù)圖 —— 其中每個 all-reduce 任務(wù)僅依賴于生成其輸入的對應(yīng) matmul 任務(wù)，從而實現(xiàn)分塊執(zhí)行與計算通信重疊。

除生成優(yōu)化任務(wù)圖外，MPK 還通過 Mirage 內(nèi)核超優(yōu)化器自動為每個任務(wù)生成高性能 CUDA 實現(xiàn)，確保任務(wù)在 GPU 流式多處理器（SM）上高效執(zhí)行。

Part 2：運行時 —— 在巨型內(nèi)核中執(zhí)行任務(wù)圖

MPK 包含內(nèi)置 GPU 運行時系統(tǒng)，可在單個 GPU 巨型內(nèi)核內(nèi)完整執(zhí)行任務(wù)圖。這使得系統(tǒng)能在推理過程中無需額外內(nèi)核啟動的情況下，實現(xiàn)任務(wù)執(zhí)行與調(diào)度的細(xì)粒度控制。

為了實現(xiàn)此機制，MPK 在啟動時將 GPU 上所有流式多處理器（SM）靜態(tài)分區(qū)為兩種角色：即工作單元（Worker）和調(diào)度單元（Scheduler）。

工作 SM 與調(diào)度 SM 的數(shù)量在內(nèi)核啟動時固定配置，且總和等于物理 SM 總數(shù)，從而徹底避免動態(tài)上下文切換開銷。

工作單元

每個工作單元獨占一個流式多處理器（SM），并維護專屬任務(wù)隊列。其執(zhí)行遵循以下高效簡潔的循環(huán)流程：

• 獲取任務(wù)：從隊列中提取下一待執(zhí)行任務(wù)。
• 執(zhí)行計算：運行任務(wù)（如矩陣乘法 / 注意力機制 / GPU 間數(shù)據(jù)傳輸）。
• 事件觸發(fā)：任務(wù)完成后通知觸發(fā)事件。
• 循環(huán)執(zhí)行：重復(fù)上述過程。

該機制既保障了工作單元的持續(xù)滿載運行，又實現(xiàn)了跨層和跨操作的異步任務(wù)執(zhí)行。

調(diào)度單元

調(diào)度決策由 MPK 的分布式調(diào)度單元處理，每個調(diào)度單元運行于單個線程束（warp）上。由于每個流式多處理器（SM）可以容納多個線程束，因此單 SM 最多可并發(fā)運行 4 個調(diào)度單元。每個調(diào)度單元維護激活事件隊列，并持續(xù)執(zhí)行以下操作：

• 事件出隊：移除依賴已滿足的激活事件（即所有前置任務(wù)均已完成）。
• 任務(wù)啟動：調(diào)度依賴該激活事件的任務(wù)集。

這種分布式調(diào)度機制在實現(xiàn)跨 SM 可擴展執(zhí)行的同時，最小化協(xié)同開銷。

事件驅(qū)動執(zhí)行

下圖 3 展示了 MPK 的執(zhí)行時間線，其中每個矩形代表一個在工作單元上運行的任務(wù)；每個圓圈代表一個事件。當(dāng)一個任務(wù)完成時，它會遞增其對應(yīng)觸發(fā)事件的計數(shù)器。當(dāng)事件計數(shù)器達到預(yù)設(shè)閾值時，該事件被視為已激活，并被加入調(diào)度單元的事件隊列。隨后，調(diào)度單元會啟動所有依賴于該事件的下游任務(wù)。

這種設(shè)計實現(xiàn)了細(xì)粒度的軟件流水線化，并允許計算與通信之間重疊，比如

• 矩陣乘法（Matmul）任務(wù)可以與來自不同層的注意力任務(wù)并行執(zhí)行。
• 一旦有部分 matmul 結(jié)果可用，即可開始 Allreduce 通信。

由于所有的調(diào)度和任務(wù)切換都發(fā)生在單一內(nèi)核上下文內(nèi)，任務(wù)間的開銷極低，通常僅需 1-2 微秒，從而能夠高效地執(zhí)行多層、多 GPU 的 LLM 工作負(fù)載。

下一步計劃

團隊對 MPK 的愿景是使巨型內(nèi)核編譯既易于使用又具備高性能。目前，你只需幾十行 Python 代碼（主要用于指定巨型內(nèi)核的輸入和輸出）即可將一個 LLM 編譯成一個巨型內(nèi)核。此方向仍有廣闊的探索空間，目前正在積極攻關(guān)的一些關(guān)鍵領(lǐng)域包括如下：

• 支持現(xiàn)代 GPU 架構(gòu)。下一個里程碑是將 MPK 擴展到支持下一代架構(gòu)，例如 NVIDIA Blackwell。一個主要挑戰(zhàn)在于如何將線程束專業(yè)化，這是新型 GPU 的一項關(guān)鍵優(yōu)化技術(shù)，與 MPK 的巨型內(nèi)核執(zhí)行模型相集成。
• 處理工作負(fù)載動態(tài)性。MPK 目前構(gòu)建的是靜態(tài)任務(wù)圖，這限制了它處理動態(tài)工作負(fù)載（如 MoE 模型）的能力。團隊正在開發(fā)新的編譯策略，使 MPK 能夠在巨型內(nèi)核內(nèi)部支持動態(tài)控制流和條件執(zhí)行。
• 高級調(diào)度與任務(wù)分配。MPK 在任務(wù)級別解鎖了新的細(xì)粒度調(diào)度能力。雖然當(dāng)前的實現(xiàn)使用簡單的輪詢調(diào)度在流式多處理器（SM）之間分配任務(wù)，但團隊看到了在高級調(diào)度策略（如優(yōu)先級感知或吞吐量優(yōu)化策略）方面令人興奮的機會，可應(yīng)用于諸如延遲服務(wù)等級目標(biāo)（SLO）驅(qū)動的服務(wù)或混合批處理等場景。

團隊相信，MPK 代表了在 GPU 上編譯和執(zhí)行 LLM 推理工作負(fù)載方式的根本性轉(zhuǎn)變，并熱切期待與社區(qū)合作，共同推動這一愿景向前發(fā)展。

該項目也在快速迭代中，非常歡迎有興趣的伙伴加入contribute。