FlashAttention-2 (2023)（フラッシュアテンションツー2023）
Tri Dao Princeton/Together AI 2023年 IO-Aware Attention 高速実装。FlashAttention v1 から 2倍高速化・GPU HBM 帯域最適化
概要
FlashAttention-2 は 2023年7月にプリンストン大学（後に Together AI 創業）の Tri Dao が arXiv:2307.08691 で発表した Transformer attention 計算の高速実装ライブラリ第2世代（論文「FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning」）。FlashAttention v1（2022年5月、Dao・Fu・Ermon・Rudra・Ré 共著）の IO-Aware 計算手法を改良し、2倍高速化・GPU HBM 帯域効率 230 TFLOPs/s 達成（A100 GPU 比 73% ピーク性能利用、PyTorch 標準実装は 25-40% のみ）。
FlashAttention は Transformer の attention 計算がメモリ律速（HBM 高速メモリ帯域がボトルネック）であることに着目し、SRAM（GPU on-chip 小容量高速メモリ）でタイル化計算する手法。GPT-3.5/Llama 2/Mistral 等 2023-2024 年の主要 LLM 推論サーバ（vLLM・SGLang・TGI・TensorRT-LLM）全てに採用された業界標準実装で、Together AI の主力プロダクト。
主な特徴・仕組み
- 理論基盤: IO-Aware Algorithm（HBM ⇔ SRAM データ転送を最小化）
- タイル化: attention 行列を 64-128 行 × 64-128 列のタイル単位で計算
- オンライン softmax: 全 attention 値が出揃わなくても softmax 計算可能な改良アルゴリズム
- 並列化改良: v1 の sequence length 並列化に加え head 並列化を導入
- 計算性能: 230 TFLOPs/s（A100 80GB）・PyTorch SDPA の 5-9倍
- メモリ削減: O(N²) → O(N) で長文 (16K-128K context) でも OOM 回避
- 対応 GPU: NVIDIA A100/H100/RTX 30/40系・AMD MI250X（ROCm 経由）
- 数値精度: FP16・BF16・FP8（H100）対応
- 後継: FlashAttention-3 (2024-07、H100 特化、740 TFLOPs/s)
- ライセンス: BSD-3-Clause（オープンソース）
スペック比較表
| 項目 | FlashAttention-2 (2023) | FlashAttention v1 (2022) | PyTorch SDPA | TensorRT-LLM |
|---|---|---|---|---|
| 性能 (A100) | 230 TFLOPs/s | 124 TFLOPs/s | 25-40 TFLOPs/s | 200 TFLOPs/s |
| 性能 (H100) | 460 TFLOPs/s | 270 TFLOPs/s | 60 TFLOPs/s | 400 TFLOPs/s |
| 長文対応 | 128K context | 32K context | 4K まで | 128K context |
| FP8 対応 | あり (H100) | なし | あり | あり |
| 採用例 | vLLM, SGLang, TGI | vLLM 旧版 | PyTorch 標準 | NVIDIA 公式 |
具体例・対応製品
- vLLM (UC Berkeley 2023): PagedAttention + FlashAttention-2、業界標準 LLM 推論サーバ
- SGLang (LMSYS 2024): RadixAttention + FlashAttention-2、トークン共有最適化
- HuggingFace TGI (2023): Text Generation Inference、Hugging Face 公式
- NVIDIA TensorRT-LLM (2023): NVIDIA 専用最適化、FlashAttention-2 互換実装
- llama.cpp + Metal (2024): macOS 向け実装、CUDA 不要
選び方・注意点
- NVIDIA GPU 必須: 2023 年時点で CUDA のみ、AMD は ROCm 経由でβ対応
- FP16/BF16 推奨: FP32 は対応外、FP8 は H100 のみ
- 長文プロンプト: 16K context 以上で v1 比優位、4K 以下は差小
- インストール: pip install flash-attn でビルド時間 10-30分、A100/H100/RTX 30/40 専用 wheel あり
- 依存関係: CUDA 11.6+ / 12.0+ 必須、Python 3.8+ + PyTorch 2.0+
- FlashAttention-3 へ移行検討: H100 環境なら FA-3 (2024-07) で更に 2倍高速化、A100 はFA-2 継続推奨
- デバッグ: 内部実装が Triton カーネルで PyTorch のスタックトレース困難
関連用語との違い
| 用語 | 違い |
|---|---|
| FlashAttention-2 (2023) | IO-Aware attention 230 TFLOPs/s |
| TensorRT-LLM (2023) | NVIDIA 専用最適化スタック |
| Continuous Batching (2022) | バッチ動的再配置 |
| GGUF Quantization (2023) | llama.cpp 標準量子化フォーマット |
| BitNet (2024) | 1.58-bit 量子化・微軽量 |
よくある質問（FAQ）
Q1: PyTorch SDPA とどちらが速い？ A: FlashAttention-2 が 5-9倍高速。PyTorch SDPA（Scaled Dot Product Attention、2.0 で標準実装）は便利だが内部実装が単純で、attention 計算で HBM 律速のまま。本気の LLM 推論なら FlashAttention-2 必須。
Q2: 学習でも使える？ A: 使える。前方計算（推論）+ 後方計算（学習）両対応。Llama 2/3 公式学習・Mistral 学習・Qwen 学習全てで採用されている。学習時間 30-50% 短縮報告。
Q3: 長文プロンプトで OOM が出る場合は？ A: O(N²) メモリの PyTorch SDPA は 16K context で 1 リクエスト 80GB 食う場合あり、FlashAttention-2 は O(N) で 16K でも 2GB 程度。長文 (32K-128K) は FlashAttention-2 必須。
まとめ
- Tri Dao 2023年7月発表・FlashAttention v1 改良版
- IO-Aware で HBM 帯域最適化・A100 で 230 TFLOPs/s
- vLLM・SGLang・TGI・TensorRT-LLM 全ての主要サーバが採用
- 長文 (16K-128K context) でも O(N) メモリで OOM 回避
- 業界標準 attention 実装・LLM 推論不可欠ライブラリ