Flash Attention 3FA3

Flash Attention 3(フラッシュアテンション スリー、FA3)は、Jay Shah・Tri Dao らが2024年7月に発表した NVIDIA Hopper アーキテクチャ(H100)向けに特化した改良版で、Hopper で導入された3つのハードウェア機能を最大限活用しています。第1に、Warp Specialization により Producer ワープ(データロード)と Consumer ワープ(計算)を分離し、TMA(Tensor Memory Accelerator)による非同期メモリ転送と Tensor Core 演算をオーバーラップさせます。第2に、Pingpong Scheduling で2つのワープグループが行列積と softmax を交互に実行することで Tensor Core を埋め続けます。第3に、FP8(8ビット浮動小数点)精度サポートを追加し、ブロックレベルのスケーリングで精度低下を抑えながら H100 の FP8 ピーク性能を活用します。これらの改良により、H100 SXM5 で BF16 では約740 TFLOPS(FA2 比1.5-2倍)、FP8 では約1.2 PFLOPS という驚異的なスループットを達成し、理論性能の75%超を引き出します。Llama 3.1 405B の128Kコンテキスト学習や、超長文推論サービングのレイテンシ削減で実用化が進んでおり、PyTorch 2.5 以降や Hugging Face Transformers の最新版で順次サポートされています。