vLLM Triton 注意力後端深度探討

vLLM 推出了基於 Triton 的注意力後端，以在多樣的 GPU 硬體上提供高效能、可移植的注意力 kernel。此後端使 vLLM 能夠維護單一的注意力運算程式碼庫，能在 NVIDIA、AMD 與 Intel GPU 上高效執行，減少硬體特定 kernel 的維護負擔，同時匹配專門實作的效能。

透過 Triton 的效能可移植性

為每種 GPU 架構（例如 NVIDIA Hopper、Blackwell 與 AMD MI300）維護數百個專門的 kernel 並不實際。vLLM 使用 Triton，這是一種領域特定語言，允許以 Python 撰寫 GPU kernel，並編譯成多平台的高效能程式碼。

Triton 採用平鋪（tiled）程式模型，開發者定義計算的邏輯平鋪。Triton 編譯器與自動調校器會將這些平鋪映射到硬體特定的執行佈局。此抽象層使後端能保持硬體無關性，同時仍可進行低階優化。

Triton 注意力後端架構

為了將注意力實作與線性層等其他元件分離，vLLM 使用注意力後端抽象層。Triton 注意力後端是 vLLM 原生的，只依賴 PyTorch 與 Triton，且全部以 Triton 實作。

使用情境

AMD GPU 的預設後端： 它是運行於 ROCm 上的 AMD GPU 的主要後端。
Intel XPU： 用於 float32 操作，因為 FlashAttention 在此平台不支援 fp32。
特定模型功能： 支援 ALiBi sqrt（用於 StepFun 音訊模型）、sink tokens 與 GPT-OSS 行為，特別是在 Hopper 前的 NVIDIA GPU（例如 A100）上。
特殊需求： 處理具有小頭尺寸、編碼器/解碼器注意力以及多模態前綴注意力的模型。
備援機制： 若 FlashAttention、FlashInfer 或其他相依項目不可用時，作為一般備援。

分頁注意力的技術實作

分頁注意力透過對 KV 快取分頁來最佳化記憶體。kernel 會處理查詢 token，遍歷查詢與 KV 的頭部，並遍歷分頁的 KV 快取以計算注意力分數。

Q 區塊最佳化

為了最大化 tl.dot（Triton 的矩陣乘法）的使用率，後端使用「Q 區塊」。由於 KV 快取的頁面大小限制了 KV 端的平鋪大小，kernel 會將多個查詢 token 與頭部合併為單一工作項目（即 Q 區塊），以提升平行度與快取重用，尤其在 Group Query Attention（GQA）中。

平行平鋪 Softmax（3D Kernel）

雖然 Q 區塊對預填（prefill）工作負載有益，但解碼（decode）工作負載僅處理單一查詢 token。為了最佳化，vLLM 使用平行平鋪 softmax 實作「3D kernel」。此方法將 KV 快取的遍歷分割到多個 kernel 實例。每個實例計算部分結果，然後透過第二次 kernel 呼叫進行彙總，以產生最終輸出。

永續 Kernel 與 CUDA Graph

CUDA graph 透過記錄固定的執行圖來降低啟動開銷。然而，標準的注意力 kernel 常根據 batch 大小與序列長度使用可變的啟動格子，導致在透過 CUDA graph 重放時產生低效率（例如工作浪費或串流多處理器未充分利用）。

為了解決此問題，vLLM 開發了 永續 kernel。不使用可變格子，而是啟動固定數量的 kernel 實例——等同於可用的計算資源。這些實例會動態從 GPU 記憶體讀取中繼資料以決定其工作負載，確保固定的啟動格子並有效重複使用 CUDA graph。

基準測試與結果

來自 2025 年底的基準測試顯示，Triton 後端在效能上表現優異，且相較於專門的替代方案，程式碼複雜度大幅降低。Triton 分頁注意力的實作約有 800 行程式碼，而 FlashAttention-3 約有 70,000 行。

效能指標（Llama 3.1 8B，500 個輸入 token，batch size 為 1）：

NVIDIA H100： 在長解碼請求中達到 FlashAttention-3 效能的 100.7%。
AMD MI300： 相較於先前的實作，提升約 5.8 倍。

未來方向：Helion

實驗性工作已開始使用 Helion，這是 PyTorch 團隊推出的新領域特定語言，被描述為更高階的 Triton 或平鋪版 PyTorch。以 Helion 實作的簡化分頁注意力 kernel 已展現出有前景的早期結果，且目前以草稿 pull request 形式於 vLLM 倉庫中提供。