Autoregressive inference and the key-value cache

Prerequisites: Matrix multiplication and attention.

Learning objectives

• Separate prefill from decoding.
• Explain what a key-value cache stores.
• Calculate cache memory under stated assumptions.

Training and generation use different schedules

An autoregressive model is trained to predict each next token from earlier tokens. A causal attention mask prevents a position from using later input positions. During training, predictions at many positions can be computed in parallel because the full target sequence is available. During generation, the next input token depends on the previous sampled or selected output, so successive token decisions are sequential.

Inference commonly has a prefill phase that processes the prompt and a decode phase that generates additional tokens. The implementation, model architecture, batching strategy, and sequence lengths determine actual latency. Tokens per second alone hides the delay before the first token and how throughput changes under concurrent load.

Cache projections, not answers

In a transformer attention layer, a query attends to keys and uses the corresponding values to form a weighted combination. During causal generation, the keys and values for earlier positions do not change when a new token is appended. A key-value cache stores them so they need not be recomputed at every step.

For a simplified model with L layers, T cached positions, H key/value heads, head dimension D, and b bytes per element, cache storage is approximately 2×L×T×H×D×b per sequence. The factor two accounts for keys and values. This excludes allocator overhead, temporary tensors, and model weights. Architectures with grouped-query attention use fewer key/value heads than query heads, so H must refer to the cached heads.

Take L=2, T=100, H=4, D=8, and b=2. The cache requires 25,600 bytes under these assumptions. Doubling T doubles this storage. The cache avoids recomputing old projections, but each new query still attends over earlier positions in ordinary full attention.

Memory and scheduling are part of correctness

A batch contains sequences of different lengths and completion times. Cache allocation must keep each sequence's state separate, release finished state, and preserve the mapping between tokens and positions. Reusing the wrong cache can silently produce incorrect output, not merely poor performance.

Longer context does not guarantee that the model uses every relevant fact correctly. Quantization, attention variants, and cache compression change resource use and may change output quality. Measure both under realistic workloads.

For deterministic debugging, fix the model version, tokenizer, input, decoding settings, and relevant random seeds. Floating-point kernels can still introduce small numerical differences. Reproducibility is a contract with specified tolerances and outputs, not a claim that every hardware execution is bit-for-bit identical.

Practice problem

Compute cache storage for 4 layers, 200 positions, 2 key/value heads, dimension 16, and 2-byte elements. What happens for eight independent sequences of that same length? Explain which inference cost this estimate omits.

Worked solution

The estimate is 2×4×200×2×16×2 = 102,400 bytes per sequence. Eight independent equal-length sequences need 819,200 bytes for these cached tensors. This excludes model weights, temporary activations, allocator overhead, and other runtime storage. It also does not measure the arithmetic or memory bandwidth required to attend over cached positions. Sharing a common prompt may permit implementation-specific reuse, but independent cache accounting should not assume that optimization without evidence.

Check your understanding

Doubling cached positions and halving KV heads, with other dimensions fixed, changes cache storage how?

• It halves.
• It stays unchanged.
• It doubles.
• It quadruples.

Explanation

Storage is proportional to positions times KV heads, so the factors two and one-half cancel.

Can all newly generated tokens usually be chosen in one ordinary autoregressive forward pass?

• Yes, because training uses parallel positions.
• Yes, if the prompt is long.
• No, each next decision depends on the previously chosen token.
• Only if the cache is disabled.

Explanation

Generation does not know the next token until the current decision is made.

Primary reading

• Language Modeling from Scratch: Stanford CS336
• Attention Is All You Need: Vaswani et al.
• Caching: Hugging Face Transformers documentation

Open the study and save your work →