The KV Cache: Build It, Break It, Measure It

Maitreya Suin — Mon, 13 Jul 2026 00:00:00 GMT

TL;DR - Autoregressive generation re-runs the transformer over the entire prefix for every new token. Causal attention makes most of that work redundant: past tokens never look at future tokens, so their keys and values never change. Cache them, and each step processes one token. We build this from scratch on GPT-2, prove it’s mathematically exact, break it in the three ways everyone breaks it, and measure a 1.4-3.2× wall-clock speedup on a free Colab T4. Then we confront the two prices: the cache becomes the memory bottleneck of modern inference, and our simple implementation quietly copies gigabytes it already had.

If you know attention already, skim to §3; if you build inference systems, §§5–8 are for you. Everything runs in one Colab notebook on a free T4.

The whole post in one plot

Run GPT-2 in the most obvious generation loop - full forward pass over the whole sequence, every token - and time each step. Then turn the KV cache on and do it again.

Without the cache, the per-token latency is a rising line: token 400 costs visibly more than token 1, because each step reprocesses a longer prefix. With the cache, the line goes flat. Same model, same weights, same output tokens - provably the same, as we’ll see - at a fraction of the cost.

The rest of this post is the why, the how, the three common bugs, and the bill.

1. The problem: generation is a loop that forgets it’s a loop

A decoder-only LLM generates one token at a time. The naive loop:

tokens = prompt
for _ in range(max_new_tokens):
    logits = model(tokens)            # forward over the WHOLE sequence
    tokens = tokens + [sample(logits[-1])]

At step 1 we run the model over tokens. At step 2, over tokens - of which are identical to what we just processed. By step 500 we’ve pushed the first prompt token through the network 500 times and gotten the exact same intermediate results 500 times.

How expensive? For hidden size and sequence length , one forward pass costs roughly in the linear layers (QKV projections + MLP) and in attention. A single full-sequence forward is quadratic in length - the cubic term below appears only because the naive loop repeats that forward for prefixes of length :

Two things to notice, because most explanations get the emphasis wrong:

The attention term is cubic in generated length. That’s the headline.
But for realistic sizes ( in the hundreds–thousands, in the thousands, so ), the linear-layer term dominates the actual FLOP count. The biggest practical waste isn’t recomputing attention - it’s re-running every MLP over the whole prefix, every step.

The KV cache eliminates the repeated prefix work in both terms. Note carefully what it doesn’t do: the new token still attends over the full history, so cached decoding still costs attention per step, total - and generation still stays sequential, because token needs token . The cache only removes redundancy. Totals:

2. What decode actually needs: the math

Causal self-attention for token (single head, scale ):

needs - the query of the current token only.
needs - keys and values of all tokens so far.
Because the mask is causal, and for were computed from a prefix that hasn’t changed. They are constants. Recomputing them is pure waste.

So the algorithm rewrites itself: keep in memory; each step, run the network on one token - compute , append to the cache, attend against the full cache. Per-step cost drops from to .

Why not cache Q too? Remember that past queries appear nowhere in any future step. A query is a search request - once token ’s attention row is computed, is dead. Causality means new tokens query the past; the past never queries the future. (In an encoder, where everyone attends to everyone and the sequence doesn’t grow, none of this machinery applies. Caching is a decoding trick.)

flowchart LR
    subgraph step["decode step t"]
        x["x_t (1 token)"] --> qkv["W_qkv"]
        qkv --> q["q_t"]
        qkv --> k["k_t"]
        qkv --> v["v_t"]
    end
    style step fill:#e3f2fd,stroke:#90caf9,stroke-width:1px
    k -->|append| K[("K cache 1..t")]
    v -->|append| V[("V cache 1..t")]
    q --> att["attention: softmax(q_t·Kᵀ)·V"]
    K --> att
    V --> att
    att --> out["o_t → MLP → logits"]

Figure 1. Cached decoding in one transformer layer. The current token produces a fresh query, key, and value. The key and value are appended to the cache; the query attends over the entire cached prefix. Past queries are never stored - nothing ever asks for them again.

3. Build it: a minimal GPT-2 with a cache

We’ll write a ~120-line GPT-2 that loads real pretrained weights from Hugging Face, with a forward that optionally consumes and returns per-layer caches. (Architecture and weight-porting follow the well-known nanoGPT pattern.) First, here are the symbols we’ll use in the code and math:

Symbol	Meaning	GPT-2 small	Where you’ll see it
	batch size	1 here	first dim of everything
	tokens in this forward call	prompt length, or 1 during decode	`x: (B, T, C)`
	embedding width	768	residual stream
	attention heads	12	`q,k,v: (B, H, T, D_h)`
	head dim	64	last dim of q/k/v
`past`	tokens already cached	grows by each call	`cache: (B, H, past, D_h)` per layer, ×2 for K and V

Cell 1 - setup

# GPU: Runtime → Change runtime type → T4
!pip -q install tiktoken transformers
import math, time, torch, torch.nn as nn, torch.nn.functional as F
from dataclasses import dataclass
torch.manual_seed(0)
device = "cuda"
print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name())

Cell 2 - the model

@dataclass
class GPTConfig:
    block_size: int = 1024   # GPT-2 max context
    vocab_size: int = 50257
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_head = cfg.n_head
        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd)  # fused QKV
        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd)

    def forward(self, x, cache=None):
        B, T, C = x.shape
        # One linear layer produces Q, K, V together:
        #   x: (B, T, C) → qkv: (B, T, 3C) → split → q, k, v: each (B, T, C)
        q, k, v = self.c_attn(x).split(C, dim=2)
        # Reshape for multi-head attention: (B, T, C) → (B, H, T, D_h).
        # Heads become a batch-like dim so each head attends independently.
        q, k, v = (t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
                   for t in (q, k, v))
        if cache is not None:
            # Decode: this call's k,v cover only the NEW T tokens.
            # Prepend history along the sequence dim (dim=2 in (B,H,T,D_h)).
            k_past, v_past = cache          # each (B, H, past, D_h)
            k = torch.cat([k_past, k], dim=2)   # (B, H, past+T, D_h)
            v = torch.cat([v_past, v], dim=2)
            # ^ Simple, correct, and secretly expensive - see §8.
        # Masking (the full story is §5, bug 2):
        #   prefill  → square T×T attention → causal mask ON
        #   decode   → 1 query vs past+1 keys → everything IS the past; no mask needed.
        # is_causal=True during decode is not just unnecessary - it's WRONG:
        # SDPA aligns non-square causal masks top-left, hiding nearly the whole cache.
        y = F.scaled_dot_product_attention(q, k, v, is_causal=(cache is None))
        # (B, H, T, D_h) → (B, T, H, D_h) → (B, T, C).
        # contiguous() because transpose changed strides and view() needs flat memory.
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y), (k, v)

class MLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.c_fc   = nn.Linear(cfg.n_embd, 4 * cfg.n_embd)
        self.c_proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd)
        self.act    = nn.GELU(approximate="tanh")   # GPT-2 uses tanh-GELU
    def forward(self, x):
        return self.c_proj(self.act(self.c_fc(x)))

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.ln_1, self.attn = nn.LayerNorm(cfg.n_embd), CausalSelfAttention(cfg)
        self.ln_2, self.mlp  = nn.LayerNorm(cfg.n_embd), MLP(cfg)
    def forward(self, x, cache=None):
        a, cache = self.attn(self.ln_1(x), cache)
        x = x + a
        x = x + self.mlp(self.ln_2(x))
        return x, cache

class GPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.config = cfg
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(cfg.vocab_size, cfg.n_embd),
            wpe = nn.Embedding(cfg.block_size, cfg.n_embd),
            h   = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer)),
            ln_f= nn.LayerNorm(cfg.n_embd),
        ))
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.transformer.wte.weight   # weight tying

    def forward(self, idx, caches=None):
        B, T = idx.shape
        # Every layer has cached the same number of tokens;
        # caches[0][0] is layer 0's key cache: (B, H, past, D_h) → read `past` off dim 2.
        past = 0 if caches is None else caches[0][0].size(2)
        assert past + T <= self.config.block_size
        # POSITION OFFSET - the #1 cache bug (§5): token t must get position
        # embedding t, not 0. During decode T=1, so without the offset every
        # generated token would sit at position 0.
        pos = torch.arange(past, past + T, device=idx.device)
        x = self.transformer.wte(idx) + self.transformer.wpe(pos)
        new_caches = []
        for i, block in enumerate(self.transformer.h):
            x, c = block(x, None if caches is None else caches[i])
            new_caches.append(c)
        x = self.transformer.ln_f(x)
        return self.lm_head(x), new_caches

Cell 3 - load the real GPT-2 weights (weight surgery; fold open if curious)

Code

Maitreya - Notes

The KV Cache: Build It, Break It, Measure It

The whole post in one plot

1. The problem: generation is a loop that forgets it’s a loop

2. What decode actually needs: the math

3. Build it: a minimal GPT-2 with a cache