KV Cache — The Memory Engine of LLM Serving

Phase 2 · Document 06 · Transformer Foundations Prev: 05 — Autoregressive Generation · Next: 07 — Prefill vs Decode

Table of Contents

1. Why This Matters
2. Core Concept
3. Mental Model
4. Hitchhiker's Guide
5. Warmup Readings
6. Deep Readings and External References
7. Key Terms
8. Important Facts
9. Observations from Real Systems
10. Common Misconceptions
11. Engineering Decision Framework
12. Hands-On Lab
13. Verification Questions
14. Takeaways
15. Artifact Checklist

1. Why This Matters

The KV cache is the single most important memory structure in LLM inference. It decides how much GPU memory serving needs, why longer contexts cost more, why concurrency is capped, and why you can run out of memory even when the model weights fit comfortably. Every marquee serving technique — PagedAttention, prefix caching, continuous batching, chunked prefill — exists to manage it. If you understand the KV cache, you understand the core trade-offs of LLM serving; if you don't, production OOMs and cost surprises will blindside you.

2. Core Concept

What it is

During attention (02), each token produces a Key (K) and Value (V). To generate token N, attention needs the K and V of all prior tokens (1…N−1).

• Without caching: recompute K/V for every prior token at every step → wasted O(N²) compute.
• With KV cache: compute each token's K/V once, store them, and reuse them on every later step → O(N) compute per step.

The KV cache trades memory for compute — almost always worth it. But that memory is the binding constraint.

What's stored, and how big

For each token, each layer, each KV head: one Key and one Value vector of head_dim floats.

kv_bytes_per_token = 2 (K&V) × num_layers × num_kv_heads × head_dim × bytes_per_element

FP16 examples (full MHA, num_kv_heads = num_heads):
  Llama-3 8B :  2 × 32 × 32 × 128 × 2  ≈ 0.5 MB / token
  Llama-3 70B:  2 × 80 × 64 × 128 × 2  ≈ 2.6 MB / token
  Llama-3 405B: 2 × 126 × 128 × 128 × 2 ≈ 8.3 MB / token

GQA/MQA shrink this dramatically (02): with num_kv_heads = 8 instead of 64, the KV cache is 8× smaller. This is the reason modern big models use GQA.

Total, and why it dominates

total_kv = kv_bytes_per_token × context_length × batch_size

Example — 8B model, 4K context, batch 10: 0.5 MB × 4096 × 10 ≈ 20 GB — more than the 16 GB of FP16 weights. KV cache, not weights, is usually what caps concurrency.

The management techniques (all KV-cache optimizations)

• PagedAttention (vLLM): manage KV like OS virtual memory — fixed-size pages allocated on demand, tracked by a page table. Eliminates the fragmentation/over-reservation of contiguous allocation, pushes utilization near 100%, and lets requests share pages for identical prefixes.
• Prefix caching: if many requests share a long prefix (e.g. a 2,000-token system prompt), compute its KV once and share it. Providers expose this as cached input tokens at a 75–90% discount (Phase 1.08).
• Continuous batching: add/remove sequences from the running batch every decode step, keeping KV slots full and GPU busy (3–5× throughput vs static batching). (Detailed in Phase 1.05.)
• Chunked prefill: split a long prompt's prefill into chunks interleaved with decodes so one huge prompt doesn't stall everyone (improves TTFT fairness).

3. Mental Model

KV CACHE = a growing notebook, one entry per token per layer per KV-head.

  generate a token →  WRITE its K,V into the notebook
                  →  READ K,V of every earlier token (that's attention)

  notebook size = per_token × context_length × batch_size
  notebook fills GPU memory → you cannot seat more/longer requests (concurrency cap)

  PagedAttention = store the notebook in tidy fixed-size PAGES (no wasted gaps, shareable)
  Prefix caching = reuse the notebook pages for a shared opening (system prompt)
  Continuous batching = refill a page slot the instant a request finishes
  Chunked prefill = write a long prompt's pages in chunks so others aren't blocked

4. Hitchhiker's Guide

What to check first on OOM: per-request context length, batch size/concurrency, and the split of VRAM between weights and KV cache. The culprit is usually KV under concurrency, not weights.

What to ignore at first: page-size tuning and CPU-offload knobs — defaults are good; revisit only when profiling says so.

What misleads beginners:

• "Weights fit, so we're fine." KV for concurrent long contexts often exceeds the weights.
• "Tokens/sec is the whole story." Concurrency is bounded by KV memory.
• Forgetting GQA — a model's KV footprint depends on num_kv_heads, not num_heads.
• Not enabling prefix caching for repeated system prompts (leaving 75–90% savings on the table).

How experts reason: they compute KV per token from the config (using num_kv_heads), multiply by target context × concurrency, and size hardware/limits accordingly; they enable PagedAttention + prefix caching + continuous batching by default.

What matters in production: a sensible max context (don't reserve 128K if users need 8K), prefix caching for shared prompts, monitoring KV utilization/preemptions, and choosing GQA models for high concurrency.

How to verify: read num_hidden_layers, num_key_value_heads, head_dim from the config; compute KV/token; watch vLLM's gpu_cache_usage_perc and num_preemptions under load.

Questions to ask: Does the model use GQA/MQA (num_kv_heads)? Does the serving stack do PagedAttention + continuous batching? Is prefix caching available and at what discount? What max context is configured?

What silently gets expensive/unreliable: long-running conversations filling KV until eviction/preemption; high batch × long context OOMs; missing prefix caching; over-reserved max context wasting memory.

5. Warmup Readings

6. Deep Readings and External References

7. Key Terms

8. Important Facts

• The KV cache stores K and V for every prior token, per layer, per KV head — one notebook entry per token.
• KV memory grows linearly with context length × batch size.
• KV cache is often the binding memory constraint, not weights (e.g. 8B/4K/batch-10 ≈ 20 GB > 16 GB weights).
• GQA/MQA shrink KV by the head-sharing ratio — compute KV/token with num_kv_heads, not num_heads.
• PagedAttention raises utilization toward 100% and enables prefix sharing.
• Prefix caching gives 75–90% input-token discounts for repeated prefixes.
• Continuous batching yields ~3–5× throughput over static batching.
• The KV cache trades memory for compute (O(N²) recompute → O(N)/step) — almost always worth it.

9. Observations from Real Systems

• vLLM exposes gpu_cache_usage_perc, num_preemptions, num_running — monitor these in Prometheus/Grafana to see KV pressure directly.
• Commercial providers don't show the KV cache, but cached input tokens in usage metadata = a prefix-cache hit (cheaper).
• llama.cpp --ctx-size and parallel slots directly allocate KV; bigger ctx → fewer parallel requests.
• Ollama holds KV per active session — idle long sessions keep memory pinned.
• GQA in Llama-3/Qwen/Mistral exists largely to shrink this cache and raise serving concurrency.
• SGLang's RadixAttention is an aggressive prefix-sharing KV scheme for agent/few-shot workloads.

10. Common Misconceptions

11. Engineering Decision Framework

Capacity planning:
  kv_per_token = 2 × layers × num_kv_heads × head_dim × bytes   (from config)
  need = weights + kv_per_token × max_context × max_concurrency + overhead
  need ≤ GPU memory?  → ship.  else → see levers below.

Levers when KV-bound:
  • Prefer a GQA/MQA model (smaller num_kv_heads).
  • Lower max_context to realistic usage.
  • Reduce max concurrency or add GPUs.
  • Enable PagedAttention + continuous batching (default in vLLM).
  • Enable prefix caching for shared system prompts (cuts cost + memory).
  • Quantize KV cache (FP8/INT8 KV) if the engine supports it.

Monitoring:
  watch gpu_cache_usage_perc & num_preemptions → preemptions mean you're over capacity.

12. Hands-On Lab

Goal

Compute KV-cache size from a model config, then observe KV usage rise with context and concurrency on a real server.

Prerequisites

• Python 3.10+; for the live part, a GPU with vLLM or llama.cpp (CPU works for small models).

Setup

pip install transformers          # for config inspection
# optional live serving:
pip install vllm                  # GPU; or build llama.cpp

Steps

A. Compute KV/token from config (no GPU needed):

from transformers import AutoConfig
c = AutoConfig.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
layers = c.num_hidden_layers
kv_heads = getattr(c, "num_key_value_heads", c.num_attention_heads)  # GQA-aware
head_dim = c.hidden_size // c.num_attention_heads
bytes_per = 2  # FP16
kv_tok = 2 * layers * kv_heads * head_dim * bytes_per
print(f"KV/token ≈ {kv_tok/1024:.1f} KB")
for ctx, batch in [(4096,1),(4096,16),(32768,16)]:
    print(f"ctx={ctx} batch={batch}: KV ≈ {kv_tok*ctx*batch/1e9:.2f} GB")

B. (Live, optional) Observe on vLLM:

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.2-1B-Instruct --max-model-len 4096 --gpu-memory-utilization 0.8
# then, while sending requests:
curl -s http://localhost:8000/metrics | grep -E "kv_cache|cache_usage|preempt"

Send (1) a 100-token prompt, (2) a 3,000-token prompt, (3) 10 concurrent 100-token prompts; record cache usage each time.

Expected output

• KV/token in KB; total KV ballooning with ctx × batch (often exceeding weights).
• Live: long context and many concurrent requests both drive gpu_cache_usage_perc up; preemptions appear when over capacity.

Debugging tips

• Forgot GQA? Using num_attention_heads overestimates KV — use num_key_value_heads.
• No GPU? The config computation alone delivers the key insight.

Extension task

Recompute KV/token assuming the model were full-MHA (kv_heads = num_heads) and report the GQA savings ratio.

Production extension

Build kv_capacity(model_config, gpu_gb, max_ctx) → max concurrent requests; this plugs into the Phase 6 memory calculator and Phase 7 capacity planning.

What to measure

KV/token; total KV vs weights at several ctx×batch; GQA savings ratio; live cache usage vs context and concurrency.

Deliverables

• A KV-size table (ctx × batch) vs model weights.
• GQA savings ratio for one model.
• (Live) KV-usage observations + the point where preemptions start.

13. Verification Questions

Basic

1. What's stored in the KV cache, and per what (token/layer/head)?
2. Why does KV memory grow with context length and batch size?
3. How does GQA reduce KV-cache size?

Applied 4. Compute KV/token for 32 layers, 8 KV heads, head_dim 128, FP16. Then total for 8K context, batch 8. 5. A 70B model uses 140 GB for weights. Roughly what context × batch makes KV exceed that?

Debugging 6. The service OOMs only under concurrent long-context load. Diagnose and give two fixes. 7. Costs are high for a bot with a fixed 2,000-token system prompt. What's the missing optimization?

System design 8. Plan KV capacity for 50 concurrent users at 8K context on a given GPU; show the formula and your levers if it doesn't fit.

Startup / product 9. Explain to your CFO how prefix caching and GQA improve gross margin for a high-volume chatbot.

14. Takeaways

1. The KV cache stores K/V for all prior tokens (per layer, per KV head) so decode is O(N)/step, not O(N²).
2. KV memory grows with context × batch and is usually the binding constraint, not weights.
3. Compute KV/token with num_kv_heads — GQA/MQA shrink it sharply.
4. PagedAttention, prefix caching, continuous batching, chunked prefill are all KV-cache optimizations.
5. Prefix caching = 75–90% input savings for shared prompts; continuous batching = 3–5× throughput.
6. Monitor KV utilization and preemptions; size max-context and concurrency to reality.

15. Artifact Checklist

• Code: KV/token + total-KV calculator from a real config (GQA-aware).
• Table: KV vs weights across ctx × batch.
• GQA savings note for one model.
• Capacity function: kv_capacity(...) → max concurrency.
• (Live) Observation log: KV usage vs context/concurrency + preemption point.
• Diagram: the growing-notebook + paging model.

Next: 07 — Prefill vs Decode