Why is attention memory-bound, and how does FlashAttention fix it?

Standard attention is slow not because of its arithmetic but because of its memory traffic. To compute softmax(QKᵀ/√d)·V it first builds the full N×N score matrix S, writes it to the GPU's high-bandwidth memory (HBM), reads it back to apply softmax, writes the N×N probability matrix P, and reads that back to multiply by V — the FlashAttention paper notes that "standard attention implementations materialize the matrices S and P to HBM, which takes O(N²) memory." Since those reads and writes dominate, and "most operations in Transformers are bottlenecked by memory accesses," the kernel spends its time waiting on memory while the tensor cores idle. FlashAttention fixes this by never writing S or P to HBM at all: it processes attention in small tiles that fit in on-chip SRAM, so the O(N²) intermediate lives only in fast memory and never round-trips to HBM.

What is online softmax and how does tiling keep it exact?

Online softmax is the trick that lets you compute a softmax over a row of scores without ever having the whole row in memory at once. The problem: a numerically safe softmax subtracts the row maximum before exponentiating, but if you only see the scores a tile at a time you do not know the true maximum yet. The fix is to carry two running numbers per query row — the maximum m seen so far and the denominator l (the running sum of exponentials) — and, each time a new tile pushes the maximum higher, rescale the accumulated denominator and output by exp(m_old − m_new) to correct for the smaller value you exponentiated against earlier. Because that rescaling exactly compensates, the final result is bit-for-bit the same as a full-matrix softmax: FlashAttention is exact attention, not an approximation.

How does FlashAttention save memory in the backward pass?

The backward pass needs the attention probability matrix P to compute gradients, and storing P for every layer would cost O(N²) memory per layer — exactly the quadratic blow-up FlashAttention is trying to avoid. So it does not store P. Instead the forward pass keeps only the tiny per-row residual it already had — the softmax statistics, which the reference kernels describe as "(maximum softmax input values, softmax denominator)" — and the backward pass recomputes each tile of the scores and probabilities on the fly from Q, K, V and those statistics. This is a classic compute-for-memory trade: it does a bit of extra arithmetic (a recomputed QKᵀ tile) to avoid reading and writing a huge matrix, and because the kernel was memory-bound to begin with, the recomputation is nearly free.

Does FlashAttention reduce the number of FLOPs?

No — and this is the honest subtlety. FlashAttention does the same number of floating-point operations as standard attention (arguably a few more, once you count the backward-pass recomputation). What it reduces is memory accesses. The paper proves it needs Θ(N²d²M⁻¹) HBM accesses versus Θ(Nd + N²) for standard attention, where N is sequence length, d is head dimension, and M is the SRAM size — so for the regime where d² is smaller than M, FlashAttention moves far less data even though it computes the same amount. The speedup comes entirely from being IO-aware: the operation was bottlenecked on memory bandwidth, so cutting memory traffic while keeping the FLOPs constant makes it run faster and use linear instead of quadratic memory.

What changed between FlashAttention 1, 2, and 3?

FlashAttention-1 (2022) introduced the tiling + online-softmax + recomputation kernel and the IO-complexity argument. FlashAttention-2 (2023) is an engineering rework for better GPU occupancy: it reduces the non-matmul FLOPs (the expensive rescaling work), parallelizes across the sequence-length dimension so more of the GPU stays busy, and improves how work is split between warps — reaching "50-73% of the theoretical maximum FLOPs/s on A100" and about 2× the throughput of FA1. FlashAttention-3 (2024) is specific to Hopper (H100) GPUs: it exploits hardware asynchrony (overlapping the matmuls with the softmax, using the tensor memory accelerator) and low-precision FP8, reaching "up to 740 TFLOPs/s (75% utilization)" in FP16 and close to 1.2 PFLOPs/s in FP8. The algorithm is the same across all three; each generation squeezes more out of the silicon.

Anatomy of a FlashAttention Kernel

Plain attention is not slow because it does too much math — it is slow because it writes an N×N score matrix out to memory and reads it back, twice. FlashAttention never writes that matrix at all: it walks the keys in tiles, keeps a running max and denominator so softmax stays exact block by block, and recomputes what it needs in the backward pass instead of storing it. We compute the materialization bill, step through the online-softmax recurrence on a real tiny attention, do the IO complexity honestly, and cover what FA2 and FA3 actually changed — drawn, computed, and animated.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Anatomy of a FlashAttention Kernel — built to be understood, not skimmed.

Questions

Why is attention memory-bound, and how does FlashAttention fix it?: Standard attention is slow not because of its arithmetic but because of its memory traffic. To compute softmax(QKᵀ/√d)·V it first builds the full N×N score matrix S, writes it to the GPU's high-bandwidth memory (HBM), reads it back to apply softmax, writes the N×N probability matrix P, and reads that back to multiply by V — the FlashAttention paper notes that "standard attention implementations materialize the matrices S and P to HBM, which takes O(N²) memory." Since those reads and writes dominate, and "most operations in Transformers are bottlenecked by memory accesses," the kernel spends its time waiting on memory while the tensor cores idle. FlashAttention fixes this by never writing S or P to HBM at all: it processes attention in small tiles that fit in on-chip SRAM, so the O(N²) intermediate lives only in fast memory and never round-trips to HBM.
What is online softmax and how does tiling keep it exact?: Online softmax is the trick that lets you compute a softmax over a row of scores without ever having the whole row in memory at once. The problem: a numerically safe softmax subtracts the row maximum before exponentiating, but if you only see the scores a tile at a time you do not know the true maximum yet. The fix is to carry two running numbers per query row — the maximum m seen so far and the denominator l (the running sum of exponentials) — and, each time a new tile pushes the maximum higher, rescale the accumulated denominator and output by exp(m_old − m_new) to correct for the smaller value you exponentiated against earlier. Because that rescaling exactly compensates, the final result is bit-for-bit the same as a full-matrix softmax: FlashAttention is exact attention, not an approximation.
How does FlashAttention save memory in the backward pass?: The backward pass needs the attention probability matrix P to compute gradients, and storing P for every layer would cost O(N²) memory per layer — exactly the quadratic blow-up FlashAttention is trying to avoid. So it does not store P. Instead the forward pass keeps only the tiny per-row residual it already had — the softmax statistics, which the reference kernels describe as "(maximum softmax input values, softmax denominator)" — and the backward pass recomputes each tile of the scores and probabilities on the fly from Q, K, V and those statistics. This is a classic compute-for-memory trade: it does a bit of extra arithmetic (a recomputed QKᵀ tile) to avoid reading and writing a huge matrix, and because the kernel was memory-bound to begin with, the recomputation is nearly free.
Does FlashAttention reduce the number of FLOPs?: No — and this is the honest subtlety. FlashAttention does the same number of floating-point operations as standard attention (arguably a few more, once you count the backward-pass recomputation). What it reduces is memory accesses. The paper proves it needs Θ(N²d²M⁻¹) HBM accesses versus Θ(Nd + N²) for standard attention, where N is sequence length, d is head dimension, and M is the SRAM size — so for the regime where d² is smaller than M, FlashAttention moves far less data even though it computes the same amount. The speedup comes entirely from being IO-aware: the operation was bottlenecked on memory bandwidth, so cutting memory traffic while keeping the FLOPs constant makes it run faster and use linear instead of quadratic memory.
What changed between FlashAttention 1, 2, and 3?: FlashAttention-1 (2022) introduced the tiling + online-softmax + recomputation kernel and the IO-complexity argument. FlashAttention-2 (2023) is an engineering rework for better GPU occupancy: it reduces the non-matmul FLOPs (the expensive rescaling work), parallelizes across the sequence-length dimension so more of the GPU stays busy, and improves how work is split between warps — reaching "50-73% of the theoretical maximum FLOPs/s on A100" and about 2× the throughput of FA1. FlashAttention-3 (2024) is specific to Hopper (H100) GPUs: it exploits hardware asynchrony (overlapping the matmuls with the softmax, using the tensor memory accelerator) and low-precision FP8, reaching "up to 740 TFLOPs/s (75% utilization)" in FP16 and close to 1.2 PFLOPs/s in FP8. The algorithm is the same across all three; each generation squeezes more out of the silicon.

Anatomy of a FlashAttention Kernel

Questions

Related explainers