Part IV · Ch. 19 — Inference, Efficiently

Part IV · Chapter 19 of 21

Inference, Efficiently

KV cache, context windows, and quantization

You’ll need:the loop: sample a token, append it, run the machine againqueries, keys, values; one row of scores per asker; the causal maskthe full stack — layers, heads, d_model, and where the parameters livesampling variation, long-run frequencies, and when errors can average out

Two speeds, one model

We trained the weights in Chapter 15. Inference uses them. The arithmetic now has two rhythms: read the prompt together, then generate one token at a time. Each chosen token must arrive before we can choose its successor.

Our KV cache keeps earlier attention keys and values. Prefill fills that cache; decoding extends it. Reusing this work speeds generation without updating the weights.

What changes between the two phases?
PhaseNew workDependency
PrefillProcess all prompt positionsKnown tokens allow parallel processing with a causal mask
DecodeProcess the newly appended tokenWait for the previous choice; reuse cached keys and values

Extend Chapter 10's invented cast with “on”: $\mathbf q_{\text{on}}=[2,1]$, $\mathbf k_{\text{on}}=[1,0]$, $\mathbf v_{\text{on}}=[0,1]$. These are supplied toy vectors; we omit the usual attention-score scaling here.

$$\mathbf{q}_{\text{on}} \cdot \mathbf{k}_{\text{cat}} = 2 \times 1 + 1 \times 1 = 3$$

In words: line on's question up against cat's key, multiply matching slots and add, for a score of 3.

Worked example: on's four scores

$\mathbf{q}_{\text{on}} \cdot \mathbf{k}_{\text{the}} = 2 \times 1 + 1 \times (-1) = 1$;  $\mathbf{q}_{\text{on}} \cdot \mathbf{k}_{\text{cat}} = 2 \times 1 + 1 \times 1 = 3$;  $\mathbf{q}_{\text{on}} \cdot \mathbf{k}_{\text{sat}} = 2 \times 0 + 1 \times 1 = 1$;  $\mathbf{q}_{\text{on}} \cdot \mathbf{k}_{\text{on}} = 2 \times 1 + 1 \times 0 = 2$.

Four scores: $(1,\, 3,\, 1,\, 2)$. No earlier vector was recomputed to get them; the keys came out of the cache.

$$(a_{\text{the}},\, a_{\text{cat}},\, a_{\text{sat}},\, a_{\text{on}}) = \mathrm{softmax}(1,\, 3,\, 1,\, 2) = \left(\tfrac{2.72}{32.92},\, \tfrac{20.09}{32.92},\, \tfrac{2.72}{32.92},\, \tfrac{7.39}{32.92}\right) = (0.08,\, 0.61,\, 0.08,\, 0.22)$$

In words: turn the four scores into four shares of on's attention. We reuse Chapter 10's symbol $a_i$ for an attention share. We use "a" rather than $w$ because these are computed on the fly. So "on" leans hard on "cat" (0.61), gives a fifth to itself (0.22), and spares a little for "the" and "sat" (0.08 each). Rounded to two decimals these shares read $0.08 + 0.61 + 0.08 + 0.22 = 0.99$, the same rounding artifact as in Chapter 10.

Worked example (optional): the value mix

The output is the shares blending the values on file, as in Chapter 10:

$$\mathbf{o}_{\text{on}} = 0.08\,[0, 2] + 0.61\,[3, 1] + 0.08\,[1, 0] + 0.22\,[0, 1] = [1.91,\, 0.99]$$

In words: scale each stored value vector by its share and add them slot by slot. Slot 1: $0.08 \times 0 + 0.61 \times 3 + 0.08 \times 1 + 0.22 \times 0 = 1.91$. Slot 2: $0.08 \times 2 + 0.61 \times 1 + 0.08 \times 0 + 0.22 \times 1 = 0.99$. Then file $\mathbf{k}_{\text{on}}$ and $\mathbf{v}_{\text{on}}$ into the cache, ready for the next token.

A 4-by-4 attention-share grid in the same style as Chapter 10's heatmap, rows labeled 'asking' (the, cat, sat, on) and columns labeled 'being scored' (the, cat, sat, on), cells brighter for larger shares. The upper-right triangle above the diagonal is blanked to dark panels with muted em-dashes (the causal mask). Rows one to three (the, cat, sat) are dimmed to about a third opacity, with a slim muted bracket beside them labeled 'computed in earlier steps — read from cache, never redone'. The bottom row (on) renders at full strength with values 0.08, 0.61, 0.08, 0.22 and is outlined by a bright rounded rectangle, annotated 'the only new work this step'.
Each new token adds just one row.
A left-to-right timeline in two phases split by a vertical hairline. Left phase, labeled 'prefill — the pause': five token chips packed tightly, with one wide bracket beneath them labeled 'processed together, one parallel pass', and an arrow down into a cache drum (a cylinder labeled 'KV cache, keys plus values per layer') that shows five tick squares inside. Right phase, labeled 'generation — the stream': three token chips appearing one at a time with tick labels t+1, t+2, t+3, each with a thin arrow down into a second cache drum showing six, seven, then eight squares, annotated 'each new token: compute 1, read the rest'. Bottom caption: 'same machine, two speeds — a big parallel gulp, then steady sips'.
Prefill is parallel; decoding is sequential.

The bill per token

A prompt of $n$ tokens has $n^2$ query–key pairs before masking. The mask removes future-looking pairs but leaves quadratic growth. Ten times the prompt length means roughly a hundred times the attention pairs:

$$n = 1{,}000 \;\rightarrow\; n^2 = 1{,}000{,}000 \qquad n = 10{,}000 \;\rightarrow\; n^2 = 100{,}000{,}000$$

In words: ten times the context means a hundred times the score grid, per layer and per head.

A line chart. X axis: context length n in tokens from 0 to 10,000. Y axis: relative cost, normalized at n equals 1,000, from 0 to about 105. A curve rising steeply as n squared is labeled 'prefill scores — grows as n squared', with two dots annotated '1,000 tokens gives 1,000,000 scores' near the bottom and '10,000 tokens gives 100,000,000 scores' near the top. A straight line rising gently is labeled 'KV-cache memory — grows as n'. A muted footnote reads 'per layer, per head — the shapes are the story, not the units'.
Scores grow quadratically; cache linearly.

One cached decode step compares its new query with $n$ stored keys: linear growth in context length. It still runs the new token through every layer. Caching removes repeated computation of old keys and values; it does not make a longer context free. Actual speed also depends on memory traffic, hardware, and implementation.

What the cache stores

Each transformer layer stores its own keys and values. For causal attention, earlier tokens cannot attend to the appended token, so their states remain valid when we extend a fixed prefix at the same positions. The query is needed now; old queries need not be saved for later attention.

Worked example: the cache bill

Take a conventional model that runs several attention heads side by side, the parallel copies from Chapter 11. With width 12,288, 96 layers and two bytes per stored number it needs $2\times12{,}288\times96\times2=4{,}718{,}592$ bytes per token for keys and values. For 2,048 tokens that is about 9.66 GB, before other memory. Architectures that share key/value heads reduce this bill.

Batching improves hardware use, but every request needs its own context state. Speculative decoding proposes several tokens with a cheaper model, then checks them with the target model. With the proper acceptance rule, it preserves the target sampling distribution.

The context window

The context window limits the prompt plus generated continuation. Model training, positional encoding, and runtime memory all matter. A larger cache alone does not teach a model to use arbitrary distances reliably.

At the limit, an application must stop, discard older material, or summarize it. That summary is another input, with possible omissions. Saved chat histories or notes are application storage, distinct from the temporary KV cache and trained weights.

Shrinking the numbers

Bits determine how many patterns we can store. Quantization replaces fine-grained values with a smaller set. Four bits provide sixteen patterns; our symmetric toy uses integers −7 through 7, leaving one unused.

Three horizontal bands. Top band: a muted number line from minus 0.7 to plus 0.7 with tick marks every 0.1, labeled minus 0.7, 0, plus 0.7, and shelf numbers minus 7 to 7 beneath each tick; a title reads 'the 4-bit ruler — 16 possible stored values, spacing = scale = 0.1'. Middle band: four dots above the line at 0.34, minus 0.12, 0.70, and 0.23 with labels, each dropping a thin arrow to its nearest shelf tick, landing on a square at 0.3, minus 0.1, 0.7, 0.2, with stored values 3, minus 1, 7, 2 beneath. Bottom band: two pill rows, '32-bit: 0.3400000035… (4 bytes each)' and '4-bit: shelf 3, ruler times 0.1 (half a byte each)', joined by an arrow.
Fewer bits give fewer possible levels.

Choose a scale from the largest magnitude, divide each weight by it, and round to the nearest supported integer. Multiply by the scale to reconstruct:

$$\text{scale} = \frac{0.7}{7} = 0.1 \qquad 0.34 \;\rightarrow\; \text{store } 3 \;\rightarrow\; 3 \times 0.1 = 0.30$$

In words: store each weight as its nearest shelf number, a small whole number that fits in 4 bits. To use it later, multiply the shelf number back by the scale. So 0.34 rounds to shelf 3, and rebuilding gives $3 \times 0.1 = 0.30$, a little under the original.

Worked example: the round-trip, all four weights

Nearest shelf numbers: $[0.34,\, -0.12,\, 0.7,\, 0.23] \rightarrow (3,\, -1,\, 7,\, 2)$. Rebuild each by multiplying back by the scale 0.1:

$3 \times 0.1 = 0.30$ (was 0.34);  $-1 \times 0.1 = -0.10$ (was −0.12);  $7 \times 0.1 = 0.70$ (exact);  $2 \times 0.1 = 0.20$ (was 0.23).

Weight storage for seven billion parameters
Bits per weightNominal weight bytesTradeoff
3228 GBFinest precision here
1614 GBCommon floating-point storage
87 GBFewer representable levels
43.5 GBMore rounding; quality needs checking
$$28\ \text{GB} \;(32\text{-bit}) \;\rightarrow\; 14\ \text{GB} \;(16\text{-bit}) \;\rightarrow\; 7\ \text{GB} \;(8\text{-bit}) \;\rightarrow\; 3.5\ \text{GB} \;(4\text{-bit})$$

In words: four bits instead of thirty-two makes the stored weights eight times smaller.

Worked example: rounding reaches the output

With input $[1,1,1,1]$, the original row sums to $0.34-0.12+0.70+0.23=1.15$. Its reconstruction gives $0.30-0.10+0.70+0.20=1.10$. The difference is 0.05, about 4.35% of the original result. Smaller storage has changed the arithmetic.

Real formats often give small groups of weights separate scales. Scales and metadata add overhead; activations and the cache need memory too. Rounding errors can accumulate, and sensitive weights need care. Measure task quality at the chosen precision.

See it move

Watch (1:47): what to notice. The counter climbs while attention redraws every earlier link from scratch, the violet shelf of stored keys and values stops that climb at one new box per token, the shelf runs into the wall of the context window, and a 32-bit number shrinks to 8 so four times as much fits.

Why the LLM cares

A card marked “7B, Q4, 8K” describes roughly seven billion parameters, four-bit weight storage, and an 8,192-token context. The nominal 3.5 GB of weights is only one memory bill. Cache size and serving overhead still matter. We can now distinguish a smaller representation from fewer parameters, and faster reuse from learning during a conversation.

What you now know

  • Prefill processes the prompt; decoding extends it one token at a time.
  • The KV cache reuses earlier keys and values at every layer.
  • Prompt attention grows quadratically; one cached query grows linearly with context.
  • Quantization reduces precision and storage, with measurable output error.
  • Weights, cache, activations, and metadata all consume memory.

Where we're headed

We have a limited workspace for context. Next we build a way to fetch the few passages that belong there: retrieval and vector search.