Key-Value Cache

A Key-Value Cache (KV Cache) is a memory optimization technique in Transformer-based models that stores the intermediate key and value tensors computed during attention, allowing previously processed tokens to be reused across autoregressive generation steps without recomputation. It is foundational to practical LLM inference, reducing the per-step computational complexity from O(n²) to O(n) and making generation of long sequences tractable on real hardware. The counterintuitive fact: the KV cache can consume more GPU memory than the model weights themselves for long-context generation — a 7B-parameter model may need 14 GB for weights but over 30 GB of KV cache at context lengths above 100K tokens.

The core insight

During autoregressive generation, the new token attends to every previous token — but those previous tokens' keys and values are identical to what was computed on the prior step. The aha moment is recognizing that attention is stateless across the sequence dimension: storing K and V matrices from past steps is lossless, and recomputing them is pure waste. The KV cache turns inference from a full forward-pass problem into an incremental one: only the newest token's query, key, and value need to be freshly computed, while the cache supplies all historical context.

How It Works

Attention Recap

In a standard multi-head self-attention layer, given an input sequence of $n$ tokens with representation matrix $X \in \mathbb{R}^{n \times d}$ , we project into query, key, and value spaces:

\quad Q = X W_Q,\quad K = X W_K,\quad V = X W_V

Explanation: Three learned linear projections transform the same input sequence into the role-specific matrices Q, K, and V used by the attention mechanism. Here $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$ are learned weight matrices, $d$ is the model dimension, and $d_k = d / H$ is the per-head dimension for $H$ heads. The attention output is then:

\quad \text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V

Explanation: Each query row scores against all key rows to produce a weighted sum of value rows, scaled to prevent vanishing gradients in the softmax.

The Autoregressive Setting

During generation, the model produces one token at a time. At step $t$ , the input is only the new token $x_t$ , giving a single query vector $q_t \in \mathbb{R}^{1 \times d_k}$ . However, it must attend to all previous tokens $1 \ldots t$ . Without caching, we would recompute the full $K, V \in \mathbb{R}^{t \times d_k}$ at every step, costing $O(t \cdot d)$ FLOPs just in projection — repeated $t$ times to generate $t$ tokens, yielding $O(t^2 d)$ total. With the KV cache, we maintain:

\quad \mathcal{K}_t = \begin{bmatrix} k_1 \\ k_2 \\ \vdots \\ k_t \end{bmatrix}, \quad \mathcal{V}_t = \begin{bmatrix} v_1 \\ v_2 \\ \vdots \\ v_t \end{bmatrix}

Explanation: The cache simply concatenates the new key and value row vectors onto the stored matrices at each step, growing the cache by one row per generated token. where $k_i = x_i W_K$ and $v_i = x_i W_V$ are computed once when token $i$ is first processed, then stored. At step $t$ we only compute $q_t = x_t W_Q$ fresh, then:

\quad \text{Attn}(q_t, \mathcal{K}_t, \mathcal{V}_t) = \text{softmax}\!\left(\frac{q_t \,\mathcal{K}_t^\top}{\sqrt{d_k}}\right) \mathcal{V}_t

Explanation: The single new query vector scores against the full cached key matrix, producing a weighted sum of cached values — no historical recomputation needed.

Memory Cost

The memory required to store the KV cache for a model with $L$ layers, $H$ heads, head dimension $d_k$ , context length $n$ , batch size $B$ , and float16 precision is:

\quad \text{Mem}_{\text{KV}} = 2 \cdot L \cdot H \cdot d_k \cdot n \cdot B \cdot 2 \text{ bytes}

Explanation: The factor of 2 accounts for both K and V tensors, and 2 bytes per float16 element; every additional token or batch element linearly increases this cost. For LLaMA-2-70B ( $L=80$ , $d=8192$ , $H=64$ , $d_k=128$ ) at $n=4096$ and $B=1$ : approximately 40 GB — exceeding the 70B weight size in float16 (~140 GB total, but the cache alone is significant).

PyTorch Implementation Sketch

```python import torch import torch.nn.functional as F class CachedAttention(torch.nn.Module): def __init__(self, d_model: int, n_heads: int): super().__init__() self.n_heads = n_heads self.d_k = d_model // n_heads self.W_q = torch.nn.Linear(d_model, d_model, bias=False) self.W_k = torch.nn.Linear(d_model, d_model, bias=False) self.W_v = torch.nn.Linear(d_model, d_model, bias=False) self.W_o = torch.nn.Linear(d_model, d_model, bias=False) def forward( self, x: torch.Tensor, # (B, T, D) — T=1 during generation kv_cache: dict | None = None, # {"k": (B,H,S,dk), "v": (B,H,S,dk)} use_cache: bool = True, ): B, T, D = x.shape H, dk = self.n_heads, self.d_k def split_heads(t): # (B, T, D) -> (B, H, T, dk) return t.view(B, T, H, dk).transpose(1, 2) q = split_heads(self.W_q(x)) # (B, H, T, dk) k = split_heads(self.W_k(x)) # (B, H, T, dk) v = split_heads(self.W_v(x)) # (B, H, T, dk)

# ── Append to cache ────────────────────────────────────────── if kv_cache is not None: k = torch.cat([kv_cache["k"], k], dim=2) # grow sequence dim v = torch.cat([kv_cache["v"], v], dim=2) new_cache = {"k": k, "v": v} if use_cache else None

# ── Scaled dot-product attention ───────────────────────────── scale = dk ** -0.5 scores = torch.einsum("bhqd,bhkd->bhqk", q, k) * scale

# Causal mask only needed for prefill (T > 1) if T > 1: mask = torch.triu(torch.ones(T, k.size(2), device=x.device), diagonal=1).bool() scores = scores.masked_fill(mask[-T:, :], float("-inf")) attn = F.softmax(scores, dim=-1) out = torch.einsum("bhqk,bhkd->bhqd", attn, v) # (B, H, T, dk)

# ── Merge heads & project ──────────────────────────────────── out = out.transpose(1, 2).contiguous().view(B, T, D) return self.W_o(out), new_cache ``` The `kv_cache` dict is passed layer-by-layer and updated at each generation step — in practice, frameworks like HuggingFace `transformers` manage a list of per-layer caches automatically via `past_key_values`.

researcher · Ashish Vaswani | b. 1983 | Inventor of the Transformer Architecture

Ashish Vaswani led the team at Google Brain that published "Attention Is All You Need" (2017), the foundational paper that introduced the Transformer and the self-attention mechanism from which KV caching directly emerges. While the original paper focused on encoder-decoder translation and did not explicitly discuss inference-time KV caching, the architectural separation of Q, K, and V projections is precisely what makes stateless incremental caching possible. Vaswani later co-founded Essential AI; the 2017 paper has accumulated over 100,000 citations, making it one of the most cited ML papers ever.

Figure

KV Cache Data Flow During Autoregressive Generation

Intuition

The Mental Model

Imagine you're a research assistant reading a long document to answer questions from a professor. Every time the professor asks a new question, you must re-read the entire document to know what's relevant — unless you've made margin notes. The KV cache is your margin notes: when you first read each page, you annotate what's important (the keys) and what each passage means (the values). From then on, a new question (query) just scans your notes — you never re-read the source pages. The query must always be freshly formulated because the question changes each step; but the document content is fixed, so its K and V representations are immortal once computed. The cost you pay is carrying a growing notebook (memory), and the benefit is never re-reading old text (computation).

Historical Development

Origins · 1 minhow this was discovered

History

2017: Vaswani et al. publish "Attention Is All You Need," introducing the Transformer with separate Q/K/V projections — the architectural prerequisite for KV caching.

2018: GPT-1 (Radford et al., OpenAI) demonstrates decoder-only autoregressive Transformers; inference still naive, but the generation setting that demands KV caching is established.

2019: GPT-2 and Megatron-LM (Shoeybi et al., NVIDIA) scale Transformers; NVIDIA's inference code introduces explicit past key-value management as `past_key_values` in early HuggingFace integrations.

2020: HuggingFace `transformers` library standardizes `past_key_values` as the KV cache API across GPT-2, CTRL, and XLNet — making KV caching accessible to the practitioner community.

2022: FlashAttention (Dao et al., Stanford) restructures attention computation for IO-efficiency; works synergistically with KV caching by accelerating the prefill phase.

2023: PagedAttention (Kwon et al.) in the vLLM system treats the KV cache as virtual memory pages, eliminating fragmentation and enabling continuous batching — a breakthrough for production LLM serving.

2023: Grouped Query Attention (GQA) introduced in Ainslie et al. and adopted by LLaMA-2-70B and Mistral-7B; shares K/V heads across multiple Q heads, reducing KV cache size by 4–8× with minimal accuracy cost.

2023–2024: Multi-Query Attention (MQA), sliding-window attention (Mistral), and quantized KV caches (KIVI, KVQuant) push cache compression; Apple, NVIDIA, and Google deploy quantized int8/int4 KV caches in production inference engines.

2024: Cross-layer KV sharing (MLA in DeepSeek-V2) and streaming KV eviction (H2O, SnapKV) emerge as frontier research, reducing cache memory by 10–20× through low-rank projection and dynamic token dropping.

Why It Matters

Real-World Impact

GPT-4 / ChatGPT (OpenAI)Uses KV caching with prompt caching across API calls — repeated system prompts are cached server-side, cutting inference latency and cost for common prefixes by up to 50%, a feature OpenAI began billing differentially in 2023.

LLaMA-2 and LLaMA-3 (Meta AI)Both models use Grouped Query Attention (GQA) specifically to shrink the KV cache footprint; LLaMA-3-70B uses 8 KV heads vs. 64 Q heads, making its cache 8× smaller than a naive multi-head design at equivalent context length.

vLLM Serving FrameworkImplements PagedAttention, which manages the KV cache as non-contiguous memory pages (analogous to OS virtual memory), increasing GPU utilization by 2–4× and throughput by up to 24× compared to HuggingFace Transformers naive serving.

Google Gemini (long context)Gemini 1.5 supports 1M-token contexts; this is only feasible because of aggressive KV cache management combined with sparse attention patterns — without caching, a single 1M-token sequence would require quadratic compute per step.

Apple On-Device Inference (Core ML)iPhone and Mac on-device LLM inference (iOS 18 "Apple Intelligence") uses 4-bit quantized KV caches to fit multi-turn conversations within the 6–8 GB memory envelope of Apple Silicon neural engine.

Speculative Decoding (DeepMind, Google)Draft-model proposals are verified in parallel against the target model's cached KV state — the KV cache enables batch verification of multiple candidate tokens at once, achieving 2–3× generation speedup.

Anthropic Claude (long context)Claude 3's 200K-token context window leverages sliding-window attention with a hybrid full-context KV cache; Anthropic's prompt caching API (released 2024) explicitly exposes cache control to users for cost reduction on repeated document analysis.

Key-Value Cache

How It Works

Attention Recap

The Autoregressive Setting

Memory Cost

PyTorch Implementation Sketch

Intuition

Historical Development

Why It Matters

Prerequisites

What This Leads To

See Also

Where this leads