FlashAttention doesn't change the math of attention at all — the output is bit-for-bit the same softmax attention from the previous piece. What it changes is where the computation lives on the GPU: it never writes the full attention matrix to slow memory. To see why that matters, we'll first build a mental model of the GPU's memory hierarchy — on-chip SRAM vs. off-chip HBM — since that's the entire reason this trick works.
A GPU has a tiny pool of extremely fast on-chip memory (SRAM) and a large pool of much slower
off-chip memory (HBM). Standard attention computes the full T×T score matrix,
writes it out to HBM, reads it back for softmax, writes the result again, and reads it once more for the final
weighted sum — several full passes over an object that's often bigger than the model's weights themselves.
FlashAttention fuses all of that into a single pass: it tiles Q, K,
V into blocks small enough to fit in SRAM, and uses a running ("online") version of softmax so it
never needs the whole row of scores in memory at once. The math comes out exactly the same — this is
not an approximation — but the number of slow memory round-trips drops from O(T²) to something
close to O(T), which is why it's faster even though it does the same FLOPs.
The previous piece established that attention's cost splits into
two very different regimes: prefill, where a whole prompt runs through in parallel and the
matmuls (projections, QKᵀ, the weighted sum) dominate wall-clock time; and decode,
where one token at a time is generated and reading the KV cache off HBM dominates instead.
Grouped-Query Attention attacked the decode side by shrinking the cache itself.
FlashAttention attacks something orthogonal: the cost of the attention computation itself —
softmax(QKᵀ/√d)V — regardless of whether you're in prefill or training a model from scratch, where
a whole batch of full-length sequences runs through every layer. That computation involves a T×T
intermediate matrix, and how you handle that matrix on the chip turns out to matter enormously once
T gets into the thousands. To see why, we need to look inside the GPU.
Every number quoted about GPUs so far — TFLOPS, HBM bandwidth — implicitly assumed data was already sitting where the compute units could reach it. In practice a GPU has several tiers of memory, and they differ by roughly two orders of magnitude in both size and speed. Think of it like a workbench versus a warehouse:
The workbench analogy is worth sitting with. A woodworker keeps the piece they're actively cutting on the workbench — walking to the warehouse for every single screw would dominate the time spent actually building anything. SRAM is that workbench: minuscule compared to HBM (megabytes vs. tens of gigabytes), but so fast that keeping data there instead of fetching it from HBM repeatedly is the difference between a kernel finishing in microseconds or in milliseconds. Model weights, the KV cache, activations — all of it lives in HBM at rest; the only question is how often a given GPU kernel has to go fetch something from there instead of reusing what's already on the workbench.
Decode being "memory-bandwidth-bound" (§4 of the previous piece) meant exactly this: every decode step has to
walk to the HBM warehouse and drag back the model's weight matrices and the KV cache, do a tiny amount of work
with them, then walk back and do it again next step. FlashAttention is the same underlying idea applied to a
different object — the T×T attention matrix — and to a different regime, since that matrix
shows up whenever T is large: training, and prefill.
A naive implementation of softmax(QKᵀ/√d)V isn't one GPU kernel — it's several, chained together,
because each operation (matmul, then softmax, then another matmul) is typically its own separate kernel launch.
Between every pair of kernels, the intermediate result has nowhere to live but HBM: it gets written out by one
kernel and read back in by the next.
T×T score matrix and writes all of it to HBM — there's nowhere else for a matrix this size to go once the kernel exits. At T=1,024, fp32: 1024²×4 bytes ≈ 4.2 MB; doubled below for the round trip.
Notice what's not different: the number of FLOPs. Both versions compute the exact same dot products,
the exact same softmax, the exact same weighted sum. The previous piece's ridge-point argument said arithmetic
intensity — FLOPs per byte moved — determines whether a step is compute-bound or memory-bound. Standard
attention's problem is that the T×T matrix has low reuse: each of its T²
entries is read only a couple of times before being discarded, so despite doing real arithmetic, the operation
as a whole is bottlenecked on shuffling that matrix in and out of HBM, not on computing it.
FlashAttention's fix has two ingredients that have to work together:
Q into row-blocks and K, V into column-blocks, small
enough that one block of each — plus a small amount of running state — fits in SRAM at once. Process one
(Q-block, K/V-block) pair at a time, entirely on-chip, and fuse the matmul → softmax → matmul
chain into a single kernel so intermediate results never have to leave SRAM.
exp across the entire row — which looks like it requires having
the whole row on hand before you can normalize anything. The trick is a numerically stable way to update a
running max and running sum incrementally, one block at a time, correcting the accumulated output every time
the running max changes. That's the only reason you're allowed to process the row in pieces at all.
Put together: instead of "compute all scores → normalize all at once → weight all values," FlashAttention does "for each chunk of keys/values, update a running estimate of the output" — and by the time it's seen every chunk, that running estimate is the exact, fully normalized answer. No approximation, just a different order of operations that happens to need far less memory traffic.
Ordinary softmax over a row of scores s1...sT first finds the row max
m = maxi si (purely for numerical stability — it keeps exp()
from overflowing) and then computes exp(si − m) / Σj exp(sj − m).
The apparent problem: both m and the sum Σ need the entire row before either can be
computed.
The fix (Milakov & Gimelshin, 2018 — FlashAttention applies this idea to fused attention specifically):
keep a running max m, a running sum l, and a
running (unnormalized) output O. Every time a new block of keys/values arrives,
update all three — and if the running max changes, rescale everything accumulated so far by a correction
factor before adding the new block's contribution:
m̃, l̃ = this block's local max and local softmax-sum ·
Ṽ = this block's contribution to the weighted sum, computed against its own local max
After the very last block B, divide once: output = O(B) / l(B).
That single division is the only normalization step in the whole algorithm — everything before it was carried
around in an unnormalized, but consistently-rescaled, running state. Crucially, this recurrence is
mathematically exact: expand it out and it telescopes to precisely the same value as computing
the max and sum over the whole row up front. Nothing here is an approximation of softmax — it's the same
softmax, computed in a different order.
Put tiling and online softmax together and the full algorithm is a nested loop: an outer loop over blocks of
Q, an inner loop over blocks of K, V (FlashAttention-2 swaps which loop is outer — more
in §8). For each (Q-block, K/V-block) pair, the kernel computes a block of scores, runs one step of
the online-softmax update above, and accumulates into that Q-block's running output — all without leaving SRAM.
Two details worth naming, since they explain why this needed a paper and not just an obvious refactor:
The FlashAttention paper's central theoretical result isn't a speed claim, it's a counting argument: how many
bytes does each approach move between HBM and SRAM, as a function of sequence length T, head
dimension d, and SRAM size M?
Since M (tens of thousands of elements) is much larger than d (64–128), the
T²d²/M term is a large factor smaller than the plain T² standard attention pays —
and unlike standard attention, FlashAttention's HBM traffic doesn't even depend on T alone, it
depends on how favorably T and M trade off. The paper proves this is
asymptotically optimal: no exact attention algorithm can do fewer HBM accesses for a given
SRAM size.
Headline numbers from the original paper (GPT-2, BERT-large, and long-sequence benchmarks) — quoted here as reported, not re-derived:
That last row matters as much as the speedup: because the T×T matrix is never materialized,
activation memory for attention no longer grows quadratically with sequence length — which is precisely what
let FlashAttention-era models push context lengths from a few thousand tokens into the hundreds of thousands,
independent of the speed gains.
FlashAttention-1 was already IO-optimal in the counting-argument sense above, but it left real GPU utilization on the table — around 25–40% of peak FLOPs, versus 80%+ for a plain dense matmul. FlashAttention-2 (2023) closes most of that gap without changing the fundamental algorithm:
Q, which parallelizes cleanly even
for a single, very long sequence (long-context prefill, or training with long sequences and small batches).
Under the original order this parallelism was harder to extract.
None of this changes the IO-complexity argument from §7 — it's the same asymptotic HBM traffic. The gain is from spending less time stalled and more time actually issuing matmul instructions, roughly a further 2× wall-clock speedup over FlashAttention-1.
FlashAttention-3 (2024) targets NVIDIA's Hopper architecture (H100) specifically, and layers on hardware features that didn't exist on Ampere (A100):
Together these get FlashAttention-3 to roughly 75% of H100's theoretical peak FLOPs for FP16, and up to ~1.2 PFLOPS with FP8 — a further 1.5–2× over FlashAttention-2 on the same hardware. The underlying algorithm — tiling plus online softmax — is unchanged from §§4–6; this is entirely about exploiting a newer chip's specific async instructions to hide latency that Hopper introduced ways to overlap.
One more place the same "compute vs. memory" trade-off shows up. Standard backprop through attention needs the
softmax output matrix P (the T×T attention weights) to compute gradients — the
obvious approach is to keep it around from the forward pass. But storing it is exactly the O(T²)
memory cost FlashAttention was built to avoid.
Instead, FlashAttention's backward pass stores only the small, O(T)-sized softmax statistics
(the running max and sum from §5, one pair per row) plus the output O — and
recomputes the blocks of P on the fly during the backward pass, using those
stored statistics to skip redoing the max/sum search. This is the same lesson as the KV-cache-vs-recompute
trade-off from the previous piece, but inverted: there, recomputing was the trap and caching won; here,
recomputing wins because the alternative (storing an O(T²) matrix) is the more expensive thing to
move through memory. Recomputing the blocks costs extra FLOPs, but those FLOPs are cheap and fast compared to
the HBM traffic that storing P would have required — the same ridge-point logic as always, just
landing on the opposite side this time.
GPUs have a tiny, fast SRAM workbench and a large, slow HBM warehouse — anything that avoids
round trips to the warehouse wins, independent of how much arithmetic it does. Standard attention
pays for several full round trips on a T×T object because each step is a separate kernel.
FlashAttention fuses those steps into one kernel and uses a mathematically exact running
softmax so the whole computation can happen tile-by-tile, entirely on the workbench, with the T×T
matrix never fully existing anywhere. FlashAttention-2 squeezes more real GPU utilization out
of the same idea; FlashAttention-3 squeezes further by using instructions specific to newer
hardware. None of it changes what attention computes — only how cheaply it gets there.