FlashAttention computes exactly the same attention as the textbook formula — softmax(QK^T / √d_k) V — and yet it runs several times faster and uses memory that grows linearly with sequence length instead of quadratically. It does this without approximating anything: the output is numerically identical (to rounding) to standard attention. The trick is not a cleverer formula but a cleverer schedule. Standard attention writes an entire N × N score matrix out to slow high-bandwidth memory, reads it back to softmax it, reads it again to multiply by V — three round trips through the slowest part of the machine. FlashAttention fuses all of that into one pass that keeps only small blocks in fast on-chip memory and never writes the big matrix at all. To make that possible it reorganizes softmax into an online, streaming form with a running maximum and a running normalizer. This piece derives that online softmax from first principles, walks the tiling scheme over Q, K and V, works a full numeric example by hand, explains the recompute-in-backward trade, and argues why a memory-IO win — not a FLOP win — is exactly what real hardware, including a CPU running a small language model, rewards.
The starting point: attention as three matmuls and a softmax
Fix the shapes so nothing is hand-waved. For one attention head with sequence length N and head dimension d, the inputs are Q, K, V, each of shape [N, d]. Standard attention is:
S = Q K^T / √d # scores, S: [N, N]
P = softmax(S) # weights, P: [N, N] (row-wise softmax)
O = P V # output, O: [N, d]Each row i of S holds the affinities between query i and every key; softmax turns that row into a probability distribution; the output row is the corresponding convex combination of value vectors. The compute is dominated by the two N × N × d matmuls, so attention is O(N^2 d) in FLOPs. The subtle cost is the middle: S and P are each N × N. At N = 8192 a single such matrix in fp16 is about 128 MB — per head, per layer. Materializing it, storing it, and streaming it back and forth is where naive attention actually loses, and it is the thing FlashAttention refuses to do.
The real bottleneck is memory traffic, not arithmetic
A modern accelerator (and, in its own way, a CPU) has a steep memory hierarchy: a large pool of high-bandwidth memory (HBM / DRAM) that is big but slow, and a tiny pool of on-chip SRAM (registers, shared memory, L1/L2 cache) that is fast but measured in kilobytes. Arithmetic throughput has grown far faster than memory bandwidth for two decades, so for many operations the chip finishes the math and then sits idle waiting for operands to arrive. Such operations are memory-bandwidth-bound.
Attention’s softmax step is exactly this kind of operation. The QK^T and PV matmuls are arithmetically heavy, but the softmax in between is elementwise and cheap in FLOPs while touching every one of the N^2 entries. Standard attention writes S to HBM, reads it back to compute row maxima and sums, writes P, reads P again for PV. Those O(N^2) reads and writes of the big matrix — not the multiplies — are the wall clock. The whole design of FlashAttention starts from measuring that the kernel is IO-bound, and IO-aware means: minimize bytes moved between HBM and SRAM, even if you must redo some arithmetic to achieve it.
Safe softmax: the numerically stable baseline
Before streaming softmax we need the stable form of ordinary softmax, because the online algorithm is built on top of it. Naively, softmax(x_i) = exp(x_i) / Σ_j exp(x_j). But scores can be tens or hundreds in magnitude, and exp(100) overflows fp16 and even fp32. The fix is the max-subtraction trick, which is exact because shifting every logit by a constant leaves the ratio unchanged:
m = max_j x_j # the row max
p_j = exp(x_j - m) # now every exponent ≤ 0, so p_j ∈ (0, 1]
l = Σ_j p_j # the normalizer (a.k.a. denominator)
softmax(x_j) = p_j / lSubtracting m guarantees the largest exponent is exactly exp(0) = 1, so nothing overflows and the smallest terms merely underflow to zero harmlessly. Two per-row scalars carry all the state: the maximum m and the sum l. FlashAttention’s entire innovation is realizing that these two scalars can be maintained incrementally as you stream the keys past in blocks — you never need the whole row of scores in memory at once.
Online softmax: updating a running max and running sum
Suppose you have already processed some keys and hold the running max m_old, running normalizer l_old, and (as we’ll see) a running output. A new block of scores arrives with its own local max m_blk. The new global max is m_new = max(m_old, m_blk). Here is the crux: everything you accumulated was exponentiated relative to the old max. To fold in the new block you must rescale the old accumulators to the new reference by the correction factor α = exp(m_old - m_new):
m_new = max(m_old, m_blk)
α = exp(m_old - m_new) # correction for the running state
p_blk = exp(S_blk - m_new) # this block, in the new reference
l_new = α · l_old + rowsum(p_blk) # rescale old sum, add newBecause m_new ≥ m_old, the factor α ≤ 1: we are always scaling the past down to match a new, larger maximum, which keeps everything in range. When the block brings no new maximum, α = 1 and the update is just an ordinary running sum. This single rescale is what makes softmax associative enough to stream: process the keys in any order, in any block size, and l_new converges to the same denominator the one-shot safe softmax would have produced.
Carrying the output vector in the same stream
Maintaining m and l would only give the normalizer; we also want the output O = P V without ever forming P. The move is to accumulate an unnormalized output O_tilde — the value-weighted sum of un-divided exponentials — and divide by l only at the very end. The same correction factor α rescales it in lockstep with l:
O_tilde_new = α · O_tilde_old + p_blk · V_blk # V_blk: [B_k, d]
# ...after the last block:
O = O_tilde_final / l_final # normalize onceTrace why this is exact. After all blocks, l_final is Σ_j exp(S_j - m_final) and O_tilde_final is Σ_j exp(S_j - m_final) V_j, both measured against the true global maximum because every earlier partial was rescaled by the running α each time the max grew. Dividing gives Σ_j softmax(S)_j V_j — identical to standard attention. The state per query row is now three small objects: the scalar m, the scalar l, and the length-d vector O_tilde. None of them is N-sized. That is the whole game.
A fully worked numeric example
Take one query row, four keys, and split them into two blocks of two. Use pre-scaled scalar scores S = [1, 3, 2, 4] and scalar values V = [10, 20, 30, 40] (so d = 1, which keeps the arithmetic visible). First the reference answer: the max is 4, exp(S - 4) = [0.0498, 0.3679, 0.1353, 1.0], summing to l = 1.5530, and the output is Σ p_j V_j / l = 51.9155 / 1.5530 = 33.429.
Block 1: S=[1,3], V=[10,20] (start m=-∞, l=0, O_tilde=0)
m1 = max(-∞, 3) = 3 α = exp(-∞ - 3) = 0
p = exp([1,3]-3) = [0.1353, 1.0]
l1 = 0·0 + (0.1353+1.0) = 1.1353
O~1 = 0·0 + (0.1353·10 + 1.0·20) = 21.353
Block 2: S=[2,4], V=[30,40]
m2 = max(3, 4) = 4 α = exp(3 - 4) = 0.3679
p = exp([2,4]-4) = [0.1353, 1.0]
l2 = 0.3679·1.1353 + (0.1353+1.0) = 1.5530
O~2 = 0.3679·21.353 + (0.1353·30 + 1.0·40) = 51.9155
Output = O~2 / l2 = 51.9155 / 1.5530 = 33.429 ✓The streamed result matches the one-shot softmax to the last digit. Notice block 2 raised the max from 3 to 4, so α = 0.3679 scaled the block-1 partials down before adding block 2 — the rescale is not bookkeeping decoration, it is what makes the two orders of accumulation agree.
The tiling scheme: two nested loops over blocks
Now lift the single row to the full matrix. Partition the rows of Q into blocks of size B_q and the rows of K and V into blocks of size B_k. FlashAttention is two nested loops:
for each Q-block Q_i (rows i): # outer
load Q_i into SRAM; init m,l = -∞,0 ; O_tilde = 0
for each K/V-block (K_j, V_j): # inner
load K_j, V_j into SRAM
S_ij = Q_i K_j^T / √d # small: [B_q, B_k]
update m, l, O_tilde with the online-softmax step
O_i = O_tilde / l ; write O_i to HBM # write onceThe only tiles that ever exist are Q_i, K_j, V_j, and the little S_ij of shape [B_q, B_k] — all block-sized, all resident in SRAM. The N × N matrix S is implicit: its tiles are computed, consumed, and discarded inside the inner loop, never assembled and never sent to HBM. Each element of Q, K, V is read a small, bounded number of times, and O is written exactly once. Block sizes are chosen so a working set of a few tiles fits the SRAM budget; that constraint, not the algebra, sets B_q and B_k.
O(N) memory instead of O(N^2)
Account for the memory the kernel actually keeps. Per Q-block it holds m and l (one scalar each per query row, so O(B_q)) and O_tilde of shape [B_q, d]. Across all query rows the persistent state is the output O: [N, d] plus the per-row statistics [N] — together O(N d), i.e. linear in sequence length. The quadratic term has vanished from HBM entirely because S and P are never stored there; only their block-sized shards briefly occupy SRAM.
Concretely: standard attention at N = 8192, one head, fp16, needs ~128 MB just for the score matrix. FlashAttention needs the O(N d) output and a couple of length-N stat vectors — a few megabytes — regardless of head count scaling the same way. This is precisely why FlashAttention is what unlocked long context: doubling N doubles its attention memory rather than quadrupling it, so 32K- and 128K-token windows stop being memory-impossible. The compute still grows as N^2, but memory no longer does — and memory was the binding constraint.
Why this is an IO win, not a FLOP win
This is the point most summaries get wrong, so state it plainly: FlashAttention does not do less arithmetic. It performs the same O(N^2 d) multiply-adds for QK^T and PV as standard attention — asymptotically identical, and the backward pass, as we’ll see, actually does more FLOPs. There is no sparsity, no low-rank approximation, no skipped work. Every attention weight is computed.
What shrinks is bytes moved between HBM and SRAM. Standard attention’s HBM traffic is O(N^2) — write S, read it, write P, read it. FlashAttention’s traffic is O(N^2 / M) where M is the SRAM size, because each byte of the implicit matrix is produced and consumed on-chip and only the O(N d) inputs and output cross the bus. Since the kernel was memory-bandwidth-bound, cutting the dominant traffic by a large constant factor cuts wall-clock time by a similar factor — the reported 2–4× speedups — even though the FLOP count is unchanged. Trading a little extra compute for a lot less memory traffic is a winning trade precisely because the compute units were starved, not saturated.
Kernel fusion and keeping the pipeline busy
The mechanism that realizes the IO saving is kernel fusion. Standard attention is typically several separate kernels — a matmul, a softmax, another matmul — and every kernel boundary is a full round trip through HBM: the previous kernel writes its result out, the next reads it back in. FlashAttention fuses the entire score-softmax-weight-value sequence into a single kernel so intermediate tiles live and die in registers and shared memory without ever touching DRAM.
Fusion does more than save bytes; it keeps the hardware pipelined. While the arithmetic units chew on the current S_ij tile, the next K_j/V_j block can be prefetched into SRAM, overlapping compute with memory movement so neither stalls waiting on the other. This is the IO-aware philosophy made concrete: the algorithm is co-designed with the memory hierarchy, choosing block sizes and a loop order that maximize on-chip reuse. The same result — identical numbers — falls out of a schedule that treats data movement, not multiplication, as the scarce resource.
The backward pass: recompute instead of store
Training needs gradients, and the gradients of attention depend on the very N × N matrices S and P that FlashAttention refused to store. A normal autograd implementation would have saved P during the forward pass for reuse in the backward pass — reintroducing the O(N^2) memory we just eliminated. FlashAttention instead saves only the O(N d) output O and the per-row statistic L = m + log(l) (the log-sum-exp), a length-N vector.
In the backward pass it then recomputes each S_ij and P_ij tile on the fly from the saved Q, K, and L — because L already encodes the normalizer, a single block matmul reconstructs the exact softmax weights without redoing the streaming reduction. Those tiles feed the gradient accumulation for dQ, dK, dV, then are discarded. This is a deliberate recomputation trade: spend extra FLOPs regenerating the attention matrix rather than pay the memory and the HBM traffic to have stored it. On a memory-bound, bandwidth-starved kernel, recompute is cheaper in wall-clock than reload — so the backward pass, too, comes out ahead despite doing strictly more arithmetic.
Exactness: why the answer is not an approximation
It is worth insisting on this because ‘efficient attention’ is a crowded field full of approximations. Linear attention, sparse attention, and low-rank methods change what is computed to beat the quadratic cost, and they trade some accuracy for it. FlashAttention changes only how and in what order the standard computation is scheduled. Every step — the max-subtraction, the running rescale by α, the deferred normalization — is an exact algebraic identity, not a bound or an estimate.
The only differences from a one-shot implementation are floating-point rounding artifacts, and even those tend to be smaller, because the running-max formulation keeps every intermediate exponential in a safe numeric range. So you can drop FlashAttention into an existing model with no retraining and no accuracy regression: the logits, the loss, and the gradients match the reference to within rounding. That combination — free speed, less memory, and bit-comparable results — is why it became the default attention kernel across the ecosystem rather than one option among many.
Where the quadratic wall still stands
Be precise about what FlashAttention does and does not solve, or you will misbudget. It removes the quadratic memory and slashes the memory traffic. It does not remove the quadratic compute: QK^T is still N^2 d multiply-adds. At modest context lengths the kernel was memory-bound, so killing the traffic is a near-total win. But push N far enough — hundreds of thousands of tokens — and the sheer O(N^2) arithmetic eventually dominates again, and no schedule of exact attention can escape it.
That is the boundary between FlashAttention and the genuinely sub-quadratic methods. When compute, not memory, becomes the wall, you reach for sliding-window attention (each token attends to a local band, O(N w)), sparse patterns, or linear-attention approximations — all of which change the math and accept some quality cost. FlashAttention and these are complementary: use the exact, IO-aware kernel as far as it carries you, and switch schemes only when the irreducible N^2 FLOPs, not the memory, are what breaks the budget.
CPU-SLM relevance: the same memory wall, closer to home
FlashAttention was born on GPUs, but its central premise — that attention is memory-bound and the fix is to cut data movement — is if anything sharper on a CPU running a small language model. A CPU has far less memory bandwidth to DRAM than a datacenter accelerator, but it has a real cache hierarchy (L1/L2/L3) that plays exactly the role SRAM plays in the original design. Materializing an N × N score matrix on a CPU blows past the cache into main memory and thrashes; a tiled, fused, block-at-a-time attention keeps the working set inside L2/L3 and streams cleanly.
For a CPU-hosted SLM this matters twice over. During prefill of a long prompt, the quadratic score matrix is exactly the thing that would not fit in cache, so a Flash-style kernel is what keeps a long context tractable on commodity hardware. And the discipline generalizes: llama.cpp-style CPU inference engines lean hard on cache-blocking, operator fusion, and quantization for the same reason — on a bandwidth-starved machine, the algorithm that moves the fewest bytes wins, even when it does the same or slightly more arithmetic. FlashAttention is the canonical worked example of that principle, and it is the principle, not the CUDA, that transfers to CPU SLM inference.
Common misconceptions to retire
A short list, because each of these leads to a wrong estimate. ‘It approximates attention.’ No — it is exact to rounding; there is no accuracy trade. ‘It reduces the FLOPs.’ No — compute is unchanged in the forward pass and higher in the backward pass; the win is bytes moved, not multiplies done. ‘It makes attention linear.’ Only the memory is linear; the compute is still O(N^2).
‘The N × N matrix is computed then thrown away.’ More precisely, only its block-sized tiles ever exist, and only in fast on-chip memory — the full matrix is never assembled anywhere. ‘It stores P for the backward pass.’ No — it stores the O(N) log-sum-exp statistics and recomputes the tiles. And ‘a faster kernel must be doing less work.’ The whole lesson is the opposite: on memory-bound hardware you can do the same work, or even more, and finish sooner — by keeping the arithmetic units fed instead of idling on the memory bus. Internalize that inversion and the rest of FlashAttention follows.
m, a running normalizer l, and an unnormalized output rescaled by α = exp(m_old - m_new) each time the max grows — so keys and values can be streamed in blocks and the full N × N score matrix is never assembled or written to memory. That drops attention memory from O(N^2) to O(N) and cuts HBM traffic by a large constant factor. The speedup is a memory-IO win, not a FLOP win: the arithmetic is identical (the backward pass even recomputes the tiles from saved log-sum-exp statistics, doing more compute to avoid storing anything), yet wall-clock drops because attention was bandwidth-bound, not compute-bound. It is exact, not approximate; it linearizes memory, not compute; and its one lesson — on starved hardware, the algorithm that moves the fewest bytes wins — is exactly why the same idea pays off for a small language model running on a CPU’s cache hierarchy.