FlashAttention（フラッシュアテンション）
FlashAttentionとは、GPUのメモリ階層（SRAM/HBM）を意識したIO-Awareなアテンション計算アルゴリズムであり、標準的なアテンション実装と数学的に同一の結果を返しながら、実行速度を2〜4倍高速化しメモリ使用量をO(n^2)からO(n)に削減する。Tri Dao（Stanford/Together AI）が2022年に提案し、FlashAttention-2（2023年）、FlashAttention-3（2024年、H100/Hopper最適化）と進化を続けている。
FlashAttention（フラッシュアテンション）とは
FlashAttentionは、標準的なSelf-Attentionの計算結果を一切変えることなく（exact attention）、GPUメモリ階層の特性を活用して実行時間を2〜4倍高速化するアルゴリズムである。Tri Dao（Stanford大学、現Together AI CTO）が2022年の論文「FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness」で発表した。
なぜFlashAttentionが必要なのか
標準アテンション実装の問題点は、n x nのアテンション行列をHBM（GPU外部メモリ）に書き出す必要があることにある：
| 問題 | 標準アテンション | FlashAttention |
|---|---|---|
| アテンション行列サイズ | O(n^2) HBMに格納 | タイル単位でSRAMに保持 |
| HBM読み書き回数 | O(n^2 * d) | O(n^2 * d^2 / M) |
| メモリ使用量 | O(n^2) | O(n) |
| 128Kトークン時のメモリ | 約64GB（fp16） | 数百MB |
| Softmax計算 | 行全体を一度にロード | オンラインSoftmaxで逐次計算 |
ここでMはSRAMサイズ（A100で20MB、H100で50MB）、dはヘッド次元である。
FlashAttentionの動作原理
FlashAttentionの核心は「タイリング」と「オンラインSoftmax」の2つの技術にある：
- タイリング（Tiling）: Q・K・V行列をSRAMに収まるブロック（タイル）に分割し、タイルごとにアテンション計算を実行。HBMへのアテンション行列の書き出しを完全に回避
- オンラインSoftmax: 行全体のSoftmax正規化をブロック単位で逐次的に計算。各ブロックの最大値と指数和を累積的に更新することで、正確なSoftmax結果を得る
- 逆伝播の再計算: 順伝播時にアテンション行列を保存せず、逆伝播時にQ・K・Vから再計算。メモリO(n)と引き換えに計算量は約20%増加するが、HBM帯域がボトルネックの実環境では全体性能が向上
バージョン別進化
| バージョン | 年 | 対応GPU | 主要改善 | 実測速度（A100比） |
|---|---|---|---|---|
| FlashAttention-1 | 2022 | A100 | タイリング + オンラインSoftmax | 2〜4x |
| FlashAttention-2 | 2023 | A100/H100 | ワープ間並列化、非因果マスク最適化 | 1.5〜2x (v1比) |
| FlashAttention-3 | 2024 | H100 (Hopper) | FP8対応、WGMMAの活用、ピンポンスケジューリング | 1.5〜1.75x (v2比、H100) |
実装とフレームワーク対応
FlashAttentionは主要なMLフレームワークすべてに統合されている：
- PyTorch 2.0+:
torch.nn.functional.scaled_dot_product_attentionにFlashAttentionバックエンドとして統合（torch.backends.cuda.flash_sdp_enabled()で切替） - Hugging Face Transformers:
model.to(device, attn_implementation="flash_attention_2")で有効化 - vLLM: デフォルトでFlashAttention-2を使用、PagedAttentionと組み合わせ
- TensorRT-LLM: NVIDIA公式最適化エンジンにFlashAttention-2カーネルを統合
- FlashInfer: FlashAttention-3ベースの高速推論ライブラリ（Together AI）
- Triton: OpenAIのTriton言語でカスタムFlashAttentionカーネルを記述可能
ベンチマーク結果
A100 80GB SXM上での性能比較（causal attention、batch=8、head=32、d=128）：
| シーケンス長 | 標準実装 (ms) | FlashAttention-2 (ms) | 速度向上 | メモリ削減 |
|---|---|---|---|---|
| 1,024 | 2.1 | 0.8 | 2.6x | 85% |
| 4,096 | 28.5 | 8.2 | 3.5x | 94% |
| 16,384 | 442 | 120 | 3.7x | 98% |
| 65,536 | OOM | 1,850 | - | 99.6% |
| 131,072 | OOM | 7,200 | - | 99.8% |
よくある質問（FAQ）
Q1: FlashAttentionを使うと計算結果が変わりますか？ A: 変わらない。FlashAttentionは「exact attention」であり、標準アテンションと数値的に同一の結果を返す（bit-for-bit identical ではないが、浮動小数点誤差の範囲内で一致）。近似手法ではないため品質への影響はゼロ。
Q2: FlashAttentionはどのGPUで使えますか？ A: FlashAttention-2はNVIDIA Ampere以降（A100、A10G、RTX 3090/4090）で動作する。FlashAttention-3はHopper世代（H100、H200）専用。AMD MI300XはROCm 6.x向けのFlashAttention対応が進行中で、2026年時点でflash-attn 2.6.x以降で実験的サポート。
Q3: FlashAttentionとPagedAttentionの違いは何ですか？ A: FlashAttentionはアテンション「計算」の高速化手法であり、PagedAttention（vLLM）はKVキャッシュの「メモリ管理」手法。両者は直交する技術で、vLLMでは両方を同時に使用している。
まとめ
- FlashAttentionはGPUメモリ階層を活用したIO-Awareアテンション、結果はexact（近似なし）
- メモリ使用量をO(n^2)→O(n)に削減し、128Kトークン以上の長文処理を実現
- FlashAttention-3はH100/Hopper向けにFP8対応で更なる高速化
- PyTorch 2.0+、Hugging Face、vLLM等の主要フレームワークに標準統合済み