Flash Attention 2FA2

Flash Attention 2(フラッシュアテンション ツー、FA2)は、Tri Dao が2023年7月に発表した Flash Attention の改良版で、初代の課題だった「非行列演算(softmax の正規化など)に対する並列化が不十分で GPU の理論性能の25-40%しか引き出せない」問題を解消しました。具体的には、(1) 非行列演算の FLOP 数を削減し softmax を逐次的に再スケーリングする方式に変更、(2) 系列長次元での並列化を追加し、Multi-Head 並列に加えて Q ブロック並列も導入、(3) ブロック内のワープ間タスク分割を最適化、という3点の改良で、A100 80GB で従来比約2倍の高速化を達成し、理論性能の72%(335 TFLOPS / 312 TFLOPS の bfloat16 ピーク超)を引き出しました。Llama 2/Llama 3 の事前学習、Mistral・Qwen・DeepSeek など主要オープン LLM の学習で標準採用されています。Hugging Face Transformers では from_pretrained に attn_implementation="flash_attention_2" を渡すだけで利用でき、推論側でも vLLM・TGI・SGLang などの主要サービングフレームワークが内部で活用します。なおバックワード(勾配計算)も最適化されており、ファインチューニング・LoRA 学習でも恩恵が大きいです。