一文要約: FlashAttentionは近似Attentionではなく、 $N^2$ 回のスコア計算も消しません。Tile分割とOnline Softmaxにより、完全な $N \times N$ 中間行列をデバイスメモリへ何度も書き戻すことを避けます。

素朴なAttentionはスコアと確率の完全な行列をHBMへ書き、FlashAttentionはチップ上のtileを出力へ累積する

21.1 ボトルネックは「計算が多い」だけではない

1つのAttention headについて、 $Q,K,V \in \mathbb{R}^{N \times d}$ とします。

O = \operatorname{softmax}\!\left(\frac{QK^T}{\sqrt d} + B_{\mathrm{mask}}\right)V

主な計算量は依然として $\Theta(N^2d)$ です。素朴な実装は $N \times N$ のスコア行列を実体化し、Softmax後に同じ大きさの確率行列も実体化する場合があります。

$N=4096$ なら、1 head、1 sampleのFP16スコア行列は：

4096^2 \times 2\ \text{bytes} = 32\ \text{MiB}

素朴な経路がbatch size 8、32 headsでこの行列を保持すれば8 GiBです。Mask、確率、Dropout、逆伝播で必要な状態はまだ含んでいません。これは中間行列を実体化する素朴な経路の計算であり、現在のPyTorch Attentionが常に8 GiBを確保するという意味ではありません。

もう半分の問題はデータ転送です。

HBMからQ/Kを読む → S=QKᵀを計算 → SをHBMへ書く
HBMからSを読む   → Mask/Softmax  → PをHBMへ書く
HBMからP/Vを読む → O=PVを計算  → OをHBMへ書く

Tensor Coreが行列を速く掛けても、データが来なければ演算器は待ちます。FlashAttentionはこのI/O経路を改善します。

21.2 Tiling：大きな紙を広げず、小さな作業台で完結する

元の本棚と机の例で考えましょう。HBMは容量の大きな本棚、registers/shared memory/L1は手元の狭い机です。机に載る分だけ取り出し、その計算を終えるか、小さな進捗状態にまとめてから次の束を取りに行きます。

FlashAttentionは $Q_i$ tileと $K_j,V_j$ のtileをチップ上へ持ち込みます。

Sᵢⱼ = QᵢKⱼᵀ / √d + mask
                 ↓
         Online Softmax state
                 ↓
           出力 Oᵢへ累積

Q tileとK/V tileからチップ上で一時スコアを計算し、完全なN×N行列はHBMへ書かない

ハードウェアの数値は混同しやすい部分です。A100全体のL2は40 MBですが、各SMのL1/texture/shared-memory複合データキャッシュは最大192 KBで、shared-memory carveoutはkernelごとに構成されます。これらは自由に共有できる「20 MBのSRAM pool」ではありません。A100 40GBのHBM2帯域幅は約1550 GB/s、80GB版は約2039 GB/sです。ハードウェアごとに値が異なるため、「SRAMは常に20倍速い」という表で永久にまとめることはできません。

原論文のAlgorithm 1で使う理想化したtile sizeは：

B_c=\left\lceil\frac{M}{4d}\right\rceil,\qquad B_r=\min\!\left(\left\lceil\frac{M}{4d}\right\rceil,d\right)

ここで $d$ はhead dimension、 $M$ は論文の抽象機械でチップ上に収まるスカラ要素数です。dtypeを無視して192 KBをそのまま代入する式ではありません。実際のkernelではdtype、register pressure、shared-memory carveout、alignment、warpの分担、GPU世代も考慮します。この式だけから「tileは必ず64」とは決まりません。

21.3 Online Softmax：1 tileずつ見ても行全体を保つ

Softmaxの分母は行全体にわたります。tileごとに独立してSoftmaxをかけ、結果を連結するのは誤りです。1つのquery行ごとに、これまで処理したkey blocksに対して次の3つを保ちます。

$m$ : 現在までの最大スコア；
$\ell$ : その最大値を基準にした指数和；
$a$ : まだ正規化していない、Valueの重み付き和。

新しいスコアblock $s$ と $V_j$ が来たら：

m'=\max(m,\max(s)),\qquad \alpha=e^{m-m'}

p=e^{s-m'},\qquad \ell'=\alpha\ell+\sum p,\qquad a'=\alpha a+pV_j

最後に $o=a/\ell$ を返します。後のtileでより大きなスコアが出ると、 $\alpha$ がこれまでの累積量を新しい数値尺度へ変換します。

Online Softmaxは現在までの最大値、分母、未正規化の出力累積量でスコアblockを統合する

次はmaskを含めないPyTorchの教材版です。漸化式の確認用であり、高速なkernelではありません。

import math
import torch


def online_attention(q, k, v, block_size):
    # q: [Nq, d], k: [Nk, d], v: [Nk, dv]
    scale = 1.0 / math.sqrt(q.size(-1))
    rows = q.size(0)
    m = torch.full((rows,), -torch.inf, device=q.device, dtype=q.dtype)
    denominator = torch.zeros_like(m)
    accumulator = torch.zeros(
        rows, v.size(-1), device=q.device, dtype=q.dtype
    )

    for start in range(0, k.size(0), block_size):
        k_block = k[start:start + block_size]
        v_block = v[start:start + block_size]
        scores = q @ k_block.transpose(-2, -1) * scale

        block_max = scores.max(dim=-1).values
        new_m = torch.maximum(m, block_max)
        correction = torch.exp(m - new_m)
        probabilities = torch.exp(scores - new_m[:, None])

        denominator = (
            correction * denominator + probabilities.sum(dim=-1)
        )
        accumulator = (
            correction[:, None] * accumulator + probabilities @ v_block
        )
        m = new_m

    return accumulator / denominator[:, None]

同じ $Q,K,V$ に対して、softmax(Q @ K.T / sqrt(d)) @ V と浮動小数点誤差の範囲内で一致します。「Exact attention」とは、疎やlow-rankの近似を入れないという意味です。有限精度で演算順序が変わる以上、bit-for-bitの一致まで約束する言葉ではありません。

21.4 計算量、追加メモリ、HBM I/Oを分ける

FlashAttentionは二次のスコア計算を保ちつつ、追加メモリを二次から一次へ減らしHBM転送を抑える

問い	中間行列を実体化する素朴なAttention	FlashAttention
主な計算量	$\Theta(N^2d)$	$\Theta(N^2d)$
追加の中間メモリ（dを固定）	$\Theta(N^2)$	$O(N)$
原論文の抽象機械におけるHBMアクセス	$\Theta(Nd+N^2)$	$\Theta(N^2d^2/M)$

第3行には条件があります。 $M$ はチップ上の要素数を表し、論文は $d \le M \le Nd$ の範囲を解析しています。これを $O(N^2d/M)$ と安易に省略することはできません。追加メモリが二次から一次に下がっても、実際のピークメモリがちょうどN分の1になるわけではありません。Q/K/V、出力、重み、optimizer state、allocator、他の層は残ります。

逆伝播では完全な確率行列を保存せず、保存済みの統計量から局所スコアのtileを再計算します。FLOPsは少し増えますが、逆伝播全体が必ず遅くなるとは限りません。HBM転送を減らすことで、kernelや学習全体が速くなることもあります。

21.5 FA1からFA4：一つの核心、ハードウェアごとのkernel

2026年7月時点のFlashAttention 1から4のアルゴリズムと対象ハードウェア

FA1 (2022): I/O-awareなtilingとOnline Softmaxを確立し、完全な行列を実体化しない実装を示しました。
FA2 (2023): 行列積以外のFLOPsを減らし、thread blockとwarpの作業分割を改善しました。論文のA100 workloadではFA1の約2倍、理論ピークの50–73%、GPT風の学習で最大225 TFLOPs/s/GPU（model FLOPs utilization 72%）と報告されています。これは論文の条件下での測定であり、任意のモデルに対する保証値ではありません。
FA3 (2024；公式repositoryでも現在はHopper beta): H100/H800のTMA、WGMMA、FP8経路を活用します。
FA4 (2026): Blackwell B200/GB200の不対称なハードウェア拡張を対象に、CuTe DSLでpipelineを協調設計します。2026-07時点で公式repositoryは flash-attn-4 のpre-releaseを提供しています。head dimension、dtype、逆伝播の対応範囲は現在のreleaseで確認し、FA2から推測しません。

バージョン番号が大きいほど、すべてのGPUで速いわけではありません。FA3はHopper、FA4はBlackwellを対象とします。Ampere/Adaでは依然としてFA2を使う場合があります。

21.6 実装：まずframeworkに選ばせる

PyTorchの scaled_dot_product_attention はdevice、dtype、shape、maskなどの制約から、利用可能なbackendを選びます。第18章の配列ではQ/K/Vを [batch, heads, time, head_dim] にします。

import torch.nn.functional as F


def attention(q, k, v, dropout_p, training):
    return F.scaled_dot_product_attention(
        q,
        k,
        v,
        attn_mask=None,
        dropout_p=dropout_p if training else 0.0,
        is_causal=True,
    )

scaled_dot_product_attention は渡された dropout_p に従って常にDropoutを適用します。外側のmoduleの training は自動で読みません。評価時は明示的に 0.0 を渡します。

ある入力がFlash backendを本当に使えるか診断するときは、backendを限定できます。

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output = F.scaled_dot_product_attention(
        q, k, v, dropout_p=0.0, is_causal=True
    )

Fused kernelが対応できない場合、PyTorchはその理由を報告します。本番コードでは、単に「Flashを使っている」と言いたいために、すべてのfallbackを無効にする必要はありません。

公式の Dao-AILab/flash-attention packageを直接使う道もあります。対応ハードウェアを「新しいNVIDIA GPUだけ」と書くのも、現在は不正確です。2026-07時点でFA2のCUDA経路はAmpere/Ada/Hopperをサポートし、公式repositoryはROCm向けのCK/Triton AMD backendsも提供しています。インストール前に、現在のPyTorch、CUDA/ROCm、GPU、dtype、head dimensionの対応表を確認します。

21.7 ベンチマークを正直に作る

少なくとも次を固定して報告します。

GPUの型番とpower mode、PyTorch、CUDA/ROCm、kernel version；
dtype、batch、heads、query length、key length、head dimension;
causal/noncausal、mask、Dropout、forwardだけかforward+backwardか；
warmup、計時方法、median/percentile、kernelのmicrobenchmarkかモデル全体のwall timeか；
ピークメモリ、出力誤差、勾配誤差。

短いsequenceや小さなbatchでは、kernel launchや配列変換のoverheadが利得を消すことがあります。Attention kernelが2倍速いことは、モデル全体や1 tokenあたりのlatencyが2倍速いことを意味しません。

FlashAttentionとKV Cacheも、別の問題を解きます。前者は1回のAttention計算の中でI/Oを減らし、後者は自己回帰decodeのstep間で過去のK/Vを再計算しないためのものです。

21.8 本章のまとめ

FlashAttentionは標準Attentionの数学的な結果を保ち、主な貢献はI/O-awareな実装にあります。
主な $\Theta(N^2d)$ 計算は残りますが、完全な $N \times N$ 中間行列は実体化しません。
Online Softmaxは分母とValueの累積量の両方を再スケーリングし、tile同士を統合できるようにします。
$d$ はhead dimension、 $M$ は抽象的なチップ上容量です。モデル幅やbyte数を無条件に代入できません。
FA1–FA4のkernelと対応ハードウェアは変化中なので、速度の数値にはworkloadとversionの境界が必要です。
PyTorch SDPAはbackendを自動的に選択できますが、評価時に dropout_p=0.0 を渡すのは呼び出し側の責任です。

章末チェックリスト

計算量、追加メモリの計算量、HBM I/Oを分けられる
$m,\ell,a$ でOnline Softmax tile mergeを導出できる
「Exact」がbit-for-bitの約束ではない理由を説明できる
PyTorch SDPAを正しく呼び、実際のbackendを確認できる
kernelの速度をモデル全体の速度に見せないベンチマークを設計できる

一次資料

次章予告

FlashAttentionは現在のAttention計算を軽くしますが、第20章の自己回帰ループは、まだstepごとに過去のKeys/Valuesを再計算します。第22章で、このstep間の冗長性を KV Cache で取り除きます。