Transformerのセルフアテンション計算をIO効率化するアルゴリズム。HBMとSRAM間のデータ転送を最小化するタイリング技法により、標準アテンションと数学的に等価な結果を2〜4倍高速に算出する。
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推論・学習の事実上の標準となっている。
標準的なセルフアテンションは以下の計算を行う:
シーケンス長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の核心は、アテンション行列を明示的にメモリ上に展開せず、ブロック単位で計算を完結させる点にある。
| バージョン | 対応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.0以降、torch.nn.functional.scaled_dot_product_attentionでFlashAttentionバックエンドが自動選択される。
vLLM 0.4以降でFlashAttention-2がデフォルト有効。--enable-flashinferオプションでFlashInfer（FlashAttention派生の高速カーネル）も選択可能。
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-70B", attn_implementation="flash_attention_2")で明示的に有効化。
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が標準。