Flash Attention

Flash Attention（FlashAttention）は、GPUのメインメモリに完全なアテンション行列を書き込まずに、Transformerモデルの正確なアテンションを計算するアルゴリズムおよびGPUカーネル群です。N個のトークンに対する標準的な自己アテンションでは、N × Nのスコア行列を作成して高帯域幅メモリ（HBM）に書き込み、ソフトマックスと値との積の計算のために読み戻します。そのため、アテンションの実行時間の大部分はデータ移動に費やされます。スタンフォード大学のTri Daoらは、2022年の論文FlashAttention: IO認識による高速かつメモリ効率の高い正確なアテンションでこの手法を発表しました。処理を高速なオンチップSRAMに収まるタイルに分割するため、メモリ使用量はシーケンス長に対して線形に増加し、出力は浮動小数点の丸め誤差の範囲内で標準的なアテンションと一致します。このカーネルはPyTorch内や、vLLMなどの大規模言語モデル（LLM）推論エンジンで実行されます。

Flash Attentionの仕組み#

GPUには、演算ユニットの近くにある少量の超高速メモリと、チップ外部にある大容量で低速なメモリがあります。FlashAttentionの論文では、HBMが40～80 GB、帯域幅が1.5～2.0 TB/sで、108基のストリーミングマルチプロセッサそれぞれに192 KBのSRAMを備え、推定帯域幅が19 TB/sのNVIDIA A100を例に挙げています。アテンションは読み込む1バイトあたりの演算数が少ないため、その速度は2つの階層間でデータを移動する頻度に左右されます。この手法ではその転送量を考慮しており、それが「IO認識」の意味です。

FlashAttentionはクエリのタイルを読み込み、キーと値のタイルをSRAMに順次流し込みます。各タイルについてスコアを計算し、クエリの各行に対する実行中の最大値と合計を更新します。最大値が大きくなった場合は部分出力を再スケーリングし、新しいタイルの寄与を加算します。このオンラインソフトマックスにより、行全体を一度に参照せずに正確なソフトマックス結果を得られます。HBMに戻るのは、最終出力と行ごとの2つの数値だけです。バックワードパスではN × Nの行列を保存せず、Q、K、Vからスコアタイルを再計算します。演算は増えますが、メモリ読み込みが大幅に少ないため、なお高速に実行できます。次に示すのは、クエリすべてを1つのキーと値のタイルに対して一度に処理するPyTorchの同じループです。

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

最大のスコアタイルは4,096 × 256で、4,096 × 4,096の行列の16分の1です。結果と標準的なアテンションとの差は最大でも4.3e-07で、FP32の丸め誤差の範囲です。実際のカーネルではクエリもタイルに分割し、GPUのスレッドブロックごとに割り当てます。また、ループ全体を単一のカーネルに融合することで、タイルがSRAMの外に出ないようにします。

FlashAttention-2、-3、-4#

各バージョンは、より新しい世代のGPUを対象としています。FlashAttention-2（2023年）はスレッド間の処理分割を再構成し、初代の約2倍の速度を達成しました。A100のピークスループットの50～73%に達します。FlashAttention-3（2024年）はH100の非同期Tensor Coresを使用し、FP8をサポートします。FP16で740 TFLOPs/s、FP8で約1.2 PFLOPs/sを達成します。FlashAttention-4（2026年）はBlackwell GPUを対象とし、B200でBF16の最大1,613 TFLOPs/s、ピーク性能の71%を報告しています。カーネルはオープンソースのflash-attentionライブラリに含まれ、FP16およびBF16入力で動作します。また、PyTorchのscaled_dot_product_attentionは、GPUと入力が対応している場合にFlashAttention-2カーネルを自動選択し、それ以外ではメモリ効率の高い実装または標準実装にフォールバックします。Ultralyticsパッケージで実行されるMetaのSAM 2は、Hiera画像エンコーダーとメモリアテンションの計算にこの関数を使います。そのため、ハードウェアと入力精度が許す場合、PyTorchはこれらのレイヤーをFlashAttentionに割り当てます。

Flash Attentionと関連用語#

• 標準的なアテンション： 同じ数学的演算ですが、スコア行列と確率行列全体をHBMに書き込んで計算します。
• スパースアテンション： 選択したクエリとキーのペアだけをスコアリングします。FlashAttentionの論文では、空のタイルをスキップするブロックスパース版も紹介されました。
• 線形アテンション： ソフトマックスを特徴マップに置き換えて計算量を長さに対して線形に増加させますが、モデルの出力は変わります。
• PagedAttention： サービング中にKVキャッシュを固定サイズのページに保存するメモリマネージャーです。FlashAttentionを含むアテンションカーネルは、これらのページからキーと値を読み込みます。
• リングアテンション： 非常に長いシーケンスを複数のデバイスに分割します。各デバイスは割り当てられた断片でブロック単位のアテンションを計算し、キーと値のブロックはリング状に受け渡されます。

学習と長いコンテキストへの影響#

メモリ使用量がシーケンス長に対して線形に増加するため、FlashAttentionによって長いコンテキストウィンドウでの学習が実用的になりました。原論文では、MLPerf 1.1の学習記録と比較して、シーケンス長512の BERT-largeでエンドツーエンドの速度が15%向上し、シーケンス長1,024のGPT-2で3倍、Long Range Arenaベンチマークで2.4倍になったと報告されています。同様の削減効果はビジョンにも当てはまります。Vision Transformer（ViT）や拡散Transformerは、高解像度の画像や動画を長いトークンシーケンスに変換するためです。シーケンスを長くすると精度も向上しました。論文では、16,384ピクセルのシーケンスに平坦化した128 × 128画像で、2点が経路によってつながっているかを判定するタスクPath-Xにおいて、初めてチャンスレベルを上回ったTransformerが報告されました。