vLLM Triton アテンションバックエンドの深掘り

vLLM Triton アテンションバックエンドの深掘り

vLLM は、単一のソースコード実装で NVIDIA、AMD、Intel の GPU 全体で最先端の性能を実現する、ポータブルな Triton ベースのアテンションバックエンドを実装しました。

vLLM Triton アテンションバックエンドの深掘り

vLLM は、さまざまな GPU ハードウェア上で高性能かつポータブルなアテンションカーネルを提供するために、Triton ベースのアテンションバックエンドを導入しました。このバックエンドにより、vLLM は NVIDIA、AMD、Intel の GPU で効率的に動作するアテンション操作用の単一コードベースを維持でき、ハードウェア固有のカーネルの保守負担を削減しつつ、専門的な実装と同等の性能を実現します。

Triton による性能ポータビリティ

すべての GPU アーキテクチャ（例: NVIDIA Hopper、Blackwell、AMD MI300）向けに数百もの専門カーネルを維持することは現実的ではありません。vLLM は、GPU カーネルを Python で記述し、複数プラットフォーム向けに効率的なコードにコンパイルできるドメイン固有言語である Triton を活用しています。

Triton はタイル化プログラミングモデルを採用しており、開発者は計算の論理タイルを定義します。Triton コンパイラと自動チューナーがこれらのタイルをハードウェア固有の実行レイアウトにマッピングします。この抽象化により、バックエンドはハードウェアに依存しないものとなりながら、低レベルの最適化も可能になります。

Triton アテンションバックエンドのアーキテクチャ

線形層などの他コンポーネントからアテンション実装を分離するために、vLLM はアテンションバックエンド抽象層を使用しています。Triton アテンションバックエンドは vLLM にネイティブで、PyTorch と Triton のみに依存し、完全に Triton で実装されています。

使用シナリオ

  • Default for AMD GPUs: AMD GPU（ROCm 上で動作）向けの主要バックエンドです。
  • Intel XPU: float32 演算に使用されます。FlashAttention はこのプラットフォームで fp32 をサポートしていないためです。
  • Specific Model Features: ALiBi sqrt（StepFun オーディオモデルで使用）や sink tokens、GPT-OSS の挙動をサポートし、特に Hopper 以前の NVIDIA GPU（例: A100）で有効です。
  • Specialized Requirements: ヘッドサイズが小さいモデル、エンコーダ/デコーダ アテンション、マルチモーダル プレフィックス アテンションを扱います。
  • Fallback Mechanism: FlashAttention、FlashInfer、またはその他の依存関係が利用できない場合の一般的なフォールバックとして機能します。

ページドアテンションの技術的実装

ページドアテンションは KV キャッシュをページングすることでメモリを最適化します。カーネルはクエリトークンを処理し、クエリと KV ヘッドを反復し、ページド KV キャッシュを走査してアテンションスコアを計算します。

Q ブロック最適化

tl.dot（Triton の行列乗算）の利用率を最大化するために、バックエンドは「Q ブロック」を使用します。KV キャッシュのページサイズが KV 側のタイルサイズを制約するため、カーネルは複数のクエリトークンとヘッドを単一の作業項目（Q ブロック）にまとめ、並列性とキャッシュ再利用を向上させます。特に Group Query Attention（GQA）で有利です。

パラレルタイル化ソフトマックス（3D カーネル）

Q ブロックはプレフィルワークロードに有利ですが、デコードワークロードは単一のクエリトークンのみを処理します。これを最適化するために、vLLM はパラレルタイル化ソフトマックスを用いた「3D カーネル」を実装しています。この手法は KV キャッシュの走査を複数のカーネルインスタンスに分割し、各インスタンスが部分結果を計算します。その後、2 番目のカーネル起動で結果を集約し、最終出力を生成します。

永続カーネルと CUDA グラフ

CUDA グラフは固定実行グラフを記録することで起動オーバーヘッドを削減します。しかし、標準のアテンションカーネルはバッチサイズやシーケンス長に基づく可変の起動グリッドを使用することが多く、CUDA グラフで再生すると非効率（作業の無駄やストリーミングマルチプロセッサの利用率低下）につながります。

この課題を解決するために、vLLM は 永続カーネル を開発しました。可変グリッドの代わりに、利用可能な計算リソースと同数のカーネルインスタンスを固定で起動します。これらのインスタンスは GPU メモリからメタデータを動的に読み取り、作業負荷を決定するため、起動グリッドが一定となり、CUDA グラフの効率的な再利用が可能になります。

ベンチマークと結果

2025 年後半のベンチマークは、Triton バックエンドが専門的な代替手段に比べてはるかに少ないコード複雑性で高効率を実現していることを示しています。Triton のページドアテンション実装は約 800 行のコードで構成されており、FlashAttention-3 の約 70,000 行と比較されます。

Performance Metrics (Llama 3.1 8B、入力トークン 500、バッチサイズ 1):

  • NVIDIA H100: 長いデコードリクエストに対して FlashAttention-3 の 100.7% の性能を達成しました。
  • AMD MI300: 以前の実装に比べて約 5.8 倍の速度向上を達成しました。

今後の方向性: Helion

実験的な取り組みとして、PyTorch チームが開発した新しいドメイン固有言語 Helion が始まっています。Helion は高レベルの Triton またはタイル化された PyTorch と表現されます。Helion で実装された簡易版ページドアテンションカーネルは有望な初期結果を示しており、現在 vLLM リポジトリのドラフトプルリクエストとして利用可能です。

Sources