FlashAttention — フラッシュアテンション
FlashAttentionは、計算精度を完全に維持しながら大規模言語モデル (LLM)の訓練と推論を大幅に高速化するために開発された、革新的なアテンション機構（attention）の計算アルゴリズムです。このアルゴリズムは、2022年にスタンフォード大学の研究者チームがTri Dao（トライ・ダオ）氏の主導のもとで初めて発表しました[1]。
FlashAttentionの重要なアイデアは、GPUのメモリ階層を考慮して計算を再構成することにあり、これにより低速なメモリへのアクセス回数を最小限に抑え、標準的なアテンション機構の主要なボトルネックを解消します。
標準的なアテンションの問題点
トランスフォーマーにおける標準的な自己アテンション機構は、次の式で計算されます。 ここで、Q、K、Vはそれぞれクエリ、キー、バリューの行列です。
このアプローチの主な問題は、シーケンス長Nに対して時間計算量とメモリ計算量が二乗のオーダー（O(N²)）になることです[1]。単純な実装では、サイズがN×Nの完全なアテンション行列Sを計算し、GPUメモリに格納する必要があり、これが2つの重大な問題を引き起こします。
- メモリ消費量が大きい: 長いコンテキストを扱う場合、N×Nの行列を格納することは不可能になります。
- 入出力（IO）操作: 主要なボトルネックは演算回数ではなく、低速なGPUメモリへの頻繁なアクセスです。
GPUのメモリ階層
この問題を理解するためには、GPUにおける2種類のメモリ（NVIDIA A100を例に）を区別することが重要です。
- SRAM（静的メモリ）: 高速なオンチップメモリで、容量は小さい（約20MB）ですが、帯域幅は非常に大きい（最大19TB/s）。
- HBM（広帯域メモリ）: 低速な大容量メモリ（40～80GB）で、帯域幅ははるかに小さい（約1.5TB/s）[2]。
この非対称性により、標準的なアテンションアルゴリズムはメモリ帯域幅に制約される（memory-bound）ことになります。なぜなら、低速なHBMから大きな行列を常に読み書きしており、これが遅延の主な原因となっているためです。
FlashAttentionの主要な革新技術
FlashAttentionはIO-aware（IOを意識した）アルゴリズムであり、HBMへのアクセスを最小限に抑えることで問題を解決します。これは、主に3つの技術によって実現されます。
タイリングとブロック処理
行列全体を一度に処理する代わりに、FlashAttentionは入力行列Q、K、Vを高速なSRAMに収まる小さなブロック（タイル）に分割します。アルゴリズムはこれらのブロックを順次ロードし、それらに対してすべてのアテンション計算を実行して最終結果を更新します。この際、完全なアテンション行列を低速なHBMに保存しません[1]。
オンラインSoftmax計算
重要な技術的ブレークスルーは、「オンライン」でのSoftmax計算でした。標準的なSoftmaxは、正規化のために入力ベクトルの全要素を知る必要があります。FlashAttentionは、Softmaxを部分的に計算できる修正アルゴリズムを使用します。このアルゴリズムは、2つの中間値（現在の最大値と指数の合計）を保持し、新しいブロックを処理するたびにこれらを更新することで、行列全体に一度にアクセスすることなく正確な結果を得ることができます[2]。
CUDAカーネルへの演算の融合
すべてのアテンション操作（行列積QKᵀ、マスキング、Softmax、Vとの乗算）は、単一の融合CUDAカーネル（fused kernel）に統合されています。これにより、HBMへの読み書き操作の回数が劇的に削減されます。アルゴリズムは、行列全体を何度も走査する代わりに、ブロックを一度SRAMにロードし、すべての計算を実行して、最終結果のみを書き込みます。
理論的および実践的な効率性
計算量と最適性
FlashAttentionは、メモリ消費量をO(N²)からO(N)に削減し、線形のスケーラビリティを実現します。このアルゴリズムのIO計算量は、2レベルのメモリ階層においてアテンションを計算するための理論的に最適なものであることが証明されています。つまり、ハードウェアを変更しない限り、正確なアテンションをこれ以上高速に実行することは不可能です[3]。
実験結果
FlashAttentionの最初のバージョンは、大幅な改善を示しました。
- 高速化:
- BERT-large（シーケンス長512）: 訓練が15%高速化。
- GPT-2（シーケンス長1K）: 3倍の高速化。
- Long-Range Arenaタスク（1K-4K）: 2.4倍の高速化[1]。
- メモリ節約: 正確なベースライン実装と比較して、最大20倍のメモリ節約。
- モデル品質の向上: より長いコンテキストを扱えるようになったことで、FlashAttentionはモデルの品質を損なうどころか、向上させました。例えば、GPT-2のパープレキシティは0.7ポイント改善し、長文文書分類タスクの精度は6.4ポイント向上しました[1]。
進化と今後の開発
FlashAttentionの成功は、ハードウェアを意識した一連のアルゴリズムの始まりとなりました。
FlashAttention-2 (2023) - FlashAttention-2（2023年）
第2のバージョンは、GPUリソースをより完全に活用することを目的としていました。オリジナルのFlashAttentionでは、NVIDIA A100での効率は最大値のわずか25～40%でした。FlashAttention-2は計算の並列化を改善し、以下のことを可能にしました[4]:
- 第1バージョンと比較して2倍の高速化を達成。
- GPU使用率を理論上の最大値の50～73%まで向上。
- サイズ256のアテンションヘッドや、Multi-Query Attention（MQA）アーキテクチャへの対応を拡大。
FlashAttention-3 (2024) - FlashAttention-3（2024年）
第3のバージョンは、GPUアーキテクチャNVIDIA Hopper (H100)に特化して最適化されました[5]。Tensor Coreの非同期性やFP8のサポートといった新しいハードウェア機能を活用し、以下のことを可能にしました:
- FlashAttention-2と比較してさらに1.5～2倍の高速化を達成。
- FP16で最大740 TFLOPS、FP8で1.2 PFLOPSに近い性能を達成。
特化型ソリューション
FlashAttentionのアイデアは、他のプロジェクトでも発展しました。
- FlashInfer（2025年）: LLMの推論タスクに特化して最適化された、カスタマイズ可能なアテンションエンジン。ストリーミング生成モードでのKVキャッシュの効率的な処理に焦点を当てています[6]。
- FlashMLA（2024年）: コンテキストキャッシュを圧縮するアテンション（latent attention）の実装であり、情報の損失を最小限に抑えながら、非常に長いシーケンスでのメモリを節約できます[7]。
業界とエコシステムへの影響
FlashAttentionは基礎的なブレークスルーとなり、LLMの効率的な訓練と推論のための業界標準へと急速に成長しました。PyTorchやHugging Faceなどの主要なライブラリに統合され、ほとんどの大規模言語モデル（LLaMA, MPT, Falcon, Claudeなど）で使用されています。
言語モデルのコンテキストウィンドウを2～4千トークン（GPT-3）から128千トークン（GPT-4）、さらには実験的なモデルでは数百万トークンにまで拡大する上で決定的な役割を果たしたのは、まさにFlashAttentionとその後のバージョンでした[8]。このアルゴリズムは、トランスフォーマーのスケーリングにおける主要な障害の1つを取り除き、長文文書の分析からマルチモーダルな理解まで、AIアプリケーションの新たな可能性を切り開きました。
外部リンク
参考文献
- Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
- Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
- Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
- Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
- Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
- Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
- Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
- Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
- Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
- Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.
脚注
- ↑ 1.0 1.1 1.2 1.3 1.4 Dao, Tri, et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” arXiv:2205.14135 [cs.LG], 28 May 2022. [1]
- ↑ 2.0 2.1 Dao, Tri, et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” OpenReview. [2]
- ↑ “We're Training AI Twice as Fast This Year as Last”. IEEE Spectrum. [3]
- ↑ Dao, Tri. “FlashAttention-2”. tridao.me. [4]
- ↑ “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision”. PyTorch Blog. [5]
- ↑ “[2501.01005] FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving”. arXiv. [6]
- ↑ “GitHub - deepseek-ai/FlashMLA: FlashMLA: Efficient MLA decoding kernels”. GitHub. [7]
- ↑ “The Evolution of Flash Attention: Revolutionizing Transformer Efficiency”. Medium. [8]