ai_ml / transformers llm / 05_kv_cache.md

KV cache

6 interview angles 5 min read source

KV cache

The single most important inference concept, and the one most often missing from otherwise-good candidates. It explains why generation is fast, why memory runs out, and why your cost model looks the way it does.

Why it exists

Generation is autoregressive: produce a token, append it, run the model again. Without caching, generating token 1000 recomputes attention over all 999 previous tokens — the whole sequence, every step. Total cost would be O(n²) forward passes.

But keys and values for previous tokens never change. Cache them, and each new token only computes its own K and V and attends against the cache:

def decode_step(x_new, cache):
    q, k, v = project(x_new)              # only the new token
    cache.k = torch.cat([cache.k, k], dim=-2)
    cache.v = torch.cat([cache.v, v], dim=-2)
    return attention(q, cache.k, cache.v)  # attend over everything so far

That turns per-token cost from O(n) to O(1) in model passes, at the price of holding the cache in memory.

Prefill and decode are different workloads

This distinction drives everything about serving, and stating it is a strong signal.

Prefill Decode
Processes the whole prompt at once one token at a time
Parallelism full — all positions in parallel none — inherently sequential
Bottleneck compute (GPU FLOPs) memory bandwidth
Scales with prompt length (quadratic in attention) tokens generated
Metric time to first token (TTFT) inter-token latency (ITL)

Decode is bandwidth-bound because each step reads the entire model weights and the entire KV cache to produce one token. The arithmetic intensity is terrible — enormous data movement, tiny computation. That’s why batching helps so much: reading the weights once to serve 32 requests amortises the dominant cost.

It also explains the pricing you see from providers: input tokens (prefill, parallel) are cheaper than output tokens (decode, sequential).

Sizing it

cache_bytes = 2 * layers * kv_heads * d_head * seq_len * batch * bytes

The 2 is K and V. Worked example — a 70B-class model, GQA with 8 KV heads, d_head 128, 80 layers, fp16, 32k context, one request:

2 * 80 * 8 * 128 * 32768 * 2 bytes ≈ 10.7 GB

For one request. Eight concurrent requests exceeds an 80GB GPU before the weights are loaded. With plain MHA (64 KV heads) it would be 8x that.

That calculation is the entire reason GQA and MLA exist — see 02_attention_mechanisms.md.

Shrinking it

Technique Saving Cost
GQA ~4x architectural, chosen at pretraining
MLA ~14x architectural, complex kernels
FP8 KV cache 2x negligible quality impact — near-free
Quantised KV (int4) 4x measurable degradation
Sliding window bounded loses long-range attention
Eviction / compression variable quality risk; keep the sink tokens

FP8 KV cache is the free win and is standard in 2026 serving stacks. Halving cache memory doubles concurrency at essentially no quality cost.

PagedAttention

The insight that made vLLM: KV cache memory suffers from the same problems as process memory — internal fragmentation and over-allocation. Naive serving reserves a contiguous block for each request’s maximum possible length, wasting most of it.

PagedAttention treats the cache like virtual memory: fixed-size blocks, a block table per sequence, non-contiguous physical placement. Blocks are allocated on demand.

Two consequences:

  • Near-zero fragmentation, so far more concurrent requests fit.
  • Blocks can be shared. Two requests with the same prefix point at the same physical blocks, copying only when they diverge.

That sharing is what makes prefix caching possible.

Prefix caching

If every request starts with the same system prompt, tools definition and few-shot examples, that prefix’s KV is identical every time. Cache it once and reuse across requests.

vllm serve MODEL --enable-prefix-caching

Reported as the single highest-value flag for typical workloads — roughly a 30% throughput gain when requests share a system prompt, and much more for long shared prefixes.

Design consequence: put the stable content first and the variable content last. A prompt structured as [system][tools][examples][user query] is cacheable; interleaving user-specific data early destroys the shared prefix. This is a concrete, actionable thing to say in a system-design interview.

Hosted providers expose the same idea as prompt caching, usually with a discount on cached input tokens.

Multi-turn conversations

Each turn re-sends the whole history. With prefix caching the previous turns’ KV is reused, so only the new message is prefilled — which is why chat feels responsive despite the context growing.

Break it by editing earlier messages: the prefix diverges from that point and everything after must be recomputed. Systems that “summarise and rewrite history” trade cache hits for context length, and it’s worth measuring both sides.

Interview angle

  • “What is the KV cache and why does it exist?” — keys and values of past tokens don’t change, so caching them makes each new token cost one incremental step instead of reprocessing the whole sequence. It trades memory for a quadratic reduction in compute.
  • “Prefill vs decode?” — prefill processes the prompt in parallel and is compute-bound; decode generates one token at a time and is memory-bandwidth-bound. Different bottlenecks, different metrics (TTFT vs inter-token latency), and it’s why output tokens cost more than input tokens.
  • “Estimate the KV cache for a 70B model at 32k context.”2 * layers * kv_heads * d_head * seq_len * bytes. With GQA that’s roughly 10GB per request; with MHA it would be several times more. That number is why cache-reduction architectures exist.
  • “What does PagedAttention do?” — manages KV cache in fixed-size blocks like OS virtual memory, eliminating fragmentation and over-allocation, and enabling block sharing between sequences with common prefixes.
  • “How would you cut LLM serving cost without changing the model?” — enable prefix caching and structure prompts so the stable part comes first; quantise the KV cache to FP8; increase batch size within the latency budget; and route easy requests to a smaller model.
  • “Why does decoding get slower as the conversation grows?” — each step attends over a longer cache, and decode is bandwidth-bound, so more cache means more bytes read per token. Cost grows with context even though the model runs once per token.