FlashAttention（フラッシュアテンション）
Transformerのセルフアテンション計算をIO効率化するアルゴリズム。HBMとSRAM間のデータ転送を最小化するタイリング技法により、標準アテンションと数学的に等価な結果を2〜4倍高速に算出する。
FlashAttentionとは
FlashAttentionは、2022年にStanford大学のTri Dao氏が発表したTransformerのアテンション計算を高速化するアルゴリズムである。GPU上のHBM（High Bandwidth Memory）とオンチップSRAM間のデータ移動をタイリング（ブロック分割）によって最小化し、標準的なアテンション実装と完全に等価な計算結果を、メモリ使用量O(N)・計算速度2〜4倍で達成する。2026年現在、vLLM・TensorRT-LLM・PyTorch 2.x・Hugging Face Transformersなど主要フレームワーク全てに統合されており、LLM推論・学習の事実上の標準となっている。
アテンション計算のボトルネック
標準的なセルフアテンションは以下の計算を行う:
- Q×K^T の行列積でアテンションスコアを計算（O(N²d)）
- Softmaxで正規化
- 正規化スコア×V で出力を計算
シーケンス長Nが4,096〜128,000に達する現代のLLMでは、N²サイズのアテンション行列をHBMに書き出す必要があり、これがメモリ帯域のボトルネックとなる。H100のHBM帯域は3.35TB/sだが、SRAM帯域は33TB/s以上であり、約10倍の差がある。
| GPU | HBM容量 | HBM帯域 | SRAM容量 | SRAM帯域 | 演算性能(FP16) |
|---|---|---|---|---|---|
| A100 80GB | 80GB | 2.0 TB/s | 20MB/SM | ~19 TB/s | 312 TFLOPS |
| H100 SXM | 80GB | 3.35 TB/s | 50MB/SM | ~33 TB/s | 989 TFLOPS |
| H200 | 141GB | 4.8 TB/s | 50MB/SM | ~33 TB/s | 989 TFLOPS |
| B200 | 192GB | 8.0 TB/s | 未公開 | 未公開 | 2,250 TFLOPS |
FlashAttentionのタイリングアルゴリズム
FlashAttentionの核心は、アテンション行列を明示的にメモリ上に展開せず、ブロック単位で計算を完結させる点にある。
動作原理
- Q, K, V 行列をブロックサイズ B_r × B_c（典型的に128×128）に分割
- 各ブロックをSRAMにロードし、そのブロック内でQ×K^T→Softmax→×Vを一気に計算
- Online Softmax（Milakov & Gimelshein, 2018）を用いて、ブロック間のSoftmax正規化を漸進的に統合
- 最終結果のみHBMに書き戻す
メモリ計算量の改善
- 標準アテンション: O(N²) の中間メモリ（アテンション行列全体を保持）
- FlashAttention: O(N) の中間メモリ（ブロックサイズ分のみSRAMに保持）
- シーケンス長128Kの場合: 標準で128K×128K×2B = 32GB → FlashAttentionで数MB
バージョン進化
FlashAttention-1（2022年）
- 初版。A100で標準アテンションの2〜4倍高速
- Tritonカーネルで実装。Causal masking対応
FlashAttention-2（2023年）
- ワークパーティショニング改善で非causal時1.7〜2倍、causal時約3倍にさらに高速化
- GPUスレッドブロック間のワーク分配を最適化
- A100で最大理論FLOPSの73%を達成（FA-1は35%）
FlashAttention-3（2024年）
- H100のTMA（Tensor Memory Accelerator）を活用した非同期データ転送
- FP8ネイティブ対応でさらに1.5〜2倍高速
- WGMMA（Warpgroup Matrix Multiply-Accumulate）命令の活用
- Hopper世代以降のGPU専用（Ampere非対応）
| バージョン | 対応GPU | 主要改善 | 対標準比速度 | 理論FLOPS利用率 |
|---|---|---|---|---|
| FA-1 | Ampere+ | タイリング導入 | 2〜4× | ~35% |
| FA-2 | Ampere+ | ワーク分配最適化 | 3〜5× | ~73% |
| FA-3 | Hopper+ | TMA/FP8/WGMMA | 5〜8× | ~85% |
主要フレームワークでの利用方法
PyTorch 2.x（組み込み）
PyTorch 2.0以降、torch.nn.functional.scaled_dot_product_attentionでFlashAttentionバックエンドが自動選択される。
vLLM
vLLM 0.4以降でFlashAttention-2がデフォルト有効。--enable-flashinferオプションでFlashInfer（FlashAttention派生の高速カーネル）も選択可能。
Hugging Face Transformers
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-70B", attn_implementation="flash_attention_2")で明示的に有効化。
FlashAttentionの制約と代替
- Ampere以前のGPU: FlashAttentionはSM80以降（A100/RTX 30xx以降）が必要。V100ではxFormers memory_efficient_attentionが代替
- 可変長シーケンス: パディングなしのvarlen APIが必要（FA-2で対応済み）
- カスタムアテンションマスク: Sliding Window Attention等の複雑なマスクは専用実装が必要（FlashAttention-2.5で一部対応）
よくある質問（FAQ）
Q1: FlashAttentionを使うとモデルの精度は変わりますか？ A: 変わらない。FlashAttentionは標準アテンションと数学的に完全等価であり、精度低下は一切ない（FP8モード除く）。速度とメモリ効率のみが改善される。
Q2: RTX 4090でFlashAttention-3は使えますか？ A: 使えない。FlashAttention-3はHopper世代（H100/H200）以降のGPU専用。RTX 4090（Ada Lovelace）ではFlashAttention-2が最速選択肢。
Q3: FlashAttentionとxFormersの違いは何ですか？ A: xFormersのmemory_efficient_attentionはFlashAttentionと同様のメモリ効率化を行うが、FlashAttention-2/3はより積極的なハードウェア最適化（ワーク分配、TMA活用）で上回る。2026年現在はFlashAttentionが標準。
まとめ
- FlashAttentionはGPUのメモリ階層を活用したIO-awareなアテンション最適化アルゴリズム
- 標準アテンションと数学的に等価で精度劣化なし、速度2〜8倍・メモリO(N²)→O(N)
- FA-3はH100以降でFP8/TMAを活用し理論性能の85%を達成
- 2026年現在、vLLM/PyTorch/HF Transformers全てに統合済みのデファクト標準