Two of the era's most elegant results trade arithmetic for memory traffic, and speed for nothing at all. Both are exact: neither changes a single output.
Consider two claims that sound like they cannot both be true. The fastest known way to compute attention does more arithmetic than the obvious way. And there is a method that makes a large model generate text two or three times faster while provably producing exactly the same distribution of outputs it would have produced anyway. Both are true, and both follow from the previous chapter's finding about what is actually scarce.
Chapter 6 counted attention's arithmetic as \( 2n^2 d \). At realistic sequence lengths, however, the binding cost is not those operations but the \( n \times n \) matrix of scores travelling to and from main memory. The output of attention is only \( n \times d \) values; the intermediate is \( n^2 \). That gap — \( \Theta(n^2) \) traffic for a \( \Theta(nd) \) result — is what an IO-aware kernel closes.
The intermediate does not fit
At \( n = 4096 \) in half precision, one head's score matrix occupies \( 2n^2 = 2 \times 4096^2 = 33.6 \) MB — larger than the reference accelerator's entire 20 MB of fast on-chip memory, and that is per head, per layer. You cannot keep it on chip; you either round-trip it through main memory at great cost, or you never form it at all.
Never forming it requires computing a softmax over values you see only a block at a time. That is possible, and exactly, because the softmax normalizer can be maintained incrementally.
Streaming softmax
To normalize safely we subtract the running maximum before exponentiating, so nothing overflows. Maintain a running maximum \( m \) and a running sum \( \ell = \sum e^{s_i - m} \) over the scores seen so far. On receiving a new block of scores \( \{s_i\} \):
\[ m' = \max\bigl(m,\ \max_i s_i\bigr), \qquad \ell' = \ell\,e^{\,m - m'} + \sum_i e^{\,s_i - m'}. \]The factor \( e^{m-m'} \) rescales the old partial sum into the new frame — it is never larger than 1, so the update is numerically safe. The same rescaling is applied to a running weighted sum of value vectors, so the attention output can be accumulated block by block without the score matrix ever existing in full.
Verify that this is exact. Let scores arrive as \( \{2, 1\} \) then \( \{3, 0\} \). After the first block, \( m = 2 \) and \( \ell = e^0 + e^{-1} = 1.368 \). After the second, \( m' = 3 \) and \( \ell' = 1.368\,e^{-1} + (e^0 + e^{-3}) = 0.503 + 1.050 = 1.553 \). Computing in one shot instead: \( e^{-1} + e^{-2} + e^{0} + e^{-3} = 0.368 + 0.135 + 1 + 0.050 = \) 1.553. Identical. The streaming form is not an approximation.
With this, attention is computed in tiles small enough to sit in fast memory: load a block of queries, stream blocks of keys and values past it, accumulate. Main-memory traffic falls by roughly an order of magnitude. The backward pass recomputes the blocks it needs rather than storing them — more arithmetic, less memory — and the kernel wins anyway. That is the roofline lesson in its purest form.
This technique is exact. It computes the same attention function, to floating-point tolerance, as the naive implementation. Approximate attention — sparse patterns, low-rank surrogates — is a different trade entirely, buying speed with fidelity. Confusing the two is the most common misreading here: an IO-aware kernel changes how a quantity is computed, not what is computed.
Scores arrive as \( \{1, 4\} \) then \( \{2, 2, 5\} \). Run the streaming update and verify against the one-shot sum.
Block 1: \( m = 4 \), \( \ell = e^{-3} + e^{0} = 1.050 \). Block 2: \( m' = 5 \), \( \ell' = 1.050\,e^{-1} + (e^{-3} + e^{-3} + e^{0}) = 0.386 + 1.100 = 1.486 \). One shot: \( e^{-4} + e^{-1} + e^{-3} + e^{-3} + e^{0} = 0.018 + 0.368 + 0.050 + 0.050 + 1 = 1.486 \). ✓
The second technique attacks the same waste from a different side. Decoding is memory-bound (Chapter 11), so the accelerator's arithmetic sits idle while weights stream past. Suppose a small, cheap draft model proposes the next \( k \) tokens; the large target model can then check all \( k \) in a single forward pass, because verifying several positions at once is a prefill-shaped, compute-bound operation — nearly free relative to \( k \) separate decode steps.
The obvious objection is that the output would then be the draft's, not the target's. The answer is a rejection-sampling scheme that makes the final distribution exactly the target's.
The obvious objection comes first, and you should feel its force before seeing it dissolved. If a small, weaker model is proposing the words, surely the output is contaminated by that weaker model — faster, perhaps, but subtly worse. It is an entirely reasonable worry. It is also, as the following shows, exactly wrong: the acceptance rule is constructed so that the weaker model's preferences cancel out of the final distribution completely.
Exactness by rejection sampling
Let the target assign probability \( p(x) \) to token \( x \) and the draft assign \( q(x) \). Sample \( x \sim q \) and accept it with probability \( \min\bigl(1, p(x)/q(x)\bigr) \). If rejected, sample instead from the normalized residual \( \operatorname{norm}\bigl(\max(0,\ p - q)\bigr) \).
Then the probability of emitting \( x \) is the accepted mass plus the resampled mass. The accepted mass is \( q(x)\min(1, p/q) = \min(q(x), p(x)) \). Summed over all tokens, the total accepted probability is \( \sum_x \min(p,q) \), so rejection occurs with probability \( 1 - \sum_x \min(p,q) = \sum_x \max(0, p-q) \). The residual distribution places mass \( \max(0, p(x)-q(x)) / \sum_x \max(0,p-q) \) on \( x \). Multiplying and adding:
\[ \min(p,q) + \max(0, p-q) = p(x). \]The emitted distribution is exactly the target's. Speculation changes speed, never samples.
How much speed? If each drafted token is accepted independently with rate \( \alpha \), a verification cycle keeps the accepted run plus one token the target supplies regardless.
Expected tokens per cycle
The run of accepted tokens has length \( i \) with probability \( \alpha^i(1-\alpha) \) for \( i < k \), and length \( k \) with probability \( \alpha^k \); each case yields \( i+1 \) tokens. Summing the truncated geometric series,
\[ \mathbb{E}[\text{tokens per cycle}] = \frac{1 - \alpha^{\,k+1}}{1 - \alpha}. \]If the draft costs \( c \) target-forwards per token, a cycle costs \( kc + 1 \) target-forwards. The speedup is the ratio.
The economics, end to end
Take \( \alpha = 0.8 \), \( k = 4 \), \( c = 0.1 \). Tokens per cycle: \( (1 - 0.8^5)/0.2 = (1-0.328)/0.2 = 3.36 \). Cycle cost: \( 4(0.1) + 1 = 1.4 \) target-forwards. Speedup: \( 3.36/1.4 = \) 2.4× — with provably identical output.
Sensitivity matters more than the headline. At \( \alpha = 0.5 \) the same configuration gives \( (1 - 0.5^5)/0.5 = 1.94 \) tokens per 1.4 forwards, a mere \( 1.38\times \). The draft's agreement with the target is everything; its standalone quality matters only through \( \alpha \).
(a) The draft assigns \( q(A) = 0.5 \), the target \( p(A) = 0.3 \). What is the acceptance probability for A? (b) With \( \alpha = 0.7, k = 3, c = 0.15 \), what is the speedup?
(a) \( \min(1, 0.3/0.5) = 0.6 \). (b) Tokens \( = (1 - 0.7^4)/0.3 = (1-0.2401)/0.3 = 2.53 \); cost \( = 3(0.15)+1 = 1.45 \); speedup \( = 1.75\times \).
Choosing \( k \) is a marginal calculation: extending the draft by one more token adds an expected \( \alpha^{k+1} \) tokens and costs \( c \) more. When those balance, you are at the optimum — and because the curve is flat near it, a range of \( k \) performs about equally.
Both techniques in this chapter trade an abundant resource for a scarce one, and both are exact — a combination worth pausing on, since it is rarer than it sounds. Streaming softmax lets attention be computed in tiles, cutting memory traffic by an order of magnitude while doing slightly more arithmetic and recomputation. Speculative decoding converts \( k \) memory-bound passes into one compute-bound verification, with a rejection rule that provably preserves the target distribution. Both illustrate the discipline this book keeps returning to: identify the binding budget, then spend freely from the other. Chapter 13 continues in the same vein but changes the object — instead of computing the same model more cleverly, it makes the model itself sparser and cheaper.
A · Drills
B · Problems
C · Challenge
Reproduce Derivation 12.1 (streaming softmax, including the value accumulator's rescaling), Derivation 12.2 (the exactness proof), and Derivation 12.3 (expected tokens per cycle) on blank paper. Verify a streaming example numerically by hand. Attempt B-2 and B-3 closed book. You pass when you can state, without hedging, that both techniques are exact and explain why.