在做大語言模型（LLM）的訓練、微調和推理時，使用英偉達的 GPU 和 CUDA 是常見的做法。在更大的機器學習編程與計算范疇，同樣嚴重依賴 CUDA，使用它加速的機器學習模型可以實現更大的性能提升。

雖然 CUDA 在加速計算領域占據主導地位，并成為英偉達重要的護城河之一。但其他一些工作的出現正在向 CUDA 發起挑戰，比如 OpenAI 推出的 Triton，它在可用性、內存開銷、AI 編譯器堆棧構建等方面具有一定的優勢，并持續得到發展。

近日，PyTorch 官宣要做「無英偉達 CUDA 參與的大模型推理」。在談到為什么要 100% 使用 Triton 進行探索時，PyTorch 表示：「Triton 提供了一條途徑，使大模型 能夠在不同類型的 GPU 上運行，包括英偉達、AMD、英特爾和其他基于 GPU 的加速器。

此外 Triton 還在 Python 中為 GPU 編程提供了更高的抽象層，使得使用 PyTorch 能夠比使用供應商特定的 API 更快地編寫高性能內核。」

在 PyTorch 博客中討論了使用流行的 LLM 模型（例如 Meta 的 Llama3-8B 和 IBM 的 Granite-8B Code）實現 FP16 推理的方法，其中計算是 100% 使用 OpenAI 的 Triton 語言執行的。

對于使用基于 Triton 內核的模型生成單個 token 的時間，PyTorch 能夠實現在英偉達 H100 GPU 上 Llama 和 Granite 的 CUDA 內核主導工作流程的 0.76-0.78 倍性能，以及在英偉達 A100 GPU 上的 0.62-0.82 倍。

圖 1. 在英偉達 H100 和 A100 上，Llama3-8B 和 Granite-8B 的 Triton 和 CUDA 變體的推理吞吐量比較。設置：批大小 = 2，輸入序列長度 = 512，輸出序列長度 = 256

也許告別英偉達的時候真要來了。

Transformer 塊的組成

PyTorch 團隊首先對基于 Transformer 的模型中發生的計算進行細分。下圖顯示了典型 Transformer 塊的「內核（kernel）」。

圖 2

Llama3 架構的核心操作總結如下：

• 均方根歸一化（RMSNorm）
• 矩陣乘法：Fused QKV
• RoPE
• 注意力
• 矩陣乘法：輸出投影
• RMSNorm
• 矩陣乘法：Fused Gate + Up Projection
• 激活函數：SiLU
• 點乘（Element Wise Multiplication）
• 矩陣乘法：Down Projection

這些操作中的每一個都是通過在 GPU 上執行一個（或多個）內核來計算的。雖然每個內核的細節在不同的 Transformer 模型中可能有所不同，但核心操作保持不變。例如，IBM 的 Granite 8B Code 模型在 MLP 層中使用偏置，與 Llama3 不同。此類更改確實需要對內核進行修改。典型的模型是這些 Transformer 塊的堆疊，這些 Transformer 塊通過嵌入層連接在一起。

模型推理

典型的模型架構代碼與 PyTorch 啟動的 python model.py 文件共享。在默認的 PyTorch Eager Execution 模式下，這些內核都是使用 CUDA 執行的。為了實現 100% Triton 進行端到端 Llama3-8B 和 Granite-8B 推理，需要編寫和集成手寫 Triton 內核以及利用 torch.compile（生成 Triton 操作）。首先，PyTorch 用編譯器生成的 Triton 內核替換較小的操作，其次，PyTorch 用手寫的 Triton 內核替換更昂貴和復雜的計算（例如矩陣乘法和閃存注意力）。

Torch.compile 自動為 RMSNorm、RoPE、SiLU 和點乘生成 Triton 內核。使用 Nsight Systems 等工具，可以觀察到這些生成的內核，它們在矩陣乘法和注意力之間表現為微小的深綠色內核。

圖 3. 使用 torch.compile 跟蹤 Llama3-8B，顯示用于矩陣乘法和閃存注意力的 CUDA 內核。

對于上面的跟蹤，PyTorch 團隊注意到，在 Llama3-8B 樣式模型中，占 E2E 延遲 80% 的兩個主要操作是矩陣乘法和注意力內核，并且兩者仍然是 CUDA 內核。因此，為了彌補剩余的差距，PyTorch 團隊用手寫的 Triton 內核替換了 matmul 和注意力內核。

Triton SplitK GEMM 內核

對于線性層中的矩陣乘法，PyTorch 團隊編寫了一個自定義 FP16 Triton GEMM（通用矩陣 - 矩陣乘法）內核，該內核利用了 SplitK 工作分解。

GEMM 內核調優

為了實現最佳性能，PyTorch 團隊使用窮舉搜索方法來調整 SplitK GEMM 內核。Granite-8B 和 Llama3-8B 具有如下形狀的線性層：

圖 4. Granite-8B 和 Llama3-8B 線性層權重矩陣形狀。

每個線性層都有不同的權重矩陣形狀。因此，為了獲得最佳性能，必須針對每個形狀輪廓調整 Triton 內核。在對每個線性層進行調整后，PyTorch 能夠在 Llama3-8B 和 Granite-8B 上實現相對于未調整的 Triton 內核 1.20 倍的 E2E 加速。

Flash Attention 內核

PyTorch 團隊使用不同的配置，對現有 Triton flash attention 內核進行了評估，包括

• AMD Flash
• OpenAI Flash
• Dao AI Lab Flash
• XFormers Flash
• PyTorch FlexAttention

PyTorch 團隊分別在 eager 模式和編譯模式下評估了每個內核的文本生成質量。下圖 5 為不同 Flash Attention 內核的比較。

上圖總結了 PyTorch 觀察到的開箱即用情況，并預計內核 2 到 5 可以在修改后滿足上述標準。不過這也表明，擁有一個可用于基準測試的內核通常只是將它用作端到端生產內核的開始。

PyTorch 團隊選擇在后續測試中使用 AMD flash attention 內核，它通過 torch.compile 進行編譯，并在 eager 和編譯模式下產生清晰的輸出。

為了滿足 torch.compile 與 AMD flash attention 內核的兼容性，PyTorch 團隊必須將它定義為 torch 自定義算子。并且封裝更復雜的 flash attention 內核遵循以下兩個步驟：

一是將函數封裝為一個 PyTorch 自定義算子。

二是向該算子添加一個 FakeTensor 內核，并在給定 flash 輸入張量的形狀（q、k 和 v）時，計算 flash 內核的輸出形狀。

在將 Triton flash 內核定義為一個自定義 op 后，PyTorch 團隊可以成功地對它進行編譯以實現端到端運行。

圖 6：在交換 Triton matmul 和 Triton flash attention 內核后，使用 torch.compile 的 Llama3-8B 軌跡。

從圖中可以看到，在集成 SplitK 矩陣乘法內核后，torch op 封裝 flash attention 內核，然后運行 torch.compile，即可實現使用 100% Triton 計算內核的前向傳遞。

端到端基準測試

PyTorch 團隊分別對運行 Granite-8B 和 Llama3-8B 模型的英偉達 H100 和 A100（單 GPU）進行了端到端測試，使用了兩種不同的配置來執行基準測試。

其中 Triton 內核配置使用了：

• Triton SplitK GEMM
• AMD Triton Flash Attention

CUDA 內核配置使用了

• cuBLAS GEMM
• cuDNN Flash Attention - Scaled Dot-Product Attention (SDPA)

在典型推理設置下，兩種 eager 和 torch 編譯模式的吞吐量和 inter-token 延遲如下圖所示。

圖 7：H100 和 A100 上 Granite-8B 和 Llama3-8B 單 token 生成延遲（批大小 = 2，輸入序列長度 = 512，輸出序列長度 = 256）。

總的來說，在 H100 上，Triton 模型最高可以達到 CUDA 模型性能的 78%；在 A100 上可以達到 82％。這些性能差距是由 matmul 和 flash attention 的內核延遲造成的。

微基準測試

下圖 8 為 Triton 和 CUDA 內核延遲比較（英偉達 H100 上運行 Llama3-8B）。輸入為一個任意 prompt（批大小 = 1，prompt 序列長度 = 44），以解碼延遲時間。

最后結果顯示，Triton matmul 內核比 CUDA 慢了 1.2 至 1.4 倍，而 AMD Triton Flash Attention 比 CUDA SDPA 慢了 1.6 倍。

以上結果凸顯了需要進一步提升 GEMM 和 Flash Attention 等核心原語內核的性能。最近的一些工作（如 FlashAttention-3、FlexAttention） 已經提出了更好地利用底層硬件和 Triton 的方法，PyTorch 希望在它們的基礎上實現更大加速。為了闡明這一點，PyTorch 團隊將 FlexAttention 與 SDPA、AMD’s Triton Flash 內核進行了比較。

PyTorch 團隊 正努力驗證 FlexAttention 的端到端性能。目前，FlexAttention 的初始微基準測試結果表明，在查詢向量較小的情況下，有望實現更長的上下文以及解碼問題形狀。

圖 9：英偉達 H100 SXM5 80GB 上 FlexAttention 內核基準測試（批大小 = 1，最大頭數 = 32，頭維數 = 128）。

未來工作

未來，PyTorch 團隊計劃探索進一步優化 matmuls 的方法，以便更好地利用硬件，并為基于 Triton 的方法實現更大的加速。

對于 flash attention，PyTorch 團隊計劃探索 FlexAttention 和 FlashAttention-3 等內核中使用到的技術，以幫助進一步縮小 Triton 與 CUDA 之間的差距。同時還將探索端到端 FP8 LLM 推理。