vLLM AMD ROCm アテンションバックエンド最適化
vLLM は AMD ROCm 向けに一連の最適化されたアテンションバックエンドを実装し、単なる移植からアーキテクチャ共同設計へと進化させました。vLLM のオーケストレーション層と AMD の AITER プリミティブを組み合わせることで、AMD Instinct MI300X、MI325X、MI355X GPU 上でマルチヘッドアテンション (MHA) およびマルチヘッド潜在アテンション (MLA) ワークロードに対して 1.2〜4.4 倍のスループット (TPS) 向上を実現しています。
混在推論ワークロードの課題
本番環境の LLM サービングは、prefill、extend、decode トークンを同時に処理する連続バッチングを伴います。各フェーズはそれぞれ異なるパフォーマンスボトルネックを持ちます:
- Prefill: 計算に依存。新しいプロンプトに対して大きなタイルサイズと最大の ALU 利用率が必要です。
- Extend: ハイブリッドワークロード。KV キャッシュが部分的に構築されたリクエストに対して、追加のプロンプト側トークンを処理します。
- Decode: メモリに依存。トークンを一度に 1 つ生成し、連続したメモリアクセスと最小限のキャッシュフェッチが求められます。
あるフェーズに最適化されたカーネルは別のフェーズで性能が劣ることが多いため、vLLM は明示的なルーティングで各リクエストタイプを専用カーネルへ振り分けます。
ROCM_AITER_FA: MHA 用の三経路オーケストレーション
ROCM_AITER_FA は、CDNA アーキテクチャ上でハードウェア利用率を最大化するために、リクエストを 3 つの専門パスに振り分ける高度なオーケストレーション層です。
技術的イノベーション
- Three-Path Routing: リクエストは動的に Prefill（
flash_attn_varlen_funcを使用してマトリクスコアを処理）、Extend（100K 以上のコンテキストに対して LSE マージを伴うチャンクドアテンションを使用）、Decode（AITER の高度に最適化されたメモリ帯域幅カーネルを使用）のいずれかに分類されます。 - Batch Reordering: vLLM Model Runner はリクエストを
[decode:extend:prefill]の順序に再配置します。これにより連続したメモリアクセスが保証され、冗長な KV キャッシュフェッチが排除されます。 - Hardware-Optimized KV Cache Layout: 事前にシャッフルされた KV キャッシュレイアウトは、AMD の CDNA アーキテクチャに合わせたメモリアクセスパターンを実現します。このレイアウトにより、デコードパスはレイアウト変換オーバーヘッドなしで AITER の
pa_fwd_asmカーネルを呼び出せ、デコードスループットが 15〜20% 向上します。 - Chunked Context Processing: 長いシーケンスは固定のイテレーショントークン予算（約 32K トークン）で処理され、数値安定性を保つために Log‑Sum‑Exp (LSE) ベースのマージが使用されます。
MHA バックエンド比較
ROCM_AITER_FA 以外にも、vLLM は以下の MHA バックエンドを提供しています:
- Unified Backends (
TRITON_ATTN,ROCM_AITER_UNIFIED_ATTN): すべてのトークンを単一カーネルパスで処理します。 - Legacy 2-Path (
ROCM_ATTN): Prefill（Triton）と Decode（HIP ページドアテンション）で別々のカーネルを使用します。このバックエンドは Radeon GPU をサポートしますが、未対応の KV ヘッドサイズの場合は遅い Triton デコードカーネルにフォールバックすることがあります。
AITER MLA バックエンド: DeepSeek 向けに最適化
DeepSeek と Kimi で使用されるマルチヘッド潜在アテンション (MLA) は KV キャッシュを 576 次元に圧縮し、パフォーマンスボトルネックを変化させます。vLLM はこの圧縮に対応するため、専用の AITER ベース MLA バックエンドを提供します。
ハイブリッド処理戦略
MLA バックエンドは処理フェーズに応じて以下のように分割されます:
- Prefill/Extend (Non-Absorbed): 圧縮されていない表現上で標準の MHA カーネルを使用してアテンションを計算します。
- Decode (Absorbed): 圧縮された 576 次元潜在空間上で直接動作する専用 MLA カーネルを使用します。
アセンブリから得られる性能向上
ROCM_AITER_MLA と ROCM_AITER_TRITON_MLA の主な性能向上は、mla_decode_fwd アセンブリカーネルによるものです。この手作業でチューニングされたカーネルは HBM3 帯域幅を最大限に活用し、TRITON_MLA ベースラインに比べて Time Per Output Token (TPOT) を 1.2〜1.6 倍高速化します。
パフォーマンスベンチマーク
ベンチマークは 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 デコードカーネルにフォールバックするためです。
MLA 結果 (DeepSeek‑R1)
AITER MLA バックエンドは TRITON_MLA に比べて最大 1.5 倍のスループット向上を提供します:
| ハードウェア | 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 は特に MI355X 上で最高の Time To First Token (TTFT) を達成するため、MLA ワークロードのデフォルトとして推奨されます。
実装とデプロイ
最適化された 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 |