天下苦英偉達久矣！PyTorch免CUDA加速推理，Triton時代要來？

機器之心報道

編輯：杜偉、小舟

在做大語言模型（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 架構的核心操作總結如下：

這些操作中的每一個都是通過在 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 內核進行了評估，包括

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 內核配置使用了：

CUDA 內核配置使用了

在典型推理設置下，兩種 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 推理。

原文鏈接：https://pytorch.org/blog/cuda-free-inference-for-llms/