Katkat
Back

Explain KV cache in Transformer inference

OpenAIMachine Learning Engineer · Onsite
MediumMachine Learning EngineerOnsiteSoftware Engineering Fundamentals

Last updated

This question evaluates understanding of KV cache mechanisms in Transformer inference, including attention-state caching, memory and latency trade-offs, and engineering optimizations for autoregressive decoding.

Question

In Transformer-based large-language-model inference, what is a key-value (KV) cache?

Give a complete, systems-level explanation that covers:

  • What gets cached — which tensors, their shapes at a high level, and at which parts of the model they live.
  • Why KV caching speeds up autoregressive decoding (the asymptotics it changes).
  • The distinction between the prefill phase (processing the prompt) and the decode phase (generating tokens one at a time), and the very different performance profile of each.
  • The main tradeoffs and pitfalls: memory growth, batched / variable-length request management, the multi-head-attention variants (MHA vs. MQA vs. GQA), positional-encoding consistency, and long-context handling.
  • At least two practical optimizations used in production serving systems (e.g. paged attention, quantized KV cache, sliding-window / streaming attention, GQA).

Constraints & Assumptions

  • Assume a standard decoder-only Transformer doing autoregressive generation (causal self-attention), served on GPU.
  • $L$ = number of layers, $B$ = batch size, $P$ = prompt length, $S$ = current sequence length, $n_q$ = query heads, $n_{kv}$ = KV heads, $d_h$ = per-head dimension.
  • The discussion is about inference, not training — no backward pass, weights are frozen.
  • "Production" means a multi-tenant serving system handling many concurrent requests of differing lengths, not a single-sequence toy script.

Clarifying Questions to Ask

A candidate would scope the question by asking:

  • Is the target a decoder-only model (the common case), or are we including encoder–decoder cross-attention caching as well?
  • Are we optimizing for time-to-first-token (TTFT), inter-token latency / throughput (TPOT, tokens/sec), or max concurrent requests — they pull the design in different directions?
  • What context lengths and batch sizes matter? (This decides whether the cache or the weights dominate HBM.)
  • What positional scheme is in use (RoPE, learned absolute, ALiBi)? It changes what must be stored and how eviction/sliding interacts with positions.
  • Is the model architecture fixed, or can we assume / choose GQA/MQA (an architecture-time decision, not an inference-time knob)?

Hints

Hint 1

Where to start: Begin from what self-attention recomputes at every decode step. For a new token's query $Q_t$, which of the per-token projections of the *earlier* tokens are functions only of already-fixed hidden states, and therefore never change once computed?

Hint 2

The key invariant: Only $K$ and $V$ of past tokens are reused across steps; $Q$ is used once for its own token and discarded. Reason about what becomes $O(1)$ per step versus what stays $O(t)$ once you stop recomputing the prefix's projections.

Hint 3

Two regimes: Separate the prompt pass from the per-token loop. One is a tall matrix–matrix multiply (many query rows at once); the other is a thin matrix–vector multiply (one query row). Think about which is limited by GPU FLOPs and which by HBM bandwidth — that dictates which one batching helps.

Hint 4

Memory and pitfalls: Write the cache size as a product of the obvious factors (layers, batch, sequence, heads, head-dim, bytes/elem, ×2 for K and V) and notice which factors grow at runtime. From there, the production fixes fall out: shrink $n_{kv}$ (GQA/MQA), shrink bytes/elem (quantization), bound the sequence term (sliding window), or stop pre-reserving `max_len` (paged/block-wise allocation).

Follow-up questions

  • Derive the KV-cache size in bytes for a concrete model (give it $L$, $n_{kv}$, $d_h$, FP16) at batch $B$ and context $S$, and compare it to the weight memory — at what context length does the cache dominate?
  • How does **paged attention** enable prefix sharing across requests with a common system prompt, and what has to happen on a copy-on-write divergence?
  • Why are **keys** more sensitive than values to low-bit quantization, and what scaling granularity (per-tensor vs. per-channel vs. per-token) mitigates it?
  • With a **sliding-window** cache, what breaks if you naively drop the oldest tokens, and how do "attention-sink" schemes keep streaming generation stable at fixed memory?
  • In **disaggregated** serving, why might you place prefill and decode on separate hardware, and what must be transferred between them?