S05Memory & KV Cache·205 lines

The KV Cache

Without a cache, generating token N costs O(N²). With one it costs O(N), and the bottleneck moves from compute to memory bandwidth.

cache layout
prefill/decode split
arithmetic intensity
bandwidth wall

Motivation

The problem

Run the S01 engine on a real prompt and profile it. Almost all the time goes into recomputing keys and values for tokens that have not changed since the last iteration. The model is deterministic and causal, so nothing about the past can change: token 500's key vector on step 501 is the vector it had on step 500.

The naive loop recomputes it anyway: 500 times, once per layer, once per request. In an unoptimised engine this is the largest single source of waste, and also the easiest to remove.

Core idea

The solution

Keep the keys and values. Attention at position t needs the K and V of every position ≤ t, so store them as they are produced and append one column per step.

cached attentionpython

def attention_with_cache(x, w, cache, layer, pos):
    q = x @ w.wq                        # [1, n_heads * head_dim]  (decode: 1 token)
    k = x @ w.wk                        # [1, n_kv_heads * head_dim]
    v = x @ w.wv
    q, k = rope(q, pos), rope(k, pos)   # rotate BEFORE storing

    cache.k[layer, pos] = k             # append — this is the whole trick
    cache.v[layer, pos] = v

    # attend over everything written so far, including what we just wrote
    k_all = cache.k[layer, : pos + 1]   # [pos+1, n_kv_heads, head_dim]
    v_all = cache.v[layer, : pos + 1]
    return grouped_query_attention(q, k_all, v_all)

Three consequences follow immediately, and they shape every remaining chapter:

The forward pass changes shape. Prefill runs with a sequence dimension of N; decode runs with a sequence dimension of 1. They are different enough that engines compile separate kernels for them.
The engine becomes stateful. An S01-style engine is a pure function. A cached engine owns per-request memory that must be allocated, tracked and freed. S06 is about doing that well.
The bottleneck moves. You have removed almost all redundant arithmetic. What remains is reading weights and cache from memory, and there is no arithmetic left to hide it behind.

Cache K after RoPE, not before

The rotation depends only on absolute position, and a cached token's position never changes. Store the rotated K and you never touch RoPE for the past again. Store the unrotated K and every step has to re-rotate the whole cache. That swaps one compute cost for a slightly smaller compute cost and buys almost nothing.

Mechanics

How it works

Diagramthe same computation, with and without a cache

On the left, ten cells are computed and nine of them duplicate the previous step. On the right only the last query row exists: a decode step attends once, over a cache built by the steps before it.

The bottleneck moves from FLOPs to bandwidth

This is the most important idea in the course, so it is worth stating carefully. During decode, for each token, the GPU must read:

every weight in the model: 16 GB for an 8B model in fp16;
every byte of the KV cache for every request in the batch. At long context that can exceed the weights.

With that data it performs roughly two FLOPs per weight. An H100 offers about 3.3 TB/s of memory bandwidth and about 1,000 TFLOP/s of fp16 compute, a ratio of roughly 300 FLOPs per byte. A decode step delivers about 1: two FLOPs against the two bytes each fp16 weight costs to read. The GPU is idle more than 99% of the time.

8B fp16, batch 1: ~5 ms; 16 GB ÷ 3.3 TB/s, a hard floor no kernel can beat
resulting ceiling: ~200 tok/s; reading fewer bytes is the only way past it
the escape hatch: ×batch; one weight read serves N requests, the subject of S09

Every optimisation after this chapter attacks that formula from one side or the other. Read fewer bytes: quantization (S08), KV compression, GQA. Amortise the read across more tokens: batching (S09), speculative decoding (S14), MoE (S16).

Explore

Try it

The top half shows the mechanism; the bottom half shows the bill. Toggle the cache and step through to see the work profile flip from quadratic to linear, then use the memory calculator to find how many concurrent requests your GPU can hold.

SimulatorKV cache: mechanism and memoryrun the top; then push concurrency until the bar turns red

part 1 · what the cache does

each step computes k,v for 1 new token and reads the rest

step 0 / 16

cache contents (one layer)

012345

cached — read only
computed this step

Exactly one new column per step. The k and v for every earlier position were computed once and are only ever read after that.

attention work per step

press run

with cache: 0
would have cost: 0

With the cache, per-step cost grows linearly (attend to n positions). Without it, cost grows quadratically (recompute all n, each attending to n).

part 2 · what the cache costs

model

context per request4kconcurrent requests32GPU memory80 GB

kv dtype

bytes / token: 131.1kB
cache total: 16.0 GB
weights: 14.9 GB
max concurrent: 126

GPU memory budget · 80 GB

weights (fixed)
activations + overhead
KV cache (grows with load)

The KV cache is 52% of the memory in use. Drag concurrency up until it turns red, then switch the KV dtype to int8 and watch the ceiling double. S08 is about that trade.

Select GPT-3 175B and note the KV cost: 96 attention heads with no grouping produce a cache larger than most models' weights. Now select Llama-3-70B, with 64 query heads and 8 KV heads. Same scale of model, 14× less cache. That is the argument for GQA, and it is an argument about memory rather than quality.

Build it

Implementation

The naive allocation strategy is one contiguous tensor per request, sized to the maximum possible length. Write that version first anyway: its failure mode is worth seeing for yourself.

code/s05_kv_cache.py (excerpt)python

class KVCache:
    """One contiguous buffer per request. Simple, and disastrously wasteful."""

    def __init__(self, layers, max_seq, n_kv_heads, head_dim, dtype=np.float16):
        shape = (layers, max_seq, n_kv_heads, head_dim)
        self.k = np.zeros(shape, dtype=dtype)
        self.v = np.zeros(shape, dtype=dtype)
        self.length = 0

    @property
    def wasted(self):
        """Every slot past self.length is reserved and unusable by anyone else."""
        return 1.0 - self.length / self.k.shape[1]

    def append(self, layer, k, v):
        self.k[layer, self.length] = k
        self.v[layer, self.length] = v

    def view(self, layer):
        return self.k[layer, : self.length], self.v[layer, : self.length]

max_seq is a promise you cannot keep

You cannot know how long a generation will run, so you allocate for the worst case. A request that stops after 40 tokens with max_seq=8192 has wasted 99.5% of its reservation, and that memory stays unavailable to anyone else for the life of the request. Measured across a real workload, 60–80% of KV memory goes this way. That measurement is what produced PagedAttention.

Run it locally

Runs the same generation twice, once with the cache disabled and once with it enabled, and asserts the outputs are token-for-token identical. Then prints a work profile, a speedup table, and a KV memory calculator for several real model configurations.

$ python code/s05_kv_cache.py

Expect: Identical output from both paths (the cache is an optimisation, not an approximation), a large reduction in attention work, and a memory table where the cache overtakes the weights past ~120k tokens, about 15 concurrent requests at 8k context.

Only NumPy is required — setup instructions.

Production notes

In production

picoLM stores the cache in fp16 to halve it, with software fp32↔fp16 conversion because the target CPUs lack hardware support. On a Raspberry Pi that decides whether the model runs at all.
quant.cpp goes further and quantizes the cache to 4 bits, keeping a full-precision window over the most recent 128 tokens. Attention concentrates on recent context, so that is where the precision belongs.
vLLM does not use contiguous caches at all. The next chapter is about what it does instead.

Exercises

1
Measure the waste yourself: run 200 requests with lengths drawn from a realistic distribution and report what fraction of reserved KV memory was ever written. Expect 60–80% waste.
2
Implement a sliding-window cache that keeps only the last W tokens. Note what breaks, and why models must be trained for windowed attention before this is sound.
3
Add fp16 and int8 cache dtypes behind a flag and measure perplexity change on a fixed passage. This is your first data point for S08.

Continue

What's next

Decode is fast now, memory is the bottleneck, and most of that memory sits in reservations nobody used. S06 fixes the reservations with the idea that made vLLM famous: treat the KV cache like virtual memory and page it.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

Adding a KV cache changes the total attention work for generating T tokens from an N-token prompt from O((N+T)³) to what?

Score 0/6