KV cache and why long context is costly
Section focus: Explain every item aloud without notes
Description
Level 1 — The Simple Intuition (ELI5)
Imagine writing a book one word at a time. If you had to re-read every single page you've already written from scratch just to decide the next word, writing page would take hours! Instead, you keep sticky notes of the key facts from earlier pages so you only process the newest word and glance back at your notes.
In an LLM, those sticky notes are the KV Cache (Key-Value Cache). During generation, past tokens' Keys () and Values () never change because causal masking prevents old tokens from looking at newer tokens. By caching and in GPU memory, generating token only computes the Query, Key, and Value for the single newest token—turning an per-step computation into !
Level 2 — How It Works (Prefill vs. Decode & Memory Math)
LLM inference happens in two distinct hardware regimes:
- Phase 1 — Prefill (Compute-Bound):
- When the user submits a prompt of tokens, the model processes all tokens in parallel in a single forward pass (matrix-matrix multiplication), populating the initial KV cache for all tokens across all layers.
- Because attention computes all pairs in the prompt, Prefill compute scales quadratically: FLOPs (this determines Time-To-First-Token / TTFT).
- Phase 2 — Autoregressive Decode (Memory-Bandwidth Bound):
- To generate token , we only feed the single latest token () into the network.
- We compute its new , append and to the cached matrices , and compute attention between and the entire KV cache.
- Compute drops from to per generated token, but for every single token generated, the GPU must stream both the model weights AND the entire growing KV cache from High-Bandwidth Memory (HBM) into SRAM (this determines Inter-Token Latency / ITL).
Exact KV Cache Memory Formula (Memorize for Interviews!)
For a batch size , sequence length , number of layers , number of KV heads , head dimension , and precision bytes ( bytes for FP16/BF16, byte for FP8/INT8):
(The leading factor of accounts for storing both K and V tensors.)
Level 3 — Advanced Depth (Why Long Context Is So Costly & How Production Engines Fix It)
Why does a or token context window cost a fortune?
- Quadratic Prefill Compute (): Doubling prompt length from to quadruples the attention FLOPs during prefill.
- Linear KV Cache Memory Explosion + HBM Bandwidth Bottleneck: Even with GQA, long sequences consume tens of gigabytes of VRAM per user request, forcing serving engines to shrink batch size , which tanks GPU throughput!
- VRAM Fragmentation PagedAttention (vLLM): Traditional serving pre-allocates a contiguous tensor of size per request, wasting of VRAM on unused padding. PagedAttention (inspired by OS virtual memory) slices the KV cache into fixed-size non-contiguous blocks (e.g., tokens per block) mapped via a block table, eliminating external fragmentation and enabling Prefix Caching (sharing system prompt KV blocks across requests).
- Architectural Compressions:
- GQA / MQA: Shrinks by .
- Multi-Head Latent Attention (MLA — DeepSeek-V2/V3): Compresses and jointly into a low-rank latent vector before caching, reducing KV cache by while outperforming standard MHA.
- Sliding Window Attention (Mistral / Gemma-2): Caps KV cache per layer to the last tokens (e.g., ).
- KV Cache Quantization (FP8 / INT4): Cuts bytes per element from to or .
Learning notes
Without KV Cache vs. With KV Cache (Decode Step )
Loading diagram…
Real-World Worked Calculation: LLaMA-2 70B vs. LLaMA-3 70B vs. Naive MHA at
Let layers, , heads, , FP16 ( bytes), Batch , Sequence ( tokens):
Architecture Variant | KV Heads () | Bytes per Token () | Total KV Cache at (Batch = 1) | Total KV Cache at Batch = 16 |
Standard MHA (70B) | (Exceeds 4x H100!) | |||
GQA 8-Head (LLaMA-3 70B, FP16) | ( smaller) | |||
GQA 8-Head + FP8 KV Cache | ( smaller) |
Example or implementation
PyTorch Autoregressive Attention with Stateful KV Cache
import math import torch import torch.nn as nn import torch.nn.functional as F class CachedCausalAttention(nn.Module): def __init__(self, d_model: int = 256, n_heads: int = 4): super().__init__() self.n_heads = n_heads self.d_k = d_model // n_heads self.q_proj = nn.Linear(d_model, d_model, bias=False) self.k_proj = nn.Linear(d_model, d_model, bias=False) self.v_proj = nn.Linear(d_model, d_model, bias=False) self.out_proj = nn.Linear(d_model, d_model, bias=False) def forward( self, x: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: B, S_cur, C = x.shape # During decode, S_cur == 1! q = self.q_proj(x).view(B, S_cur, self.n_heads, self.d_k).transpose(1, 2) k = self.k_proj(x).view(B, S_cur, self.n_heads, self.d_k).transpose(1, 2) v = self.v_proj(x).view(B, S_cur, self.n_heads, self.d_k).transpose(1, 2) # Concatenate new k, v onto cached history along sequence dimension (dim=2) if kv_cache is not None: k_prev, v_prev = kv_cache k = torch.cat([k_prev, k], dim=2) v = torch.cat([v_prev, v], dim=2) new_kv_cache = (k, v) # Compute attention: q has length S_cur (1 in decode), k/v have total length S_total scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k) if S_cur > 1: # Prefill phase needs causal mask mask = torch.tril(torch.ones(S_cur, S_cur, device=x.device)) scores = scores.masked_fill(mask == 0, float("-inf")) attn = F.softmax(scores, dim=-1) out = (attn @ v).transpose(1, 2).contiguous().view(B, S_cur, C) return self.out_proj(out), new_kv_cache # Demo: Prefill 10 tokens, then Decode 3 tokens one-by-one layer = CachedCausalAttention(d_model=256, n_heads=4) prompt = torch.randn(1, 10, 256) out, cache = layer(prompt, kv_cache=None) # Prefill: processes 10 tokens print("After Prefill -> Cache K shape:", cache[0].shape) for step in range(3): next_tok = torch.randn(1, 1, 256) # Decode: pass ONLY 1 new token! out, cache = layer(next_tok, kv_cache=cache) print(f"Decode step {step+1} -> Cache K shape: {cache[0].shape}")
Interview-ready explanation
90-Second Spoken Pitch for Interviews
"During autoregressive generation, predicting token only requires the hidden state at the final position . Because causal masking prevents earlier tokens from attending to later ones, the Key and Value projections of tokens through never change as new tokens are generated. Instead of recomputing them at every step—which would cost per token and to generate an -token sequence—we store past Keys and Values in GPU memory as the KV Cache, reducing per-step compute to .
This splits inference into two phases: Prefill, which processes the user's prompt in parallel and is compute-bound with attention FLOPs, and Decode, which generates one token at a time and is memory-bandwidth bound because we must stream both model weights and the entire KV cache from HBM for every single token.
The KV cache memory scales as . At long contexts like , a single request on a 70B GQA model takes of VRAM just for the KV cache—almost as much as the weights themselves! That is why production serving uses Grouped-Query Attention (GQA), PagedAttention (vLLM) to eliminate memory fragmentation, Prefix Caching, and FP8 KV quantization."