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

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

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

3文要約

1. 標準Attentionはscoreとsoftmax後の確率というN×N行列を作り、段階ごとにHBMへ保存するとデータ移動が大きくなります。
2. FlashAttentionはQ/K/Vをtileに分け、online softmaxで統計を更新しながらオンチップSRAM上で計算するため、Attentionの追加メモリを系列長に対して二乗から線形へ抑えます。
3. 効果は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\)として次の順で計算します。

系列長を\(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の解説を参照してください。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1. \(Q\)のtileと\(K\)、\(V\)のtileをHBMからオンチップSRAMへ読む
2. tile内のscoreを計算する
3. causal maskなどを適用し、online softmaxの統計を更新する
4. 確率tileと\(V\)を掛け、出力の部分和を更新する
5. 次の\(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では、最大値と分母が次の値になります。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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が支配的になりやすくなります。

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

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

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

量子化は数値表現を低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要件を確認します。

Benchmark前のチェックリスト

1. 使用中のframeworkとruntime versionを記録する
2. GPU、driver、CUDAまたはROCm versionを記録する
3. dtype、batch size、head数、head dimension、query/key lengthを固定する
4. causal mask、任意mask、dropout、training modeを揃える
5. warm-up後にprefillとdecodeを分けて測る
6. latencyだけでなくpeak memoryとthroughputも測る
7. profilerやbackend限定で、実際に選ばれたkernelを確認する
8. 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を扱います。

参考文献

• Attention Is All You Need
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• PyTorch scaled_dot_product_attention
• PyTorch sdpa_kernel
• Dao-AILab/flash-attention