KV Cache From First Principles

Oct 13, 2025PyTorch · LLM

KV Cache From First Principles

How transformer inference reuses Key and Value state to avoid recomputing the past.

01

The need for KV caching

KV caching is one of the most impactful advances for large language model inference. It lowers latency and avoids repeating work while a model generates a response.

LLMs are autoregressive: each new token is predicted from the tokens that came before it. Attention therefore grows quadratically with sequence length, O(n²), if every decoding step recomputes the whole sequence. At production scale, that repeated work quickly becomes expensive.

02

A decoder-only LLM

At inference time, a prompt becomes a sequence of n tokens. The model samples one new token, appends it to the sequence, and repeats until it produces <EOS>. Without a cache, every generated token causes the model to recompute the Key and Value representations for all earlier tokens.

Vanilla attention

Attention derives scores from Q × Kᵀ / √dₖ. During prefill, the model processes the entire prompt. From the next token onward, the earlier states already have valid Key and Value vectors, but a naïve implementation calculates them again.

Illustration of vanilla attention

Repeated work without a cache

StepInput sequenceKeys and Values computedRedundant work
1Once upon a timeK/V for tokens 1–4None
2Once upon a time thereK/V for tokens 1–5Tokens 1–4
3Once upon a time there livedK/V for tokens 1–6Tokens 1–5
4Once upon a time there lived aK/V for tokens 1–7Tokens 1–6
03

The solution

A cache stores the Key and Value matrices already calculated for the active sequence. This trades memory for compute: each decoding step only derives K and V for its new token, then appends them to the cached matrices.

Illustration of KV caching

Procedure

  1. Prefill. Compute Query, Key, and Value matrices for the prompt, then retain the Key and Value matrices as the cache.
  2. Generation. Look up the next token embedding, compute its Q, K, and V vectors, append K and V to the cache, then attend with softmax(Q × Kᵀ / √dₖ) × V.
  3. Repeat until the model produces <EOS>.
  4. Discard the cache before starting an unrelated sequence.

Complexity

Prefill remains O(n²), because it processes the full input once. After that, KV caching changes decoding from repeated quadratic work to linear growth across generated tokens.

Trade-off

Caching avoids redundant compute, but its memory use grows linearly with context length. Long-context serving is often bandwidth-bound.
04

KV cache memory formula

For each token, every transformer layer stores a Key vector and a Value vector for each KV head:

KV cache size = 2 × nlayers × nheads × dhead × nbytes × ntokens
  • 2 stores both Keys and Values.
  • n_layers is the number of transformer layers.
  • n_heads is the count of attention heads used for KV.
  • d_head is each head's dimension.
  • n_bytes is bytes per element, such as 2 for BF16.
  • n_tokens is the context length.

Qwen3 0.6B example

  • n_layers = 28
  • d_head = 64
  • n_heads (KV) = 8 for grouped-query attention
  • n_bytes (BF16) = 2
  • n_tokens = 32,768

2 × 28 × 8 × 64 × 2 × 32,768 ≈ 1.75 GB. At full context, that cache can outweigh the roughly 1.4 GB model parameters.

05

Advantages and disadvantages

Advantages

  • Reduced computation. Previously computed states are reused.
  • Lower latency and higher throughput. Less repeated work shortens each decoding step.

Disadvantages

  • Memory overhead. Cache size grows with sequence length.
  • Bandwidth pressure. As contexts grow, moving cached state can become the bottleneck rather than arithmetic.
06

Mitigation strategies

  1. Sliding window. Retain only the most recent tokens rather than the full cache.
  2. Grouped-query attention. Let multiple Query heads share KV heads, reducing cache size.
  3. Multi-head latent attention. Compress Key and Value state into a lower-dimensional representation.
  4. Quantization. Store the cache at lower precision, such as FP8, to reduce its footprint.
  5. Offloading. Move cache state from GPU memory to system RAM or disk when context grows.