FlashAttentionとは？仕組み・メモリ削減・KV Cacheとの違い

FlashAttention（GPUのメモリ階層を意識してAttentionを計算する手法）は、Attentionの式やmodel weightを変えず、計算順序とデータ移動を組み替える技術です。巨大なAttention行列をGPUのHBMへ何度も書き戻さず、小さなblockに分けてオンチップメモリ上で処理します。

重要なのは、FlashAttentionがdense Attentionの演算量を線形にする技術ではないことです。token同士の組み合わせを扱う計算は系列長の二乗に応じて増えますが、N×Nの中間行列を丸ごと保持しないため、追加メモリとHBMへの読み書きを大きく抑えられます。

標準Attentionの巨大な中間行列を小さなtileへ分けて処理するFlashAttentionのアイキャッチ

3文要約

標準Attentionはscoreとsoftmax後の確率というN×N行列を作り、段階ごとにHBMへ保存するとデータ移動が大きくなります。
FlashAttentionはQ/K/Vをtileに分け、online softmaxで統計を更新しながらオンチップSRAM上で計算するため、Attentionの追加メモリを系列長に対して二乗から線形へ抑えます。
効果はtrainingやprefillで現れやすく、1-token decodeではKV Cacheの読出しが別のbottleneckになるため、実際のruntimeがどのbackendを選んだか確認して測定します。

代表的な論文

FlashAttentionの代表資料は、Tri DaoらのFlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenessです。論文はGPUの演算回数だけではなく、HBMとオンチップSRAMの間を何byte移動するかに注目し、tiling（行列を小さなblockへ分割する方法）によってexactなAttentionを計算しました。

後続のFlashAttention-2は、IO-aware（メモリ階層間のデータ移動を意識する設計）という基本原理を保ちながら、sequence方向の並列化、thread blockとwarp間の仕事分担、行列積以外の演算削減を進めています。

本記事はLLMシリーズの基礎編として、標準Attentionとのデータフローの違い、online softmaxの具体計算、training・prefill・decodeでの役割を扱います。原論文のIO complexityや実験結果を詳しく読みたい場合は、FlashAttention論文の詳しい解説も参照してください。

標準AttentionでN×N行列が生まれる理由

Scaled Dot-Product Attentionは、queryを\(Q\)、keyを\(K\)、valueを\(V\)、head dimensionを\(d\)として次の順で計算します。

S=\frac{QK^{\mathsf{T}}}{\sqrt{d}}

P=\operatorname{softmax}(S)

\[O=PV\]

系列長を\(N\)とすると、1 headの\(Q\)、\(K\)、\(V\)はそれぞれ\(N\times d\)です。各queryと各keyの組み合わせを計算するため、score \(S\)とattention確率\(P\)は\(N\times N\)になります。Self-Attentionの式自体を先に確認したい場合は、Self-AttentionとScaled Dot-Product Attentionの解説を参照してください。

各段階の中間値を実体化する実装を単純化すると、次の流れになります。

\(QK^{\mathsf{T}}\)を計算し、\(S\)をHBMへ書く
HBMから\(S\)を読み、softmaxを計算し、\(P\)をHBMへ書く
HBMから\(P\)と\(V\)を読み、\(O\)を計算する

GPU kernel（GPU上で実行される処理単位）の境界ごとに大きな中間行列をHBMへ置くと、計算器が速くてもデータの往復が待ち時間になります。

標準Attentionでscoreと確率のN×N中間行列がHBMへ書き戻される流れ

系列長が4倍になるとscoreは16倍になる

score行列1枚の要素数は\(N^2\)です。1要素をFP16の2 bytesとして、1 head分だけを概算します。

系列長\(N\)	scoreの要素数	FP16のscore 1枚
1,024	1,048,576	約2.1 MB
4,096	16,777,216	約33.6 MB
16,384	268,435,456	約536.9 MB

これはbatch数、head数、softmax後の\(P\)、gradient、その他のactivationを含まない10進MBの概算です。\(N\)を4,096から16,384へ4倍にすると、scoreだけで16倍になります。

系列長4096から16384への4倍化で1 headのFP16 score容量が33.6 MBから536.9 MBへ16倍になる比較

FLOPsだけでは速度を説明できない

GPU処理の所要時間は、FLOPs（浮動小数点演算回数）だけでは決まりません。演算に必要な値をどこから読み、結果をどこへ書くかも重要です。

GPUメモリを単純化すると、次の2階層に分けられます。

比較軸	HBM	オンチップSRAM
位置	GPU package上で演算器から離れた大容量memory	GPU chip内のshared memoryやregisterに近い領域
容量	大きい	小さい
データ転送	相対的に遅く、帯域が有限	相対的に速い
Attentionでの役割	Q/K/V、出力、通常の中間値を保持	計算中の小さなtileを保持

論文では高速なオンチップ領域をまとめてSRAMと表現しています。実際のkernelではshared memoryやregisterなどへ配置され、具体的な使い分けは実装とGPUに依存します。

大容量だが遠いHBMと小容量だが計算器に近いオンチップSRAMの間でtileを運ぶ図

標準AttentionとFlashAttentionの違いは、3つの軸を分けると理解しやすくなります。

比較軸	標準的なmaterialize実装	FlashAttention
数学的なAttention	dense Attention	同じdense Attention
pairwise scoreの演算	系列長に対して二乗	系列長に対して二乗の側面を残す
Attentionの追加メモリ	N×N中間行列を保持	出力と行ごとの統計を中心に線形
HBM IO	\(S\)や\(P\)を段階間で読み書き	tile内で融合し、中間行列の往復を回避

したがって「FlashAttentionでAttentionが\(O(N)\)になった」という説明は不正確です。線形へ抑えられるのは主に保存する追加メモリであり、denseな全query-key pairを扱う計算そのものではありません。

FlashAttentionは小さなtileの中で処理を完結させる

FlashAttentionは、巨大な\(S\)と\(P\)を一度に作る代わりに、\(Q\)、\(K\)、\(V\)をオンチップSRAMへ入る大きさのtileへ分けます。概念的な処理は次のとおりです。

\(Q\)のtileと\(K\)、\(V\)のtileをHBMからオンチップSRAMへ読む
tile内のscoreを計算する
causal maskなどを適用し、online softmaxの統計を更新する
確率tileと\(V\)を掛け、出力の部分和を更新する
次の\(K/V\) tileへ進み、最後に出力と小さな統計だけをHBMへ書く

score tileはオンチップで使った後に捨てられます。全体の\(S\)や\(P\)をHBMへmaterialize（中間結果を完全な配列として実体化）しない点が核心です。行列積、mask、softmax、valueとの積を一つの流れへ融合するため、kernel間の不要な読み書きも減らせます。

ただし、softmaxは行全体の最大値と分母を必要とします。まだ見ていないtileに、より大きなscoreがあるかもしれません。この問題を解くのがonline softmaxです。

Online softmaxを具体値で追う

1行のscoreが\([1,2,3,4]\)で、\([1,2]\)と\([3,4]\)の2 blockに分かれている例を考えます。softmaxを安定して計算するため、最大値\(m\)を引いた指数の和\(l\)を保持します。

block 1では、最大値と分母が次の値になります。

\[m_1=2\]

l_1=e^{1-2}+e^{2-2}=e^{-1}+1\approx1.3679

block 2では次のとおりです。

\[m_2=4\]

l_2=e^{3-4}+e^{4-4}=e^{-1}+1\approx1.3679

全体の最大値は\(m=4\)です。block 1は\(m_1=2\)を基準にしていたため、新しい基準4へ合わせて\(e^{2-4}\)倍します。統合後の分母は次の式になります。

l=e^{m_1-m}l_1+e^{m_2-m}l_2

l=e^{2-4}\times1.3679+e^{4-4}\times1.3679\approx1.5530

これは全scoreを一度に見た次の計算と一致します。

e^{1-4}+e^{2-4}+e^{3-4}+e^{4-4}\approx1.5530

出力の部分和も同じ考え方で古い最大値から新しい最大値へrescale（基準変更に合わせて倍率補正）します。これにより、過去のscore tileを保存せず、新しいtileを読むたびに最大値、分母、出力を正しく更新できます。

scoreを1、2と3、4の2 blockで処理し最大値4へ合わせてsoftmax分母1.5530を得る計算

exactはbitwise identicalという意味ではない

FlashAttentionがexactであるとは、sparse化や近似式によってAttentionの定義を変えないという意味です。数学的には通常のdense softmax Attentionと同じ出力を計算します。

一方、GPUの浮動小数点演算には丸め誤差があります。加算や乗算の順番を変えると最下位bit付近の結果が変わるため、異なるbackendの出力がbit単位で完全一致するとは限りません。PyTorchのscaled dot product attention documentも、融合backendによって数値出力が異なる可能性を説明しています。

表現	意味	FlashAttention
数学的にexact	dense Attentionを近似しない	該当する
bitwise identical	実装が違っても全bitが一致する	保証されない
決定論的	同条件の実行順やalgorithmを固定できる	runtime設定とbackendに依存

品質確認では、bitwise一致だけを要求するのではなく、用途に応じたtolerance、loss、生成品質、再現性要件を決めます。

backwardでは保存せずに再計算する

trainingのbackwardでは、gradientを求めるためにforward時の中間値が必要です。素朴な実装は\(S\)や\(P\)を保存するため、系列長が長いほどactivation memoryが大きくなります。

FlashAttentionは、出力\(O\)とsoftmaxの行ごとの統計など小さな情報を保存し、backward時に必要なscoreや確率をtile単位で再計算します。recomputation（保存量を減らす代わりに必要時に再計算する方法）によりFLOPsは増えますが、HBMへ巨大な中間値を書いて読み直す量を減らせます。

Attentionの中間行列を保存する方法と小さな統計だけを保存してbackwardでtileを再計算する方法の比較

「計算を増やしたのに速くなる」可能性があるのは、GPUで行列積の計算能力よりHBM IOが先に上限へ達していた場合です。FLOPsの最小化と実行時間の最小化は、常に同じではありません。

training・prefill・decodeで効果は同じではない

LLMの利用段階によって、queryの長さと主なbottleneckが変わります。

段階	query length	Attentionの特徴	FlashAttentionの主な役割
Training	複数token	forwardとbackwardを実行	N×N中間値の保存・IOを抑え、backwardで再計算
Prefill	prompt全体	多数のqueryをまとめて処理	大きなAttention中間値のmaterializeを避ける
1-token decode	通常1	新tokenのqueryから過去K/Vを参照	N×N中間値回避より、KV読出しやkernel効率が重要

autoregressive decode（1 tokenずつ出力する生成）では、過去tokenのK/VをKV Cacheへ保存します。batchとheadの次元を省けば、1-token decodeのscore shapeは概ね\(1\times N\)であり、trainingやprefillと同じN×N行列は作りません。この段階では、長くなったKV Cacheを読むmemory bandwidthが支配的になりやすくなります。

trainingとprefillでは複数queryのAttention、1-token decodeではKV Cache読出しが中心になる比較

FlashAttention系の最適化kernelがdecodeで役立たないという意味ではありません。query length、batch size、head dimension、mask、KV Cache layoutに応じて最適なkernelが変わるため、「FlashAttentionを有効にした」という名前だけで全段階の速度を説明しないことが大切です。

KV Cache・量子化・GQA/MQAとの違い

LLMの省メモリ・高速化技術は、減らす対象が異なります。

技術	何を変えるか	主に減らすもの	Attentionの式
FlashAttention	Attention内部の計算順序とデータ移動	N×N中間値とHBM IO	dense Attentionのまま
KV Cache	過去tokenのK/Vを保存して再利用	decode時のK/V再計算	変えない
量子化	weightやactivationなどのbit幅	保存容量と転送byte数	原則として構造は変えない
GQA/MQA	複数query headでK/V headを共有	KV Cache容量とdecode時のKV読出し	head構成を変える

量子化は数値表現を低bit化する技術です。FlashAttentionと排他的ではなく、runtimeとkernelが対応すれば組み合わせられます。次回扱うGQA/MQAはK/V head数そのものを減らすため、とくにdecode時のKV Cacheへ効く点が異なります。

PyTorch SDPAでFlash backendを確認する

PyTorchのscaled_dot_product_attentionは、利用可能な融合kernelとC++ math実装から入力条件に合うbackendを選びます。現行documentではCUDA向け候補としてFlashAttention-2、Memory-Efficient Attention、C++ math implementationが説明されていますが、API、対応GPU、dtype、head dimension、maskなどの条件はversionにより変わります。

通常は自動選択に任せます。動作確認時だけFlash backendへ限定すると、非対応の場合にwarningやerrorから理由を確認できます。

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

query = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
key = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
value = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)

# backendを固定する理由は、benchmark前にFlash kernelの利用可否を診断するためです。
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        dropout_p=0.0,
        is_causal=True,
    )

sdpa_kernelの使い方はPyTorchの公式documentで確認できます。直接flash-attn packageを導入する場合は、Dao-AILabの公式repositoryで対象versionのGPU、CUDAまたはROCm、dtype、head dimension、build要件を確認します。

PyTorch SDPAが入力条件に応じてFlash、Memory-Efficient、Math backendを選び利用不能理由を確認する流れ

Benchmark前のチェックリスト

使用中のframeworkとruntime versionを記録する
GPU、driver、CUDAまたはROCm versionを記録する
dtype、batch size、head数、head dimension、query/key lengthを固定する
causal mask、任意mask、dropout、training modeを揃える
warm-up後にprefillとdecodeを分けて測る
latencyだけでなくpeak memoryとthroughputも測る
profilerやbackend限定で、実際に選ばれたkernelを確認する
output差、loss、生成品質が許容範囲か確認する

入力shapeが小さい場合は、kernel launchやdispatchのoverheadが相対的に大きくなります。非対応条件では別backendへfallbackするため、すべての環境で同じ倍率の高速化が得られるわけではありません。論文やrepositoryのbenchmark値は、そのGPU、dtype、shape、versionの結果として読みます。

よくある質問

FlashAttentionは近似Attentionですか

いいえ。標準的なdense softmax Attentionと同じ数学的結果を計算するexactな手法です。ただし、浮動小数点演算の順序が変わるため、異なるbackendとのbitwise一致を意味しません。

Attentionの計算量はO(N)になりますか

なりません。保存する追加メモリは系列長に対して線形へ抑えられますが、dense Attentionでquery-keyの組み合わせを扱う演算には二乗の側面が残ります。

KV Cacheの代わりになりますか

なりません。FlashAttentionは一回のAttention演算内部のデータ移動を減らし、KV Cacheはdecodeで過去K/Vの再計算を避けます。両者は別の問題を解き、組み合わせて使えます。

有効にすれば必ず速くなりますか

いいえ。GPU、dtype、shape、mask、runtime version、fallbackの有無で変わります。training、prefill、decodeを分け、実際のbackendとpeak memory、latency、throughputを測定してください。

まとめ

FlashAttentionは、Attentionの数学を近似せず、GPU上の計算順序をメモリ階層に合わせる技術です。Q/K/Vをtileに分け、online softmaxで最大値・分母・出力を更新することで、N×Nのscoreと確率行列をHBMへ丸ごと保存しません。

理解の要点は、演算量、追加メモリ、HBM IOを混同しないことです。FlashAttentionは追加メモリとIOを減らしますが、dense Attentionのpairwise計算を線形にはしません。また1-token decodeではKV Cache読出しが別のbottleneckになるため、trainingやprefillの効果をそのまま当てはめず、実行backendを確認して測ります。

次回は、K/V headを共有してKV Cacheとdecode時の転送量を減らすGQA・MQAを扱います。

FlashAttentionとは？Attentionを高速化する技術