llm/curriculum
Module 04 · ready

Inference, KV cache and context

Why is the first token slow, and what eats memory in long chats?

Qwen/Qwen3-0.6B (measured on CPU)configs of 9 models

Two phases: prefill and decode

Generating an answer happens in two very different phases:

  1. Prefill: the model reads your whole prompt in one go. All prompt tokens go through every layer in parallel, as one big matrix multiply. This sets the time to first token (TTFT).
  2. Decode: the model produces the answer one token at a time. Each new token needs a full pass through every layer, but only for that single token.

Prefill is limited by compute (FLOPs), because there’s lots of work per weight read. Decode is limited by memory bandwidth, because each token re-reads all active weights to do very little math (module 01). That’s why a Mac reads prompts slower than a GPU but generates at a decent speed.

Measured on this curriculum’s build machine

Qwen3-0.6B, FP32, PyTorch on 8 CPU cores (labs/inside_qwen.py). Prefill: 16 tokens at 54.8 tok/s, 64 tokens at 54.9 tok/s, 256 tokens at 92.7 tok/s, 1024 tokens at 139.7 tok/s. Longer prompts get faster per token because the hardware gets bigger matrices to chew on. Decode with KV cache: 8.1 tok/s at 64 context, 5.1 tok/s at 512 context, 5.49 tok/s at 1024 context. A CPU is a slow machine for this, but the ratios are the point.

The KV cache: why decode doesn’t start from scratch

In attention (module 03), each new token needs the keys and values of every earlier token. Those don’t change once they’re computed, so engines save them: that’s the KV cache. Without it, every new token would recompute the whole conversation.

Measured with the same script, generating 32 tokens:

context with KV cache without (recompute everything)
| 64 tokens | 8.1 tok/s | 1 tok/s | | 512 tokens | 5.1 tok/s | 0.41 tok/s | | 1024 tokens | 5.49 tok/s | 0.3 tok/s |

The longer the context, the bigger the gap, because without a cache the work per token grows with the context.

How big it gets

Per token, each attention layer stores one key and one value vector per KV head. For Qwen3-0.6B: 28 layers × 2 × 8 KV heads × 128 numbers × 2 bytes (BF16) = 112 KiB per token. We checked it against the real cache: the K tensor of each layer had shape [1, 8, 512, 128] for 512 tokens, 29,360,128 numbers in total, exactly the formula.

At Qwen3-0.6B’s full 40,960-token context that’s about 4.7 GB, roughly 4× the size of the model’s own weights (1.2 GB). For long contexts and parallel users, the KV cache, not the weights, is what fills the GPU.

KV cache size, from each model's own config.jsonper sequence · weights not included

Formula: tokens × layers × 2 (key and value) × KV heads × head_dim × bytes. MLA models store one compressed latent instead of per-head K and V. Linear-attention layers (Qwen3.5, Qwen3-Next, Ornith) keep a constant-size state that is not counted here. For Qwen3-0.6B the formula matched the real cache tensors exactly (29,360,128 numbers for 512 tokens, measured). Engines add a little overhead for paging and alignment.

This is why architectures fight over KV size:

  • GQA cuts it by sharing K/V heads (Qwen3-Next: 16 query heads, only 2 KV heads).
  • MLA stores one compressed latent per token: DeepSeek-V3 needs 61 × (512 + 64) numbers per token. Plain multi-head attention with 128 full heads would need about 70× more.
  • Sliding windows: gpt-oss alternates full-attention layers with layers that only look at the last 128 tokens, so half its layers have a capped cache.
  • Linear attention: Qwen3.5, Ornith and Qwen3-Next keep a full KV cache in only 1 of every 4 layers.

Context length

Context length is how many tokens the model can attend to at once: prompt, history, tool outputs and its own answer together. max_position_embeddings in config.json is the trained limit (40,960 for Qwen3-0.6B, 262,144 for Qwen3.5-9B).

Three things limit what you get in practice:

  1. Training: RoPE positions beyond what the model saw in training produce garbage. YaRN and similar methods rescale the rotations to stretch the window (Peng et al., 2023). Qwen documents enabling it with a rope_scaling block.
  2. Memory: the KV cache figure above.
  3. Quality: models get worse at using information buried deep in long contexts (“lost in the middle”, Liu et al., 2023). Advertised context ≠ usable context.

KV cache quantization

The cache can be stored in fewer bits, just like weights (module 07): vLLM has --kv-cache-dtype fp8, and llama.cpp has --cache-type-k q8_0 --cache-type-v q8_0. 8-bit halves the memory with usually small quality loss. 4-bit keys tend to hurt more. Measure it with the KLD method from module 07 before trusting it for agent work.

Serving many users: batching and paged attention

  • Batching: during decode, the GPU reads all the weights for each token anyway. Processing 32 users’ next tokens in the same pass costs barely more than one user. Throughput goes up, latency per user stays roughly the same.
  • Continuous batching: new requests join the running batch at any step instead of waiting for the current batch to finish.
  • PagedAttention (vLLM, Kwon et al., 2023): the KV cache is stored in small fixed-size blocks, like memory pages in an operating system, so no memory is wasted on reserved-but-unused context. It also enables prefix caching: sharing the cache of a common system prompt between requests.
  • FlashAttention (Dao et al., 2022) computes attention in tiles that fit in fast on-chip memory, never materializing the full [tokens × tokens] score matrix.

Speculative decoding and MTP

Decode wastes compute: the hardware could process many tokens per pass, but only gets one. Speculative decoding (Leviathan et al., 2022) uses a small fast draft model to guess the next few tokens. The big model then checks all the guesses in one parallel pass and keeps the ones it agrees with. The output follows exactly the same distribution as the big model alone, just faster when the guesses are good.

Multi-token prediction (MTP) builds the draft into the model: DeepSeek-V3 has an extra prediction module (num_nextn_predict_layers: 1 in its config) trained to predict the token after next, and engines can use it for speculative decoding (DeepSeek-V3 report).

Where to go deeper