在 AI 領域，英偉達開發的 CUDA 是驅動大語言模型（LLM）訓練和推理的核心計算引擎。

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

近日，CMU 助理教授賈志豪（Zhihao Jia）團隊創新玩法，推出了一個名為「Mirage Persistent Kernel（MPK）」的編譯器，可以自動將 LLM 轉化為優化的巨型內核（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 毫秒，逼近基于內存帶寬計算得出的 10 毫秒理論下限。

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

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

引入 MPK 的必要性

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

在這種設計中，系統僅啟動一個 GPU 內核來執行整個模型 —— 從逐層計算到 GPU 間通信 —— 整個過程無需中斷。這種方法提供了以下幾個關鍵的性能優勢：

• 消除內核啟動開銷：通過避免重復的內核調用，即使是在多 GPU 環境下，也能消除內核啟動開銷；
• 實現跨層軟件 pipeline 允許內核在計算當前層的同時，開始為下一層加載數據；
• 重疊計算與通信：由于巨型內核可以同時執行計算操作和 GPU 間通信，從而隱藏通信延遲。

盡管有這些優勢，將 LLM 編譯成巨型內核仍然極具挑戰性。

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

這種碎片化使得將整個推理 pipeline 整合進一個單一的、統一的內核變得非常困難。

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

MPK 的優勢

MPK 的一個關鍵優勢在于：通過消除內核啟動開銷，并最大程度地重疊跨層的計算、數據加載和 GPU 間通信，實現了極低的 LLM 推理延遲。

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

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

MPK 的工作原理

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

• Part 1：MPK 編譯器，其將 LLM 的計算圖轉化為優化的任務圖；
• Part 2：MPK 運行時系統，該系統在單個巨型內核內執行任務圖，以實現高吞吐量與低延遲。

編譯器 —— 將 LLM 轉化為細粒度任務圖

LLM 的計算過程通常表示為計算圖，其中每個節點對應一個計算算子（如矩陣乘法、注意力機制）或集合通信原語（如 all-reduce），邊表示算子間的數據依賴關系。現有系統通常為每個算子啟動獨立的 GPU 內核。

然而，這種「單算子單內核」的執行模型難以實現 pipeline 優化，因為依賴關系是在整個內核的粗粒度層面強制執行的，而非實際數據單元層面。

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

下圖 2 展示了 MPK 編譯器將 PyTorch 定義的 LLM 計算圖轉化為優化細粒度任務圖，最大化暴露并行性。右側展示次優方案 —— 其引入不必要的數據依賴與全局屏障，導致跨層流水線優化機會受限。

為了解決此問題，MPK 引入的編譯器可將 LLM 計算圖自動轉化為細粒度任務圖。該任務圖在子內核級別顯式捕獲依賴關系，實現更激進的跨層流水線優化。

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

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

任務圖使 MPK 能夠發掘計算圖中無法實現的 pipeline 優化機會。例如，MPK 可以構建優化任務圖 —— 其中每個 all-reduce 任務僅依賴于生成其輸入的對應 matmul 任務，從而實現分塊執行與計算通信重疊。

除生成優化任務圖外，MPK 還通過 Mirage 內核超優化器自動為每個任務生成高性能 CUDA 實現，確保任務在 GPU 流式多處理器（SM）上高效執行。

Part 2：運行時 —— 在巨型內核中執行任務圖

MPK 包含內置 GPU 運行時系統，可在單個 GPU 巨型內核內完整執行任務圖。這使得系統能在推理過程中無需額外內核啟動的情況下，實現任務執行與調度的細粒度控制。

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

工作 SM 與調度 SM 的數量在內核啟動時固定配置，且總和等于物理 SM 總數，從而徹底避免動態上下文切換開銷。

工作單元

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

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

該機制既保障了工作單元的持續滿載運行，又實現了跨層和跨操作的異步任務執行。

調度單元

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

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

這種分布式調度機制在實現跨 SM 可擴展執行的同時，最小化協同開銷。

事件驅動執行

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

這種設計實現了細粒度的軟件流水線化，并允許計算與通信之間重疊，比如

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

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

下一步計劃

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

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

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

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