高速アテンションカーネルと演算子融合がaiモデルを加速

トランスフォーマーの層は、線形射影、アテンション、正規化、活性化関数、残差加算といった一般的な演算に依存している。素直な実装では、各演算ごとに別個の計算カーネル（「カーネル」）を起動し、中間テンソルをアクセラレータのメモリに書き込む。計算は正しく行われるが、実行時間の多くが値の移動やカーネルの起動に費やされ、有用な計算に割かれる時間が少なくなる可能性がある。

演算子融合

多くのデコーダー専用モデルで使用されるSwiGLUブロックはこの問題をよく表している。素直な実装では、射影のための行列乗算、SiLU関数のためのカーネル、要素ごとの乗算のための別のカーネル、そして最後の行列乗算をそれぞれ起動する。各境界で、中間テンソルが高帯域幅メモリ（HBM）に書き込まれ、その後再読み込みされる可能性がある。融合は、これらのステップを単一のカーネルに組み合わせることを意味する。これは、メモリ移動と起動のオーバーヘッドが大きい場合に時間を節約するが、大きな行列乗算が既に実行時間の大部分を占めている場合には効果は限定的である。カーネルの選択は経験則に基づく：ワークロード全体をプロファイリングする必要がある。

トランスフォーマーでは他にも、残差接続の加算に続くRMSNormやLayerNormの正規化、または脱量子化と行列乗算の組み合わせなど、融合が一般的である。ただし、融合には限界もある：データ配置が互換性のない演算や、非常に大きな中間状態はうまく融合できない。カーネルが大きすぎると、レジスタへの負荷が増加し、占有率が低下し、複数の適切に調整されたカーネルよりも遅くなる可能性がある。

FlashAttention型アテンションとページングされたアテンション

従来のアテンションは、シーケンス長Tに対してヘッドごとにサイズT²のスコア行列を計算する。GPU上では、これを生成し、グローバルメモリに書き込み、その後再読み込みするのに多くの帯域幅を消費し、長いプリフィル（「prefill」）時にボトルネックとなる可能性がある。FlashAttentionは、クエリ、キー、値をブロック単位で処理する。それぞれに対して、スコアの一部を計算し、softmaxの統計情報を更新し、対応する出力をオンチップメモリ内で累積し、完全な行列を保存することはない。PyTorchはこの最適化をscaled_dot_product_attention()関数を通じて公開している。チュートリアルでは、バックエンドをハードコーディングすると、入力やハードウェアをサポートするバックエンドがない場合にコードが失敗する可能性があると指摘している。

計算済みのキーと値を保持するKVキャッシュは、最終的な長さが予測不能な状態で、トークンごとに成長する。したがって、リクエストごとに連続した領域を割り当てると、予約された未使用のスペースが生じる。ページングされたキャッシュは、ストレージを固定サイズのブロックに分割し、論理位置と物理ブロックを関連付けるブロックテーブルに依存する。アテンションカーネルは、このテーブルを通じて各位置を変換しなければならない。ページングは論理要素の数を減らすものではないが、フラグメンテーションと予約容量を制限する。

提供された例では、vLLMはモデルQwen/Qwen2.5-1.5B-Instructを16トークンのブロック（block_size=16）とバックエンド「TRITON_ATTN」を使用している。チュートリアルによれば、vLLMは初期化時にバックエンドとモデルおよびハードウェアの互換性をチェックし、非互換性がある場合はエラーを報告する。

torch.compileによる融合とCUDA Graphs

関数torch.compile()は、TorchDynamoで一連の演算を捕捉し、その後デフォルトのバックエンドであるTorchInductorでコンパイルする。TorchInductorは、標準モードでは別々に起動される互換性のある演算に対して、単一のTritonカーネルを生成できる可能性がある。データ依存のフロー制御やテンソルからPythonスカラーへの変換などによるグラフの中断は、コンパイラが最適化できる範囲を制限する。

「reduce-overhead」モードでは、捕捉されたグラフが適格な場合にTorchInductorに対してCUDA Graphsの使用を要求する。CUDA GraphはGPU演算のシーケンスを記録し、CPUへの単一のサブミッションで再生する。追加の融合を作成するものではなく、TorchInductorによる融合を補完する。安定したメモリアドレスと互換性のある形状を要求し、より多くのメモリを消費する可能性がある。チュートリアルでは、データ型に適した許容誤差で結果を検証し、実行をプロファイリングし、エンドツーエンドのパフォーマンスを測定することを推奨している。