FlashAttention

Definition

FlashAttention (Dao, Fu, Ermon, Rudra, Ré, 2022; arXiv:2205.14135) is an IO-aware, exact attention kernel that computes standard scaled-dot-product attention without ever materializing the $N\times N$ score matrix in HBM. It fuses the whole matmul → softmax → matmul chain into one kernel by tiling $Q,K,V$ into SRAM and using online softmax, collapsing HBM traffic from $O(N^2)$ to $O(N)$. It is the canonical example of algorithm–hardware co-design.

How it works

Naive attention runs $S=Q{K}^{\top }/\sqrt{d}$, $P=\text{softmax}(S)$, $O=PV$ as separate kernels, round-tripping the $N\times N$ matrix through HBM ~4× (write $S$, read $S$, write $P$, read $P$) — $\Theta (N^2)$ traffic with tiny attached math, so arithmetic intensity collapses to $\sim d$, far below the ridge. The softmax is the smoking gun: an elementwise/reduction kernel that forces $S$ to be materialized. At $N=8192$ one BF16 head’s scores are 128 MiB; at $128\text{k}$, 32 GiB — the capacity wall that capped context.

FlashAttention tiles: for each query block $Q_i$, stream over K/V blocks, forming each score tile $S_{ij}$ only in SRAM, updating running softmax state $(m,\ell ,\mathbf{o})$ via online softmax, and discarding $S_{ij}$. Only $Q,K,V$ (in) and $O$ (out) cross HBM: $\Theta (Nd)$, linear in $N$. The paper’s bound is $\Theta (N^2{d}^2/M)$ accesses ($M$ = SRAM size) vs. $\Theta (N^2+Nd)$ — measured ~9× fewer HBM accesses, 2–4× wall-clock on GPT-2. The backward uses recomputation (store only per-row logsumexp $L_i=m_i+\log {\ell }_i$, $O(N)$; recompute $P_{ij}=\exp (S_{ij}-L_i)$ on the fly), trading FLOPs for the $N^2$ memory.

The FA-1 → FA-4 lineage — same math, each generation retargeting the then-binding hardware constraint:

• FA-1 (2205.14135): kill HBM traffic (IO-awareness). ~30–50% A100 peak.
• FA-2 (2307.08691): kill idle warps/blocks — defer output rescaling, parallelize over seq-len, split Q (not K) across warps. ~2×, 50–73% A100.
• FA-3 (2407.08608): Hopper async — warp-specialized TMA/wgmma producer-consumer overlap, ping-pong matmul/softmax, FP8. 1.5–2× over FA-2, ~740 TFLOP/s FP16, ~1.2 PFLOP/s FP8.
• FA-4 (arXiv:2603.05451, 2026): Blackwell — software-emulated exp on FMA units (SFUs didn’t scale), conditional rescaling, TMEM for backward, 2-CTA MMA, CuTe-DSL. ~1.6 PFLOP/s BF16 (~71% util), 1.3× cuDNN.

Why it matters

Attention is where LLMs spend a growing fraction of time as context grows; FlashAttention removed the quadratic memory wall without approximation, making million-token context tractable and turning a memory-bound op near-compute-bound. It is exposed via F.scaled_dot_product_attention and backs PyTorch FlexAttention — the default attention kernel in modern training/inference.

Taught in

• 06-memory-wall-and-flashattention — §5–§9 the full derivation, IO analysis, and the FA-1→FA-4 lineage.
• 07-triton-and-modern-kernels — the fused-attention Triton kernel.

See also

• online-softmax — the streaming recurrence at its heart
• io-aware-algorithms — the design principle
• kernel-fusion — attention as extreme vertical fusion
• shared-memory-tiling — tiling $Q,K,V$ into SRAM
• recomputation — the backward pass
• memory-wall — the wall it tears down
• arithmetic-intensity · gpu-memory-hierarchy · gemm
• gpu-systems-for-llms