FlashAttentionを1から — 計算を増やして速くする逆説

FlashAttentionは逆伝播でわざと同じ計算をやり直します。演算回数は増えるのに速い——GPUでは掛け算よりデータの運搬のほうが高いからです。HBMとSRAMの往復を数えるところから、タイリングを阻むsoftmaxの壁、それを壊すオンラインsoftmaxの漸化式、そして現場でバックエンドが静かにフォールバックする落とし穴までを前提知識ゼロから追います。

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

最適化なのに、計算が増える

高速化と聞けば、ふつうは無駄な計算を削る話を想像します。FlashAttentionは逆です。逆伝播（誤差を戻す計算）で、順伝播とほぼ同じ計算をもう一度やり直します。演算回数は明確に増える。それでも素朴な実装より速い。

種明かしは一行で済みます。いまのGPUでは、掛け算をするより、数をメモリから運んでくるほうが高い。 だから「運ぶ回数が減るなら、計算はやり直したほうが安い」という交換レートが成立します。しかもFlashAttentionは近似ではありません。間引きも疎化もせず、素朴な実装と数学的に同じ答えを出します。変えたのは答えではなく、答えにたどり着く経路です。

比喩: 採点表を紙に書かない

資料が地下の倉庫にあり、机はとても狭い図書館を想像してください。倉庫との往復は遅く、机の上の作業は一瞬で終わります。

仕事は「1000人の応募者を1000項目で採点し、点数の割合に応じて意見を混ぜる」こと。素朴なやり方では、まず100万マスの採点表を作り、机に載らないので倉庫にしまい、行ごとの合計を出すために出し直し、割合に直してまたしまい、意見を混ぜるのに三度目を出してきます。

FlashAttentionは採点表を紙に書きません。応募者を50人ずつの束に分け、束ごとに机の上で採点し、答えの途中経過だけを机に置いたまま更新し続ける。最後の束を終えたとき、机には完成した答えが残っています。

倉庫がGPUのHBM（大容量・低速な外部メモリ）、机がSRAM（小容量・高速なオンチップメモリ）です。落差が実際に何桁あるかはGPUのメモリ階層で扱っています。

おさらい: 途中だけ N×N に膨らむ

Attention機構を1から理解するで見たとおり、自己注意は3行で書けます。系列長（トークン数）を $N$、ヘッド次元を $d$ とします。

$Q, K, V$ はどれも $N \times d$ の行列で、各行が1トークンの「質問」「名札」「発言内容」です。$S$ は全員と全員を突き合わせた $N \times N$ の点数表、$P$ はそれを行ごとに合計1の割合へ直したもの、$O$ が出力（$N \times d$）。つまり入力も出力も $N \times d$ なのに、途中だけ $N \times N$ に膨らむ。ここがすべての元凶です。

遅いのは計算ではなく、往復である

数字にすると効きます。$N = 4096$、$d = 128$ なら、$Q$ は約52万要素なのに $S$ は約1678万要素——32倍です。系列長を4倍にすれば $Q$ は4倍にしかならないのに $S$ は16倍になり、比は128倍へ開きます。

素朴な実装は、この巨大な表をHBMに何度も往復させます。$S$ を書き出す、softmaxのために読み直す、$P$ を書き出す、$V$ との積でまた読む。運ぶバイト数は $N^2$ に比例し、演算量は $N^2 d$ に比例するので、1バイトあたりの演算回数は $d$ 程度。完全に帯域律速の領域です。

しかも $S$ も $P$ も、最後には捨てられる中間結果です。最終出力より大きな中間結果を、わざわざ倉庫に往復させている。ここに手を入れられる、というのが出発点でした。

タイリングを阻む、softmaxという壁

「大きな行列を小さなブロックに切り、速いメモリに載る分だけ処理する」——タイリングは行列積では昔からの定石です。ではなぜattentionでやっていなかったのか。

softmaxが邪魔をするからです。

$$\mathrm{softmax}(s)_j = \frac{e^{s_j - m}}{\sum_{k} e^{s_k - m}}, \qquad m = \max_k s_k$$

$s_j$ は点数表の1マス、$m$ はその行の最大値です。最大値を引いてから指数を取るのは $e^x$ が桁あふれしないための定番の安定化で、分子と分母に同じ定数がかかるだけなので答えは変わりません。

問題は $m$ と分母の合計です。どちらも行を最後まで見ないと確定しません。ブロックを1つ処理した時点では、後からもっと大きな点数が現れるかもしれない。だから「表を出し切ってから正規化する」実装になり、$N \times N$ を持たざるを得なかった——長らくそう思われていました。

オンラインsoftmax: 途中経過を、後から縮める

この壁を壊すのがオンラインsoftmaxです。発想は単純で、「最大値が更新されたら、それまでに足した分をまとめて縮める」。

行を左から右へブロック単位で見ていきます。手元に持つのは3つだけです。

• $m$: いままで見た中の最大点数
• $\ell$: いままでの指数の合計（分母の途中経過）
• $\tilde{O}$: いままでの重み付き和の途中経過（$d$ 次元、まだ割っていない）

新しいブロックから局所的な最大値 $\tilde{m}$、指数和 $\tilde{\ell}$、重み付き和 $\tilde{P}\tilde{V}$ を求めたら、次の3本で更新します。

言葉にすると、新しい最大値を決め、古い蓄積には $e^{m - m'}$ を、新しいブロックには $e^{\tilde{m} - m'}$ を掛けて同じ土俵に揃えてから足す。最大値が更新されなければ $e^{m-m'} = 1$ で、古い分はそのままです。

核心は、縮小率 $e^{m - m'}$ がスカラー1つだという点です。過去に足し込んだ何千項をひとつずつ直す必要はなく、蓄積したベクトル全体に1回掛ければ済む。だから「後で最大値が更新される」ことを恐れずに前から足していけます。全ブロックを見終わったら、最後に一度だけ $O = \tilde{O}/\ell$ と割る。得られる $O$ は素朴な実装と数学的に等しく、近似は一切入りません。

この漸化式自体は2018年のオンラインsoftmaxの論文が示したもので、FlashAttentionはそれを $V$ との積まで巻き込み、GPUのメモリ階層に合わせて設計し直しました。

擬似コードで書くと

import numpy as np

def flash_attention(Q, K, V, Bc=128):        # Bc = 机に載るブロック幅
    N, d = Q.shape
    O = np.zeros((N, d))                      # 重み付き和の途中経過（未正規化）
    m = np.full(N, -np.inf)                   # 行ごとの最大値
    l = np.zeros(N)                           # 行ごとの指数和

    for j in range(0, N, Bc):                 # K/V を1ブロックずつ「机」に運ぶ
        Kj, Vj = K[j:j+Bc], V[j:j+Bc]
        S = Q @ Kj.T / np.sqrt(d)             # ここは N×Bc。N×N は作らない
        m_new = np.maximum(m, S.max(axis=1))
        alpha = np.exp(m - m_new)             # 古い蓄積を縮める率
        P = np.exp(S - m_new[:, None])
        l = alpha * l + P.sum(axis=1)
        O = alpha[:, None] * O + P @ Vj       # まとめて縮めて、足す
        m = m_new

    return O / l[:, None]                     # 最後に一度だけ割る

N × N の配列がどこにも現れないことを確かめてください。生きているのは N × Bc のブロックと、行ごとのスカラー2本だけです。実物のCUDAカーネルでは $Q$ 側もブロックに切り、ブロックの組を1つのSMに割り当てて $S$ をSRAMから出しません。

HBMの往復はどれだけ減るのか

論文はこれをIO計算量として定式化しています。素朴な実装のHBMアクセスは $\Theta(Nd + N^2)$、FlashAttentionは $\Theta(N^2 d^2 M^{-1})$。$M$ はSRAMの大きさ（要素数）です。

つまりどちらも $N^2$ に比例するが、FlashAttentionには $d^2/M$ という小さな係数が掛かる。$d$ が64〜128、$M$ が100KB相当という現実的な設定では $d^2$ は $M$ よりずっと小さく、往復が桁で減ります。オーダーが同じでも係数で桁が変わる——最適化の効き方としてはむしろ典型です。

逆伝播で、わざと計算し直す

冒頭の逆説の正体はここです。逆伝播には $P$ が要ります。素朴な実装は順伝播でそれをHBMに保存し、逆伝播で読み直します。FlashAttentionは保存しません。保存するのは出力 $O$ と、行ごとの $m$ と $\ell$ だけ——$N \times N$ ではなく $N$ に比例する量です。

逆伝播では $Q, K, V$ と $m, \ell$ から、必要なブロックの $P$ をSRAM上でその場で作り直します。$S$ を計算する分の演算がまるまる増える。それでも全体は速い。増えたのは机の上の作業で、減ったのは倉庫との往復だからです。帯域律速の世界では、FLOPsは通貨ではありません。 通貨はバイトです。

論文が報告している値としては、GPT-2（系列長1K）の学習で約3倍、BERT-large（系列長512）ではMLPerf 1.1の学習速度記録に対して実時間で15%短縮、といったものが挙げられています。効き方は系列長・ヘッド次元・GPU世代に強く依存するので、自分のワークロードで測るのが前提です。

その後の版と、推論での事情

FlashAttention-2（2023）は行列積以外の演算を減らし、並列化の軸を系列長方向にも広げ、ワープ間の共有メモリ経由のやり取りを削りました。FlashAttention-3（2024）はHopper世代の非同期性——転送と行列積を重ねて走らせる仕組み——を使い、FP8のような低精度でも精度を保つ工夫を入れています。

推論では事情が一段変わります。生成中（デコード）のクエリは1トークン分しかないので、系列長方向で並列化する前提が崩れ、コアが遊びます。そこで $K, V$ 側を分割して部分結果を並列に求め、最後に結合する方式が使われます。この結合に使うのが、まさに上の3本の式です——部分ごとの $m, \ell, \tilde{O}$ を持ち寄って揃えて足せば、分割しても答えは変わりません。

なお、ここで減るのはスコア行列のメモリであって、KVキャッシュではありません。長文でメモリが足りなくなる主因はたいていKVキャッシュ側で、見積もり方はKVキャッシュを1から理解するにあります。この混同は現場で本当によく起きます。

現場ではこう使う

誰が、いつ触るか。 LLMの学習基盤・推論基盤のエンジニア、長文入力を売りにするプロダクトの担当者、そして独自のマスク形状や位置バイアスを実装する研究寄りのエンジニアです。前2者は「使う側」、最後は「既存カーネルが使えなくなる側」で、事故の種類が違います。

実際に触る名前。 PyTorchなら torch.nn.functional.scaled_dot_product_attention（SDPA）が入口で、内部にMATH／FLASH_ATTENTION／EFFICIENT_ATTENTION／CUDNN_ATTENTIONといったバックエンドを持ちます。どれが選ばれたかを固定・検証するには torch.nn.attention.sdpa_kernel を使います。より直接には flash-attn の flash_attn_func、可変長バッチ向けの flash_attn_varlen_func。推論スタックならvLLMやTensorRT-LLMが内部で呼んでいます。

知らないと事故になる落とし穴。

• マスクを渡した瞬間、黙ってフォールバックする。 任意の加算マスク（attn_mask にテンソルを渡す形）はflashバックエンドの対象外になりやすく、MATHへ落ちます。例外は出ません——$N \times N$ を実体化する実装に戻るので、速度が数倍落ちるかOOMするかという形で初めて気づきます。因果マスクは is_causal=True で表現し、sdpa_kernel でバックエンドを明示して、外れたら失敗させること。
• dtypeと形の制約。 flashカーネルはfp16／bf16が前提でfp32では動きません。ヘッド次元にも上限があり、版とGPU世代で変わります。fp32でデバッグしていると「本番だけ速い」という不可解な差になります。
• ビット一致しないのはバグではない。 足し算の順序が変わるので、素朴な実装と出力が完全一致することは期待できません。厳しい許容誤差の回帰テストは、正しい移行なのに落ちます。相対誤差での比較に直すのが筋です。
• 短い系列で測ると効果が見えない。 往復そのものが小さいので、系列長512程度のベンチでは差が出ません。本番の系列長・バッチ・ヘッド数で測ること。
• パディングしたまま流さない。 長さの違う文を最長に合わせて詰めると、その分の計算と往復をまるごと払います。可変長API（cu_seqlens で区切りを渡す形）で消えます。

設計レビューで問われること。 「FLOPsが増えるのになぜ速いのか」——帯域律速なので通貨はバイトであり、SRAM上の再計算はHBM往復より安いから。「順伝播で保存するものは何か」——出力 $O$ と行ごとの $m, \ell$、合わせて $N$ に比例する量だけ。この2問に自分の言葉で答えられれば、「速いらしいから使う」から「なぜ効くか分かって使う」に移れます。

まとめ

• FlashAttentionは近似ではない。素朴な実装と同じ答えを、$N \times N$ の表を作らずに出す
• タイリングを阻んでいたsoftmaxの壁は、オンラインsoftmax——古い蓄積をスカラー1つで縮めてから足す漸化式——で壊せる
• 逆伝播では $P$ を保存せず作り直す。演算は増えるが、帯域律速の世界ではバイトのほうが高いので勝つ
• 現場の事故は性能ではなく静かなフォールバックから起きる。バックエンドを固定し、本番サイズで測る