Explain KV cache in Transformer inference

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).

Hint: 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: 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: 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: 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).

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)?

What a Strong Answer Covers

A strong answer is judged on these dimensions (not the answers themselves):

• Precision of the cached object: that per-layer $K$ and $V$ for every past position are stored, $Q$ is not cached, and the cache lives only in attention layers (not embeddings/MLP/LM head).
• Correct asymptotics: identifies the redundant recomputation a cache removes (per-step projection cost $O(t) \to O(1)$, eliminating the $O(N^2)$ projection work over $N$ tokens) while being honest that the attention scan over cached keys stays $O(t)$.
• Prefill vs. decode as distinct regimes: compute-bound prompt pass → TTFT vs. bandwidth-bound per-token loop → TPOT, and why batching amortizes the weight read during decode.
• The memory model: a correct size formula, the observation that it is linear in $S \times B$, and the consequence that throughput is usually KV-memory-limited.
• Attention variants and their cache effect: MHA vs. MQA vs. GQA and how $n_{kv}$ scales the footprint.
• At least two concrete production optimizations named with their mechanism, not just buzzwords (e.g. paged attention's block table + prefix sharing; quantization's bytes/elem reduction; sliding window's bounded sequence term).
• Awareness of the sharp edges: ragged/variable-length batching, per-sequence causal masking when packing, positional consistency under eviction, and beam/multi-sample cache duplication.

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?

Solution

What a KV cache is

In a decoder-only Transformer, generation is autoregressive: token $t$ is produced by attending over all previous tokens $1\dots t-1$. Self-attention at each layer computes three projections of the hidden state $x$:

$Q = xW_Q,\quad K = xW_K,\quad V = xW_V$

and then $\text{Attn} = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}} + M\right)V$, where $M$ is the causal mask and $d_h$ is the per-head dimension.

The key observation: when you append a new token, the $K$ and $V$ vectors of all previous tokens do not change — they depend only on those tokens' hidden states, which are already fixed. Only the new token's $Q$ needs to attend over them. So instead of recomputing $K$ and $V$ for the whole prefix at every step, you store them.

A KV cache holds the per-layer, per-head $K$ and $V$ tensors for every position seen so far. At each decode step you compute $Q,K,V$ only for the single new token, append its $K,V$ to the cache, and run attention of that one query against the full cached $K,V$.

Note: $Q$ is not cached — each query is used once (for its own token) and then discarded. Only $K$ and $V$ are reused across future steps.

What gets cached and the shapes

For each of the $L$ transformer layers you keep two tensors:

K_cache: [batch, n_kv_heads, seq_len, head_dim]
V_cache: [batch, n_kv_heads, seq_len, head_dim]

• seq_len grows by 1 every decode step (and starts at the prompt length after prefill).
• n_kv_heads is the number of KV heads, which may be fewer than the number of query heads under grouped-query attention (see below).
• The cache lives in every attention layer; the embedding layer, MLPs, and final LM head are not cached because they have no cross-token reuse.

Exact memory cost

Total KV-cache bytes:

$\text{bytes} = 2 \cdot L \cdot B \cdot S \cdot n_{kv} \cdot d_h \cdot \text{bytes\_per\_elem}$

where the leading $2$ is for $K$ and $V$, $L$ = layers, $B$ = batch, $S$ = sequence length, $n_{kv}$ = KV heads, $d_h$ = head dim. The cache is linear in sequence length and batch size — this is what makes it the dominant memory consumer at long context. For a single sequence it only rivals the weights at extreme context length, but across a large batch the aggregate cache routinely exceeds the weights, and that is the regime that caps how many requests fit in HBM.

---

Why it speeds up decoding

Without a cache, generating each new token re-runs attention over the entire prefix, recomputing $K$ and $V$ projections for all $O(t)$ past tokens. Over $N$ generated tokens that is $O(N^2)$ redundant projection work.

With a cache:

• Projection work per step is $O(1)$ — you project only the new token.
• Attention for the new query still reads $O(t)$ cached keys/values (the dot products can't be avoided), but you no longer recompute the projections of the past.

So the cache turns the per-step projection cost from $O(t)$ to $O(1)$. The remaining $O(t)$ attention scan is cheap relative to the projections and MLP, and is memory-bandwidth-bound rather than compute-bound. The net effect is a large, roughly constant-time-per-token decode loop instead of one that degrades with prefix length.

---

Prefill vs. decode

These are two operationally distinct phases with different performance profiles.

	Prefill	Decode
Input per forward pass	Whole prompt, $P$ tokens	1 token
Cache action	Populate all $P$ positions, all layers	Append 1 position
Matrix shape	matrix–matrix (tall $Q$)	matrix–vector (single-row $Q$)
Bottleneck	Compute-bound (GPU FLOPs)	Memory-bandwidth-bound
Parallelism	All prompt positions in parallel	Inherently sequential

Prefill runs the full prompt through the model once. Because $Q$ has $P$ rows, the attention and FFN are large dense matmuls that saturate the GPU's compute units — it's throughput-oriented and batches well. Its latency determines time-to-first-token (TTFT).

Decode processes one token at a time, so each step is a thin matrix–vector multiply. The arithmetic intensity is low: you stream the entire model weights (and the growing KV cache) from HBM to do very little compute, so you're limited by memory bandwidth, not FLOPs. Decode latency determines inter-token latency (TPOT, time per output token) — the per-token latency whose reciprocal is the tokens-per-second throughput. This is why serving systems batch many concurrent decodes together — it amortizes the one-time weight read from HBM across many requests.

This split is why the two phases are often scheduled, and even placed, on separate hardware ("disaggregated" prefill/decode) in large deployments — see the follow-up below.

---

Tradeoffs and pitfalls

1. Memory growth dominates. As shown above, KV memory is linear in $S \times B$. At long context the cache can dwarf the weights, capping how many concurrent requests (batch size) fit in HBM. Throughput in production is usually KV-memory-limited, not FLOP-limited.

2. Variable-length / ragged batching. Sequences in a batch have different prompt lengths and finish at different times. Naively reserving max_seq_len per slot wastes most of the cache. Finished sequences leave holes; you need bookkeeping to reclaim and reuse that space (this is the core motivation for paged attention, below).

3. Multi-head attention variants change the cache size. The cache scales directly with $n_{kv}$:

• MHA (multi-head): $n_{kv} = n_q$ — largest cache.
• MQA (multi-query): a single shared KV head, $n_{kv} = 1$ — smallest cache, some quality loss.
• GQA (grouped-query): $n_{kv}$ between $1$ and $n_q$ (query heads share KV heads in groups) — the common modern compromise. GQA shrinks the cache by the grouping factor with little quality cost, which is why it is now standard.

4. Positional-encoding consistency. With RoPE, the rotation is applied to $K$ as a function of absolute position before caching, so cached keys already carry their positional phase — straightforward as long as you rotate at the correct index. The pitfall appears when positions shift (sliding window or cache eviction): naive eviction misaligns the position math unless you re-anchor or use a relative scheme. With learned absolute positions you must add the embedding at the right index.

5. Masking with packed batches. The causal mask must keep each token attending only to earlier positions of its own sequence. When multiple sequences are packed into one tensor, you need per-sequence (block-diagonal) masks so tokens don't leak across sequence boundaries.

6. Beam search / multi-sample. Beam search or parallel sampling multiplies the cache by the beam width / number of samples unless the shared prompt prefix is stored once and only the branches are duplicated (prefix sharing).

---

Practical production optimizations

1. Paged attention (block-wise KV management). Store the cache in fixed-size blocks ("pages") and keep a per-sequence block table mapping logical positions to physical pages, analogous to OS virtual memory. Benefits: near-zero internal fragmentation, growth one page at a time instead of pre-reserving max_len, and prefix sharing — multiple requests with a common system prompt point at the same physical pages, with copy-on-write on divergence. This typically allows substantially larger effective batch sizes from the same HBM. (Popularized by vLLM.)

2. Quantized KV cache. Store $K$ and $V$ in lower precision (e.g. INT8 or FP8) while keeping matmul accumulation in BF16/FP16. Roughly halves or quarters both the cache footprint and the bandwidth needed to read it during decode, raising batch capacity and decode throughput. Requires per-channel/per-token scaling — keys are more sensitive than values because their outlier channels, once corrupted, distort every attention score through the $QK^\top$ dot product (see follow-up).

3. GQA / MQA at the architecture level. Reducing $n_{kv}$ is the most direct lever — it shrinks the cache proportionally essentially for free at inference time (decided at training/architecture time). This is why nearly all recent large models ship with GQA.

4. Sliding-window / streaming attention. Cap the cache to the most recent $W$ tokens, giving bounded memory regardless of total length. Pure truncation loses long-range dependencies; "attention-sink" variants keep the first few tokens plus the recent window, which preserves stability for streaming/infinite generation at fixed memory (see follow-up).

5. KV offloading / hierarchical cache. Spill cold KV blocks to CPU RAM (or NVMe) and page them back when needed — trades latency for capacity, useful for very long contexts or for caching reusable prefixes across requests. Often combined with chunked/retrieval-style attention so you don't read the whole offloaded cache every step.

---

Worked answers to the follow-ups

KV-cache size vs. weights — when does the cache dominate?

Take a model with $L = 32$ layers, $n_{kv} = 8$ KV heads (GQA), $d_h = 128$, FP16 (2 bytes). Per token per sequence:

$2 \cdot L \cdot n_{kv} \cdot d_h \cdot 2 \text{ bytes} = 2 \cdot 32 \cdot 8 \cdot 128 \cdot 2 = 131{,}072 \text{ bytes} \approx 128\text{ KiB/token}.$

• One sequence at $S = 8{,}000$ tokens: $\approx 1.0$ GiB.
• Batch of $B = 64$ at $S = 8{,}000$: $\approx 64$ GiB — comparable to or larger than the weights of a ~7–13B model in FP16, and this is the regime that fills an 80 GB GPU.

The crossover is governed by $\text{cache} = 2 L\,n_{kv}\,d_h \cdot \text{bytes} \cdot B\,S$ versus a fixed weight term, so the cache overtakes the weights once $B \cdot S$ is large enough — pushed sooner by long context, large batch, and (had we used MHA instead of GQA) more KV heads.

How paged attention enables prefix sharing

Each request's logical token positions are mapped, through its block table, to physical KV pages. Two requests that begin with the same system prompt prefill that prefix once; both block tables point their leading entries at the same physical pages (refcounted). They diverge only at the first differing token. On a write that would mutate a shared page (the point of divergence, or any later token), the server does copy-on-write: allocate a fresh page, copy the shared block, redirect that request's block-table entry, and decrement the shared page's refcount. The shared prefix's compute and memory are paid once across all requests that share it.

Why keys are more quantization-sensitive than values

Keys enter the model through the $QK^\top$ dot product before the softmax, so a quantization error on a key perturbs the pre-softmax logits, which the softmax can amplify into a large shift in the attention distribution. Values enter after the softmax as a weighted average, so value errors are averaged and damped. Empirically, key activations also carry large per-channel outliers; per-channel (or grouped) scaling for keys and per-token scaling for values, rather than a single per-tensor scale, contains the error and lets INT8/FP8 KV match BF16 quality closely.

Sliding window — why naive truncation breaks, and attention sinks

Naively dropping the oldest tokens removes the model's ability to refer back beyond $W$ — but more acutely, models trained with full attention place a disproportionate amount of attention mass on the first few tokens (an "attention sink"). When sliding eviction removes those initial tokens, the softmax has no low-information landing spot to dump excess attention, the distribution destabilizes, and generation degrades or diverges. StreamingLLM-style schemes fix this by always retaining the first few "sink" tokens plus the recent $W$-token window, giving stable, fixed-memory streaming over effectively unbounded length (at the cost of true long-range recall, which sliding inherently sacrifices).

Disaggregated prefill/decode

Prefill is compute-bound (saturates FLOPs) and decode is bandwidth-bound; co-locating them lets a long prefill stall the latency-sensitive decode loop ("decode interference"). Disaggregation runs prefill on one pool of GPUs and decode on another, each tuned for its bottleneck and scaled independently. The cost is that the prefilled KV cache must be transferred from the prefill node to the decode node (over NVLink/InfiniBand) before the first token is generated — a bandwidth and TTFT tradeoff that pays off at scale because it removes interference and improves overall goodput.

---

How to land this in an interview

• Define it precisely: cached $K,V$ per layer per position; $Q$ is not cached; lives only in attention layers.
• State the cost formula and call out that it's linear in sequence length × batch — the central scaling constraint.
• Frame the core tradeoff as latency vs. memory, and connect prefill → TTFT (compute-bound) and decode → TPOT (bandwidth-bound).
• Name a concrete optimization with its mechanism (paged attention for fragmentation/sharing, quantization or GQA for footprint, sliding window for bounded context) rather than just listing buzzwords.