An Indian LLM startup — the kind now building Hindi- and Tamil-fluent assistants for legal and financial documents — wants its model to read an entire High Court judgment or a full RBI monetary policy report in one context window: 8,000 to 30,000 tokens, not the 512-token snippets a first transformer chapter usually assumes. The model architecture does not change. The attention formula does not change. What breaks is the hardware. A GPU that trains a small transformer in minutes suddenly stalls, runs out of memory, or both, the moment the sequence length N grows past a few thousand tokens — even though the number of floating-point operations the GPU is theoretically capable of has not gone down at all. FlashAttention is the fix, and understanding why it works requires looking past the attention formula itself and into how a GPU actually moves bytes around.
Attention is not slow because of arithmetic
Recall the self-attention computation from the transformers chapter: given query, key and value matrices Q, K, V ∈ ℝ^(N×d) for a sequence of length N and head dimension d,
S = QKᵀ / √d (N × N score matrix)
P = softmax(S) (row-wise, N × N)
O = P·V (N × d output)The number of floating-point operations here is Θ(N²d) — quadratic in sequence length, same as it has always been described. But a modern GPU does not have one speed of memory; it has a hierarchy. An NVIDIA A100 ships with 40–80 GB of High Bandwidth Memory (HBM) — the "GPU RAM" that holds your tensors between kernel calls — reachable at roughly 1.5–2 TB/s. Sitting directly on each of its 108 streaming multiprocessors is a sliver of on-chip SRAM, about 192 KB per multiprocessor, reachable at an estimated 19 TB/s — an order of magnitude faster, but four to five orders of magnitude smaller. A standard implementation of attention (the one every deep learning framework used before 2022) computes S in one kernel, writes the full N×N matrix out to HBM, reads it back for the softmax kernel, writes P back to HBM, reads it back again to multiply by V, and finally writes O to HBM. Every one of those N×N matrices — which does not fit in 192 KB of SRAM for anything but a tiny toy sequence — round-trips through the slow memory.
This is the crux of the whole optimization: attention's arithmetic intensity — floating-point operations performed per byte moved from memory — is low relative to what the GPU's compute units need to stay busy. A rough "roofline" calculation makes this precise. The A100's dense FP16 tensor cores deliver about 312 TFLOP/s; its HBM delivers roughly 1.5–2 TB/s. The ratio, 312 × 10¹² ÷ 2 × 10¹² ≈ 156 (or ≈ 208 at 1.5 TB/s), is the "ridge point": an operation needs at least this many FLOPs per byte moved to be compute-bound rather than memory-bound. The elementwise steps of attention — exponentiating S, normalizing rows, the extra write-then-read of two full N×N matrices — do only a handful of FLOPs per byte touched, nowhere near 150+. So even though the matrix multiplications QKᵀ and PV are individually efficient, the full attention block spends most of its wall-clock time waiting on HBM, not computing. This is why doubling the GPU's peak FLOP rate barely speeds up attention, while cutting its HBM traffic speeds it up directly and substantially.
The fix: tile the computation, never materialize the full matrix
FlashAttention (Tri Dao et al., 2022) restructures the same mathematics — it does not change the softmax formula or approximate it — so that the N×N matrix is never written to HBM at all. The trick is to split K and V into blocks along the sequence dimension, load one block of K and V into SRAM at a time, compute only that block's slice of scores, and fold it into a running estimate of the output. The obstacle is that softmax needs a normalization constant (the sum of exponentials across the *entire* row) before any individual weight can be finalized — you cannot normally compute softmax one piece of the row at a time without seeing the whole row first.
The resolution is the online softmax recurrence, which keeps three small running values per query row instead of the full row: a running maximum m (for numerical stability, exactly like the max-subtraction trick in ordinary softmax), a running sum of exponentials l, and a running weighted accumulator of the output acc. Each new block of scores contributes a local max, local sum, and local weighted value; these are folded into the running totals using a rescaling factor whenever the running max changes:
for each new block with local scores {s}, local values {v}:
m_block = max(s in block)
p_block = exp(s - m_block) # relative to this block's own max
l_block = sum(p_block)
acc_block = sum(p_block * v)
m_new = max(m_old, m_block)
scale_old = exp(m_old - m_new) # correction for the running total
scale_new = exp(m_block - m_new) # correction for this block
l_new = l_old * scale_old + l_block * scale_new
acc_new = acc_old * scale_old + acc_block * scale_new
m_new is carried forwardAfter the last block, dividing acc / l gives exactly the same output as computing the full softmax over all N scores at once — not an approximation, an algebraically equivalent regrouping of the same sum. Because each block only needs a Q tile, a K tile, a V tile and these three small running values in SRAM at once, the full N×N matrix never has to exist anywhere except implicitly, block by block, for a few microseconds before being discarded.
Worked example: online softmax by hand
Take one query row against N = 4 keys, split into two blocks of size 2. Let the (already √d-scaled) scores be s = [1, 3, 2, 5] and the corresponding values (using scalars instead of d-dimensional vectors, purely to keep the arithmetic visible) be v = [10, 20, 30, 40].
Direct softmax (what a non-tiled implementation would compute): subtract the global max 5, exponentiate, normalize.
exp(1-5)=0.01832 exp(3-5)=0.13534 exp(2-5)=0.04979 exp(5-5)=1.00000
sum = 1.20345
weights = [0.01522, 0.11246, 0.04137, 0.83095]
O = 0.01522·10 + 0.11246·20 + 0.04137·30 + 0.83095·40 = 36.8806Tiled online softmax, block size 2: process block 1 = {1, 3} first.
Block 1: local max m₁ = 3
p = exp([1-3, 3-3]) = [0.13534, 1]
l₁ = 1.13534
acc₁ = 0.13534·10 + 1·20 = 21.3534
running state: m = 3, l = 1.13534, acc = 21.3534Block 2 = {2, 5} arrives. Its local max (5) exceeds the running max (3), so the running total must be rescaled down before the new block is added — this rescaling is the entire trick that makes tiling exact.
Block 2: local max m₂ = 5
p = exp([2-5, 5-5]) = [0.04979, 1]
l₂ = 1.04979
acc₂ = 0.04979·30 + 1·40 = 41.4937
m_new = max(3, 5) = 5
scale_old = exp(3-5) = 0.13534
scale_new = exp(5-5) = 1
l = 1.13534·0.13534 + 1.04979·1 = 0.15366 + 1.04979 = 1.20345
acc = 21.3534·0.13534 + 41.4937·1 = 2.8897 + 41.4937 = 44.3834
O = acc / l = 44.3834 / 1.20345 = 36.880636.8806, identical to the direct computation to four decimal places. Nothing was approximated — the rescale-and-merge step exactly undoes the earlier normalization relative to the wrong (too-small) max and replaces it with the correct global one, block by block, without ever holding all four scores in memory simultaneously.
Tracing the algorithm as code
import numpy as np
def flash_attention_row(scores, values, block_size):
n = len(scores)
m = -np.inf # running max
l = 0.0 # running sum of exponentials
acc = 0.0 # running weighted sum (unnormalized)
for start in range(0, n, block_size):
block_s = scores[start:start+block_size]
block_v = values[start:start+block_size]
m_block = np.max(block_s)
p_block = np.exp(block_s - m_block)
l_block = np.sum(p_block)
acc_block = np.sum(p_block * block_v)
m_new = max(m, m_block)
scale_old = np.exp(m - m_new) # 0.0 on the very first block, since m = -inf
scale_new = np.exp(m_block - m_new)
l = l * scale_old + l_block * scale_new
acc = acc * scale_old + acc_block * scale_new
m = m_new
return acc / l
scores = np.array([1.0, 3.0, 2.0, 5.0])
values = np.array([10.0, 20.0, 30.0, 40.0])
print(flash_attention_row(scores, values, block_size=2))
# 36.880565891447...On the first loop iteration m is -inf, so scale_old = exp(-inf - 3) = exp(-inf) = 0.0 — NumPy evaluates this to exactly zero with no special-casing needed, correctly discarding a running total that does not exist yet. This is exactly the recurrence traced by hand above, and running it (verified independently, not asserted) reproduces 36.8806. A real GPU kernel does the same three-variable bookkeeping for every query row and every d-dimensional value vector at once, with the block loop implemented so that K_j and V_j are loaded into SRAM once and reused for every query block before being evicted.
Sizing a tile, and why the backward pass recomputes instead of storing
SRAM capacity puts a hard ceiling on block size B. A tile for one block needs roughly four B×d arrays in SRAM at once (Q, K, V blocks and the output accumulator), each element a 4-byte float: 16·B·d bytes. With a typical head dimension d = 128 and 192 KB of SRAM per streaming multiprocessor (196,608 bytes), B ≤ 196608 / (16·128) = 96 — consistent with the block sizes of 64–128 that real FlashAttention kernels use.
The same IO-aware logic extends to training. Backpropagation through attention normally needs the probability matrix P again to compute gradients — a standard implementation stores all of P from the forward pass for this purpose, which is exactly the O(N²) memory that tiling was trying to avoid. FlashAttention instead stores only the small per-row statistics (m and l, size O(N) total) from the forward pass and recomputes the needed P blocks on the fly during the backward pass, using saved Q, K, V tiles. This spends extra FLOPs — recomputation is not free — but FLOPs are cheap on a GPU relative to HBM bandwidth, which is precisely the imbalance this whole chapter has been about. Trading compute for memory traffic is a net win whenever the operation is memory-bound, which attention is.
Correcting a common misconception
The most common misunderstanding is that FlashAttention is a faster-but-approximate attention mechanism — grouped in students' minds with sparse attention (Longformer, BigBird, which skip computing most of the N×N matrix) or linear-attention approximations (which replace softmax with a kernel trick to avoid quadratic cost algebraically). It is neither. FlashAttention computes the exact same softmax attention function, with output numerically identical to standard attention up to ordinary floating-point rounding — the same kind of rounding difference you'd get from summing a list of numbers in a different order. Sparse and linear attention change what is computed to make it cheaper; FlashAttention changes only where and in what order the same computation happens in the memory hierarchy. That distinction is why it is safe to drop into any existing transformer — pretrained or not — without retraining or accuracy loss, whereas swapping in a sparse or linear attention variant is an architecture change with real accuracy trade-offs.
Why it matters at the scale the hook described
Two consequences follow directly from never materializing the N×N matrix. Peak memory for attention drops from O(N²) to O(N), since only O(d) of running state is needed per query row instead of an entire row of N scores — this is what makes 16K- or 32K-token context windows fit on the same GPU that used to run out of memory around 2K tokens. And because HBM traffic — not FLOPs — was attention's real bottleneck, the original paper reports substantial wall-clock speedups from the reorganization alone: on the order of 2–3× faster training on GPT-2-scale models and roughly 15% faster end-to-end training on BERT-large, with no change in the model's output. For an assistant meant to reason over an entire judgment or an entire policy report in one pass, that is the difference between a workable product and a GPU that runs out of memory before finishing the first document.
Active recall
Attempt each question before reading its answer.
1. Scores s = [2, 4, 1, 3] and values v = [5, 15, 25, 35], block size 2. Compute the output using tiled online softmax, showing whether a rescale is needed at the second block.
2. Why is standard attention described as "memory-bound" rather than "compute-bound" on a GPU? Use the roofline idea (FLOPs per byte) in your answer.
3. True or false, with justification: FlashAttention produces a slightly different, approximate output because it changes the order in which the softmax sum is computed.
4. SRAM per streaming multiprocessor is 192 KB and each tile needs roughly 16·B·d bytes for head dimension d = 64. Estimate the largest block size B that fits.
5. Why does FlashAttention recompute the probability matrix P during the backward pass instead of storing it from the forward pass, and why is this a good trade on a GPU specifically?
6. In the loop order that keeps one K,V block resident in SRAM across the whole inner loop over Q blocks, that K,V block is read from HBM exactly once, but each Q block gets reread once per outer iteration. For N = 4096 and block size B = 128 (32 blocks), how many total Q-block reads happen across the whole algorithm, versus the ideal minimum of reading each Q block once?
Answers.
1. Block 1 = {2,4}: local max 4, p = [exp(-2),1] = [0.1353,1], l₁ = 1.1353, acc₁ = 0.1353·5+1·15 = 15.6767. Block 2 = {1,3}: local max 3, which is less than the running max 4, so m_new = max(4,3) = 4 stays unchanged — scale_old = exp(4-4) = 1 (no rescale of the old total needed), while the new block is scaled by scale_new = exp(3-4) = 0.3679. Local block values relative to its own max 3: p = [exp(-2),1] = [0.1353,1], l_block = 1.1353, acc_block = 0.1353·25+1·35 = 38.3835. Merge: l = 1.1353·1 + 1.1353·0.3679 = 1.5531, acc = 15.6767 + 38.3835·0.3679 = 29.7971. O = 29.7971/1.5531 = 19.187, matching a direct softmax over all four scores exactly.
2. Attention's matrix-multiply steps (QKᵀ and PV) are compute-efficient, but the softmax and the extra HBM write/read of the full N×N matrices perform only a handful of FLOPs per byte moved. An A100 needs roughly 150–200 FLOPs per byte moved (312 TFLOP/s ÷ ~1.5–2 TB/s) to keep its compute units the bottleneck rather than its memory bus. Attention's effective intensity, once you count the wasted round-trips of S and P through HBM, falls well under that ridge point, so the GPU spends most of its time waiting on memory transfers, not computing — the definition of memory-bound.
3. False. The rescale-and-merge recurrence exactly undoes the normalization against the "wrong" (too-small) local max and replaces it with the correct global one at every step — it is an algebraic identity, not an approximation, and the worked example above reproduces the direct softmax result exactly (36.8806 both ways). Any tiny discrepancy that shows up in practice is ordinary floating-point rounding, the same kind you'd see from summing any list of numbers in a different order — not a modeling approximation like sparse or linear attention.
4. 16·B·64 ≤ 196608 → B ≤ 192. Halving the head dimension from 128 to 64 doubles the usable block size, since each tile element takes half as much SRAM.
5. Storing P from the forward pass would need O(N²) memory again — exactly what tiling was built to avoid — while storing only the small per-row statistics (the running max and sum, O(N) total) and recomputing P block-by-block during the backward pass costs extra FLOPs but no extra HBM traffic. Since attention is memory-bound, GPU compute is comparatively cheap and idle time is the expensive resource, so spending more FLOPs to save HBM bandwidth is a favorable trade specifically on this hardware.
6. Each of the 32 K,V blocks is held resident while all 32 Q blocks are streamed through it, so a Q block is reloaded once per outer K,V iteration: 32 (K,V blocks) × 32 (Q blocks) = 1024 total Q-block reads, versus an ideal single pass that would read each of the 32 Q blocks only once. This 32-fold redundant rereading of Q (cheap, since Q is only O(N·d) total data) is exactly the tradeoff accepted to avoid ever materializing the O(N²) matrix — and it is also the specific inefficiency that a later version of the algorithm, FlashAttention-2, restructures the loop order to reduce.
Think About It
Think about this: How would you explain flash attention optimization to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind flash attention optimization, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.