The Arithmetic of Inference

ch03 · The Arithmetic of Inference¶

The problem¶

Chapter 2 established that the naive engine saturates at a particular throughput. It did not explain why that number, rather than one ten times higher or lower.

Without a model of where the time goes, every optimisation is guesswork — and guesswork in this field is expensive, because the obvious answer is usually wrong. The rest of this book is a sequence of specific interventions, and you should be able to predict roughly what each will buy before implementing it.

The idea¶

Two operations, two different bottlenecks¶

Recall the split from ch01 · What an Inference Server Actually Does. Prefill processes the whole prompt at once; decode produces one token at a time. The crucial fact is that these are limited by different resources.

Prefill is compute-bound. The prompt’s tokens all exist already, so they go through the model together as a matrix multiplication with plenty of parallel work. The cost is roughly 2 × params FLOPs per prompt token — one multiply and one add per parameter.

Decode is memory-bandwidth-bound. To produce a single token, the model must read every one of its weights out of memory, and then does barely any arithmetic with each. The work is trivial; the data movement is not. Decode speed is therefore set by how fast the device can stream weights, not by how fast it can multiply.

That asymmetry explains almost everything that follows, so it is worth stating as a rule:

Prefill is limited by FLOPs. Decode is limited by bytes moved.

Two consequences that will keep recurring:

• Making decode faster usually means moving fewer bytes — smaller weights (ch14 · Quantisation for Serving [DRAFT]), a smaller KV cache (ch12 · Attention at Speed [DRAFT]), or reading the same weights for more sequences at once (ch07 · Continuous Batching).

• Adding sequences to a decode batch is nearly free. The weights are read once per step whatever the batch size, so the second sequence costs only its own KV cache. This is the single most important economic fact in LLM serving, and continuous batching exists to exploit it.

The two formulas worth memorising¶

KV cache per token. Each token’s keys and values are stored for every layer:

kv_bytes_per_token = 2 × n_layers × n_kv_heads × head_dim × bytes_per_element

The 2 is for K and V. Note n_kv_heads, not n_heads — grouped-query attention shrinks this term directly, which is why every modern model uses it.

The decode ceiling. Per step the device must read the weights, plus the KV cache of every sequence in the batch:

decode_tok_per_s ≤ bandwidth / (weight_bytes + kv_bytes_per_token × context × batch)

No kernel beats this. It is a ceiling, not a forecast.

The build¶

Both formulas, written out so predictions cannot drift from the model they describe — the same ModelConfig object builds the model and feeds the arithmetic:

def kv_bytes_per_token(model: ModelConfig, kv_dtype: str | None = None) -> float:
    """Bytes of KV cache a single token occupies.

    The factor of 2 is because we cache both K and V. ``kv_dtype`` is separate from the model's
    dtype because quantising the cache independently of the weights is a real technique
    (chapter 14) and usually the bigger win at long context.
    """
    element = BYTES_PER_DTYPE[kv_dtype] if kv_dtype else model.bytes_per_element
    return 2 * model.n_layers * model.n_kv_heads * model.head_dim * element

def decode_bytes_per_step(model: ModelConfig, context_length: int, batch_size: int = 1) -> float:
    """Bytes that must be read from memory to produce one token per sequence in a batch.

    Every weight is read once regardless of batch size — which is the entire reason batching
    works — plus each sequence's KV cache.
    """
    weights = model.n_params * model.bytes_per_element
    return weights + kv_bytes_for(model, context_length, batch_size)

That the prediction and the model share a source is worth more than it looks: a test asserts that the predicted parameter count equals the built model’s, so the arithmetic cannot quietly describe a model we are not running.

The measurement¶

Apply it to the reference model, then check it against chapter 2’s measurements:

Conditions: TinyGPT (reference, random weights) (5,838,080 params), 4x x86_64 CPU, torch 2.14.0+cu130, poisson, prompt 32-96, output 16-48, seed 7, arrival rate 1.0/s, measured 2026-09-13T15:26:06+00:00.

Three things to take from that table.

The KV cache is small here, and that is a property of this model, not of serving. Weights are 98.8% of the bytes read per decode step. Scale to a 7B model in fp16 and the picture inverts at long context: weights are fixed at about 14 GB, while KV grows with every token of every concurrent sequence. That crossover is why Part III is mostly about memory.

The implied bandwidth is a useful number to keep. Dividing bytes-per-step by measured time-per-token gives roughly 5 GB/s achieved on this machine. That figure now predicts things — ch06 · Static Batching and Its Limits uses it to forecast batched throughput before implementing batching.

The FLOPs model over-predicts by an order of magnitude, and the last two rows of that table are the most instructive lines in this chapter. Counting arithmetic alone, removing the KV cache should be enormously expensive: without one, generating token n redoes all n previous tokens. Measured, the penalty is a small single-digit factor.

The explanation is the rule above. At this model size, a decode step is dominated by reading 23 MB of weights, and that cost is paid whether we process one token or ninety-six. The extra arithmetic rides along in time the machine was spending on memory anyway. The FLOPs model counted the work; it did not count what was actually scarce.

This is worth dwelling on because it generalises. Any optimisation that reduces FLOPs without reducing bytes moved will disappoint you during decode. Most of the ones that work in this book — batching, quantisation, GQA, prefix caching — reduce bytes, or amortise them over more useful output.

The cost¶

The arithmetic is a model, and it ignores:

• Attention’s quadratic term. Fine at 96 tokens, badly wrong at 32k. ch21 · RAG and Long Context [DRAFT] is where this stops being safe.

• Kernel launch and framework overhead. At this model size a meaningful share of each step is Python and dispatch, not memory traffic. That is one reason the predicted penalty did not materialise.

• Cache hierarchy. “Bandwidth” is one number standing in for registers, several cache levels and DRAM, each an order of magnitude apart.

• Everything except the model. Tokenisation, scheduling and HTTP are all invisible here.

Use it to predict orders of magnitude and to decide which of two optimisations is worth trying. Do not use it to predict a percentage. Then measure, and when the measurement disagrees, the disagreement is the interesting part — as it was above.

Key takeaways¶

• Prefill is compute-bound; decode is memory-bandwidth-bound. Every technique in this book attacks one of those two, and knowing which tells you when it will help.

• KV cache per token is 2 × n_layers × n_kv_heads × head_dim × bytes. It is the budget all of Part III competes over.

• Decode throughput is bounded by bandwidth / bytes read per step. Weights are read once per step regardless of batch size, which is why batching is close to free and why ch07 · Continuous Batching works.

• A FLOPs count predicts prefill reasonably and decode badly. When the two disagree, the scarce resource is bytes, not arithmetic.

Looking ahead¶

We can now predict, measure, and explain the gap. Time to fix something. ch04 · Generating Tokens Correctly takes ownership of the decode loop itself — sampling, stop conditions, and the streaming-detokenisation bug that makes non-ASCII output look like model corruption — so that every optimisation afterwards can be tested for producing identical output.

Further reading¶

The roofline model is the general form of the argument here, and reading the original Williams, Waterman and Patterson paper will make the prefill/decode split feel inevitable rather than particular to transformers. For the same arithmetic applied to production-scale models, the FlashAttention paper (Appendix E) opens with an unusually clear statement of why attention is IO-bound rather than compute-bound.