Making Inference Cheap: KV Cache, Batching and Quantization

To generate a token, a decoder-only transformer takes the whole sequence so far, runs it through every layer, and reads the probability distribution off the last position. Then it appends the chosen token and does it again. That description is correct, and it is what the slow run does.

It is also enormously wasteful, and the reason is the causal mask from Chapter 9. Position 7's key and value vectors are computed from position 7's input and the positions before it. When position 8 arrives, position 7 cannot see it — that is what causal means — so position 7's key and value are exactly the same numbers as before. The slow run recomputes them anyway, at every step.

So store them. That store is the key-value cache, the single most consequential optimisation in language model serving:

out = model(prompt_ids, use_cache=True)          # prefill: the whole prompt
past = out.past_key_values                        
nxt = out.logits[:, -1].argmax(-1, keepdim=True)

for _ in range(n - 1):
    out = model(nxt, past_key_values=past, use_cache=True)   
    past = out.past_key_values                                
    nxt = out.logits[:, -1].argmax(-1, keepdim=True)

Look at what is fed to the model inside the loop: nxt, one token. Not the sequence. The new token's query attends against every cached key, and the cached keys were never going to change. This is not an approximation — the identical-output check above is the point. The cache does not trade quality for speed; it deletes redundant arithmetic.

To see the scaling cleanly, strip the transformer away and time a single attention head with $d = 64$, one step of generation computed both ways:

The right-hand column is the cause. Recomputing builds the full $n \times n$ attention matrix every step — the $O(n^2)$ from Chapter 9's asymptotic-notation box, paid once per token. With the cache you build a $1 \times n$ row instead: at 4,096 tokens, 67 MB of scores against 16 KB.

Counting multiply-accumulates instead of milliseconds removes the machine from the argument. To generate $T$ tokens from a cold start:

Per step the cached version is linear in the context and the uncached one quadratic; summed over a generation, $O(T^2)$ against $O(T^3)$, with the ratio growing without limit. The nine-fold difference in the opening was measured over 48 tokens — short of that table's first row.

The cache also changes what has to be in memory. On an 8 GB laptop GPU generating 256 tokens in fp16, taking the allocator's peak and subtracting the resident weights:

8.3 times more memory, spent to produce the same tokens more slowly. This is the promise made in Chapter 5, arriving from an unexpected direction: there, reverse-mode autodiff had to keep every intermediate alive for the backward pass, and activations dominated training memory. At inference there is no backward pass and nothing to retain for it — so what dominates memory instead is the cache, and it is a deliberate choice rather than an unavoidable cost.

Look again at the fast run: its first token behaved unlike the other forty-seven.

prefill, 40 prompt tokens : 1.0224 s   ->  25.6 ms per token
decode,  47 steps         : 0.1665 s mean per step

The prompt cost 25.6 ms per token and each generated token cost 166 ms. Same model, same hardware, same weights, a six-fold difference per token — and it goes the way most people do not expect. The prompt is the cheap part. Generation splits into two phases with genuinely different physics:

Both halves are measurable. Prefill, one pass over $P$ tokens:

Decode, one token against a cache of $C$:

Read the second table twice. Going from 16 tokens of context to 1,024 — sixty-four times more history to attend over — changed the cost of a step by nothing measurable. Attention against the cache is real work, but it is dwarfed by the fixed cost of dragging half a billion weights through the memory bus to produce one vector. That fixed cost is the reason for everything in the next section.

These two phases are the origin of the two numbers every serving system reports. Time to first token is essentially prefill, and it grows with the prompt, which is why a long conversation feels slow to start. Tokens per second is $1/\text{decode step}$, and it is roughly constant, which is why the reply then flows evenly. A chat that starts slow and then streams smoothly is not a rendering trick. It is these two tables.

The cache trades arithmetic for memory, and the memory it wants is not small. For every token in the context, every layer holds one key vector and one value vector per key-value head:

The 2 is for keys and values; everything else is the architecture. For the model measured throughout this chapter — 24 layers, 14 query heads, 2 key-value heads, head dimension 64 — in fp16 that is $2 \times 24 \times 2 \times 64 \times 2 = 12{,}288$ bytes per token.

Formulae in this field have a habit of being off by a factor of two, so check it against the allocator rather than believing it:

KV cache tensors per layer: (1, 2, 295, 64) float16
measured: 3,624,960 bytes for 295 tokens = 12,288 bytes/token
formula : 2 * 24 * 2 * 64 * 2                = 12,288 bytes/token

Exact, and it stays exact across every shape tried:

The last three rows deserve a second look. Thirty-two users with 2,048 tokens each, sixty-four with 1,024, one hundred and twenty-eight with 512 — the cache is 768 MB in every case, because all three hold 65,536 tokens. The cache depends only on the total number of tokens resident, not on how they are distributed among users. That fact is the foundation of the batching section.

Chapter 9 introduced multi-query and grouped-query attention and deferred the reason to this chapter. The reason is that formula, and specifically the $H_{kv}$ in it.

Standard multi-head attention gives every query head its own key and value heads. The model here has 14 query heads; with full multi-head attention its cache would be $2 \times 24 \times 14 \times 64 \times 2 = 86{,}016$ bytes per token — 84 KB instead of 12 KB, exactly seven times more, the ratio of query heads to key-value heads.

Multi-query attention¹ takes this to the limit: all query heads share a single key-value head. Grouped-query attention² is the compromise that won — a handful of key-value heads, each shared by a group of query heads — because MQA's quality loss was real and GQA's is not. Neither buys any arithmetic. They exist to divide that formula by an integer, and they spread across the industry the moment long contexts made the cache the binding constraint.

Which it does, quickly. For a 7B-class model with 32 layers and 8 key-value heads of dimension 128, the cache is 128 KB per token in fp16:

That model's own weights are 13.0 GB in fp16, the figure in the table at the end of this chapter. So at a 128,000-token context, one user's cache is larger than the model. This is the arithmetic Chapter 16 turns into money, and it is why a long conversation is not merely a slow one — it occupies a fixed slice of a machine for as long as the request is alive.

Decode is memory-bound: the weights are dragged through the bus to produce one token, and the arithmetic units idle. So put more work in the same step. Run several requests at once, and the weights, read once, serve all of them. Measured on the same model, each request holding a 64-token cache and decoding one token:

Read the two right-hand columns against each other, because they are the whole point. Going from one request to sixteen multiplies throughput by 6.0 and multiplies the wait for any individual request by 2.67. The batch made the server better and every user worse.

That is not a bug to be tuned away; it is the trade itself, and it has a name on each side. Latency is what a person waiting for a reply experiences. Throughput is what the invoice is divided by. No setting improves both.

Note also where it stops. From 16 to 32, throughput gains 9 % while latency nearly doubles: the step has stopped being memory-bound and become compute-bound, and past that knee the batch buys nothing. Every deployment has such a knee; its location has to be measured on yours, but its existence does not.

The naive way to batch is to collect $B$ requests, run them together, and return when all are done. But they do not finish together: some replies are twenty tokens and some five hundred. A fixed batch runs until its longest member finishes, and every finished request keeps occupying its slot, contributing padding, until then.

Take 64 requests with a realistic skew of output lengths — median 18 tokens, longest 231, 1,874 in total — and simulate both policies at the measured per-step cost for eight slots:

Throughput improves by 1.6x. Mean latency improves by more than ten times, because under static batching a request that finished in four steps still waits for a 231-token neighbour before anyone hears about it.

Continuous batching³ is the fix, and it is as simple as it sounds: the batch is not a group but a set of slots, and a slot that frees admits the next queued request on the very next step. The scheduler works at the granularity of one token rather than one request. Every serving stack in production does this now.

It has a second half, which is the cache. Slots that come and go leave cache memory fragmented, and reserving each slot its maximum possible context wastes most of the reservation. PagedAttention⁴ borrows the answer from operating systems: store the cache in fixed-size blocks with a block table per sequence, so a sequence's cache can be physically scattered while remaining logically contiguous — which also lets two sequences with a shared prefix share the blocks holding it. That is what vLLM is built on, and why a serving engine is a memory allocator with a transformer attached.

The other half of the bill is the weights themselves. Half a billion parameters at four bytes each is 1.98 GB; at two bytes, 0.99 GB; at one byte, 0.49 GB. Fewer bits per weight shrinks the model on disk, shrinks it in memory, and — because decode is bandwidth-bound — makes each step faster, since there are fewer bytes to move.

The simplest scheme is symmetric absolute-maximum quantization, and it fits in three lines:

qmax  = 2 ** (bits - 1) - 1
scale = W.abs().max() / qmax                        
Wq    = torch.round(W / scale).clamp(-qmax - 1, qmax)
W_hat = Wq * scale                                  # dequantized

Pick a scale so the largest weight maps to the largest integer, divide, round, store the integers and the scale. Reconstruct by multiplying back. Nothing about it is clever, and it works — right up until it does not.

Measured on the real weights of the model: all 168 projection matrices, 357.8 million parameters, relative error $\lVert W - \hat{W}\rVert / \lVert W \rVert$:

The fourth row is the collapse. A relative error of 0.99 on the worst matrix means the reconstruction retains essentially nothing of the original — the matrix has been replaced by noise of about the right magnitude. The cause is visible in the same experiment on a single matrix:

model.layers.12.mlp.down_proj.weight   (896 x 4864)
mean |w| 0.01386   std 0.01822   max |w| 0.43945   max/std 24.1
weights beyond 6 sigma: 692 of 4,358,144   (0.016 %)

One weight in six thousand sits beyond six standard deviations, and the largest is 24 out. With a single scale for the whole matrix, that one weight sets the step size for all 4.3 million of them. At 8 bits there are 256 steps and the typical weight still lands on a meaningful one. At 4 bits there are 16, the outermost reserved for a value almost nothing has, and the ordinary weights — which is all of them — round to two or three distinct levels.

Everything after that row is the same repair at different granularities: give the scale a smaller territory. Per output row divides the error by 3.4; per group of 128 consecutive weights divides it again. The cost is bookkeeping — a 16-bit scale per group of 128 is $4 + 16/128 = 4.125$ bits per weight instead of 4 — and it buys most of the gap back.

NF4 goes at it from the other side.⁵ The levels do not have to be equally spaced. Weights within a block are approximately normally distributed, so choose the sixteen levels as the quantiles of a normal distribution: dense near zero where the weights actually are, sparse in the tails where they are not. Same four bits, same block scaling, at a smaller block — 4.25 bits per weight against group-128's 4.125 — and the measured error drops from 0.1323 to 0.0952, 28 % lower. Part of that is the finer block and the rest is putting the levels where the mass is, and separating the two would need a third row.

Chapter 2's floating-point box ended with a promise: that this chapter would quantize weights to 8 and 4 bits and find a handful of outlier features refusing to be squeezed. Here they are, and they explain why "just round the numbers" was never going to work on activations.

The weights above were badly behaved. The activations are in a different league. Take an ordinary 84-token prompt, capture the residual stream at each layer, and measure the largest magnitude each of the 896 dimensions reaches:

Dimension 62 reaches 1,579.6 while the median dimension never exceeds 1.6. It is not a fluke of one token or one layer: the same dimension is there at layer 4 and still there at layer 20, with almost the same value. These are the outlier features,⁶ and they are systematic — a property of the trained model, not of the input.

The histogram of those 896 per-dimension maxima at layer 16 makes the shape unmistakable:

     0 -      1 | ######################################## 254
     1 -      2 | ######################################## 283
     2 -      4 | ######################################## 226
     4 -      8 | ######################################## 93
     8 -     16 | ##################                       18
    16 -     32 | #########                                9
    32 -     64 | #######                                  7
    64 -    128 | #####                                    5
   128 -    256 |                                          0
   256 -    512 |                                          0
   512 -   1024 |                                          0
  1024 -   4096 | #                                        1

Nine hundred dimensions in a tidy pile below 8, nothing at all for three octaves, then one dimension alone at the far end. Now quantize that tensor to INT8 and count what happens:

Fourteen levels out of 256. The scale was set by 1,579.6, so every step is 12.44 wide, and the typical activation — median magnitude 0.26, ninety-ninth percentile 2.51 — has nowhere to land. Per dimension it is starker:

single tensor-wide scale = 12.4378
  dim 826 (max |h| = 4.77):  1 distinct level out of 256
  dim 336 (max |h| = 1.62):  1 distinct level out of 256
  dim  96 (max |h| = 0.69):  1 distinct level out of 256

after excluding the top 4 dimensions, scale = 0.5749  (22x smaller)
  dim 826: 8 levels    dim 336: 4 levels    dim  96: 3 levels

One level. The whole dimension, every token, quantized to the same number. Eight bits were allocated and roughly zero were used, and the model reading those activations is handed a constant.

That measurement is the justification for every technique people actually use:

Almost every article about quantization stops at the previous section: it explains the method, quotes a compression ratio, and asserts that quality is "largely preserved". Chapter 4 was about not fooling yourself, so let us find out.

Same model, weights quantized in place with each scheme, then three measurements: perplexity on 2,048 tokens of held-out English prose — here, the draft of this course, which is why the repository substitutes a fixed public-domain book and prints a table of the same shape with different numbers — a battery of 16 short factual questions with known answers under greedy decoding, and the fraction of tokens on which the quantized model agrees with the full-precision one given identical context.

Four things in that table are worth stating plainly.

INT8 done properly is free. Per-row INT8 scores 22.96 against the reference's 23.08 — a gap of one part in two hundred, which is noise and should be read as "identical". Which way the noise points is not stable: on the repository's public-domain corpus the same two schemes come out 22.24 against 22.18: half that distance, and pointing the other way. It agrees with the full-precision model on 142 of 144 generated tokens. A quarter of the memory against the fp32 reference, half against the fp16 you would actually deploy, and no detectable cost. INT8 done carelessly is nearly free too: one scale per matrix costs 0.5 perplexity points and no battery answers. Eight bits is forgiving enough that granularity barely matters, which is exactly why people generalise from INT8 to INT4 and get hurt.

INT4 with one scale per tensor destroys the model. Perplexity 365 million: not degraded, annihilated. Granularity is then the whole game — per-tensor 365,416,000, per-row 46.18, per-group-of-128 31.08, NF4 24.55. Same four bits per weight, a factor of fifteen million between worst and best.

Perplexity is a coarse instrument and the battery a coarser one. Between NF4 and group-128 INT4 the perplexity gap is 6.5 points and the battery differs by one question — and Chapter 4's confidence interval says one question of sixteen distinguishes nothing whatsoever. There is a sharper demonstration than the interval: run the same battery with the model's stock repetition penalty switched off, which is what greedy decoding actually means, and those two rows swap places. One question of sixteen is not a small effect, it is no effect. Chapter 8's warning applies too: perplexity is comparable only between models sharing a tokenizer, so a number from someone else's write-up cannot be compared with yours.

The agreement column is the sharpest of the three, and nearly free: run the full-precision model greedily, then ask the quantized one, at every position, what it would have chosen given the same prefix. It has 144 independent observations instead of 16, needs no ground truth, and degrades smoothly where the battery degrades in jumps. It is also exactly the quantity the next section needs.

This is the promise Chapter 1 made about this chapter, arriving on schedule: the mathematics says a 4-bit model is possible, and the engineering decides whether it is usable.

Chapter 12 announced this and left the bill here.

The idea comes straight out of the prefill/decode split. Verifying a proposed sequence of $\gamma$ tokens costs one forward pass over $\gamma$ positions — a matrix-matrix product, barely more expensive than the pass over one. So:

The output distribution is unchanged. With greedy decoding that is obvious — a token is accepted only if the target would have produced it. With sampling it requires a modified acceptance rule, and Leviathan et al. prove the resulting distribution is exactly the target's.¹⁰ This is the second exact optimisation in this chapter.

Everything therefore hinges on the acceptance rate $\alpha$, which is measurable — it is the agreement column above, which is why it was computed there. Using each quantized model as a draft for the full-precision target, over 144 generated positions:

The expected tokens accepted per verification pass, at draft length $\gamma$, is

and the net speedup divides that by the draft's own cost, a fraction $c$ of the target per token:

The bold entry is the one to remember: speculative decoding can make generation slower. At 30 % acceptance with a draft costing a fifth of the target, you pay for five forward passes and keep 1.4 tokens. The last column is the other trap — a longer draft only helps when acceptance is high, because the tail of a $\gamma$-token guess is almost never reached. At 90 % acceptance $\gamma = 8$ is worth 3.40x and at 30 % it is worth 0.79x: the same configuration, a win or a loss depending on a number measured on your traffic.

Quantization shrinks a model by storing the same function in fewer bits. Distillation shrinks it by training a smaller model to imitate a larger one¹¹ — an idea that predates deep learning by nearly a decade.¹²

The subtle part is what the student learns from. Not the correct answer: it could have been trained on that directly. What the teacher adds is the whole distribution. Ask the model what follows a phrase and look past the argmax:

"She poured the milk into the"
  ' jug' 0.1355   ' cup' 0.1051   ' bowl' 0.0605   ' large' 0.0380   ' milk' 0.0360

The hard label says jug and nothing else. The soft label says jug, and also that cup was nearly as good, bowl plausible, and large — an adjective, a completely different grammatical continuation — still live. That is the original argument: this is a 7, but it looks quite a lot like a 1, and the resemblance is information the hard label throws away.

It is also why distillation uses a temperature. Dividing the logits by $T$ before the softmax flattens the distribution and raises the relative weight of the runners-up: on this phrase, the ratio between the top token and the third falls from 2.24 at $T = 1$ to 1.50 at $T = 2$ — the square root of the first, which is what dividing the logits by two does to a ratio. Same ordering, more of the loss's attention on the near misses. The student's gradient carries the teacher's uncertainty and not only its verdict.

Everything in this chapter is now one sum:

where $T$ is the total tokens resident across all concurrent requests. Applying it: the 7B and 70B rows assume 8 key-value heads of dimension 128, the 13B row full multi-head attention with 40 heads, which is how those generations of model were built — and it shows.

8 GB

16 GB

24 GB

Look at the 13B row in the 8 GB table. The weights fit — 6.2 GB of 8 — so by the usual way of talking, a 13B model "runs on an 8 GB card". It has 337 tokens of context, which is not a conversation but barely a prompt. "Does it fit" is the wrong question. The right one is "with how much context, and for how many users at once".

Look also at the two 16 GB int8 rows. The 7B gets 65,378 tokens and the 13B gets 3,136 — a twenty-fold difference from 5.6 GB of extra weights, because the 13B here has multi-head attention and its cache costs 800 KB per token against the 7B's 128 KB. Two models of similar size, one unusable for long context, for a reason that appears in no model card's headline.

Thirteen chapters ago this was a perceptron with two weights and a bias. It is now a transformer that has been designed, trained, aligned, taught to spend compute on hard questions, and served at a measured cost per token — with no box left in it unopened.

That ends here, and it ends on purpose.

Chapter 14 begins with the model somewhere else. Not in your process, not in your memory, not in a variable you can print: on a machine you do not administer, behind an API key, a port and a bill. Everything measured here is still happening — the prefill still runs before the first token, the cache still grows with the conversation, the batch you are in still belongs to somebody else and still decides your latency — but from now on you observe it through a stream of Server-Sent Events, a finish_reason, and an HTTP 429 with a Retry-After header. The questions change with the vantage point: not how is this gradient computed but why did my invoice triple. So does the language, and Chapter 14 explains that rule rather than announcing it — up to here the code held weights, gradients, logits and tokenizer bytes; from there on it holds a connection, a retry, a cancellation and accumulated state. The thirteen chapters behind you are not discarded by the crossing. They are the description of what is running on the other side of the port.

Two omissions are deliberate. FlashAttention (Dao et al., arXiv:2205.14135) is not a different attention — it computes the same function by tiling the operation so the $n \times n$ score matrix is never written to memory, which is why the 67 MB in this chapter's second table is smaller in practice than the arithmetic suggests. And the kernels themselves are delegated: lecture 10 of Stanford's CS336 covers inference systems in the depth this does not attempt, and the llama.cpp repository and the GGUF specification are the primary sources for the CPU side.