FlashAttentionとは
GPUのメモリ階層を意識した計算順序により、通常のAttentionと同じ結果を高速かつ省メモリに計算するアルゴリズム
ひとことで言うと
AIの内部計算のやり方だけ工夫して、Attention計算を同じ結果のまま速く軽くする技術。
概要
FlashAttentionは、Self-Attentionの計算結果自体は変えずに、GPU内の高速だが容量の小さいメモリ(SRAM)と低速で容量の大きいメモリ(HBM)の間のデータ転送量を減らす計算順序を工夫することで、Attentionの計算を高速化しGPUメモリ使用量も削減するアルゴリズム。 通常のAttention実装は、系列長の2乗サイズの巨大な中間行列をHBMへ何度も書き出し・読み込みする必要があり、この転送がボトルネックになっていた。 FlashAttentionは、入力をタイル状の小さなブロックへ分割し、各ブロックの計算をSRAM上だけで完結させ、HBMとのやり取りを最小限に抑える。近似ではなく、数学的に厳密な計算を保ったまま高速化を実現する。 PyTorchやTransformersなど主要なライブラリに組み込まれており、長いコンテキストを扱うLLMの学習・推論の両方で標準的に使われる最適化技術になっている。
背景
Self-Attentionの計算量とメモリ使用量は系列長の2乗に比例して増加し、特に長いコンテキストを扱う際にGPUメモリの帯域幅がボトルネックになっていた。 FlashAttentionは、演算量自体は変えずメモリアクセスパターンを最適化することで、この帯域幅のボトルネックを解消する目的で考案された。
歴史
2022年: DaoらがFlashAttentionを提案する論文「FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness」を発表。 2023年: 並列化とワークロード分割をさらに改善したFlashAttention-2が発表され、GPU利用効率が向上。 2024年: FlashAttention-3が発表され、最新世代GPUの性能をさらに引き出す改良が加えられる。
アーキテクチャ
入力のクエリ・キー・バリューを、GPUのSRAMに収まる大きさのブロックへ分割する。 各ブロックの組み合わせについてAttentionスコアと出力を逐次的に計算し、オンライン方式のソフトマックス正規化を用いて最終的な出力を段階的に積み上げていく。 この過程で系列長の2乗サイズの巨大な中間行列をHBMへ書き出す必要がなくなり、メモリ転送量が大きく削減される。
コード例
transformersでのFlashAttention有効化例
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-7b-hf",
attn_implementation="flash_attention_2",
torch_dtype="auto",
)利点
- 通常のAttentionと数学的に同一の結果を保ちながら、計算速度とメモリ効率を高められる
- 長いコンテキストウィンドウを扱う際のGPUメモリ使用量を抑えられる
- 主要な深層学習フレームワークに組み込まれており、モデル側の変更なしに利用できることが多い
欠点
比較
- Self-Attention — FlashAttentionはSelf-Attentionの計算結果を変えず、計算順序とメモリアクセスだけを最適化する実装手法
- KVキャッシュ — FlashAttentionが計算そのものを高速化するのに対し、KVキャッシュは過去の計算結果を再利用して生成を高速化する
- vLLM — vLLMなどの推論エンジンは、FlashAttentionのようなカーネル最適化を内部で組み合わせて高速な推論を実現する
関連用語
よくある質問
FlashAttentionは精度に影響する？
計算順序を変えるだけで、数学的には通常のAttentionと同一の結果を出すよう設計されている。近似手法ではないため、精度への悪影響は基本的にない。
FlashAttentionはどこで使われている？
PyTorchの標準Attention実装やHugging Face Transformers、vLLMなど多くのLLM学習・推論フレームワークに組み込まれ、標準的な最適化技術として広く使われている。