FlashAttention: GPU別バックエンドとアテンションAPIを読む

FlashAttentionで見る対象と設計上の焦点

このリポジトリは、README で引用されている論文による FlashAttention と FlashAttention-2 の公式実装です。FlashAttention の論文は Tri Dao、Daniel Y. Fu、Stefano Ermon、Atri Rudra、Christopher Ré によるもので、FlashAttention-2 は Tri Dao によるものです。FlashAttention は IO を意識した高速でメモリ効率の良い正確なアテンションと説明され、FlashAttention-2 はより良い並列性とワーク分割を目指しています。README には、Hopper GPU 向けに最適化された FlashAttention-3 beta と、Hopper と Blackwell 向けに CuTeDSL で書かれた FlashAttention-4 のセクションもあります。リポジトリにはソース、テスト、完全な GPT モデル実装、トレーニングスクリプトが含まれます。

FlashAttentionは正確なアテンションを計算しながら入出力の扱いを意識した実装として公開されています。FlashAttention-3はHopper向けのベータ版、FlashAttention-4はCuTeDSLでHopperとBlackwellを狙う別系統です。ベータ版を安定版の機能と混同しない切り分けが必要です。 flash-attentionでは第1節の確認対象を独立させ、入力、結果、失敗時の記録を同じ単位で残します。READMEに明記された値と、手元で観測した値を分けて記録することで、flash-attention固有の判断材料になります。

FlashAttentionで見る導入条件と依存関係

README は要件として CUDA または ROCm ツールキット、PyTorch 2.2 以上、packaging、psutil、ninja を挙げ、Linux を想定しています。標準のインストールは `pip install flash-attn --no-build-isolation`、またはソースから `python setup.py install` です。96GB 未満の RAM で CPU コアが多いマシンでは、`MAX_JOBS` を設定して並列コンパイルジョブ数を制限できます。例: `MAX_JOBS=4 pip install flash-attn --no-build-isolation`。FlashAttention-3 は `hopper` ディレクトリ内に個別のインストール手順があり、FlashAttention-4 は `pip install flash-attn-4` でインストールし、CUDA 13 向けの `cu13` extra があります。README は Windows が v2.3.2 以降で動作する可能性があるが、加えてテストが必要と述べています。

CUDAではPyTorch 2.2以上、CUDA 12.0以上、Linuxが基本条件です。AMD側はROCm 6.0以上で、CKとTritonの2つの経路があります。READMEがWindowsについて「動く可能性はあるが追加テストが必要」としている点からも、OSを先に固定して検証する構成です。 flash-attentionでは第2節の確認対象を独立させ、入力、結果、失敗時の記録を同じ単位で残します。READMEに明記された値と、手元で観測した値を分けて記録することで、flash-attention固有の判断材料になります。

FlashAttentionで見る処理経路と設定項目

NVIDIA 向けでは、CUDA 版 FlashAttention-2 は Ampere、Ada、Hopper GPU、fp16 と bf16、ヘッド次元 256 までをサポートします。flash-attn 2.5.5 以降、dropout なしであればヘッド次元 256 の backward はコンシューマ GPU でも動作します。Turing GPU はこのリポジトリの対象外で、README は別の flash-attention-turing リポジトリを指しています。AMD 向けの ROCm バージョンには2つのバックエンドがあります。composable_kernel バックエンドがデフォルトで、MI200x、MI250x、MI300x、MI355x、RDNA 3/4 GPU をサポートします。Triton バックエンドは CDNA と RDNA GPU を fp16、bf16、fp32 でサポートし、MQA/GQA、回転埋め込み、ALiBi、ページ化アテンション、FlashAttention v3 インターフェースによる FP8 などを含みます。Triton バックエンドのスライディングウィンドウアテンションは作業中とされています。FlashAttention-3 には H100 または H800 と CUDA 12.3 以上が必要で、CUDA 12.8 が推奨されています。

READMEのベンチマーク値はA100やH100の特定条件による自称値です。A100の例は系列長512から16k、ヘッド次元64または128、隠れ次元2048です。自分のGPU、dtype、系列長を同じ表にそろえない限り、225 TFLOPs/secや3から5倍という数値をそのまま期待値にはできません。 flash-attentionでは第3節の確認対象を独立させ、入力、結果、失敗時の記録を同じ単位で残します。READMEに明記された値と、手元で観測した値を分けて記録することで、flash-attention固有の判断材料になります。

FlashAttentionで見る出力と運用記録

Python インターフェースは `flash_attn_qkvpacked_func` と `flash_attn_func`、および推論用の `flash_attn_with_kvcache` を提供します。これらの関数はスケーリングドット積アテンション、softmax(Q @ K^T * softmax_scale) @ V を計算します。2つのコア関数は `dropout_p`、`softmax_scale`、`causal`、`window_size`、`alibi_slopes`、`deterministic` を受け取ります。qkvpacked 版は Q、K、V が1つのテンソルにまとめられている場合に、backward で勾配の明示的な結合を避けます。`flash_attn_func` は Q のヘッド数より少ない KV ヘッドを受け取ることでマルチクエリおよびグループクエリアテンションをサポートしますが、Q のヘッド数は KV のヘッド数で割り切れる必要があります。`flash_attn_with_kvcache` は k_cache と v_cache をインプレースで更新し、`rotary_cos` と `rotary_sin` が与えられた場合は回転埋め込みを適用し、backward はサポートしません。README は `flash_attn/modules/mha.py` のマルチヘッドアテンションモジュールも参照しています。

FlashAttentionは正確なアテンションを計算しながら入出力の扱いを意識した実装として公開されています。FlashAttention-3はHopper向けのベータ版、FlashAttention-4はCuTeDSLでHopperとBlackwellを狙う別系統です。ベータ版を安定版の機能と混同しない切り分けが必要です。 flash-attentionでは第4節の確認対象を独立させ、入力、結果、失敗時の記録を同じ単位で残します。READMEに明記された値と、手元で観測した値を分けて記録することで、flash-attention固有の判断材料になります。

FlashAttentionで見る対応範囲の読み方

README には A100 と H100 のベンチマーク図があり、FlashAttention と PyTorch 標準アテンションを forward と backward を合わせて比較しています。A100 のテスト構成はヘッド次元 64 または 128、隠れ次元 2048、シーケンス長 512 から 16k、バッチサイズは 16k をシーケンス長で割った値です。高速化の図は画像であり、本文は構成とメモリ削減のみを示しています。シーケンス長 2K でメモリが10倍、4K で20倍削減されると報告されています。FlashAttention のメモリ使用量はシーケンス長に線形ですが、標準アテンションは2次です。H100 のセクションにも同様の高速化図があります。README はまた、完全な GPT モデル実装が Huggingface ベースラインと比較してトレーニングを3〜5倍高速化し、A100 あたり最大225 TFLOPs/sec、72% のモデル FLOPs 利用率に相当すると述べています。アクティベーションチェックポイントは不要です。

CUDAではPyTorch 2.2以上、CUDA 12.0以上、Linuxが基本条件です。AMD側はROCm 6.0以上で、CKとTritonの2つの経路があります。READMEがWindowsについて「動く可能性はあるが追加テストが必要」としている点からも、OSを先に固定して検証する構成です。 flash-attentionでは第5節の確認対象を独立させ、入力、結果、失敗時の記録を同じ単位で残します。READMEに明記された値と、手元で観測した値を分けて記録することで、flash-attention固有の判断材料になります。

FlashAttentionで見る制約を切り分ける確認

変更履歴は動作の変更と新機能を示しています。バージョン 2.0 では unpadded 関数が varlen 関数に改名され、完全な書き直しで2倍高速と説明されています。2.1 では query と key の長さが異なる場合に因果マスクを右下に合わせるよう変更されました。2.2 では query のシーケンス長が短い推論を最適化し、`flash_attn_with_kvcache` を追加しました。2.3 ではスライディングウィンドウアテンションを追加し、Mistral AI に謝意を示しています。2.4 では ALiBi と決定論的 backward を追加し、2.5 ではページ化 KV キャッシュ、2.6 では Gemma-2 と Grok で使われる softcapping、2.7 では torch compile との互換性を追加しました。FlashAttention 1.x からアップグレードする場合に役立ちます。

READMEのベンチマーク値はA100やH100の特定条件による自称値です。A100の例は系列長512から16k、ヘッド次元64または128、隠れ次元2048です。自分のGPU、dtype、系列長を同じ表にそろえない限り、225 TFLOPs/secや3から5倍という数値をそのまま期待値にはできません。 flash-attentionでは第6節の確認対象を独立させ、入力、結果、失敗時の記録を同じ単位で残します。READMEに明記された値と、手元で観測した値を分けて記録することで、flash-attention固有の判断材料になります。

FlashAttentionで見るライセンスと採用判断

README によると、テストは FlashAttention の出力と勾配が参照実装と数値許容範囲内で一致することを確認し、最大数値誤差はベースライン実装の誤差の最大2倍です。テストコマンドは CUDA 用の `pytest -q -s tests/test_flash_attn.py`、composable kernel バックエンド用の `pytest tests/test_flash_attn_ck.py` です。Triton AMD テストスイートは `FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pytest tests/test_flash_attn_triton_amd.py` で実行し、数時間かかるとされています。リポジトリは BSD-3-Clause ライセンスで、条件付きで再配布と変更が許可され、ライセンス文には保証と責任の免責が含まれます。README は利用時に FlashAttention の論文を引用するよう求めています。正式なサポートプロセスについては、GitHub issue を開くこと以外は言及されていません。

FlashAttentionは正確なアテンションを計算しながら入出力の扱いを意識した実装として公開されています。FlashAttention-3はHopper向けのベータ版、FlashAttention-4はCuTeDSLでHopperとBlackwellを狙う別系統です。ベータ版を安定版の機能と混同しない切り分けが必要です。 flash-attentionでは第7節の確認対象を独立させ、入力、結果、失敗時の記録を同じ単位で残します。READMEに明記された値と、手元で観測した値を分けて記録することで、flash-attention固有の判断材料になります。

編集部の結論

このプロジェクトは、READMEが示すFast and memory-efficient exact attention. FlashAttention This repository provides the official implementation of FlashAttention and FlashAttention-2 from the following papers.という対象と、記載された環境・権限・入力条件が一致するチームに向きます。採用前にはREADMEのベンチマーク値はA100やH100の特定条件による自称値です。A100の例は系列長512から16k、ヘッド次元64または128、隠れ次元2048です。自分のGPU、dtype、系列長を同じ表にそろえない限り、225 TFLOPs/secや3から5倍という数値をそのまま期待値にはできません。を実際の小さな構成で確認し、READMEにない性能や互換性を前提にしないでください。

公式情報源

1. Official README
2. Project repository
3. Release notes