在 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。
