FlashAttention-2（フラッシュアテンションツー）
Tri Dao氏らが開発したTransformerのAttention計算を高速化するアルゴリズム。GPUのSRAMを活用してメモリアクセスを最小化し、標準Attentionより2〜4倍の高速化とメモリ削減を実現する。LLM推論・学習の事実上の標準最適化技術。
FlashAttention-2 とは
FlashAttention-2（フラッシュアテンションツー）は、Tri Dao氏（スタンフォード大学・Together AI）らが2023年に発表したTransformerアーキテクチャのSelf-Attention計算を高速化するアルゴリズム。2022年発表のオリジナルFlashAttentionからさらに改良され、GPU上でのAttention計算を標準実装より2〜4倍高速化し、メモリ使用量をO(N²)からO(N)に削減する。
なぜAttentionがボトルネックなのか
TransformerのSelf-Attentionは系列長Nに対してO(N²)の計算量とメモリを必要とする。LLMの長コンテキスト（4K→32K→128K tokens）化が進むにつれ、この二次的なスケーリングが深刻なボトルネックとなる。
標準Attentionの問題点:
- Q, K, V行列のHBM（High Bandwidth Memory）への読み書きが頻繁に発生
- HBMの帯域幅（A100=2TB/s）がSRAM（A100=19TB/s相当）より約10倍遅い
- N²サイズのAttention行列全体をHBMに書き込む必要がある
FlashAttention-2の仕組み
タイリング（Tiling）
Q, K, V行列をGPUのSRAM（共有メモリ・L2キャッシュ）に収まる小ブロックに分割し、Attention計算をブロック単位で実行する。HBMへのアクセス回数をO(N)に削減する。
Online Softmax
Softmax計算は全スコアが揃わないと実行できないが、FlashAttentionは「数値安定なIncremental Softmax」を使い、ブロックを逐次処理しながら正確なSoftmax値を算出する。
並列化の改善（FlashAttention-2 固有）
- Warps間の並列化: Q行列方向に並列実行するよう変更（FA-1はK/V方向）
- シーケンシャル削減: オリジナルに比べて非行列積（non-GEMM）演算を削減
- 因果的マスキングの最適化: Causal Attention（GPT系）でのマスク処理を約2倍高速化
パフォーマンスベンチマーク
- A100 80GB での Attention FLOP 利用率: 標準=約25% / FA-1=約60% / FA-2=約72%
- LLaMA 2 70B 学習速度: FA-2採用で約3倍高速（HuggingFace公式）
- 推論スループット (Llama 3 8B, A100): FA-2で約2.5倍向上
対応フレームワーク・エンジン
- PyTorch: ○（SDPA経由） / バージョン2.0+
- vLLM: ○（デフォルト） / バージョン0.2.0+
- TensorRT-LLM: ○ / バージョン0.5.0+
- SGLang: ○ / 全バージョン
- HuggingFace Transformers: ○（
attn_implementation="flash_attention_2"） / バージョン4.34+ - Axolotl: ○ / 全バージョン
FlashAttention-3（2024年）
H100 GPU向けに最適化されたFlashAttention-3が2024年に発表。FP8精度・非同期WGMMA命令・Persistent Kernel等でFA-2比さらに1.5〜2倍の高速化を達成。FA-2は依然としてA100/V100/RTX 4090等で有効だが、H100/H200環境ではFA-3が推奨。
まとめ
FlashAttention-2はLLMの推論・学習インフラにおいて事実上の標準となったAttention最適化技術。長コンテキストモデル（Claude 3.5/GPT-4/Gemini 1.5等）の実現を支える基盤技術であり、vLLM・TensorRT-LLM・SGLangといった主要推論エンジンがすべてデフォルト採用する。LLM推論エンジニアにとって必須の知識だ。