KV Cache Mechanics

What Is the KV Cache?

In transformer attention, every token computes Q (query), K (key), and V (value) matrices. On each decoding step, the model must compute attention over all previous tokens and the current token. Without caching, this means recomputing K and V for every token that has already been processed — O(n²) computation for a sequence of length n.

The KV cache stores pre-computed K and V matrices from all previous tokens. When a new token arrives, the model only computes its K and V, then attends over the cached KV from all past steps. This reduces the second and subsequent decoding steps from O(n) complexity to O(1) per token, making the amortized cost O(n).

The cache lives in GPU VRAM and is evicted when the context window fills or memory runs out. It is the primary bottleneck for long-context inference, often consuming more memory than the model weights themselves.

Cache Growth and Memory Budget

The KV cache grows linearly with sequence length. The formula for cache memory is:

cache_bytes = 2 × layers × heads × head_dim × seq_len × 2 (for fp16)

The 2× factor accounts for both K and V. The final 2 is the byte size of fp16. For a 7B model with 32 layers, 32 heads, and 128-dim heads, caching a 2K-token sequence costs ~0.5 GB. At 128K tokens, it's 32 GB — larger than the model weights.

This is why long-context requests are expensive and why serving systems enforce max_batch_tokens. A single batch of 16 requests × 128K-token context for a 70B model requires ~3 TB of GPU VRAM just for cache. Even sharding across multiple GPUs becomes challenging.

Grouped Query Attention and MLA

Grouped Query Attention (GQA) reduces cache size by sharing K and V heads across multiple Q heads. Instead of one K/V head per Q head (Full Multi-Head Attention), GQA uses one K/V head per group of Q heads — typically a group size of 8. This shrinks cache by 8×.

Multi-Query Attention (MQA) is the extreme case: a single K/V head shared by all Q heads. Early models like Falcon and PaLM used MQA and saw modest quality drops but dramatic cache savings.

MLA (Multi-Head Latent Attention, DeepSeek) takes compression further by projecting K/V into a low-rank latent space before caching. Only the compressed latent is stored, then decompressed at attention time. This can achieve 2–4× additional cache reduction compared to GQA.

Cache Eviction Strategies

When context exceeds the window or cache fills up, something must be evicted. The strategy matters: drop the wrong tokens and the model loses crucial context (like the system prompt). Drop the right tokens and quality barely changes.

Eviction Strategy Comparison

PagedAttention (vLLM) is now the industry standard. Its virtual memory approach eliminates fragmentation and enables sharing of cache blocks across concurrent requests — crucial for batching efficiency.

Prefix Caching and Prompt Reuse

When many requests share an identical prefix (system prompt, few-shot examples, shared document), compute K/V once and reuse it across all requests. This amortizes the prefill cost and can reduce total tokens processed by 50–90% depending on prefix length.

SGLang, vLLM, and TensorRT-LLM all implement prefix cache (SGLang calls it "radix cache"). At scale, the savings are dramatic: a 100-token system prompt at 1000 requests/sec = 100K free tokens/sec that never need recomputation.

Prefix Cache Implementation

The cache is keyed by a hash of the token sequence (e.g., SHA256 of concatenated embeddings). When a new prefix arrives, compute the hash; if it matches a prior request, retrieve the cached K/V blocks directly. Prefix caches are stored in a distributed KV store (Redis, Memcached) or local GPU cache depending on scale.

Speculative Decoding and Cache Interaction

Speculative decoding uses a smaller, faster draft model to propose k tokens. The target model then verifies all k+1 tokens (k proposed + 1 from the target) in a single forward pass. Tokens that match the draft are accepted; those that don't are discarded and resampled.

The KV cache interacts subtly: during the verification pass, the target model processes k+1 tokens in parallel, building cache for all of them. Accepted tokens' cache gets written; rejected tokens' cache is thrown away. This means you compute forward passes on tokens you ultimately discard, but the verification batch processes multiple tokens at once (higher GPU utilization).

Speculative decoding trades a small increase in compute (wasted on rejected tokens) for much lower latency. It's effective when the draft model is fast (e.g., a 1B model as draft, 70B as target) and when the acceptance rate is high (>70%, which happens with well-calibrated models). Streaming applications benefit significantly.

from vllm import LLM, SamplingParams
import time

llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
          enable_prefix_caching=True,   # prefix cache = reuse shared prefixes
          max_model_len=4096)

params = SamplingParams(temperature=0.0, max_tokens=64)

# Simulate workload with shared system prompt (benefits from prefix caching)
SYSTEM_PROMPT = "You are a helpful assistant. " * 50  # ~50 tokens of shared prefix
queries = [
    SYSTEM_PROMPT + f"Question {i}: What is {i} + {i}?"
    for i in range(20)
]

# First batch — cold cache
t0 = time.perf_counter()
outputs_cold = llm.generate(queries, params)
cold_time = time.perf_counter() - t0

# Second batch — warm cache (same system prompt prefix)
t1 = time.perf_counter()
outputs_warm = llm.generate(queries, params)
warm_time = time.perf_counter() - t1

print(f"Cold cache: {cold_time:.2f}s")
print(f"Warm cache: {warm_time:.2f}s")
print(f"Speedup from prefix caching: {cold_time/warm_time:.1f}×")
# Typical: 1.5–3× speedup on workloads with long shared prefixes

Production Knobs

In production serving systems, KV cache behavior is controlled by a handful of key parameters. Understanding these lets you trade off throughput, latency, and quality for your workload.

Key Tuning Parameters

--max-model-len: The maximum sequence length. Larger values allow longer contexts but consume more cache. Set based on your longest expected request + prefill buffer.

--gpu-memory-utilization: Fraction of GPU VRAM to allocate to cache (typically 0.8–0.95). Higher = more concurrent requests, but risks OOM on bursts.

--enable-prefix-caching: Enable automatic caching and reuse of identical prefixes. Essential for workloads with shared system prompts.

--num-gpu-blocks-override: Manually specify cache block count. Useful for tuning to your VRAM constraints.

Common Pitfalls

• Setting max_model_len to 128K when your GPU has only 80GB VRAM and model is 70B. Result: OOM on first long request.
• Using FIFO eviction (default in some old systems) and losing system prompts in long sessions.
• Not enabling prefix caching when 80% of requests share the same system prompt.
• Batching requests with vastly different sequence lengths (e.g., 100 tokens + 8K tokens). The cache pads to the longest, wasting VRAM on short requests.

References

• Paper Zhou, Y. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180. — arxiv:2309.06180 ↗
• Paper Dao, T., Fu, D., et al. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691. — arxiv:2307.08691 ↗
• Paper Xiao, G., et al. (2024). Streaming LLM: Efficient Streaming Language Models with Attention Sinks. arXiv:2309.17453. — arxiv:2309.17453 ↗
• Paper Zhang, W., et al. (2024). H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models. arXiv:2306.14048. — arxiv:2306.14048 ↗
• Paper Leviathan, Y., et al. (2023). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:2211.15297. — arxiv:2211.15297 ↗

• Docs vLLM PagedAttention. vllm.ai ↗
• Blog SGLang RadixAttention: Efficient Prefix Caching. lmsys.org ↗
• Docs HuggingFace Text Generation Inference KV Cache. huggingface.co ↗
• Blog Together AI: KV Cache Quantization for Efficient Serving. together.ai ↗

• Guide vLLM GitHub. github.com/vllm-project ↗
• Guide SGLang GitHub. github.com/hiyouga ↗
• Guide llama.cpp. github.com/ggerganov ↗

Learning Path

KV cache is an inference optimisation. Understand it in this order: