FlashAttention 4：LLM向けのより高速でメモリ効率の良いアテンション

はじめに

スケールド・ドットプロダクト注意（SDPA）は、大規模言語モデル（LLM）の推論時間とエネルギー予算を支配しています：ほぼすべての演算は単一のプリミティブ内で行われます。注意はクエリ（Q）を取り、それをキー（K）と掛け合わせます。その結果をソフトマックス操作で正規化し、さらに値（V）と掛け合わせて出力を生成します。これはメモリ制約に依存しています：高帯域幅メモリからQ/K/Vを何度も読み込み、中間タイルを書き出す操作を繰り返します。

FlashAttentionシリーズは、GPUのストリーミングマルチプロセッサ（SM）を活用して、より多くのデータを長時間チップ上に保持することで、このボトルネックを軽減します。2022年に初めてオープンソース化されたFlashAttentionは、注意機構の処理を再設計し、ソフトマックスをその場で計算するようにしました。任意の時点でチップ上に格納されるQ × Kの行はごくわずかです。

この記事は、実務者向けにFlashAttention 4（FA4）を詳しく解説するものです。以前のバージョンからの変更点をまとめ、導入やベンチマークに関する指針を提供します。私たちの目標は、LLMインフラエンジニア、カーネルエンジニア、MLプラットフォームチームに、FA4を自分たちのスタックでいつ – そして使用すべきかどうか – 判断するために十分な情報を提供することです。

重要ポイント

• A4は、SM100 GPUへの展開を目的としたBlackwell初のアテンションカーネルとして設計されています。従来のFlashAttention世代と比べて、重なりとオンチップ再利用を向上させる、ワープ専用の5段階パイプラインを特徴としています。
• ソフトマックスの効率は主に2つの方法で改善されます：1つは、SFUへの競合を減らすためにソフトウェアの exp2() を CUDA コア上で直接実行すること、もう1つは、安定性を保ちながら不要な再スケーリング計算を避けるための適応オンライン再スケーリングです。
• まだいくつかの機能が欠けています：FA4はフォワード優先であり、現在はバックワードパス、可変長、GQA/MQAの実装が欠けているため、トレーニングや一部のモデルアーキテクチャでの使用が制限されています。
• FA4を実用的に採用する：まずBlackwell推論に使用し、ベンチマークや正確性チェック、フォールバックの後に制御して導入し、広範な本番環境のニーズにはFA3（Hopper）/FA2（Ampere/Ada）を安定したデフォルトとして維持する。

なぜアテンションカーネルがまだLLMのコストを支配しているのか

GPUが高いFLOPSを提供しているにもかかわらず、トランスフォーマーモデルは注意機構でQ、K、Vをグローバルメモリから何度も読み込む必要があるため、しばしばメモリ帯域幅に制約されます。各クエリトークンの注意出力を計算するには、それぞれが最大で数千のキー トークンとドット積を取り、この結果を正規化および再スケーリングし、最後に値を取得する必要があります。単純な実装では、Q×K行列全体を計算してメモリに格納します。これによりオンチップメモリの容量をすぐに超えてしまい、オフチップ読み出しによるパフォーマンス低下を招きます。

FlashAttentionは、アテンション行列をタイル状に分割することでこの問題を緩和します。Q×Kを一度に生成するのではなく、QとKの小さなタイルを共有メモリに逐次読み込み、そのタイル上でソフトマックスを計算し、各クエリ行についての実行中の最大値と合計のみを保持します。これによりメモリトラフィックが大幅に削減され、カーネルが計算バウンドになることを可能にします。

FlashAttentionの進化の概要

FlashAttentionが時間とともにどのように進化してきたかを理解することは有用です。各バージョンは、新しいハードウェア機能を活用して速度や効率を向上させています。FA4は特にBlackwell世代のGPU向けに最適化されています。以下は、各FlashAttentionバージョンとそれぞれの主な改善点、およびパフォーマンスへの影響を示した表です。

FlashAttention 4の新機能

FlashAttention 4 は、以前の FlashAttention リリースで確立されたタイル型計算パラダイムを継続しています。Q/K/V テンソルは、チップ上のデータ局所性を最大化し、グローバルメモリへのトラフィックを最小化するために、タイルごと（ブロックごと）に処理されます。個々のカーネルインスタンス（協調スレッド配列、CTA と呼ばれる）は協力して、グローバルメモリから対応する Q、K、V のタイルをストリーミングし、それらを最終的なアテンション出力へ投影することで、1つ以上の出力タイルを計算します。

FA4の内部構造は、高度にパイプライン化されたスケジュールで構成されており、注意計算の複数の段階を重ね合わせて潜在的な遅延を隠蔽し、GPUリソースの使用効率を最大化します。FlashAttention 3が基本的に2段階のロード/計算の重ね合わせにとどまるのに対し、FA4は計算を複数の同時段階のセットに分割し、パイプラインの異なる部分を別々のワープグループで処理します。これらのワープグループはそれぞれ、タイルが注意関数を通過する「ライフサイクル」の間において単一の役割に特化しています。

このワープ特化モデルは、32スレッドのワープを次のような特定のタスクに割り当てます:

• データ移動（ロード）：クエリ、キー、バリューのタイルをグローバルメモリから高速なオンチップストレージに取得すること。
• 計算ステージ：テンソルコアを用いた行列乗算-加算ハードウェアで部分的なアテンションスコアを計算する。
• ソフトマックスと正規化のロジック: 注意ソフトマックスの重みに対する指数関数と正規化の操作。
• リスケーリングと縮小ロジック: 蓄積前に数値調整を行ったり、部分結果を統合したりすること。
• エピローグ／ストア操作: 出力テンソルの完了したタイルをグローバルメモリに書き戻します。

基本的な実行モデルの直感的な考え方は、各ワープグループに非同期パイプラインを維持することで、GPUのワープスケジューラがソースオペランドが利用可能になるとすぐにワープグループ間を切り替えられるようにすることです。

CUDAコアによるソフトウェアの指数演算

FA4で特に目を引く新しいトリックのひとつは、ソフトマックスで使用される指数関数の計算方法です。通常、カーネルはGPU上の専用SFUを利用して、exp()のような高価な数学演算を行います。しかし、SMごとのSFUの数は比較的少ないため、簡単にキューのボトルネックになってしまうことがあります。

FA4は、通常のCUDAコアがソフトウェアで近似的な指数関数を実行することで、この制限を回避します。具体的には、FA4にはハードウェア精度でexp2(x)多項式近似を実装したカスタム機能が含まれています。FA4は三次多項式を使って指数計算をシミュレートし、それを少数のSFUに集中させるのではなく、多くのCUDAコアで並列に実行できるようにします。

結論：もはやSFUがソフトマックスを処理するのを待つ必要はありません。通常の汎用コアでexp()を処理することで、FA4はソフトマックスを計算するワープがパイプラインの進行によりうまく追従できるようにします。

よりスマートなオンラインソフトマックス再スケーリング（FA4で新登場）

FA4のもう一つの革新は、ソフトマックスの累積が数値的安定性をどのように維持するかに関係しています。すべてのFlashAttentionのバージョンが行うように、FA4も「オンライン」ソフトマックスを行います。これは、シーケンスを通して部分的なmax()値や合計を逐次的に累積することで（オーバーフローを防ぐため）、安定性を保つ方法です。以前のバージョンでは、ランニングスケールファクター（通常はこれまでの最大ロジット）を常に更新していました。新しい最大値を見つけるたびに再スケーリングを行い、ソフトマックスの安定性を維持していたのです。FA4では、適応型再スケーリングポリシーが導入されています。これは、新しい最大値が以前の最大値よりも十分に大きい場合にのみ再スケーリングを行うというものです。この方法により、再スケーリング操作の頻度が約10倍減少すると報告されています。次の図を考えてみましょう:

この図は、FlashAttentionで使用される2つのリスケーリングのバージョンを比較しており、注意機構の計算を安定かつメモリ効率良く行えるようにしています。元の戦略では、各ブロックの出力で最大値が現れた場合にリスケーリングを行い、寄与をステップごとに正規化することで、小さな和の値が正確に保たれるようにしています。適応リスケーリングでは、局所的および全体的な最大値を追跡し、値が飛躍した場合やブロックが大きくなった場合にのみリスケーリングを行います。これにより、オーバーヘッドを削減しつつ、オーバーフローを防ぎ、少数値の項の数値的に安定した合計を可能にしています。

互換性と現状

この表は、FlashAttention 4（FA4）の期待される互換性と「今日動作するもの」のスタック全体での互換性の概要を示しています。

FA4の順方向パスはすでに一般公開用のソースツリーに広く提供（コミット済み）されています。LabのFlashAttentionリポジトリのようなプロジェクトでは、ユーザーが逆方向パスや可変長、グループ、GQAの完全サポートを求めるオープンな投稿があります。これは、まだ公開されている課題であることを示しています。

可変長シーケンスのサポート（パディングなしで異なる長さのシーケンスを効率的にバッチ処理すること）やGQA（Grouped-Query Attention）は、まだバックワードパスではサポートされていません。GQAは複数のアテンションヘッドが同じキー/バリュープロジェクションを使用するアテンショントリックで、LLMでメモリ使用量を削減するために使われることがあります。以前のバージョンのFlashAttentionでは、MQA/GQAスタイルの設定をサポートしていました。しかし、FA4の主要なフォワードカーネルは、おそらく各クエリに同じ数のヘッドが使われることを前提としています。そのため、現時点ではグループヘッドが正しくサポートされていない可能性があります。グループドクエリアテンションを使用するモデル（Llama2の一部のバリアントなど）は、この機能が実装されるまでFA4ではサポートされません。

エコシステムの統合: FA4 はまだ新しいため、多くの上位フレームワークやレイヤーでは採用されていません。例として、vLLM サービングエンジン（LLM 推論向けに高度に最適化されている）が挙げられますが、現在 FA4 はサポートされておらず、FlashAttention 2 と 3 のみがサポートされています。vLLM で FA4 を選択すると「サポートされていないバージョン」のエラーが表示されます。これは、現時点でバージョンセレクターが 2 または 3 のバージョンのみを選択可能だからです。

FlashAttentionの導入方法 – 判断ガイド

この表を使って、あなたのニーズに最適なFlashAttentionのバージョンを確認してください。

よくある質問

FlashAttention 4（FA4）はどのような問題を解決しようとしているのですか？

これは、注意機構の重いHBMメモリトラフィックを、注意計算をチップ上でタイルごとに行うことで削減し、QK⊤行列全体を具体化することを避けます。

FA4とFA3の主な建築上の変更点は何ですか？

FA4は、Blackwell（SM10.x）上で重なりとオンチップの再利用を増やすために、ワープに特化した約5段階のパイプライン（ロード、計算、ソフトマックス／正規化、再スケール／縮小、ストア）を使用しています。

なぜFA4はSFUではなくCUDAコアで指数関数を計算するのですか？

SFUは限られており、ソフトマックスを計算する際にボトルネックになりやすいため、FA4は代わりにCUDAコア上でexp2()のソフトウェア近似を計算します。

「適応型オンラインリスケーリング」はソフトマックスで何を変えるのですか？」

FA4は新しい最大値が現れるたびに再スケーリングするのではなく、その最大値が大幅に増加したときにのみ再スケーリングを行います。これにより、数値の安定性を維持しながら再スケーリングのオーバーヘッドを減らすことができます。

エンジニアは、今日の生産でいつFA4を採用すべきですか？

まずは、機能フラグの背後でBlackwell 推論専用、固定長、標準注意機構を使用します。cuDNN/SDPA（および古いGPUではFA3/FA2）は、バックワード/可変長/GQAサポートが成熟するまでフォールバックとして残します。

結論

FlashAttention 4 は、注意機構のメモリ帯域幅の要求を削減するという FlashAttention の目標を基に、Q/K/V の処理をより多くオンチップで行うように設計されています。しかし、以前のバージョンとは異なり、FA4 は Blackwell 専用のワープに特化した 5 段階パイプラインと、2 つの主要なソフトマックス改善（CUDA コア上でのソフトウェア指数関数を使用して SFU ボトルネックを軽減する方法、および冗長な安定性計算を排除する適応オンライン再スケーリング）を組み合わせています。FA4 は長いシーケンス長において cuDNN に対する前向きパスのスループット改善を示していますが、依然として前向き優先で、Blackwell 専用です。 さらに、FA4 は本日、重要な機能に欠けており（後方互換、可変長、GQA などが不足）、エコシステムのサポートも十分には整っていません。現時点での要点としては、Blackwell 推論において、厳密なベンチマーク、正確性、およびフォールバックを前提とした採用が必要なものの、潜在的に高いメリットのある選択肢として FA4 の利用を検討するということです。同時に、トレーニング作業負荷やより広範なハードウェアについては、引き続き本番環境向けに成熟した FA3/FA2 のコードパスを使用するべきです。

参考文献

• FlashAttention-3：非同期かつ低精度で高速かつ高精度なアテンション
• FlashAttention 4：GPUアテンションカーネルでペタフロップの壁を突破
• B200にFA4を適用する方法は？
• BlackwellでのFA4バックワードパスのサポート

Linux-Console.netコミュニティで学んでいただきありがとうございます。