FlashAttention: IOを意識した高速かつメモリ効率の高い厳密なアテンション

論文解説

問題設定と動機

Transformerは系列長$N$に対して2乗の時間およびメモリ計算量を持ちます。近似アテンション手法は理論的なFLOP数を削減しますが、現代のGPU計算は演算性能ではなくメモリ転送（IO）が主なボトルネックとなっているため、実時間の高速化が達成できないことがよくあります。

標準的なアテンションにおいて、$\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$、$\mathbf{P}=\mathrm{softmax}(\mathbf{S})$、$\mathbf{O}=\mathbf{PV}$の計算は、低速な高帯域幅メモリ（HBM）にサイズ$N\times N$の中間行列を書き込み、それを複数回読み出します。これにより$\Theta (Nd+N^2)$のHBMアクセスと$O(N^2)$のメモリフットプリントが発生し、実用的なコンテキスト長が制限されていました。

主要なアイデア

• IO-Awareness（IOへの意識）: GPUのメモリ階層を明示的に考慮します。高速かつ容量の小さいオンチップSRAM（例: A100では192 KB/SM、約19 TB/sの帯域幅）と、大容量だが低速なHBM（40–80 GB、1.5–2.0 TB/sの帯域幅）を区別します。
• タイリング（ブロックごとのソフトマックス計算）: 入力$\mathbf{Q},\mathbf{K},\mathbf{V}$をSRAMに完全に収まるブロックに分割します。実行中のソフトマックス行最大値と正規化統計量を追跡することで、HBM上に完全な$N\times N$行列を実体化することなく、ブロック間でインクリメンタル（逐次的）にソフトマックスのリダクションを計算します。
• 逆伝播における再計算: 逆伝播のために中間値である$N\times N$のアテンション行列を保持することを回避します。代わりに、出力$\mathbf{O}$とソフトマックス統計量$(\ell ,m)$のみを保持し、逆伝播の実行中にSRAM内でブロックごとに$\mathbf{S}$と$\mathbf{P}$をオンチップで再計算します。
• カーネル融合（Kernel Fusion）: 行列積、マスキング、ソフトマックス、ドロップアウトを単一のCUDAカーネルに融合し、HBMとの無駄な往復転送を排除します。

手法（主要な数式を含む）

$\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$が与えられたとき、FlashAttentionはブロックサイズを$B_c=\lceil M/(4d)\rceil$および$B_r=\min (\lceil M/(4d)\rceil ,d)$に設定します。ここで$M$はSRAMの容量です。

インクリメンタルなソフトマックス追跡

部分ベクトル$x^{(1)}$と$x^{(2)}$に分解されたベクトルに対し、実行統計量は以下のように更新されます: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

キー/バリューのブロック$j$とクエリのブロック$i$に対して:

1. SRAM上で未正規化スコアを計算: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$。
2. ブロックの局所統計量を計算: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$、${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$、${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$。
3. グローバルの実行正規化統計量を更新: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. 出力ブロック${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$の実行値を更新: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

計算量の解析

• IO計算量: 標準的なアテンションは$\Theta (Nd+N^2)$のHBMアクセスを必要とします。FlashAttentionは$\Theta (N^2{d}^2{M}^{-1})$のHBMアクセスを必要とします。著者らは、あらゆる$M\in [d,Nd]$に対して漸近的に$o(N^2{d}^2{M}^{-1})$のアクセスを達成できる厳密アテンションアルゴリズムは存在しないことを証明しています。
• FLOP数とメモリ: FlashAttentionは$O(N^{2}d)$のFLOPを実行し（逆伝播での再計算によりFLOP数はわずかに増加）、入力と出力を除いて$O(N)$の追加メモリしか必要としません。

拡張: Block-Sparse FlashAttention

事前定義されたスパースマスク$\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$に従ってゼロでないブロックのみに評価を制限することで、$\Theta (Nd+N^2{d}^2{M}^{-1}s)$のIO計算量を達成します（$s$は非ゼロブロックの割合）。

実験結果（主要な数値）

• 学習の高速化:
  • BERT-large（系列長512）: MLPerf 1.1の学習速度記録を15%上回る（$8\times$ A100で20.0分に対し17.4分）。
  • GPT-2（系列長1K）: HuggingFaceに対して最大3.5倍、Megatron-LMに対して最大1.8倍の高速化を達成。
  • Long-Range Arena (LRA): 標準アテンションに対して2.4倍の高速化。Block-sparse FlashAttentionは2.8倍の高速化を達成。
• 長コンテキスト機能:
  • 4KコンテキストのGPT-2: Megatron-LMの1Kコンテキスト時よりも高速に学習し、検証パープレキシティが0.7ポイント向上（18.2から17.5へ）。
  • 長文分類: コンテキストを512から16Kへ拡張することで、MIMIC-IIIでmicro-$F_1$スコアが4.3ポイント向上（52.8から57.1へ）、ECtHR（8K時）で最大8.5ポイント向上（72.2から80.7へ）。
  • Path-XおよびPath-256: FlashAttentionは、Path-X（系列長16K）を61.4%の精度で解いた初のTransformerとなりました。Block-sparse FlashAttentionはPath-256（系列長64K）で63.1%の精度を達成しました。

限界点と未解決の課題

• CUDAエンジニアリングのオーバーヘッド: IOを意識したカーネルの実装には、標準的なPyTorchやコンパイラの抽象化ではなく、カスタムの低レベルCUDAプログラミングが必要であり、異なるハードウェアアーキテクチャ間での移植が困難です。
• マルチGPUアテンション: 単一ノードにおける最適なIOスケジュールは、インターコネクト（NVLink、InfiniBand）をまたぐ分散マルチGPUアテンションへ自動的に拡張できるわけではありません。
• 厳密アテンションの下界: 一般的な$M$に対して最適である一方、$M$に関してより詳細にパラメータ化された計算量の下界を証明することは未解決の課題のままです。