近日，來自 CMU 的 Catalyst Group 團隊發布了一款 PyTorch 算子編譯器 Mirage，用戶無需編寫任何 CUDA 和 Triton 代碼就可以自動生成 GPU 內核，并取得更佳的性能。

隨著 GPU 加速器的不斷發展以及以大語言模型為代表的生成式 AI 應用的不斷推廣，通過開發高性能 GPU 內核來優化 PyTorch 程序的計算效率變得越來越重要。目前，這項任務主要由專門的 GPU 專家來完成。在 NVIDIA CUDA 或 AMD ROCm 中編寫高性能 GPU 內核需要高水平的 GPU 專業知識和大量的工程開發經驗。目前的機器學習編譯器（如 TVM、Triton 和 Mojo）提供了一些高級編程接口，以簡化 GPU 編程，使用戶可以使用 Python 而非 CUDA 或 ROCm 來實現 GPU 內核。

然而，這些語言仍然依賴用戶自行設計 GPU 優化技術以達到更高的性能。例如，在 Triton 中實現一個 FlashAttention 內核大約需要 700 行 Python 代碼（在 CUDA 中需要大約 7,000 行 C++ 代碼）。在這些程序中，用戶需要手動劃分線程塊之間的工作負載，組織每個線程塊內的計算，并管理它們之間的同步與通信。

用 Triton 實現的 FlashAttention 算子

能否在不使用 CUDA/Triton 編程的情況下就獲得高效的 GPU 內核呢？基于這一動機，來自卡內基梅隆大學的 Catalyst Group 團隊發布了 Mirage 項目，基于 SuperOptimization 技術（https://arxiv.org/abs/2405.05751），為 PyTorch 自動生成高效 GPU 內核算子。例如，對于一個 FlashAttention 算子，用戶只需編寫幾行 Python 代碼來描述注意力（Attention）的計算過程而無需了解 GPU 編程細節，如下所示：

# Use Mirage to generate GPU kernels for attention
import mirage as mi
graph = mi.new_kernel_graph ()
Q = graph.new_input (dims=(64, 1, 128), dtype=mi.float16)
K = graph.new_input (dims=(64, 128, 4096), dtype=mi.float16)
V = graph.new_input (dims=(64, 4096, 128), dtype=mi.float16)
A = graph.matmul (Q, K)
S = graph.softmax (A)
O = graph.matmul (S, V)
optimized_graph = graph.superoptimize ()

Mirage 會自動搜索可能的 Attention GPU 內核實現，搜索空間不僅包括現有的手動設計的注意力內核（如 FlashAttention 和 FlashDecoding），還包括在某些場景中比目前的手寫版本快多達 3.5 倍的其他實現。Mirage 生成的 GPU 內核可以直接在 PyTorch 張量上操作，并可以在 PyTorch 程序中直接調用。

import torch
input_tensors = [
    torch.randn (64, 1, 128, dtype=torch.float16, device='cuda:0'),
    torch.randn (64, 128, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn (64, 4096, 128, dtype=torch.float16, device='cuda:0')
]
# Launch the Mirage-generated kernel to perform attention
output = optimized_graph (input_tensors)

Why Mirage?

與使用 CUDA/Triton 編程相比，Mirage 提供了一種新的編程范式，包含三個主要優勢：

更高的生產力：隨著 GPU 架構日新月異，現代 GPU 編程需要持續學習大量的專業知識。Mirage 的目標是提高機器學習系統工程師的生產力 —— 他們只需在 PyTorch 層面描述所需的計算，Mirage 便會自動生成適用于各種 GPU 架構的高性能實現。因此，程序員不再需要手動編寫 CUDA/Triton 或特定架構的低級代碼。

更好的性能：目前手動設計的 GPU 內核由于無法充分探索和利用各種 GPU 優化技術，往往只能達到次優性能。Mirage 可以自動搜索與輸入的 PyTorch 程序功能等價的潛在 GPU 實現，探索并最終發現性能最優的內核。在多個 LLM/GenAI 基準測試中的測試結果顯示，Mirage 生成的內核通常比 SOTA 的專家人工編寫或編譯器生成的替代方案快 1.2 至 2.5 倍。

更強的正確性：手動實現的 CUDA/Triton GPU 內核容易出錯，而且 GPU 內核中的錯誤難以調試和定位，而 Mirage 則利用形式化驗證（Formal Verification）技術自動驗證生成的 GPU 內核的正確性。

LLaMA-3-8B 和 Chameleon-7B 端到端推理延遲對比（NVIDIA A100，batch size=1，context length=4K），相比于 CUDA/Triton 的實現，Mirage 可以實現 15-20% 的加速

GPU 架構與 Mirage 中的

GPU 計算的內核函數以單程序多數據（SPMD）方式在多個流處理器（SM）上同時運行。GPU 內核（Kernel）借助由線程塊（Thread Block）組成的網格結構來組織其計算，每個線程塊在單個 SM 上運行。每個塊進一步包含多個線程（Thread），以對單獨的數據元素進行計算。GPU 還擁有復雜的內存層次結構，以支持這種復雜的處理結構。每個線程都有自己的寄存器文件（Register File），以便快速訪問數據。線程塊內的所有線程可以訪問一個公共的共享內存（Shared Memory），這有助于它們之間高效的數據交換和集體操作。最后，內核內的所有線程可以訪問分配給整個 GPU 的大型設備內存（Device Memory）。

GPU 計算架構和編程抽象示意圖

Mirage 使用來描述 GPU 內核，包含多個層次，代表內核、線程塊和線程級別的計算。大體上，Kernel Graph、Thread Block Graph 和 Thread Graph 分別代表整個 GPU、一個流處理器（SM）和一個 CUDA/tensor 核心上的計算。

對細節感興趣的讀者可以參考：

• https://mirage-project.readthedocs.io/en/latest/mugraph.html

Mirage 工作流示意圖

上圖展示了 Mirage 的工作流程：對于輸入的 PyTorch 程序，Mirage 的生成器自動搜索與輸入程序功能等價的其他，搜索空間涵蓋了內核、線程塊和線程級別的各種 GPU 優化。所有生成的都被發送到等價性驗證器，該驗證器自動檢查每個是否與所需程序等價。最后，轉譯器將所有經過驗證的轉譯為 CUDA 內核。最后，Mirage 會從中返回性能最佳的 CUDA 內核。

• 項目成員：Mengdi Wu (CMU), Xinhao Cheng (CMU), Shengyu Liu (PKU), Chuan Shi (PKU), Jianan Ji (CMU), Oded Padon (VMWare), Xupeng Miao (Purdue), Zhihao Jia (CMU)
• 項目地址：https://github.com/mirage-project/mirage

為什么 Mirage 生成的內核更高效？

在多個 LLM/GenAI 基準測試中的測試結果顯示，Mirage 生成的內核通常比現有的手寫或編譯器生成的內核快 1.2 至 2.5 倍。接下來，本文以 LLM 中的 Transformer 架構為例，展示現有系統中缺失的幾項 GPU 程序優化技術：

Transformer 架構示意圖

Case 1: Normalization + Linear

歸一化（Normalization）操作，如 LayerNorm、RMSNorm、GroupNorm 和 BatchNorm，廣泛應用于當今的機器學習模型。當前的機器學習編譯器通常在獨立的內核中啟動歸一化層，因為歸一化涉及到歸約和廣播，難以與其他計算融合。然而，Mirage 發現，大多數歸一化層可以通過進行適當的代數變換，與后續的線性層（如 MatMul）融合。

Normalization + Linear 現有內核 v.s. Mirage 發現的內核

Mirage 發現的自定義內核利用了 RMSNorm 中的除法和 MatMul 中的乘法的可交換性，將除法移到 MatMul 之后。這一變換保持了功能等價性，同時避免了中間張量 Y 的實例化。該內核的性能比單獨運行這兩個操作快 1.5 到 1.7 倍。

Normalization + Linear 內核性能對比

Case 2: LoRA + Linear

LoRA 廣泛用于預訓練模型的微調場景，以適配到特定領域和任務。這些 LoRA 適配器通常會被插入到模型的線性層中，引入額外的矩陣乘法?，F有系統通常為原始矩陣乘法和 LoRA 中的兩個矩陣乘法啟動獨立的內核，從而導致較高的內核啟動開銷。

LoRA+Linear 現有內核 v.s. Mirage 發現的內核

如上圖所示，Mirage 發現了一個將三個矩陣乘法和隨后的加法融合為單個內核的內核。這是通過將計算重組為兩個線程塊級別的矩陣乘法實現的，利用了以下代數變換：W×X+B×A×X=(W|B)×(X|(A×X))，其中的兩個拼接操作不涉及任何計算，而是通過在 GPU 共享內存中更新張量偏移量來完成。Mirage 發現的內核比現有系統中使用的內核快 1.6 倍。

LoRA+Linear 內核性能對比

Case 3: Gated MLP

Gated MLP 層目前在許多 LLM 中使用（如 LLAMA-2、LLAMA-3 及其變體），它的輸入張量 X 與兩個權重矩陣相乘，輸出結果被組合以產生最終結果。Mirage 發現了一個內核，該內核執行兩個矩陣乘法、SiLU 激活以及隨后的逐元素乘法，從而減少了內核啟動開銷和對設備內存的訪問。

Gated MLP 現有內核 v.s. Mirage 發現的內核

Gated MLP 內核性能對比

Case 4: Attention Variants

如今的大多數 LLM 基于注意力及其變體，雖然現有系統通常提供高度優化的注意力實現，如 FlashAttention、FlashInfer 和 FlexAttention，但支持注意力變體通常需要新的自定義內核。下面用兩個例子來展示 Mirage 如何為非常規注意力計算發現自定義 GPU 內核。

Case 4.1: Attention with Query-Key Normalization

許多最近的 LLM 架構（包括 Chameleon、ViT-22B 等）在 LLaMA 架構中引入了 QK-Norm 來緩解訓練過程中的數值發散問題。QK-Norm 在注意力之前對 Query 和 Key 向量應用 LayerNorm 層?，F有注意力實現中并不支持這些額外的歸一化層，并且它們還需要作為獨立內核啟動。

QK-Norm 注意力現有內核 v.s. Mirage 發現的內核

對于在注意力之前和 / 或之后引入計算的注意力變體，這些計算可以與注意力融合以提高 GPU 性能，而這需要自定義內核。對于帶有 QK-Norm 的注意力，Mirage 發現了上述內核來融合計算，從而避免在 GPU 設備內存中實例化中間結果。這個自定義內核還對注意力進行了現有的 GPU 優化，實現了 1.7 至 2.5 倍的性能提升。

QK-Norm 注意力內核性能對比

Case 4.2: Multi-Head Latent Attention

MLA 的現有內核 v.s. Mirage 發現的內核

另一個常用的注意力變體是 MLA（Multi-Head Latent Attention），它將注意力的 KV Cache 壓縮為一個向量，以減少存儲 KV Cache 的內存開銷。這一變化還在注意力之前引入了兩個線性層，如下圖所示。與 QK-Norm 類似，現有注意力實現中并不支持這些額外的歸一化層，同樣需要作為獨立內核啟動，而 Mirage 可以將線性層和注意力融合為一個單獨的自定義內核。

長期愿景

Mirage 項目的長期目標是希望能夠讓未來的 AI 開發者無需學習 CUDA 或者 Triton 等復雜的 GPU 編程語言，只需指定所需的數學操作，就能在 GPU 上輕松實現 AI 模型。通過利用 Mirage 的 SuperOptimization 技術，各種計算任務可以自動轉換為高度優化的 GPU 實現。隨著 LLM 和其他生成式 AI 應用的飛速發展，在各種實際部署場景都需要高效的 GPU 支持，降低 GPU 編程門檻并提高程序效率也愈發重要。