KV Cache & PagedAttention

KV Cache is the inference memory structure that caches the Key/Value vectors of already-generated tokens during autoregressive decoding, avoiding recomputing attention over the full history at every step. Each decoding step only computes the new token's Q, attends against cached K/V, and appends new K/V to the cache. Memory formula (FP16, 2 bytes/element): $\text{KV Bytes} = 2 \cdot 2 \cdot n_{layers} \cdot n_{kv\_heads} \cdot d_{head} \cdot s \cdot b$ — the first 2 = one K and one V, the second 2 = 2 bytes per FP16 element, $s$ = sequence length, $b$ = batch size. Worked example (LLaMA-3 70B: 80 layers, 8 GQA KV heads, $d_{head}=128$, 32K context, batch 8): $2 \times 2 \times 80 \times 8 \times 128 \times 32768 \times 8 \approx 85.9\text{ GB} \approx 86\text{ GB}$; the same model with MHA (64 KV heads) needs about $687\text{ GB}$ — exactly $64/8 = 8\times$ more, the direct payoff of GQA. PagedAttention (vLLM): instead of allocating KV as one contiguous array sized for worst case, KV is sliced into fixed-size blocks (default 16 tokens per block, i.e. $16 \times n_{kv\_heads} \times d_{head}$ elements) and a block table maps logical positions to physical blocks, paging like an OS. This eliminates internal fragmentation (static pre-allocation exceeding actual generated length) and external fragmentation (interleaved sequences of different lengths leaving unusable gaps), pushing effective memory utilization from roughly 60-80% for contiguous allocation toward ~100%.

standard in every LLM serving engine (vLLM/SGLang/TensorRT-LLM internals); the dominant memory bottleneck for long-context and high-concurrency batched decode; interview staples: "hand-compute KV cache memory", "why PagedAttention saves memory", "how GQA/MLA compress it".

without caching, each step's attention is $O(s)$ and total decode compute is $O(s^2)$ with massive recomputation; KV caching makes each step $O(1)$ incremental, but memory grows linearly with $s \times b$ (tens of GB at long context / high concurrency, 30-60% of a GPU's memory). PagedAttention removes fragmentation and approaches 100% utilization, letting a 70B model serve 16K-context requests in batches; GQA/MLA shrink $n_{kv\_heads}$ by 8× or more, cutting KV memory linearly.

Hand-compute KV cache memory: write $\text{KV Bytes} = 2\cdot2\cdot n_{layers}\cdot n_{kv\_heads}\cdot d_{head}\cdot s\cdot b$ and explain what each 2 means; for 80 layers, 8 KV heads, $d_{head}=128$, 16K context, batch 16, compute per-sequence and total GB.

Why PagedAttention removes memory fragmentation: 16-token blocks plus a block table; how internal and external fragmentation arise and disappear; utilization improvement from ~60-80% to ~100%.

How GQA/MLA compress KV memory: which term in the formula shrinks when MHA → GQA shares KV heads; LLaMA-3 70B (64→8) gives an 8× reduction; MLA's low-rank projection idea.

KV write patterns in prefill vs decode: prefill writes all s KV entries at once, decode appends one per step; why the memory peak sits at the end of prefill.

Precision trade-offs of KV quantization: why quantizing KV (e.g. INT8) is more sensitive than weight quantization; how it composes with FlashAttention and streaming/sliding-window KV to cap memory.