Prefill, decode and KV-cache sizing

Why it matters

LLM inference has two phases with opposite bottlenecks. Prefill processes the whole prompt in parallel and is usually compute-bound; decode produces one token per sequence per step and is usually bandwidth-bound, because each step reads the weights and every active sequence's key/value cache. KV-cache size determines how many requests fit, how long contexts can be, and why multi-query and grouped-query attention exist. These estimates come up in nearly every inference-system discussion.

Core ideas

Why a cache. In autoregressive generation each new token attends to all previous tokens. Recomputing keys and values for the whole prefix at every step would repeat work, so they are stored per layer and reused.

Sizing. Per token, the cache holds one key and one value vector for every KV head in every layer:

bytes per token = 2 × layers × kv_heads × head_dim × bytes_per_element

Total memory is that times the number of cached tokens across all active sequences. Illustrative arithmetic: 32 layers, 8 KV heads, head dimension 128 and 16-bit values give 2 × 32 × 8 × 128 × 2 = 131,072 bytes (128 KiB) per token, or 512 MiB for one 4,096-token sequence. With 32 KV heads, as in standard multi-head attention, it would be four times larger.

Prefill. Prompt tokens are processed together, so the matmuls have many rows and high arithmetic intensity, and attention over the prompt costs time quadratic in its length. Prefill plus queueing dominates time to first token.

Decode. Each step multiplies a few vectors (one per sequence in the batch) by every weight matrix. At batch size b the weight matmuls have arithmetic intensity of roughly b FLOPs per byte, far below the ridge point for small b, so a lower bound on step time is (weight bytes + KV bytes read) / memory bandwidth. Batching amortizes weight reads across sequences, which is why throughput rises almost linearly with batch size while the step stays bandwidth-bound. KV reads, however, belong to individual sequences and grow with batch size × context length, so at long contexts decode attention stays bandwidth-bound no matter how you batch.

MQA and GQA. Multi-query attention shares one key/value head across all query heads; grouped-query attention shares each KV head within a group of query heads. Both shrink the cache and the bytes read per decode step by the ratio of query heads to KV heads. The cost is a quality risk, larger for MQA, which GQA moderates; the GQA paper also shows how to convert multi-head checkpoints by uptraining.

Concurrency. Available KV memory ≈ GPU memory − weights − activation and workspace reserve. Divide by bytes per token for total cached tokens, then by the typical context per request for concurrent sequences. Engines such as vLLM preallocate the KV pool from a configured fraction of GPU memory and expose its usage as a metric, which gives you something concrete to check your estimate against.

Other levers. KV-cache quantization (for example 8-bit) halves cache bytes relative to 16-bit; sliding-window attention bounds cache growth; prefix sharing avoids duplicate caches; speculative decoding spends extra compute to reduce sequential steps; separating prefill and decode onto different GPUs removes their interference.

Common misconceptions

• "The KV cache is small next to the weights." At long contexts and high concurrency it can exceed the weights; it grows linearly with tokens, layers and KV heads.
• "Batching hurts decode latency, so avoid it." While decode is bandwidth-bound, batching raises throughput nearly linearly at a modest per-token latency cost; SLOs set the right batch size.
• "GQA mainly saves FLOPs." It mainly saves KV memory and bandwidth; query-side attention FLOPs are similar.
• "Prefill and decode benefit from the same optimizations." Prefill gains from compute efficiency (fused attention, tensor cores); decode gains from bandwidth efficiency (batching, smaller KV, weight or KV quantization).
• "Context length only affects memory." It also raises prefill compute quadratically in attention and increases bytes read per decode step.

How to approach it

Practice the arithmetic until it is fast: KV bytes per token, per sequence and per batch, then a bandwidth bound on decode step time. Then validate on a real model: predict memory at several context lengths and batch sizes and compare with measured allocations or the serving engine's reported KV capacity. In interviews, state assumptions (precision, KV heads, context mix) before computing.