FlashAttention：メモリ効率的なアテンション計算

アテンションの真のボトルネックは計算ではない

Self-Attention の計算 で見たとおり、アテンションの中身は次の式です。

Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V

系列長を n、ヘッド次元を d とすると、QKᵀ は n × n の スコア行列 S になります。これを n が大きいほど巨大化します。素朴な実装は、この S を作り、softmax をかけて P（同じく n × n）にし、P·V を計算します。FLOPs は確かに O(n²·d) ですが、実機の遅さの主因は FLOPs ではありません。本当のボトルネックは、この n × n の中間行列を GPU の主記憶（HBM）に書き出し、また読み戻す メモリ往復にあります。

GPU メモリ階層：SRAM と HBM の決定的な差

最適化の理由は GPU のメモリ階層にあります。容量と帯域が大きく違う2階層を区別するのが出発点です。

ポイントは、SRAM が HBM より おおよそ一桁速い一方、容量は HBM の数万分の一 しかないことです。n × n のスコア行列は系列長が数千ともなれば SRAM に収まりません。だから標準実装は S と P を HBM へ書かざるを得ず、softmax 行ごとに「書く・読む」を繰り返します。この HBM 入出力が律速する 状態を memory-bound と呼びます。アテンションはまさにこれに当たり、演算ユニットは HBM 待ちで遊んでいます。

オンライン softmax：全列を見ずに正規化する

中間行列を HBM に書かないには、K・V を縦に小ブロックへ割り、ブロックを順に SRAM へ載せて処理する必要があります。ところが素朴な softmax は分母に 全列の指数和 が要るため、本来は全 K を見終えるまで正規化を確定できません。これを解くのが オンライン softmax です。

数値安定化のため softmax は通常、行の最大値 m を引いてから指数を取ります。ブロックを順に見ていく過程で、それまでの最大値 m と指数和 l、そして部分的な出力 o を持ち回り、新ブロックごとに 補正しながら更新 します。

各クエリ行について、ブロックを順に処理:
  m_new = max(m, ブロック内スコアの最大値)
  α     = exp(m - m_new)          ← 旧最大値が更新された分の補正係数
  l     = l*α + Σ exp(s_j - m_new)        ← 指数和を新基準にスケールし直して加算
  o     = o*α + Σ exp(s_j - m_new) * v_j  ← 出力も同じ係数で補正して加算
  m     = m_new
最後に: 出力 = o / l                ← 全ブロック走査後に一度だけ正規化

肝は補正係数 α = exp(m - m_new) です。後から大きいスコアが出て基準の最大値が更新されたとき、過去に積み上げた l と o を α 倍してスケールを揃え直します。この再スケールにより、ブロックを分割して逐次処理しても、全列を一度に見て softmax した結果と厳密に一致 します。FlashAttention は近似ではなく、出力が数学的に等価である点が決定的です。

タイル化：HBM 往復を消すループ構造

オンライン softmax を使えば、Q・K・V をブロック（タイル）に分け、二重ループで処理できます。外側を Q のブロック、内側を K・V のブロックで回します。

for each Q ブロック Qi (SRAM に載せる):
    o, l, m を初期化
    for each K,V ブロック Kj, Vj (SRAM に載せる):
        Sij = Qi · Kjᵀ / √d_k        ← SRAM 上で計算。HBM に書かない
        オンライン softmax で o, l, m を更新（上の手順）
    Oi = o / l
    Oi だけを HBM に書き戻す            ← 出力 n×d のみ。中間 n×n は書かない

このループでは、SRAM に載るのは小さなタイルだけで、巨大な S や P は 一度も HBM に実体化されません。HBM へ書き戻すのは最終出力 O（n × d、入力と同サイズ）だけです。結果としてアテンションの 追加メモリは O(n²) から O(n) へ下がり、HBM 入出力量も大幅に減ります。memory-bound だった処理が HBM 待ちから解放され、実測でおおむね 2〜4倍 の高速化が報告されています（モデル・系列長・ハードに依存）。

逆伝播：保存しない代わりに再計算する

学習時の逆伝播は一見やっかいです。勾配計算には n × n の P が要りますが、FlashAttention はそれを保存していません。ここで取る戦略が 再計算（recomputation） です。順伝播では各行の正規化に使った m と l（各 O(n)）だけを保存しておき、逆伝播で必要になったタイルの S・P を その場で作り直します。

一見すると無駄な再計算ですが、演算をやり直すコストより、巨大な P を HBM に置いて読み書きするコストの方が高い のが現代 GPU の実情です。だから「保存せず再計算」が正味で得になります。これは 自動微分 の文脈で言う勾配チェックポイントと同じ発想を、アテンション内部に特化して埋め込んだものと捉えられます。

KV キャッシュとの関係：別レイヤーの最適化

混同しやすいのが KV キャッシュ との違いです。両者は 別レイヤーの最適化 であり、競合せず併用します。

FlashAttention は softmax の計算過程 を効率化し、KV キャッシュは K/V の再計算 を省く。前者は段取りの最適化、後者は結果の使い回しで、レイヤーが違うため両立します。実際の LLM サービングでは、PagedAttention・連続バッチングと並んで FlashAttention が標準的に組み込まれています。

まとめ：データ移動で考えるという発想

FlashAttention が示したのは、同じ数式・同じ結果でも、メモリ階層をどう使うかで実速度は桁で変わる という事実です。

この「IO-aware」――処理を演算量ではなくデータ移動量で捉える視点は、アテンションに限らず GPU 上の高速化全般に通じます。アテンション本体の数理は Self-Attention の計算 と Multi-Head Attention で、推論時のメモリ設計は KV キャッシュ で押さえると、長文 LLM の速度とメモリの全体像が線でつながります。