The KV cache, by hand.

A language model writes one word-piece at a time, and to choose each one it looks back over everything written so far. The KV cache is the model keeping notes on those earlier words, so it never has to reread them from scratch.

It is the difference between a model that slows to a crawl as an answer gets longer and one that keeps an even pace. On a laptop it turns a three-minute GPT-2 run into twelve seconds, with the exact same words out. Almost every serving trick that follows (fitting many users on one GPU, reusing the block of instructions every conversation starts with, guessing several words at once) is really a way of managing these notes. And “what does the KV cache store?” is about the most likely inference question you will get in an interview.

You built GPT-2 from scratch, and you will add the notes to that very code. Why ~17 tokens per second? gave you the two phases of generation, reading the prompt and then writing word by word. How much memory does a model need? sized the notes; here you build them. Your server from the baseline has quietly used them all along. After this come batching, prefix caching and speculative decoding, each of which bends these notes to a new purpose.

Today

• Say exactly what the model keeps between steps: which vectors, what shape, in every layer, and why it keeps the two vectors each past word is looked up by (its key and value) but throws away the one it used to do its own looking (its query).
• Explain why the rule that a word can only look backwards is what makes keeping the notes safe, and why a model like BERT, where every word sees every other word, cannot do it.
• Count the work a generation loop does with and without the notes, and say how the cost of each new word, and of the whole answer, grows with length.
• Add the cache to your own GPT-2 by hand, prove it writes exactly the same words, and plot time and memory per word against position.
• Explain the designs that make the notes smaller (letting attention heads share one set, remembering only the most recent few thousand words, storing them in fewer bits) and compute what each one saves.

You have built GPT-2 or Llama from scratch, so you know what one pass through the model does. From Why ~17 tokens per second? you want the two phases of generation and the chart of when memory or math is the bottleneck; from Baseline: my Spark vs Fireworks, the server and benchmark you will poke at the end. About one evening on your laptop and an hour on the Spark.

The Problem

Here is the generation loop almost everyone writes the first time they finish their GPT-2. It is short, it is obviously correct, and you may well have written exactly this:

@torch.no_grad()
def generate(model, idx, n_new):
    for _ in range(n_new):
        logits = model(idx)                          # run the WHOLE sequence
        nxt = logits[:, -1].argmax(-1, keepdim=True) # keep only the last position
        idx = torch.cat([idx, nxt], dim=1)
    return idx

On a MacBook (an M4 Pro, CPU only, GPT-2 small in fp32) that loop takes about 170 seconds to write 1,000 tokens. The first token takes 17 milliseconds. Token 500 takes about 160. Token 1,000 takes about 360. The model did not get bigger. The text did not get harder. Each step just got slower, in a straight line, until it was twenty times slower than the first.

Then you change about twenty lines, and the same 1,000 tokens take 12 seconds. Every token is the same token as before. Token 1,000 costs about what token 1 did.

You already know the fix by name, the KV cacheKV cacheMemory that holds the key and value vectors of every past token, in every layer, so a generation step only has to compute them for the newest token., and your description in the diagnostic was right: store the keys and values for past tokens instead of recomputing them. This lesson goes one level down. Why exactly keys and values, and not queries? What makes it legal to reuse them? What shape is the thing, how big does it get, and what are all the tricks people use to make it smaller? Then you build it into your own GPT-2 and measure it.

The Idea

A lazy librarian

Picture a librarian answering a stream of questions about a growing pile of notes. Each time a new note arrives, she writes an index card for it: a short label (what this note is about) and a summary (what it says). To answer a question, she compares it with every label, picks the relevant cards, and blends their summaries.

The no-cache loop is a librarian who, every time one new note arrives, throws away all her index cards and rewrites every one from scratch, then answers one question. The cards come out identical every time, because the notes have not changed. A sensible librarian keeps the cards in a box and writes one new card per note.

In attention, the labels are keyskeysFor each token, a vector that says “here is what I contain”, used to decide how much attention it gets. Computed as k = W_k · h from that token’s hidden state h. and the summaries are valuesvaluesFor each token, a vector holding the information it hands over when attended to. Computed as v = W_v · h. The attention output is a weighted average of values.. The question is the new token’s queryqueryFor each token, a vector that says “here is what I am looking for”. Its dot product with each key gives a score; softmax of the scores gives the attention weights.. The box of cards is the KV cache. The rest of this section is about why the cards really do never change, and why nobody bothers to keep the old questions.

What one new token needs

Take one attention head in one layer. Write the hidden state of position i as h_i. The head computes q_i = W_q h_i, k_i = W_k h_i and v_i = W_v h_i for every position, then for each position t:

score(t, i) = q_t · k_i / sqrt(d_head)        for i = 0 .. t    (the causal mask hides i > t)
a(t, i)     = softmax over i of score(t, i)
out_t       = sum over i of a(t, i) · v_i

Now ask what you need in order to produce the output for the newest position t, and nothing else.

Why the old cards never change

Keeping old keys and values is only allowed if they are truly the same numbers every step. Are they? In a decoder like GPT-2, yes, and the reason is the causal maskcausal maskThe rule in a decoder’s attention that position t may only look at positions 0 to t, never later ones. Implemented by setting the scores for later positions to minus infinity before the softmax..

Follow one position i up through the model:

1. Layer 0 input. h_i is the embedding of token i plus the position embedding for slot i. It depends only on token i and the number i.
2. Attention. The output at position i mixes values from positions 0 to i. Nothing later. So it depends only on tokens 0 to i.
3. Everything else. Layer norm, the MLP and the residual add all work on one position at a time. They never mix positions.
4. Next layer. So the next layer’s h_i also depends only on tokens 0 to i. Repeat for every layer.

By induction, every hidden state at position i, in every layer, is a function of tokens 0 to i and nothing else. Appending token 501 cannot reach back and change it. So k_i = W_k h_i and v_i = W_v h_i in every layer are frozen the moment they are first computed. They are immutableimmutableNever changes after it is created. Here: once position i’s keys and values exist, no later token can alter them..

Two conditions are hiding in that argument, and each one breaks a cache if you violate it. The mask has to be causal (an encoder like BERT lets every position see every other, so adding a word changes every hidden state, and there is nothing to cache). And a position’s encoding has to depend only on its own index. GPT-2’s learned position table and Llama’s RoPERoPERotary position embedding: positions are encoded by rotating each query and key vector by an angle that depends on its position. Llama applies it to q and k before attention, so the cached keys are already rotated. both qualify. Note that with RoPE you cache the keys after rotation, which is what Hugging Face’s Llama does.

One more thing falls out of step 2. The cache is per layer, and it has to be. Layer 5’s keys are computed from layer 5’s input, which is layer 4’s output. There is no single “key of token i”; there are 12 of them in GPT-2 small (one per layer), each in 12 heads.

Prefill fills it, decode appends to it

With a cache, generation has the same two phases you met in the bandwidth lesson, and now you can say precisely what each does to the cache.

Prefill. Run the whole prompt through the model in one forward pass, exactly like training (causal mask and all). As a side effect, every layer computes K and V for every prompt position. Keep them. The cache now holds P rows per layer. The output at the last position gives the first new token.

Decode. Feed only the newest token, as a sequence of length 1, telling the model its position is P (then P+1, and so on). In each layer, compute its q, k, v. Append k and v to that layer’s cache. Score the one query against every cached key, softmax, mix every cached value. One row out, through the MLP, on to the next layer. At the top, one row of logits, one token.

The shape of each layer’s cache is the shape you already use inside attention:

K cache, one layer:  (batch, n_kv_heads, seq_len, head_dim)
V cache, one layer:  (batch, n_kv_heads, seq_len, head_dim)

GPT-2 small:         (1, 12, T, 64)   for each of 12 layers, K and V
Llama 3.1 8B:        (1,  8, T, 128)  for each of 32 layers, K and V

Prefill writes T = P rows in one go. Each decode step grows T by exactly one, along dimension 2. Nothing else ever changes size.

Two small details in that decode step are where everyone’s first implementation goes wrong, so name them now.

• The position. The new token is at position P + (steps so far), not position 0. Your forward pass usually builds positions with torch.arange(T). With T = 1 that gives 0, every step. You need arange(past_len, past_len + T).
• The mask. A single new query is allowed to see every cached key, so it needs no mask at all. But if you ever feed more than one new token on top of a cache (a second chunk of prompt, or checking several draft tokens in speculative decoding), the mask is no longer square. Query j of the new chunk sits at absolute position past_len + j and may see keys 0 to past_len + j. In PyTorch that is torch.ones(T, S).tril(diagonal=S - T), where S is the total key length. Plain tril() on a T×S matrix is wrong by exactly past_len columns.

Step through it

Here is one head of one layer, generating a 12-token sequence from a 4-token prompt. Step through it with the cache on. Then reset, switch the cache off, and step again. Watch the counter of wasted vectors, and watch the attention triangle on the right.

Three things to notice. First, prefill is the same picture either way: a cache does not speed up the prompt, it speeds up everything after it. Second, with the cache off, the amount of waste per step grows with the length, so the total grows with the square of the length. Third, look at the Q row with the cache on: past queries go blank after their step. Nobody keeps them, because nobody needs them.

Counting the work

Now put numbers on it. Let n be the current sequence length and N_p the number of parameters in the model’s matrices (about 124 million for GPT-2 small, counting the output projection). A matrix-vector multiply costs about 2 FLOPs per weight per token, so:

                    linear layers per step     attention per step (all layers)
no cache, step n    2 · N_p · n                ~ 4 · L · d · n²   (whole triangle)
with cache, step n  2 · N_p                    ~ 4 · L · d · n    (one row)

L = layers, d = model width (768 for GPT-2 small)

So without a cache, step n costs about n times a cached step in the linear layers, and attention is worse still. Add it up over a generation of n tokens: the no-cache loop does about n²/2 token-passes through the weights where the cached loop does n. Per step, cost goes from growing like n (and n² inside attention) to roughly constant plus a small term that grows like n. Total cost goes from about n² to about n.

That small growing term with the cache is real, and worth seeing. At 1,000 tokens in GPT-2 small, the cached step’s attention is about 4 × 12 × 768 × 1,000 ≈ 37 million FLOPs, against about 250 million for the matrices. Small, but it grows, and at long context it wins. In bytes it is the same story you saw at the end of the bandwidth lesson: every decode step has to read the entire cache, every layer, every head. For Llama 3.1 8B at 128K context that read is bigger than the weights.

Why the measured speedup is not 500x

The FLOP count says a no-cache loop at 1,000 tokens does around 500 times too much matrix work. The laptop measured 14x in total time. Where did the rest go?

Back to the roofline. A step with no cache at length n is a prefill of n tokens: one read of the weights, n tokens of math on it. While n is small, that is still far below the machine’s ridge pointridge pointPeak FLOP/s divided by memory bandwidth. Work with fewer FLOPs per byte than this is limited by memory; more, by math. From the bandwidth lesson., so the step is bound by reading the weights (and by fixed per-step overhead), not by math. Doing 10 tokens of math instead of 1 is nearly free, because the math units were idle anyway. Only when n gets big enough to make the math the bottleneck does the no-cache step start growing.

Here is what that looks like on real hardware, GPT-2 small in fp32, greedy, 1,000 new tokens after a 12-token prompt:

Notice the cached GPU run is no faster than the cached CPU run. At 124 million parameters, a decode step on this GPU is dominated by fixed costs (launching dozens of small kernels from Python, about 10 ms worth) rather than by bytes or math. That is the “overhead-bound” corner of the roofline, and it is why the cache’s speedup looks modest on a tiny model and short text and becomes enormous on a big model with long text.

The price: memory

The cache trades compute for memory. Its size has a formula you can do in your head, and you saw a version of it in the last lessons:

KV bytes per token = 2 (K and V) × layers × kv_heads × head_dim × bytes per number
KV bytes total     = that × tokens in context × sequences in the batch

The GPT-2 small row is not just arithmetic: the cache in the run above held exactly 73,728 bytes per position (1011 positions, 74.5 MB).

It grows linearly: one more token, one more row in every layer. Weights are a fixed cost; the cache is a per-token, per-user cost. That makes it the thing that decides how many users and how much context fit on a machine, which is the whole subject of memory math. And because each step reads the whole cache, it is also bytes per token on the bandwidth side of the ledger. So every trick for making the cache smaller buys you two things at once: more users or context in memory, and faster steps.

Shrinking the cache: share the keys

The formula has four knobs: layers, KV heads, head size, and bytes per number. Layers and head size are mostly about model quality. The other two are where the clever designs live.

In the original Transformer, multi-head attentionmulti-head attentionAttention split into several heads, each with its own query, key and value projections, so each head can look for different things. “MHA” when every head has its own keys and values. gives every head its own queries, keys and values. With 32 heads, that is 32 sets of keys and values per layer to cache.

Shazeer’s 2019 idea, MQAMQAMulti-query attention: every query head shares one single key head and one value head. Cuts the cache by the number of heads. Some loss in quality., keeps 32 query heads (so the model can still ask 32 different questions) but gives them a single shared key head and value head. The cache shrinks 32-fold. The catch is quality: one set of keys has to serve every kind of question, and models trained this way tend to lose a little accuracy. Several big models used it anyway (Google’s PaLM and TII’s Falcon 7B, for example).

GQAGQAGrouped-query attention: query heads are split into groups, and each group shares one key head and one value head. Llama 3.1 8B has 32 query heads in 8 groups of 4., from Joshua Ainslie and colleagues at Google in 2023, is the middle ground: split the query heads into groups and give each group its own key and value head. With 8 groups, the cache is 4x smaller than MHA, and the paper reports quality close to full MHA with speed close to MQA. Llama 2 used it only for the 70B model; every Llama 3 model and most new open models use it. In code it is a small change: the K and V projections output n_kv_heads × head_dim instead of n_heads × head_dim, and before attention each KV head is repeated (or broadcast) to line up with its group of query heads. The cache stores the small, unrepeated version. That is why the Llama shape above has 8 in the heads dimension, not 32.

Things to try: start at GQA-8, 8K context, 1 user. The cache is a sliver next to the weights. Switch to MHA: Llama 2 7B’s design, at 512 KiB per token. Now drag users to 16. With MHA the caches outweigh the model; with GQA they fit comfortably; with MQA they barely register. Then push context to 128K with 8 users and watch the memory bar run out.

DeepSeek went further in 2024 with MLAMLAMulti-head latent attention, from DeepSeek-V2: instead of caching full keys and values, cache one small compressed vector per token per layer and expand it back into per-head keys and values when needed.: cache one small compressed vector per token per layer and reconstruct keys and values from it with a learned matrix. Same goal, different bargain: a bit more math per step for many fewer bytes, a good trade on a bandwidth-bound machine.

Shrinking the cache: forget the distant past

Another knob is “tokens in context”. What if a layer only attends to the last W tokens? Then it only needs to cache the last W tokens, and the cache stops growing at W. That is sliding-window attentionsliding-window attentionEach position attends only to the previous W positions instead of all of them. The cache for that layer never holds more than W tokens.. The idea was popularized for long documents by Longformer (2020); Mistral 7B (2023) used it in every layer with W = 4,096.

The cache becomes a ring buffer, which Mistral’s paper calls a rolling buffer: token i goes into slot i mod W, overwriting the token from W steps ago. Memory per sequence is capped no matter how long the conversation runs. (For Mistral 7B, with 8 KV heads of 128 in 32 layers, that cap is 128 KiB × 4,096 = 512 MiB per sequence in bf16.)

Is information older than W tokens simply lost? Not entirely. Layer 2’s view of position t includes positions back to t − W, but each of those already absorbed positions back another W in layer 1. Stack L layers and information can in principle travel back about L × W positions, in the way a rumor travels down a line of people. In practice it is lossy, so newer models mix the two: Google’s Gemma 2 alternates sliding-window and full-attention layers, and per its technical report Gemma 3 uses five local layers (window 1,024) for every global one. Only the global layers pay for the full context.

Shrinking the cache: fewer bits

The last knob is bytes per number. Store keys and values in 8 bits instead of 16 and the cache halves, in memory and in bytes read per step. FP8 KV caches are a standard option in serving engines (vLLM exposes it as a cache dtype setting; check its docs for the current name and supported formats). Going lower is an active research area. The KIVI paper (Liu et al., 2024) found that keys have a few channels with huge values, so it quantizes keys per channel and values per token, and gets down to about 2 bits with small accuracy loss. The cost is a little error in every attention score and a little extra work to unpack. The quantization lesson does the same trade for weights, with measurements.

How real code stores it

Your build will grow the cache with torch.cat, which is the simplest thing that works and also slightly wasteful: each cat allocates a new tensor and copies the whole old cache into it. That is a second full pass over the cache’s bytes, every step, every layer. Real implementations avoid it in one of three ways:

• Preallocate. Allocate (batch, heads, max_len, head_dim) once and write row t into slot t. Hugging Face calls this a static cache; it also lets the model be compiled, since shapes never change.
• Grow in chunks. Reserve room for many tokens at a time, like a Python list does.
• Pages. vLLM’s PagedAttentionPagedAttentionStore the KV cache in fixed-size blocks (say 16 tokens each) scattered in memory, with a per-sequence table of which blocks it owns, like virtual memory pages. No wasted reservation, and blocks can be shared. splits every sequence’s cache into small fixed blocks, borrowed from the operating system’s virtual memory idea. That is the next lesson but one: paged attention and prefix caching.

Build It

You will work on your laptop first (any machine with PyTorch; CPU is fine for GPT-2 small), then on the Spark.

kvcache/
  gpt2.py          # your model (or the boilerplate below)
  generate.py      # no-cache and cached generation, the equality check
  bench_pos.py     # latency vs position, memory vs position
  plot.py
  results/

python3 -m venv .venv && . .venv/bin/activate
pip install torch transformers matplotlib

# gpt2.py: a minimal GPT-2 that loads Hugging Face's weights. No cache yet.
import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class Config:
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    block_size: int = 1024
    vocab_size: int = 50257


class Attention(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.n_head, self.head_dim = cfg.n_head, cfg.n_embd // cfg.n_head
        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd)  # x -> q, k, v in one matmul
        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.c_attn(x).split(C, dim=2)
        # (B, T, C) -> (B, n_head, T, head_dim)
        q, k, v = (t.view(B, T, self.n_head, self.head_dim).transpose(1, 2) for t in (q, k, v))
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)          # (B, H, T, T)
        mask = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
        att = att.masked_fill(~mask, float("-inf")).softmax(dim=-1)
        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)


class MLP(nn.Module):
    def __init__(self, cfg: Config):
        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)

    def forward(self, x):
        return self.c_proj(F.gelu(self.c_fc(x), approximate="tanh"))


class Block(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.ln_1, self.attn = nn.LayerNorm(cfg.n_embd), Attention(cfg)
        self.ln_2, self.mlp = nn.LayerNorm(cfg.n_embd), MLP(cfg)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        return x + self.mlp(self.ln_2(x))


class GPT(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.wpe = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.h = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer))
        self.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.wte.weight  # GPT-2 ties input and output embeddings

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.wte(idx) + self.wpe(pos)
        for block in self.h:
            x = block(x)
        return self.lm_head(self.ln_f(x))                                 # (B, T, vocab)

    @classmethod
    def from_pretrained(cls, name: str = "gpt2"):
        from transformers import GPT2LMHeadModel
        sizes = {"gpt2": (12, 12, 768), "gpt2-medium": (24, 16, 1024),
                 "gpt2-large": (36, 20, 1280), "gpt2-xl": (48, 25, 1600)}
        L, H, C = sizes[name]
        model = cls(Config(n_layer=L, n_head=H, n_embd=C))
        hf = GPT2LMHeadModel.from_pretrained(name).state_dict()
        sd = model.state_dict()
        for k in sd:
            src = hf["transformer." + k if not k.startswith("lm_head") else k]
            # HF stores these four as Conv1D, i.e. transposed relative to nn.Linear
            if k.endswith(("c_attn.weight", "c_proj.weight", "c_fc.weight")):
                src = src.t()
            assert src.shape == sd[k].shape, k
            sd[k].copy_(src)
        return model.eval()

class Attention(nn.Module):
    def forward(self, x, past=None):
        """x: (B, T, C), the NEW positions only.
        past: None, or (k_past, v_past), each (B, n_head, P, head_dim).
        Returns (y, (k_all, v_all)) where k_all, v_all are (B, n_head, P + T, head_dim)."""


class Block(nn.Module):
    def forward(self, x, past=None):
        """Returns (x_out, present) where present is this layer's (k_all, v_all)."""


class GPT(nn.Module):
    def forward(self, idx, past=None):
        """idx: (B, T) new token ids. past: None, or a list of n_layer (k, v) pairs.
        Returns (logits, presents): logits (B, T, vocab), presents a list of n_layer (k, v)."""

@torch.no_grad()
def generate_nocache(model, idx, n_new) -> torch.Tensor: ...
    # the loop from The Problem: full sequence every step

@torch.no_grad()
def generate_cache(model, idx, n_new) -> tuple[torch.Tensor, list]: ...
    # prefill: logits, past = model(idx)
    # then n_new steps: pick the token from logits[:, -1], feed ONLY it with past

a = generate_nocache(model, prompt, 500)
b, past = generate_cache(model, prompt, 500)
assert torch.equal(a, b), first_mismatch(a, b)

def timed_steps(step_fn, n_new, device) -> list[float]:
    """Run step_fn() n_new times; return seconds per step."""
    # synchronize before reading the clock if on a GPU:
    #   torch.cuda.synchronize() on the Spark, torch.mps.synchronize() on a Mac
    ...

def cache_bytes(past) -> int:
    """Sum of nbytes over every K and V tensor in the cache."""
    ...

out = model(**inputs, use_cache=True)
pkv = out.past_key_values
print(type(pkv).__name__)          # a Cache object, e.g. DynamicCache, in recent versions
# transformers 5.x: pkv.layers[0].keys.shape. Older versions: pkv.key_cache[0] or pkv[0][0].

Measure

Fill these in. The Mac column is what this lesson measured, so you have something to compare against; the rest is yours.

Correctness

Latency per token (ms), GPT-2 small

Memory

Llama 3.1 8B on the Spark, Hugging Face generate

What to expect, and why:

• The no-cache curve on the Spark should be a hockey stick, not a straight line. On the Mac CPU, whose ridge point is low, recomputing more tokens costs more almost immediately. The Mac GPU stayed flat for the first hundred-odd tokens. The Spark has far more compute per byte than either, so with GPT-2 small expect a long flat stretch before the bend, maybe hundreds of tokens; with an 8B model, expect the bend near your measured ridge point (see the prediction above).
• The cached curve is “flat-ish”, not flat. It creeps up because each step reads a longer cache and computes a longer score row. On the Mac CPU it went from about 8 ms to 16 ms over 1,000 tokens. With torch.cat part of that creep is copying; t7 tells you how much.
• For Llama at short lengths, use_cache=False may barely hurt. At 64 new tokens, each no-cache step is a tiny prefill and the step is still dominated by reading 16 GB of weights. The ratio should grow with length, and by 1,024 tokens it should be large. If you want a number to beat: if the bend is near 200 tokens, the no-cache run of 1,024 tokens should take a couple of times longer than the cached one, and the gap keeps widening the longer you go. Work it out from your measured FLOP/s and bandwidth before you run it.
• Memory should be a perfectly straight line from cache_bytes, and a slightly noisier one from the allocator. The slope is the formula. If the allocator line jumps in steps, that is PyTorch’s caching allocator reserving memory in blocks, from the stack and heap story.

Break It

Forget the position offset. Use arange(T) instead of arange(P, P + T) in the cached path. Every new token now thinks it is at position 0. The output stays grammatical for a few words, then wanders, and t3 fails within the first couple of steps. This is the most common KV-cache bug in the wild, and it hides well, because the text still looks like English.

Remove the causal mask. Delete the mask in both paths and run t3 again. The two generations diverge almost immediately. Without the mask, position 3’s hidden state depends on tokens that arrive later, so the keys you cached at step 3 are stale by step 4. This is exactly why encoder models like BERT have no KV cache. Then put the mask back and break it subtly instead: use plain tril() for a T×S mask, and feed the prompt in two chunks (prefill the first half, then the second half on top of the cache). Single-token decode still works, because one query needs no mask at all; the two-chunk prefill silently produces different logits. That is the bug that bites later, in chunked prefill and in speculative decoding.

Cache too much. Keep the query vectors too, and the hidden states. Everything still works. Measure memory: you have more than doubled the cache to store things nobody reads. A good way to feel why the cache holds K and V and nothing else.

Run it in bf16. Cast the model to bf16 on the Spark (model.to("cuda", torch.bfloat16)) and rerun t3 for 1,000 tokens. There is a good chance the two loops eventually disagree on a token, and after the first disagreement they write different text. Nothing is wrong with your cache: bf16 keeps only about 3 significant digits, the cached and uncached paths add numbers in different orders, and somewhere a near-tie between two tokens tips the other way. Print the gap between the top two logits at the first mismatch to prove it to yourself. The lesson: “identical” is a statement about exact math; in low precision the right test is “same tokens until a near-tie, and tiny logit differences everywhere”.

Explain It Back

Answer each out loud, as if in an interview, before opening the answer.

Go Deeper

• “Attention Is All You Need”, Ashish Vaswani et al., 2017. The decoder, the causal mask and multi-head attention, in the original. Re-read the decoder section with the cache in mind.
• “Fast Transformer Decoding: One Write-Head is All You Need”, Noam Shazeer, 2019. Short, and the first clear statement that cached decoding is bound by reading keys and values from memory. Introduces MQA.
• “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, Joshua Ainslie et al., 2023. The middle ground every Llama 3 model uses, and how to convert an existing MHA model to it.
• “Efficiently Scaling Transformer Inference”, Reiner Pope et al. (Google), 2022. The KV cache as the memory cost that limits batch size and context, worked through with real numbers, including why MQA matters at scale.
• “Mistral 7B”, Albert Jiang et al., 2023. Sliding-window attention and the rolling buffer cache, in about two pages.
• “Efficient Memory Management for Large Language Model Serving with PagedAttention”, Woosuk Kwon et al., SOSP 2023. The vLLM paper: what goes wrong when you store caches as contiguous tensors at scale, and the paging fix. Sets up the prefix caching lesson.
• “KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache”, Zirui Liu et al., 2024. Why keys and values want different quantization, with the plots that show the outlier channels.
• Andrej Karpathy, nanoGPT (github.com/karpathy/nanoGPT) and the Hugging Face transformers documentation on caching (the “KV cache strategies” page). The first is the model you just extended; the second shows the dynamic, static and quantized caches that real code uses, and the current API.

Connections

Built on: Why ~17 tokens per second? (prefill vs decode, the roofline, and the cache as bytes read per step) and Baseline: my Spark vs Fireworks (the server and scoreboard you just ran an ablation on).

Looking back: How much memory does a model need? turns the cache formula into a capacity plan. Next: Batching puts many users’ caches into one step, which is where the per-user cache becomes the limit. Paged attention and prefix caching stores caches in pages and shares them between requests with the same prompt. Speculative decoding feeds several tokens on top of a cache at once, which is where the T×S mask from this lesson earns its keep. And FlashAttention speeds up the attention that reads all of this.