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$ を持たざるを得なかった——長らくそう思われていました。