// シス担のミカタ
LLM推論・最適化

Flash Attention ふらっしゅあてんしょん

Flash Attentionアテンション高速化GPUSRAMメモリ効率長文対応
Flash Attentionについて教えて

簡単に言うとこんな感じ！

アテンション計算を「GPUの速い小さなメモリ（SRAM）」を賢く使って超高速化したアルゴリズムだよ！従来は遅い大容量メモリを何度も読み書きしてたのを、なるべくSRAM内で完結させることで、速くなってメモリも節約できるんだ！

Flash Attentionとは

Flash Attention とは、スタンフォード大学のTri Dao氏らが2022年に提案したアテンション計算のIO効率を最適化したアルゴリズムです。Transformerのアテンション計算はシーケンス長の二乗（O(N²)）のメモリを必要とするのがボトルネックでしたが、Flash Attentionはタイリング（分割処理） の技術でGPUのSRAM内に収めながら計算することで、GPUとHBM（高帯域幅メモリ）間のデータ転送量を劇的に削減します。

通常のアテンションが「計算して → 大きいメモリに書いて → 読み直して → また計算」を繰り返すのに対し、Flash Attentionは「小さなブロックに分けてSRAM内で計算を完結させる」ことでアクセス回数を減らします。結果として2〜4倍の速度向上と、メモリ使用量の大幅削減を実現します。

従来のアテンション vs Flash Attention

項目従来のアテンションFlash Attention
メモリ使用量O(N²)（系列長の2乗）O(N)（系列長に比例）
HBMアクセス多い（何度も読み書き）少ない（SRAM内で完結）
速度ベースライン2〜4倍高速
勾配計算中間値を保持再計算で省メモリ
長文サポートメモリ不足になりやすい長文対応が現実的に

歴史と背景

  • 2022年6月：Flash AttentionをNeurIPS 2022で発表（スタンフォード大・Tri Dao）
  • 2022年後半：HuggingFace・Mosaic MLなどが採用、実用化が進む
  • 2023年7月：Flash Attention 2が登場、並列化をさらに改善
  • 2024年：Flash Attention 3がH100 GPU向けに最適化され公開
  • 現在：LLaMA、GPT-4o、Claude等の主要LLMほぼすべてで採用

Flash Attentionが重要な理由

コンテキスト長4K → Flash Attentionで速度2倍
コンテキスト長32K → Flash Attentionで速度4倍以上
コンテキスト長100K → 従来手法ではメモリ不足で実行不可
                     Flash Attentionなら実行可能

つまり「長文対応LLM」の実現には Flash Attention が不可欠

関連用語