TL;DR: Standard attention materializes the N by N score matrix in HBM and reads it back for softmax and again for the value product. At 8k context that is 134 MB per head per layer and about 544 MB of traffic for 34 GFLOPs of math, an arithmetic intensity of 63 FLOP/B against an H100 ridge of 295, so the kernel sits on the bandwidth side of the roofline and runs far below peak. FlashAttention tiles Q, K and V into blocks that fit in SRAM, carries a running row max and row sum so softmax can be finished incrementally, and never writes the scores to HBM. Traffic per head falls to about 8 MB, intensity rises into the thousands, and the kernel becomes compute-bound. IO-aware means the algorithm was designed around the memory hierarchy rather than the FLOP count: it does slightly more arithmetic and vastly less movement. The memory footprint drops from O(N squared) to O(N), which is what makes 128k context possible at all.
How to approach it
Say what standard attention does to memory first, with the number, because that is the whole motivation. Then give the two ideas that fix it, tiling and online softmax, and why neither works without the other. Then do the traffic arithmetic before and after and place both on the roofline. Close with the memory argument and the backward pass, which is where most candidates stop early.
A strong answer
A typical situation: a team moves a 70B model from 4k to 32k context and prefill starts failing with out-of-memory before it gets slow; the first fix they try is a smaller batch, which halves throughput and still does not fit, because the thing that grew was the score matrix, not the weights.
Start by naming the object whose size decides everything. Attention gives every query position a weighted average of the value vectors, and the weights come from a softmax over the dot products between that query and every key. Those dot products form the score matrix S: one row per query position, one column per key position. So S is square in the sequence length, and the sequence length is the number a product team keeps raising.
The traffic, for one attention head at sequence length N = 8,192, head dimension d = 128, bf16:
scores S = Q Kt is N x N: 8,192 x 8,192 x 2 B = 134 MB
standard implementation, HBM traffic per head:
read Q, read K 2 x (N x d x 2 B) = 2 x 2 MB = 4 MB
write S = 134 MB
read S, write P (softmax) = 268 MB
read P, read V, write O 268 + 2 + 2 = 138 MB
total ~ 544 MB
FLOPs per head: 2 x N x N x d (QKt) + 2 x N x N x d (PV) = 4 x 6.71e7 x 128 = 3.4e10
arithmetic intensity = 3.4e10 FLOP / 5.44e8 B = 63 FLOP/B
H100 ridge = 989e12 / 3.35e12 = 295 FLOP/B
sanity: 63 is well under 295, so the kernel is memory-bound by about 4.7x and cannot exceed
3.35 TB/s / 544 MB = 6.2 ms per head no matter how fast the tensor cores are
FlashAttention, same head:
read Q, K, V once, write O once: 4 x (N x d x 2 B) = 8 MB (K and V re-reads mostly hit L2)
intensity = 3.4e10 / 8.4e6 = ~4,000 FLOP/B
sanity: 4,000 is far above 295, so the kernel is now compute-bound and the tensor cores set
the time; that is the whole point of the redesign
Memory Coalescing and Memory-Bound vs Compute-Bound Kernels are the frame; the Roofline Model is where those two intensity numbers sit, on opposite sides of the ridge.
The obstacle is softmax. Softmax needs the row maximum and the row sum before any output can be produced, which is why the naive implementation materializes the whole row. The online form removes that dependency by carrying two scalars per row and rescaling whenever the max moves:
state per row: m (running max, init -inf), l (running sum of exp, init 0), O (accumulator, init 0)
for each K/V block j:
S_ij = Q_i . K_j^T computed in SRAM, never stored to HBM
m_new = max(m, rowmax(S_ij))
alpha = exp(m - m_new) the correction factor for everything accumulated so far
l = l * alpha + rowsum(exp(S_ij - m_new))
O = O * alpha + exp(S_ij - m_new) @ V_j
m = m_new
after the last block: O = O / l
sanity: with one block the loop reduces to the textbook softmax, and alpha = exp(m - m) = 1,
so the identity is exact rather than an approximation
Tiling and the online form need each other. Tiling alone cannot work, because a block of scores is useless until the row's max is known. The online form alone does not help, because without tiles the scores still land in HBM. Together they give one pass over Q, K and V with the N by N intermediate living only in SRAM. FlashAttention Internals has the block-size arithmetic: the tiles must fit the 228 KB of shared memory per SM on an H100, which is what fixes B_r and B_c at 64 to 128 rows for d = 128.
The backward pass is where the trade shows most clearly. It needs the scores again, and storing them would undo everything, so it recomputes them from Q, K and V inside the same tiled loop. That costs roughly 30% more FLOPs than a stored-score backward and still wins, because those FLOPs are free on a kernel that has moved to the compute side and the alternative is 134 MB per head of traffic.
The memory argument is the one that matters in production. Peak activation memory for attention goes from O(N squared) per head to O(N), so the 32k context in the scenario above stops being an allocation problem. Speed follows, but a team can work around a slow kernel and cannot work around an allocation that does not fit.
How you check that it is actually running, which is the practical half of the answer: in PyTorch, attention routed through F.scaled_dot_product_attention dispatches to a fused backend only when the dtype, head dimension and mask shape are supported, and falls back to the materializing path with no error when they are not. Two signals settle it. torch.cuda.max_memory_allocated() around one forward pass drops by the N-squared term when the fused path is taken. In an Nsight Systems timeline the fused path shows a single long kernel per attention call, with a name containing flash or fmha, in place of a GEMM, a softmax and a second GEMM. Profiling with Nsight covers reading that timeline.
The reversal condition: the gain is proportional to how memory-bound the baseline was. At short sequences (N of a few hundred) the score matrix fits in L2, the naive kernel is already fast, and the fused version wins only the launch overhead. At decode time the query is a single token, the score matrix is a vector, and the kernel is bandwidth-bound on the KV cache rather than on scores; that regime needs a different kernel that splits the work across KV blocks and combines partial softmax states, which is what FlashDecoding does.
What interviewers probe next
- "Why not keep the scores in L2 instead?" An H100 has 50 MB of L2 and one head at 8k context needs 134 MB, so the row you want has already been evicted by the next head. L2 helps with the K and V re-reads, not with the scores.
- "Does it change the numerics?" The rescaling identity is exact in exact arithmetic and the max subtraction keeps every exponent at or below zero. The floating-point result differs from the naive summation order, in the same way any reordered reduction does, and it is not less accurate.
- "Where do the block sizes come from?" Shared memory per SM divided by the tiles that must be resident at once (Q block, K block, V block, accumulator), then rounded to keep the tensor-core fragment shapes whole.
- "Is it an approximation like sparse or linear attention?" No. The output is the same attention, computed in a different order. Sparse attention changes which scores exist; this changes only where they live.
Common mistakes
- Calling FlashAttention an approximation, or confusing it with linear attention.
- Claiming it reduces FLOPs. It slightly increases them, and it reduces bytes moved by two orders of magnitude.
- Quoting a speedup without a sequence length, when the whole effect scales with N.
- Describing the tiling and forgetting the online softmax, which is the part that makes tiling legal.
- Treating the speedup as the headline when the O(N) memory is what unlocks long context.
Key takeaways
- At N = 8,192 the score matrix is 134 MB per head; standard attention moves about 544 MB per head for 34 GFLOPs, an intensity of 63 FLOP/B against a 295 ridge.
- FlashAttention moves about 8 MB per head instead, pushing intensity into the thousands and the kernel across the ridge to compute-bound.
- The online identity is m_new = max(m, rowmax(S)), then rescale O and l by exp(m minus m_new) before accumulating.
- Memory drops from O(N squared) to O(N) per head, and the backward pass recomputes scores for about 30% more FLOPs rather than storing them.