FlashAttention Explained: Tiling, SRAM & the Online Softmax
Attention isn’t slow because of the math — it’s slow because of the memory traffic. How tiling and an online softmax make exact attention fit in fast on-chip memory.
Explain it like I’m five
Imagine doing homework at a tiny desk next to a huge, slow warehouse. The warehouse holds every book you own, but each trip there takes forever. The desk fits three books — but everything on it is instant.
Naive attention hauls every book out of the warehouse, spreads them across the floor, then puts them all back. FlashAttention works one shelf at a time: carry a small stack to the desk, do the math, and keep a running answer on a sticky note. The sticky note is the clever part — it tracks the best score seen so far and a running total, so when a bigger score shows up later, the earlier answers get rescaled instead of recomputed. The final answer is exactly the same as if you’d spread everything on the floor.
Intuition: the bottleneck is memory, not math
Your training run sits at 35% GPU utilization, and your 32K-context model dies with an out-of-memory error on a single GPU. The arithmetic isn’t the problem — a modern GPU does hundreds of teraflops. The problem is that naive attention writes a 32K×32K score matrix to slow memory and reads it back, and the GPU spends its time waiting on data instead of computing. In interviews, nobody asks “what is FlashAttention?” They ask “why is attention memory-bound, and what would you change about how it moves data?” By the end of this guide you’ll explain tiling, derive the online-softmax rescaling, and know exactly when the trick stops helping.
Naive attention makes three heavy passes over slow memory: build the N×N score matrix S = QKᵀ in HBM, read it back to apply softmax, read it again to multiply by V. At N = 16K with 64 heads, S alone is over 30 GB per layer at fp16 — written once, read twice, with almost no data reuse.
A GPU has two memories that matter here. HBM is the big, slow main memory — tens of gigabytes at roughly 2 TB/s on an A100. SRAM is the on-chip scratchpad next to the compute cores — about 20 MB total on an A100, but around 10× the bandwidth. The score matrix will never fit in SRAM. A single tile of it will.
Same FLOPs as naive attention — the bytes are what changed.
FlashAttention never materializes the N×N matrix: it tiles Q, K, V into SRAM-sized blocks, fuses the whole computation into one kernel, and repairs the softmax with running statistics.
How the trick works
The memory wall, quantified
Put numbers on the A100. There are 108 streaming multiprocessors, each with 192 KB of shared memory — about 20 MB of SRAM in total. HBM is 40–80 GB at roughly 2 TB/s; SRAM bandwidth is close to 19 TB/s. Now take N = 16K, d = 128, 64 heads: the score matrix holds 16K² = 268M entries per head, 512 MB at fp16, over 32 GB across heads, per layer. Softmax and the V-multiply stream those bytes at roughly one FLOP per byte — deep in the memory-bound region of the roofline model. Wall-clock time is set by HBM traffic; the FLOPs were never the constraint.
Tiling
Pick block sizes Br (query rows) and Bc (key/value rows) so that a Br×d query block, a Bc×d key block, a Bc×d value block, the Br×Bc score block, and the Br×d output block all fit in SRAM at once. Loop over query blocks outside, key/value blocks inside; each output block accumulates in SRAM and is written back to HBM exactly once.
This is the same idea as the naive attention you’d write in PyTorch — same FLOPs, same math — except the N×N intermediate never exists in slow memory.
The online softmax — exact, not approximate
Softmax needs each row’s max m and sum l to normalize. Tiling sees one block of the row at a time, so FlashAttention keeps running statistics per row: the max so far, the sum so far, and the partial output.
When a new block arrives with local max m_b and local sum l_b, merge with two rescaling identities:
- m_new = max(m, m_b)
- l_new = l·e^(m − m_new) + l_b·e^(m_b − m_new)
The partial output gets the same rescaling, so the weighted average is preserved. By induction over blocks, the final (m, l, output) equal the full-row values exactly. Every step is algebra — exp(a − c) factors cleanly — so there is no approximation anywhere. This is also what keeps it numerically safe: the running max holds every exponent at or below zero, the incremental version of the standard max-subtraction trick.
The backward pass recomputes
A standard backward pass needs the N×N attention probabilities to compute dQ, dK, and dV. FlashAttention stores only the per-row max and sum — O(N) — plus the output, and recomputes each score block from Q and K inside the backward kernel, rebuilding probabilities from the stored statistics. The recompute is cheap, compute-bound matmul work; what it avoids is the quadratic HBM traffic. The classic rematerialization trade — and here it wins decisively.
Training long-context models: the FlashAttention paper
The original FlashAttention paper turned tiling plus the online softmax into a fused kernel and made long-context training practical on existing GPUs — the same kernel family now ships in PyTorch, xFormers, and inference servers.
Q, K, V stream through SRAM in blocks. The N×N score matrix is never written to HBM — memory drops from quadratic to linear in sequence length.
The forward pass — and a recomputation-based backward pass — run as fused kernels instead of many separate memory-bound passes over giant intermediates.
Linear memory means the same GPU trains on far longer sequences, and the freed bandwidth goes directly into tokens per second.
The paper reports roughly 3× faster training than a standard attention implementation and 4× longer sequences on the same hardware. This is a research result from the paper itself, not a vendor benchmark.
Inference serving uses the same kernels — the KV cache attacks the decode memory problem while FlashAttention attacks the attention compute problem. They compose: tiled attention over a paged cache.
Work the rescaling by hand
One attention row, split into two blocks: block A holds scores [1, 3], block B holds [2, 5].
Same weights as the true softmax. That is the whole trick, in four numbers.
Tiling in twelve lines
import math
import torch
def tiled_attention(Q, K, V, Br=128, Bc=128):
# Q, K, V: [N, d], one head. Illustrative — the real
# FlashAttention is a fused CUDA kernel, not Python loops.
N, d = Q.shape
O = torch.zeros(N, d)
for i in range(0, N, Br):
Qi = Q[i:i+Br]
m = torch.full((Qi.shape[0],), float("-inf"))
l = torch.zeros(Qi.shape[0])
Oi = torch.zeros_like(Qi)
for j in range(0, N, Bc):
Kj, Vj = K[j:j+Bc], V[j:j+Bc]
S = Qi @ Kj.T # Br x Bc block: lives in "SRAM"
m_b = S.max(dim=-1).values
P = torch.exp(S - m_b[:, None])
l_b = P.sum(dim=-1)
m_new = torch.maximum(m, m_b) # running max
a = torch.exp(m - m_new) # rescale old stats
b = torch.exp(m_b - m_new) # rescale new block
l = l * a + l_b * b
Oi = Oi * a[:, None] + (P * b[:, None]) @ Vj
m = m_new
O[i:i+Br] = Oi / l[:, None]
return OThe real kernel fuses these loops into one launch, keeps every block in SRAM, and stores only the per-row (m, l) for the backward pass. Same math — the memory movement is the product.
Six questions that test the real understanding
What to say (≈90 sec): “Standard attention is IO-bound, not FLOP-bound. On an A100, SRAM is about 20 MB total with roughly 19 TB/s of bandwidth, while HBM is tens of gigabytes at roughly 2 TB/s. Naive attention materializes the N×N score matrix — at N = 16K with 64 heads that’s over 30 GB per layer at fp16 — writes it to HBM, then reads it back for softmax and the V-multiply, doing roughly one FLOP per byte moved. On the roofline model that’s deep in the memory-bound region, so the GPU sits idle waiting on data; wall-clock time is set by HBM traffic, and adding more FLOPs changes nothing.”
Likely follow-up: “So why does adding FLOPs not help here?” → Because the bottleneck is bytes moved, not operations: the kernel is memory-bound, so time is set by bandwidth, not by compute throughput.
The answer that sinks you: “Because attention is quadratic in sequence length.” Why it fails: quadratic FLOPs would make it compute-bound, not memory-bound — the interviewer wants the IO argument, not the complexity class.
What to say (≈90 sec): “Each row needs its max m and sum l of exponentials. We track running statistics as we tile: the max seen so far, the sum seen so far, and the partial output. A new block arrives with local max m_b and local sum l_b. The merged max is the larger of the two; both sums must be shifted to that new max before adding, which is just the max-subtraction trick: multiply the old sum by exp(m − m_new) and the new block’s sum by exp(m_b − m_new). The partial output gets the same rescaling factors, so the weighted average is preserved. By induction over blocks, the final (max, sum, output) equal the full-row values exactly — every step is algebra, so there is no approximation anywhere. This is also what keeps it numerically safe: the running max holds every exponent at or below zero, the incremental version of the standard max-subtraction trick.”
Likely follow-up: “What breaks if you skip the rescaling?” → The old block’s statistics were computed against a stale max, so the merged sum and output are silently wrong — a numeric bug with no shape error, the worst kind.
The answer that sinks you: “You just average the two softmax outputs.” Why it fails: softmax is nonlinear — averaging block softmaxes ignores the partition-function mismatch between blocks and gives the wrong answer.
What to say (≈90 sec): “The only ‘change’ from naive attention is computation order. Tiling reorders the same Br×Bc matrix products over the same entries — floating-point rounding order may differ at the last bit, but mathematically every product is identical. The online softmax is provably identical to full softmax by induction: the rescaling identities are exact algebra on exp, with no truncation or sampling anywhere. The output is therefore the same weighted average of V that naive attention computes — exact, up to floating-point association order, which is true of any kernel rewrite. That’s the sharp contrast with Linformer, Performer, and other linear-attention variants: those change the math; FlashAttention only changes the IO.”
Likely follow-up: “So it doesn’t reduce FLOPs at all?” → Correct — the asymptotic FLOP count is unchanged and still quadratic; the savings are memory traffic and auxiliary memory.
The answer that sinks you: “It’s exact enough for practical purposes.” Why it fails: hedging makes the interviewer doubt you know the difference — FlashAttention is algebraically exact, not ‘exact enough.’
What to say (≈90 sec): “A standard backward needs the N×N attention probabilities to compute dQ, dK, and dV. FlashAttention stores only the per-row max and sum — O(N) — plus the output, and recomputes each score block from Q and K inside the backward kernel, rebuilding the block probabilities from the stored statistics. The recompute is cheap, compute-bound matmul work, while what it avoids is the quadratic HBM traffic of writing and re-reading the probability matrix. It’s the classic rematerialization trade — spend extra FLOPs to save memory bandwidth — and here it wins decisively because attention is memory-bound: the ‘extra’ compute was idle bandwidth anyway.”
Likely follow-up: “When would recomputation be a bad trade?” → When the kernel is compute-bound rather than memory-bound — then the extra FLOPs land on the bottleneck instead of filling idle capacity.
The answer that sinks you: “It recomputes the whole forward pass.” Why it fails: sloppy — it recomputes score blocks tile-by-tile inside the backward kernel, not the entire forward computation.
What to say (≈90 sec): “The constraint is SRAM: each streaming multiprocessor has 192 KB of shared memory, and the working set of one Br×d query block, one Bc×d key block, one Bc×d value block, the Br×Bc score block, and the Br×d output accumulator must fit in it. Given that, you want blocks as large as possible: bigger Br amortizes the kernel-launch overhead across more output rows and increases data reuse, and bigger Bc gives better parallelism in the inner loop. The real tradeoff is occupancy versus reuse — huge blocks starve the SM of resident warps, tiny blocks re-read K and V too often. In practice this is autotuned per GPU and per (N, d), not derived by hand — but the interview answer is the working-set equation and the two sides of the tradeoff.”
Likely follow-up: “What happens if blocks are too big?” → SRAM overflows — the kernel fails or the compiler spills to local memory, which is HBM, destroying the whole point.
The answer that sinks you: “Bigger blocks are always faster.” Why it fails: ignores the SRAM working-set constraint and the occupancy-vs-reuse tradeoff — the ceiling is physical.
What to say (≈90 sec): “When the sequence is short. At N = 512 or 1K, the score matrix is tiny — the naive implementation’s memory traffic is negligible, and attention is a small fraction of a forward pass dominated by the MLPs anyway. The speedup also shrinks when the kernel is compute-bound rather than IO-bound — huge batch × head dimensions can push the tiled matmuls onto the compute side of the roofline, where FlashAttention’s extra recomputation is real work, not free. And it doesn’t help the decode phase of serving: per-token generation is already memory-bound on the KV cache and weight reads, which is a different bottleneck solved by the KV cache and paged attention, not by tiling attention. The diagnostic is always the roofline: FlashAttention pays off exactly where attention is IO-bound with a large intermediate.”
Likely follow-up: “So how do you decide whether to enable it?” → Profile first — check whether attention is a meaningful fraction of runtime and whether the kernel sits on the memory-bound side of the roofline before rewriting anything.
The answer that sinks you: “FlashAttention is always faster — that’s why everyone uses it.” Why it fails: no free lunch — at short sequences or in decode-bound serving, the tiling overhead buys nothing and can cost.
Key takeaways
- Naive attention is memory-bound: the N×N score matrix is written to HBM and read back, at roughly one FLOP per byte.
- Tiling streams Q, K, V through SRAM in blocks — the quadratic intermediate never exists in slow memory.
- The online softmax is exact algebra: running max/sum rescaling reproduces full softmax by induction.
- The backward pass stores O(N) statistics and recomputes score blocks — a rematerialization trade that wins.
- Same asymptotic FLOPs as naive attention; the gains are memory traffic and auxiliary memory.
- Profile before enabling: it pays off where attention is IO-bound with a large intermediate.