vLLM AMD ROCm 注意力後端最佳化

vLLM 為 AMD ROCm 實作了一系列最佳化的注意力後端，已超越單純移植，進入架構共同設計階段。透過將 vLLM 的協調層與 AMD 的 AITER 原語結合，系統在 AMD Instinct MI300X、MI325X 與 MI355X GPU 上，對 Multi-Head Attention (MHA) 與 Multi-Head Latent Attention (MLA) 工作負載達成 1.2‑4.4 倍的吞吐量提升（TPS）。

混合推論工作負載的挑戰

生產環境的 LLM 服務需要持續批次處理，同時處理 prefill、extend 與 decode 令牌。每個階段都有不同的效能瓶頸：

  • 預填：計算受限；需要大型 tile 大小與最大 ALU 使用率以處理新提示。
  • 延伸：混合工作負載；為具有部分 KV 快取的請求處理額外的提示端令牌。
  • 解碼：記憶體受限；一次產生一個令牌，需要合併的記憶體存取與最小化快取抓取。

由於針對某一階段調校的核心在其他階段常表現不佳，vLLM 採用明確的路由機制，將每種請求類型導向專門的核心。

ROCM_AITER_FA：MHA 的三路協調

ROCM_AITER_FA 是一個複雜的協調層，將請求透過三條專門路徑導向，以在 CDNA 架構上最大化硬體利用率。

技術創新

  • 三路路由：請求會動態分類為 預填（使用 flash_attn_varlen_func 於矩陣核心）、延伸（使用分塊注意力與 LSE 合併處理 10 萬以上的上下文）、以及 解碼（使用 AITER 高度最佳化的記憶體頻寬核心）。
  • 批次重新排序：vLLM Model Runner 會將請求重新排序為 [decode:extend:prefill] 序列。此舉確保連續的記憶體存取，並消除多餘的 KV 快取抓取。
  • 硬體最佳化的 KV 快取布局：預先洗牌的 KV 快取布局使記憶體存取模式與 AMD 的 CDNA 架構對齊。這使得解碼路徑能以零布局轉換開銷呼叫 AITER 的 pa_fwd_asm 核心，提升解碼吞吐量 15‑20%。
  • 分塊上下文處理：長序列以固定的每次迭代令牌預算（約 32K 令牌）處理，使用基於 Log-Sum-Exp (LSE) 的合併以維持數值穩定性。

MHA 後端比較

除了 ROCM_AITER_FA，vLLM 還提供其他 MHA 後端：

  • 統一後端（TRITON_ATTN、ROCM_AITER_UNIFIED_ATTN）：透過單一核心路徑處理所有令牌。
  • 傳統雙路（ROCM_ATTN）：為 prefill（Triton）與 decode（HIP 分頁注意力）使用不同核心。此後端支援 Radeon GPU，但對於不支援的 KV 頭尺寸可能退回較慢的 Triton decode 核心。

AITER MLA 後端：為 DeepSeek 最佳化

Multi-Head Latent Attention（MLA）在 DeepSeek 與 Kimi 中使用，將 KV 快取壓縮至 576 維，改變了效能瓶頸。vLLM 提供專門的基於 AITER 的 MLA 後端，以處理此壓縮。

混合處理策略

MLA 後端根據處理階段採用分割策略：

  • 預填/延伸（未吸收）：注意力使用標準 MHA 核心在未壓縮的表示上計算。
  • 解碼（已吸收）：專門的 MLA 核心直接在壓縮的 576 維潛在空間上運作。

來自組合語的效能提升

ROCM_AITER_MLA 與 ROCM_AITER_TRITON_MLA 的主要效能提升來自 mla_decode_fwd 組合語核心。此手工調校的核心最大化 HBM3 頻寬，較 TRITON_MLA 基線快 1.2‑1.6 倍的每輸出令牌時間（TPOT）。

效能基準測試

基準測試使用 ROCm 7.0.0 在 Qwen3-235B（MHA）與 DeepSeek-R1（MLA）模型上執行。

MHA 結果（Qwen3-235B）

ROCM_AITER_FA 在所有測試硬體上於輸出吞吐量（TPS）方面顯著優於傳統後端：

硬體 ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN TRITON_ATTN ROCM_ATTN
MI300X (64 req) 1.00x 1.05x 1.30x 3.82x
MI325X (64 req) 1.00x 1.02x 1.19x 4.36x
MI355X (64 req) 1.00x 0.95x 1.08x 3.61x

ROCM_AITER_FA 在此模型上比 ROCM_ATTN 提升 2.7‑4.4 倍的吞吐量，部分原因是 ROCM_ATTN 會因不支援的頭尺寸退回至 Triton decode 核心。

MLA 結果（DeepSeek-R1）

AITER MLA 後端提供最高 1.5 倍於 TRITON_MLA 的吞吐量提升：

硬體 ROCM_AITER_MLA ROCM_AITER_TRITON_MLA TRITON_MLA
MI300X (64 req) 1.00x 0.98x 1.33x
MI325X (64 req) 1.00x 0.98x 1.41x
MI355X (64 req) 1.00x 1.03x 1.52x

建議將 ROCM_AITER_MLA 設為 MLA 工作負載的預設，尤其在 MI355X 上可取得最佳的首令牌時間（TTFT）。

實作與部署

若要啟用最佳化的 AITER 後端，使用者應設定以下環境變數：

export VLLM_ROCM_USE_AITER=1

設定後，vLLM 會自動為 MHA 模型（例如 Llama、Qwen、Mistral）選擇 ROCM_AITER_FA，為 MLA 模型（例如 DeepSeek、Kimi）選擇 ROCM_AITER_MLA。

硬體支援矩陣

GPU 記憶體 架構
MI300X 192GB HBM3 gfx942
MI325X 256GB HBM3e gfx942
MI355X 288GB HBM3e gfx950

Sources