Flash Attention(高速アテンション)FA

Flash Attention(フラッシュアテンション)は、Tri Dao らが2022年に提案した Self-Attention の I/O-aware な厳密計算アルゴリズムで、GPU の HBM(High Bandwidth Memory)と SRAM(オンチップ・シェアードメモリ)の帯域差を最大限に活用するタイリング戦略により、N×N のアテンション行列を一度に具体化することなくブロック単位で逐次計算します。これにより、メモリ消費が系列長 N に対して O(N^2) から O(N) に改善され、A100 GPU で2-4倍、長文ではさらに大きな高速化が得られます。重要なのは数値的に厳密(exact)であり近似ではない点で、ドロップアウトや任意のマスクとも互換性があります。Online Softmax の手法を再帰的に適用することで、すべてのスコアを保持せずに正規化を実現しています。PyTorch 2.0 の scaled_dot_product_attention で透過的に呼ばれるようになり、Hugging Face Transformers の attn_implementation="flash_attention_2" 指定で簡単に使えます。GPT-4 を含む主要 LLM の事前学習・ファインチューニング・推論で標準的に採用され、Flash Attention 2(2023)、Flash Attention 3(2024、H100 最適化)へと進化しています。