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