KV caching and the generation loop

Folklore, and the papers that fixed it

Why inference is affordable. It's a direct consequence of the causal mask, and nobody explains it properly.

Download kv_cache_from_scratch.py — everything below runs from it. NumPy only, random weights, under a second.

This one isn't a paper. It's the engineering consequence of one line in Attention Is All You Need, and it lives in folklore: everyone who has served a model knows it, nobody writes it up. I shipped an inference stack before I could have told you, on a whiteboard, why the cache is allowed to exist. The papers here are the ones that later went and optimised it.

The whole thing is one observation. Generating token n means running attention over all n tokens. Do that naively for a whole sequence and you pay roughly n3. But the causal mask guarantees that token i's key and value never change once computed. So you cache them. The loop becomes n2, and each step is one new row instead of a full matrix.

The KV cache is memoisation

You've written this a hundred times:

@lru_cache
def expensive(x):        # safe to cache, because it is PURE:
    return f(x)          # same x always gives the same answer

And you know the rule: you may only memoise a pure function. Cache something whose answer depends on hidden mutable state and you get the worst class of bug — right answer on Tuesday, wrong answer under load.

A decoder's per-token key and value are that pure function. Token i's k and v are computed from token i's embedding and position, and from nothing else that arrives later. So: 

cache = {"K": [], "V": []}

def decode_step(token, pos, cache):
    q, k, v = project(embed(token) + pos_enc(pos))   # q is per-step, throwaway
    cache["K"].append(k)                             # append ONE row
    cache["V"].append(v)
    scores  = q @ stack(cache["K"]).T                # (1, n), not (n, n)
    return softmax(scores) @ stack(cache["V"])

That's it. That's the KV cache. The interesting part is not the code, it's the licence to write it — and the licence comes from the mask.

TermMeansProgrammer's analogy
PrefillRun the whole prompt through once, filling the cacheWarming the cache on startup
DecodeOne token in, one token out, cache grows by a rowThe steady-state request loop
KV cacheEvery layer's keys and values for every token so farThe memo table
Causal maskRow i may only read columns ≤ iThe proof that the function is pure
EvictionDropping cache entries when you run out of roomLRU on a bounded cache

The one formula

Not an attention formula — you already have that. The formula that decides what you can actually serve:

bytes = 2 × layers × kv_heads × head_dim × seq_len × dtype_bytes × batch

Read it as terms you control:

TermTypicalWhat it is
22K and V. Never Q — Q is used once and thrown away
layers32Every layer keeps its own cache. They are not shared
kv_heads32 → 8 → 1The only term the architecture papers attack. MHA → GQA → MQA
head_dim128dmodel / nheads
seq_len4k…128kGrows every single decode step. This is the one that hurts
dtype_bytes2fp16 or bf16. Quantising the cache to 8 bits halves it
batch1…256Every concurrent request pays the whole thing again

Two loops, and a 34× gap

The script runs both loops over a toy decoder and counts multiply-accumulates analytically, so the number doesn't depend on how fast your laptop is. Per token, per layer: 4·d·d for the four projections, 2·d·n for attention over n cached keys, 2·d·ff for the feed-forward.

Don't take my word for it. Here's the actual output:

toy model: 2 layers, 4 heads, d_model=32, seq=64

  naive  loop total MACs:   62,566,400
  cached loop total MACs:    1,839,104
  ratio                 :         34.0x

  per-step cost (MACs), naive grows, cached is flat:
    step |        naive |      cached
       1 |       24,704 |      24,704
       8 |      204,800 |      25,600
      16 |      425,984 |      26,624
      32 |      917,504 |      28,672
      48 |    1,474,560 |      30,720
      64 |    2,097,152 |     32,768 ✓

Look at the two columns. The naive step cost multiplies by 85 between step 1 and step 64. The cached step cost grows by a third — and that residue is real, not rounding. Attention still has to scan every cached key. What went flat is the projections and the feed-forward, which is most of the work.

Be careful with the big-O in interviews. The n3 → n2 claim is about the attention part of a decode loop. Caching does nothing at all to the feed-forward cost per token: you still run the FFN once per generated token either way. What caching removes is re-running the prefix.

Same tokens, to float precision

A cache that changes the answer is not a cache, it's an approximation. So the script generates the same 12 tokens both ways and diffs the logits. 

  naive  tokens: [3, 17, 42, 8, 16, 16, 16, 16, 44, 44, 16, 16, 16, 16, 44, 16]
  cached tokens: [3, 17, 42, 8, 16, 16, 16, 16, 44, 44, 16, 16, 16, 16, 44, 16]
  same token sequence? True
  max |naive logits - cached logits| = 1.887e-15 ✓

1.9 × 10−15 is float64 noise from doing the additions in a different order. The weights are random and untrained, so the model happily loops on token 16 — that's fine, this page is about compute and memory, not learning. What matters is that both loops loop identically.

The mask is the proof

Here is the part nobody writes down. Why is it legal to keep those rows? Because with a causal mask, appending a token cannot change any earlier attention row. Row i has −∞ in every column past i, so the new column contributes exactly nothing to its softmax denominator.

Same layer, same head, before and after appending a 5th token:

  WITH causal mask -- layer 0, head 0 attention weights
    before appending a 5th token (4x4):
      [1. 0. 0. 0.]
      [0.671 0.329 0.    0.   ]
      [0.53  0.235 0.235 0.   ]
      [0.553 0.162 0.159 0.126]
    after  appending a 5th token, first 4 rows (4x5):
      [1. 0. 0. 0. 0.]
      [0.671 0.329 0.    0.    0.   ]
      [0.53  0.235 0.235 0.    0.   ]
      [0.553 0.162 0.159 0.126 0.   ]
    max change in the first 4 rows: 0.000e+00 ✓

Exactly zero. Not "small". Now delete the mask — turn the decoder into a BERT-style bidirectional encoder — and run the identical experiment: 

  NO mask -- layer 0, head 0 attention weights
    before appending a 5th token (4x4):
      [0.386 0.182 0.305 0.127]
      [0.417 0.204 0.261 0.117]
      [0.47  0.209 0.208 0.113]
      [0.553 0.162 0.159 0.126]
    after  appending a 5th token, first 4 rows (4x5):
      [0.362 0.17  0.285 0.119 0.064]
      [0.394 0.193 0.247 0.111 0.054]
      [0.448 0.199 0.198 0.107 0.048]
      [0.506 0.149 0.146 0.115 0.084]
    max change in the first 4 rows: 4.631e-02

  now the consequence for a cache:
    unmasked model, cached vs recomputed logits, max abs diff: 0.4017
    argmax token: recomputed=44  cached=16

Every earlier row moved, because softmax renormalises over the new column. The cached rows are now stale, the logits are off by 0.40, and the model picks a different token. Same code, one boolean flipped.

This is why BERT has no KV cache and GPT does. It was never a choice about performance. Bidirectional attention makes K and V depend on the future, so there is nothing pure to memoise. The mask is the invariant the cache is built on — and once you have it, you don't even need the −∞ matrix at decode time, because the cache only ever holds the past.

The cache outgrows the model

Now the arithmetic that decides your hardware bill. A Llama-2-7B-shaped model: 32 layers, 32 heads, head_dim 128, fp16. 

  bytes = 2 * layers * kv_heads * head_dim * seq_len * dtype_bytes

  weights (fp16, 6.74B params):  12.55 GiB
  cache per token: 524,288 bytes = 512 KiB

      seq_len |   batch 1 |  batch 32 | batch 32 vs weights
        1,024 |   0.50 GiB |   16.0 GiB |    1.3x the weights
        4,096 |   2.00 GiB |   64.0 GiB |    5.1x the weights
       32,768 |  16.00 GiB |  512.0 GiB |   40.8x the weights
      131,072 |  64.00 GiB | 2048.0 GiB |  163.2x the weights

Half a megabyte per token. One 128k-token conversation is 64 GiB of cache on its own — five times the weights, and it doesn't fit on an 80 GB card alongside the model. Batch 32 at a mundane 4k context already outweighs the checkpoint. When someone says a model "needs an A100", they're usually quoting the weights, which is the small number.

MQA and GQA: shrink one term

Look at the formula again. layers and head_dim are the model. seq_len and batch are your users. That leaves kv_heads, and that is exactly where the papers went.

Multi-Query Attention (Shazeer, 2019): keep all 32 query heads, but give them one shared key/value head. Grouped-Query Attention (Ainslie et al., 2023) is the middle ground: 8 KV heads, four query heads sharing each. Q is never touched, in either — because Q is never stored. 

      variant           kv_heads |  4k, b32 | 32k, b32 | reduction
      multi-head (2017)       32 |   64.0GiB |  512.0GiB |      1x
      GQA, 8 groups            8 |   16.0GiB |  128.0GiB |      4x
      MQA, 1 kv head           1 |    2.0GiB |   16.0GiB |     32x

  same table as bytes per token, batch 1:
      multi-head   524,288 bytes/token
      GQA-8        131,072 bytes/token
      MQA           16,384 bytes/token ✓

The reduction is exactly n_heads / n_kv_heads. No subtlety. MQA takes 32× off and the GQA paper's ablations show it losing quality; GQA-8 takes 4× off and lands close to multi-head. That is why Llama-2 70B uses GQA-8, and why almost every model released since has too. It's a memory decision that happens to be architectural.

Prefill is compute-bound, decode is memory-bound

Take a 512-token prompt. Process it in one batched pass, or as 512 sequential decode steps. Same total arithmetic — and completely different machines. 

  Llama-2-7B-shaped, 512-token prompt (approximate, attn + FFN):
    prefill, one batched pass : 3,332,961,730,560 MACs, weights read 1x
    512 sequential decode steps: 3,332,961,730,560 MACs, weights read 512x
    one decode step           :    6,576,668,672 MACs

  arithmetic intensity (MACs per byte moved off memory):
    prefill :    242.5 MACs/byte ✓
    decode  :      0.48 MACs/byte

A modern accelerator needs somewhere around 100–300 MACs per byte of memory traffic to keep its arithmetic units busy. Prefill clears that. Decode misses it by more than two orders of magnitude: each step drags all 12.55 GiB of weights across the memory bus to do six billion MACs, which the chip finishes while it's still waiting.

Two consequences you will be asked about. Batching helps throughput, not latency — the weights are in flight anyway, so 32 concurrent requests ride along for nearly free, and each user's token still arrives at the same time. And buying more FLOPs doesn't speed up decode; memory bandwidth and cache size do. Every trick in serving — continuous batching, speculative decoding, quantisation, GQA — is aimed at that one imbalance.

Where PagedAttention comes in

One more problem the formula hides. You don't know how long a reply will be, so the naive implementation reserves a contiguous slab for the maximum length, per request. Most of it is never used, and what's left is fragmented — you have the free bytes, just not in one piece.

PagedAttention (Kwon et al., 2023, the paper behind vLLM) does the obvious thing once you name it: treat the cache like virtual memory. Fixed-size pages, a block table per sequence, allocate on demand, and share pages between sequences with a common prefix. Nothing about attention changes; the storage layout does. That's what made high-throughput serving practical, and it is the reason "KV cache" is a phrase you now hear in ops meetings.

The code

Pure NumPy. No torch, no training — a two-layer decoder with random weights, because the point is compute and memory, not learning. Runs in about a quarter of a second.

Download the full script — 481 lines plus 7 runnable demos. Then: 

python3 kv_cache_from_scratch.py
def forward_full(model, cfg, ids, use_mask=True):
    """Run the whole prefix through every layer. Returns (n, vocab) logits."""
    n = len(ids)
    x = model["emb"][ids] + model["pos"][:n]
    mask = np.tril(np.ones((n, n), dtype=bool)) if use_mask else None
    for L in model["layers"]:
        h = layernorm(x)
        Q = split_heads(h @ L["Wq"], cfg)      # recomputed every step
        K = split_heads(h @ L["Wk"], cfg)      # recomputed every step -- waste
        V = split_heads(h @ L["Wv"], cfg)      # recomputed every step -- waste
        scores = Q @ K.transpose(0, 2, 1) / np.sqrt(cfg.head_dim)
        if mask is not None:
            scores = np.where(mask, scores, -np.inf)
        x = x + merge_heads(softmax(scores) @ V) @ L["Wo"]
        h = layernorm(x)
        x = x + np.maximum(h @ L["W1"], 0) @ L["W2"]
    return layernorm(x) @ model["head"]


def new_cache(cfg):
    """Per layer: K and V, each (n_heads, seq_so_far, head_dim). Never Q."""
    return [{"K": np.zeros((cfg.n_heads, 0, cfg.head_dim)),
             "V": np.zeros((cfg.n_heads, 0, cfg.head_dim))}
            for _ in range(cfg.n_layers)]


def decode_step_cached(model, cfg, token_id, pos, cache, append=True):
    """One token in, one row of logits out. `pos` indexes the position table."""
    if pos >= cfg.max_ctx:
        raise IndexError(f"position {pos} past context window {cfg.max_ctx}")
    x = model["emb"][[token_id]] + model["pos"][[pos]]     # (1, d_model)
    for L, c in zip(model["layers"], cache):
        h = layernorm(x)
        q = split_heads(h @ L["Wq"], cfg)      # (H, 1, hd) -- NOT cached
        k = split_heads(h @ L["Wk"], cfg)      # (H, 1, hd) -- appended
        v = split_heads(h @ L["Wv"], cfg)      # (H, 1, hd) -- appended
        if append:
            c["K"] = np.concatenate([c["K"], k], axis=1)
            c["V"] = np.concatenate([c["V"], v], axis=1)
        # No mask needed: the cache only ever holds the past. The mask has
        # become a data-structure invariant instead of a matrix of -inf.
        scores = q @ c["K"].transpose(0, 2, 1) / np.sqrt(cfg.head_dim)
        x = x + merge_heads(softmax(scores) @ c["V"]) @ L["Wo"]
        h = layernorm(x)
        x = x + np.maximum(h @ L["W1"], 0) @ L["W2"]
    return (layernorm(x) @ model["head"])[0]               # (vocab,)


def macs_per_token(cfg, n_ctx):
    d, ff = cfg.d_model, cfg.d_ff
    return cfg.n_layers * (4 * d * d + 2 * d * n_ctx + 2 * d * ff)


def macs_cached_step(cfg, n_cached):
    """One new token attending over n_cached keys. Flat in the projections."""
    return macs_per_token(cfg, n_cached)


def macs_naive_step(cfg, n_ctx):
    """Recompute all n_ctx tokens from scratch. Linear in n_ctx, per step."""
    return n_ctx * macs_per_token(cfg, n_ctx)


def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, dtype_bytes=2,
                   batch=1):
    return 2 * layers * kv_heads * head_dim * seq_len * dtype_bytes * batch

Now break it

This is the part that actually makes it stick. Reading won't do it. You debug for a living — and every row below is a bug that has shipped to production somewhere, in a real serving stack, written by someone competent.

Break thisWatch for
Pass use_mask=False to forward_full in demo 3 Earlier attention rows shift by 4.631e−02, cached logits drift 0.4017, and the argmax flips from token 44 to 16. The cache was never valid without the mask.
Run generate_cached(..., append=False) — skip c["K"] = np.concatenate(...) The model attends only to the prompt and freezes: it generates token 16 eight times, 1 distinct token instead of 2. It never sees what it just wrote.
Also cache q and reuse it next step Nothing to reuse — ||q|| is 5.414, 5.678, 5.368, 6.161, a new value every step, read exactly once. K and V get read again at every future step; Q does not. That asymmetry is the whole reason the cache holds two tensors, not three.
Generate past cfg.max_ctx with no eviction IndexError: position 128 past context window 128. In a real stack with a wrapping position table you get no error at all — just positions silently colliding and output turning to mush.
Hand sequence B's warm cache to sequence A: generate_cached(..., cache=warm) Prompt [3, 17, 42, 8] generates [44, 44, 44, 16, 16, 44, 44, 23] instead of [16, 16, 16, 16, 44, 44, 16, 16] — it is attending to the other user's tokens. This is the classic serving bug, and it leaks context between requests.
Then read the papers — but narrowly. In the transformer paper, only the masked-attention part of the model section; that mask is the entire licence for everything above. In the GQA paper, the method section and the ablation table comparing MHA, GQA and MQA. In the vLLM paper, the section on the block table and memory sharing. Skip every throughput benchmark: the hardware has moved, the idea hasn't.

On the maths anxiety

There is almost no maths here, and that's the point. You need matrix shapes, and you need to notice that softmax normalises over a row — which is the one fact that makes the −∞ columns contribute nothing. That's the entire proof of correctness for the KV cache.

Everything else on this page is counting. Multiply-accumulates per token. Bytes per token. Bytes per byte moved. It's the same arithmetic you already do when you size a Redis instance or work out whether a query fits in the page cache — and it is genuinely more useful, day to day, than being able to derive backprop.

What this page skips: the actual attention kernel. I compute the full (1, n) score row in NumPy and materialise it. Real implementations use FlashAttention-style tiling so the scores never hit memory at all, which is a separate and orthogonal optimisation — that one is about the working set, this one is about the history. I also skip quantised caches, sliding-window and streaming attention, and cross-attention caches in encoder-decoder models.

Where to go next

This page gave you one thing: the mask as a purity proof, plus the arithmetic that follows. These three go where I stopped.

Order I'd take them in: this page → the vLLM blog post (the memory problem made visual) → the scaling book's inference chapter (the arithmetic, done right) → Weng's survey as a reference you come back to. And read attention first if the mask isn't yet obvious to you, then GPT for the generation loop this page is optimising.