The KV cache and the systems around it

In the previous blog, we built GPT-2 from scratch and made it generate text. At the very end, I showed that storing the keys and values of the earlier tokens made generating 60 tokens almost 4× faster. In this one, I will build that cache properly, check that it gives exactly the same answers, see how fast it actually is, and then look at what happens when you have to keep it for thousands of users at once. That last part is where most of the interesting work in inference serving is.

Everything here runs on my laptop: the NumPy GPT-2 from the last blog on the CPU, and a small PyTorch version on its GPU (an RTX 3060 with 6 GB of memory). All the code is in the blog itself, so you can follow along without anything else.

The waste in our generation loop

Here’s the generation loop from the last blog, in short:

for _ in range(n_new):
    logits = gpt2(ids)            # run the WHOLE sequence through the model
    ids.append(int(np.argmax(logits[-1])))

Every step runs every token through all 12 layers again, just to get logits for the last position. When the text is 500 tokens long, we recompute 499 tokens that we had already computed in the previous step, and the step before that, and so on.

How bad is that? Here is the time for one decode step, depending on how many tokens are already there, without and with a cache:

tokens already there     one step without cache     one step with cache
        16                       37 ms                    20 ms
       128                      167 ms                    23 ms
       256                      435 ms                    24 ms
       512                      795 ms                    22 ms
     1,000                    2,400 ms                    24 ms

(Laptop CPU, NumPy with 8 threads. Laptops are noisy, so every number is the fastest of 7 runs, and I ran the whole thing twice to check that the numbers agree.)

Without the cache, each step gets slower as the text grows, and at 1,000 tokens one step takes 2.4 seconds. With the cache, every step takes about the same time, no matter how long the text is. That’s 100× faster at 1,000 tokens. The 4× from the last blog was only for 60 tokens. The longer the text, the bigger the difference.

Notice that the time without the cache grows even faster than the number of tokens (2× the tokens from 512 to 1,000 gives 3× the time). That’s attention: every token is compared with every token before it, so that part grows with the square of the length.

Why the keys and values never change

The whole idea rests on one claim: once a token has been processed, its keys and values in every layer stay the same forever, however many tokens come after it. Let’s convince ourselves, layer by layer:

The input to layer 0. Token i’s vector is wte[token] + wpe[i]. That depends only on the token itself and its position.
Attention. Because of the causal mask, token i only looks at tokens 0 to i. Tokens that come later are invisible to it.
Layer norm and the MLP work on each token on its own, so they only depend on token i’s own vector.
So after layer 0, token i’s vector depends only on tokens 0 to i. The same reasoning works for layer 1, then layer 2, and all the way up to layer 11.

So everything computed for token i, including its key and value in every layer, is fixed the moment it’s computed. Recomputing it at the next step gives exactly the same numbers, which is pure waste.

Then why not cache the queries too?

A query is only used once: to compute its own token’s output, by comparing it with the keys before it. Once token i’s output is computed, its query is never needed again.

When a new token comes in, the only thing to compute is the new token’s output. For that, we need the new token’s query, and the keys and values of every token so far:

new token's query  ──compared with──►  K of tokens 0, 1, 2, ..., n
                                       V of tokens 0, 1, 2, ..., n  ──► new token's output

The old queries play no part in it. That’s why it’s called a KV cache and not a QKV cache.

This also shows why the cache only works for models with a causal mask, like GPT-2 and every chat model. In a model like BERT, every token looks at every other token, including the later ones. So adding a token changes everything before it, and nothing can be reused.

Adding the cache to GPT-2

We need somewhere to keep the keys and values. For every layer and every head, that’s one row of 64 numbers per token. GPT-2 can handle at most 1,024 tokens, so we can just reserve space for all of them up front:

# (layers, heads, max tokens, head size)
k_cache = np.zeros((12, 12, 1024, 64), dtype=np.float32)
v_cache = np.zeros((12, 12, 1024, 64), dtype=np.float32)

Now attention gets one more argument, start: how many tokens are already in the cache. The new tokens are written right after them:

def attention(h, layer, start):
    T = h.shape[0]                                     # number of NEW tokens
    qkv = h @ W_attn + b_attn
    q, k, v = map(split_heads, np.split(qkv, 3, axis=-1))   # (12, T, 64) each

    # write the new keys and values after the cached ones
    k_cache[layer, :, start:start + T] = k
    v_cache[layer, :, start:start + T] = v
    K = k_cache[layer, :, :start + T]                  # (12, start + T, 64): all keys so far
    V = v_cache[layer, :, :start + T]

    scores = q @ K.transpose(0, 2, 1) / 64 ** 0.5      # (12, T, start + T)
    mask = np.arange(start + T)[None, :] <= (start + np.arange(T))[:, None]
    scores = np.where(mask, scores, -np.inf)

    out = softmax(scores) @ V                          # (12, T, 64)
    out = out.transpose(1, 0, 2).reshape(T, 768)
    return out @ W_proj + b_proj

Compared with the last blog, only three things changed:

The keys and values come from the cache, not just from the current input. The new ones are written first, so K and V hold everything so far.
The scores are no longer square. There are T new queries but start + T keys, so the scores are (12, T, start + T).
The mask has to know where the new tokens are. New token number r is really at position start + r, so it can look at keys 0 to start + r.

Let’s look at the shapes when we generate the 7th token. The 6 prompt tokens are already in the cache, so start = 6 and T = 1:

q        (12, 1, 64)      one query per head, only for the new token
K, V     (12, 7, 64)      6 cached + 1 new
scores   (12, 1, 7)       the new token against all 7 keys
mask     [[1 1 1 1 1 1 1]]

The mask is all ones. The newest token is allowed to look at everything, so during decode the mask actually does nothing. It matters during the prefill, where it’s the same triangle as last time. It also works for the case in between, where some tokens are cached and a few new ones come in together. For example, with start = 4 and T = 2:

            keys: 0  1  2  3  4  5
new token 0 (position 4)  [[1  1  1  1  1  0]
new token 1 (position 5)   [1  1  1  1  1  1]]

This case comes back later in the blog, when parts of a prompt are already cached.

The model itself barely changes. The positions now start at start, and we only need the logits of the last position:

def gpt2(ids, start):
    x = wte[ids] + wpe[start:start + len(ids)]
    for layer in range(12):
        x = x + attention(ln(x, f"h.{layer}.ln_1"), layer, start)
        x = x + mlp(ln(x, f"h.{layer}.ln_2"), layer)
    return ln(x[-1], "ln_f") @ wte.T                   # (50257,): last position only

The wpe[start:start + len(ids)] part is easy to get wrong. If you keep np.arange(len(ids)) from the last blog, every decoded token gets the positional embedding of position 0. The model still produces text, but it’s quietly wrong, which is the worst kind of bug.

And generation becomes a prefill, followed by decode steps that each send in just one token:

def generate(prompt_ids, n_new):
    ids = list(prompt_ids)
    logits = gpt2(ids, start=0)                        # prefill: the whole prompt at once
    for _ in range(n_new):
        nxt = int(np.argmax(logits))
        ids.append(nxt)
        logits = gpt2([nxt], start=len(ids) - 1)       # decode: only the new token
    return ids

That’s it! Around 15 lines changed, and every decode step now does the work of one token instead of the whole sequence.

Is it actually correct?

A cache is an optimization, so it must not change the answers. With greedy sampling, the model with the cache has to pick exactly the same tokens as the model without it.

So I generated 100 tokens with greedy sampling and, at every step, also ran the full sequence through the model without the cache and compared the two:

100 greedy tokens: identical
largest logit difference over all 100 steps: 0.0027

The tokens are the same at every one of the 100 steps. The logits are not bit-for-bit equal, and that’s expected: multiplying one row by a matrix and multiplying 100 rows by it take different paths inside NumPy, which add the numbers up in a different order. As in the last blog, that only changes the last few digits.

And here’s the output, if you’re curious:

Michael Jordan played the sport of basketball for the University of Michigan, where he was named
the 2011-12 Player of the Year. He was also named the 2011-12 Defensive Player of the Year.

Jordan was named the 2011-12 Defensive Player of the Year. Jordan was named the 2011-12 Defensive
Player of the Year. Jordan was named the 2011-12 Defensive Player of the Year. ...

Greedy decoding with a small model loves to get stuck in a loop like this lmao. That’s one reason why we use temperature and top-p, which I covered in the last blog.

Preallocating vs concatenating

You might have written the cache differently: start with an empty array, and add the new key and value with np.concatenate at every step. That works too, and it gives the same answers. But np.concatenate creates a new array and copies the entire old cache into it, every step, in every layer:

tokens already there     one step, concatenating     one step, preallocated
       128                       25 ms                        23 ms
       512                       33 ms                        22 ms
     1,000                       36 ms                        24 ms

At 1,000 tokens, every step copies the whole 70 MB cache, and that makes the step 50% slower. It gets worse as the text grows. Real engines never do this. They reserve the memory up front, like we did, and write into it. But reserving memory up front has its own problem, which we will get to soon.

How big does it get?

For every token, the cache keeps one key and one value, in every layer, for every head:

\[\text{bytes per token} = 2 \times \text{layers} \times \text{KV heads} \times \text{head size} \times \text{bytes per number}\]

For GPT-2 small in float32, that’s 2 × 12 × 12 × 64 × 4 bytes = 72 KiB per token, and 72 MiB for the full 1,024 tokens. You can check it against the timing table above: at 1,000 tokens, the cache held 70.3 MiB.

That doesn’t sound like much. But here are some real models, all in 16-bit numbers (2 bytes):

model              layers   KV heads   head size   per token    one long request
Llama 2 7B           32        32         128       512 KiB      4K tokens:    2 GiB
Llama 3.1 8B         32         8         128       128 KiB      128K tokens: 16 GiB
Llama 3.1 70B        80         8         128       320 KiB      128K tokens: 40 GiB

A single 128K-token request on Llama 3.1 8B needs 16 GiB of cache, which is about as much memory as the model’s weights. And a server doesn’t have one request. It has dozens or hundreds of them at the same time, each with its own cache. 32 users with 8K tokens each on the same model is another 32 GiB.

So for long texts and many users, the problem isn’t really the model anymore. It’s the cache.

(You might wonder why Llama 3.1 8B has only 8 KV heads when Llama 2 7B has 32. That’s one of the tricks to shrink the cache, and we will get to it near the end.)

The cache decides how many users you can serve

In the first blog, I said decode is memory bound: to generate one token, the GPU reads all of the model’s weights, and does very little math with them. The fix was batching: if many users are served together, the weights are read once and used for all of them.

But there’s a catch. The weights are shared between the users. The caches are not. Every user has their own cache, and every decode step has to read all of it.

To see this, I ported the decode step to PyTorch (the code is in the box below) and ran it on my laptop’s GPU in float16, with every sequence already holding 512 tokens. Then I increased the number of sequences in the batch:

sequences    time per step    tokens per second    KV cache size
    1           3.2 ms               315               18 MiB
    8           3.4 ms             2,400              144 MiB
   32           6.7 ms             4,800              577 MiB
   64          12.0 ms             5,300            1,154 MiB
  128          23.7 ms             5,400            2,308 MiB
  256          45.4 ms             5,600            4,617 MiB

There’s a lot going on here:

From 1 to 8 sequences, the time per step barely changes, so we get almost 8× more tokens for free. This is batching doing its job. (GPT-2 small is so tiny that the 3.2 ms is mostly the time to launch about 200 small GPU operations from Python, but the effect is the same on big models.)
After 32 sequences, the time grows with the batch. Doubling the batch doubles the time, and the tokens per second stop growing at around 5,400. The weights are read once, but every sequence brings 18 MiB of its own cache, and reading those caches becomes most of the work.
At 256 sequences, the cache is 4.6 GiB, on a GPU with 6 GB in total. The next doubling just doesn’t fit. On my first try it actually crashed with an out of memory error, because the previous batch’s cache was still in memory.

This is the main reason inference engines care so much about the cache. The memory you have for it decides how many sequences you can batch together, and for long texts that decides your throughput. Every byte you waste is fewer users per GPU.

The batched decode step in PyTorch

It’s the same model as before, but in PyTorch, with a batch of B sequences, and with fake caches that are already filled with 512 tokens (we only want to time the decode step). F.scaled_dot_product_attention does the scores, softmax and weighted sum in one call.

import torch, torch.nn.functional as F

dev, dt = "cuda", torch.float16
p = {name: torch.tensor(w).to(dev, dt) for name, w in weights.items()
     if not name.endswith(".attn.bias")}          # skip the stored masks
P = lambda l, n: p[f"h.{l}.{n}"]
ln = lambda x, l, n: F.layer_norm(x, (768,), P(l, n + ".weight"), P(l, n + ".bias"))

@torch.no_grad()
def decode_step(ids, K, V, pos):                  # ids: (B,), K/V: (12, B, 12, max, 64)
    B = ids.shape[0]
    x = p["wte.weight"][ids] + p["wpe.weight"][pos]
    for l in range(12):
        h = ln(x, l, "ln_1") @ P(l, "attn.c_attn.weight") + P(l, "attn.c_attn.bias")
        q, k, v = (t.view(B, 12, 1, 64) for t in h.split(768, -1))
        K[l, :, :, pos], V[l, :, :, pos] = k[:, :, 0], v[:, :, 0]
        a = F.scaled_dot_product_attention(q, K[l, :, :, :pos + 1], V[l, :, :, :pos + 1])
        x = x + a.reshape(B, 768) @ P(l, "attn.c_proj.weight") + P(l, "attn.c_proj.bias")
        h = F.gelu(ln(x, l, "ln_2") @ P(l, "mlp.c_fc.weight") + P(l, "mlp.c_fc.bias"), approximate="tanh")
        x = x + h @ P(l, "mlp.c_proj.weight") + P(l, "mlp.c_proj.bias")
    return F.layer_norm(x, (768,), p["ln_f.weight"], p["ln_f.bias"]) @ p["wte.weight"].T

pos = 512
for B in [1, 8, 32, 64, 128, 256]:
    K = torch.randn(12, B, 12, pos + 8, 64, device=dev, dtype=dt)
    V = torch.randn_like(K)
    ids = torch.full((B,), 5, device=dev)
    for _ in range(5):
        decode_step(ids, K, V, pos)               # warm up
    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(20):
        decode_step(ids, K, V, pos)
    end.record()
    torch.cuda.synchronize()
    ms = start.elapsed_time(end) / 20
    print(f"B={B:4d}  {ms:6.2f} ms/step  {B / ms * 1000:6.0f} tokens/s")

GPU work runs in the background, so a normal Python timer would only measure how long it took to queue the work. CUDA events are timestamps the GPU writes itself when it gets to them, which is why they’re used here.

The systems built around it

So the cache is big, it grows with every token, every user has their own, and how well you manage it decides how many users you can serve. Most of what a serving engine like vLLM or SGLang does is about this. Here’s a short map of the main ideas. I will go deeper into each of them in future blogs.

Not wasting memory: PagedAttention

When a request comes in, you know how long its prompt is, but not how long the answer will be. The model decides when to stop. Our GPT-2 dealt with that by reserving space for the maximum of 1,024 tokens. Older serving systems did the same, and it wastes a lot:

Space reserved for tokens that will come later, which sits empty until they do.
Space reserved for tokens that never come at all, because the answer stopped early.
Gaps between requests, when finished requests leave holes that are too small for the next one.

A small simulation shows how bad it is. With 2,000 requests, prompts of around 390 tokens, answers of around 144 tokens, and a limit of 4,096 tokens:

reserve the maximum (4,096 tokens) for every request:      20.9% of the memory holds real tokens
reserve the prompt + the max answer length (1,024):        51.5%
paged, in blocks of 16 tokens:                             99.1%

PagedAttention, from the vLLM paper, borrows the idea from how operating systems manage memory. The cache is split into small blocks of, say, 16 tokens each. A request gets a new block only when its last one is full, and the blocks don’t have to be next to each other in memory. Each request keeps a block table: a list of which blocks hold its tokens, in order. The attention kernel follows the table to find the keys and values.

So the only waste is the unused part of each request’s last block, at most 15 tokens. Since memory decides the batch size, this is a big deal: the paper reports 2–4× more throughput than the systems before it.

Reusing it: prefix caching

Many requests start with exactly the same text: the same system prompt, the same instructions, the same document with different questions about it. Their keys and values for that shared beginning are identical. That’s the same argument we used at the start: a token’s keys and values depend only on the tokens before it. So why compute them again?

Prefix caching keeps the blocks around after a request finishes. Each full block gets a hash computed from its tokens and all the tokens before it. When a new request comes in, the engine hashes its blocks the same way, and every block whose hash already exists is reused instead of computed. Only the rest goes through the model. This is the start > 0, T > 1 case from our mask earlier: some tokens cached, a few new ones.

I tried this with vLLM, serving Qwen2.5-1.5B on an RTX 5070 Ti. I sent the same ~5,900-token prompt three times, changing only the question at the end, and measured the time to the first token:

                          request 1    request 2    request 3
prefix caching on           311 ms        30 ms        27 ms
prefix caching off          351 ms       344 ms       341 ms

With caching on, the second and third requests are more than 10× faster. vLLM’s metrics show exactly what happened: each of them reused 5,920 tokens. That’s 370 full blocks of 16 tokens. Only the remaining 17 tokens were actually computed. (Only full blocks are reused, because a partly filled block can still change.)

The remaining 30 ms isn’t the model anymore. It’s everything around it: the HTTP request, turning 5,900 tokens of text into IDs on the CPU, looking up 370 blocks, and one small forward pass that still has to attend over all the cached tokens. Prefix caching removes the prefill work, and what’s left is the fixed overhead.

This is also why API providers charge less for “cached input tokens”. And it’s why you should put the parts of your prompt that don’t change at the start, and the parts that do at the end. A single different token early in the prompt changes the hash of every block after it.

Making it smaller

The other way to fit more is to make the cache itself smaller:

Fewer KV heads (MQA and GQA). In GPT-2, each of the 12 heads has its own keys and values. With grouped-query attention, several query heads share one key and value head. Llama 3.1 8B has 32 query heads but only 8 KV heads, which is exactly why its cache is 4× smaller than Llama 2 7B’s in the table above, with almost no loss in quality. Multi-query attention takes it all the way down to a single KV head.
Compressing it (MLA). DeepSeek’s models store a small compressed vector per token instead of full keys and values, and expand it when needed. DeepSeek-V3 has 61 layers and 671 billion parameters, but its cache is about 69 KiB per token. That’s less than our tiny GPT-2 in float32!
Fewer bytes per number. Storing the cache in 8-bit floats (FP8) instead of 16-bit halves it, usually with very little effect on the output.
Not keeping all of it. Some models only attend to the last few thousand tokens in some of their layers (sliding window attention), so those layers only need to keep that many. Remember the attention sink from the last blog, where many heads put a lot of weight on the first token? It turns out that if you throw away old tokens, you have to keep those first few, or the model falls apart. That’s the idea behind StreamingLLM.

When it still doesn’t fit

Even with all of that, a busy server runs out of cache memory. Then it has to choose:

Preemption. Pause some requests and free their blocks. When there’s space again, either recompute their cache from the text (a prefill, which is fast), or swap it out to the CPU’s memory and copy it back later.
Offloading. Keep the cache that isn’t needed right now, like old conversations, in CPU memory or even on disk, so it can be brought back instead of recomputed.
Splitting prefill and decode. Big deployments run prefill on some GPUs and decode on others, because one is compute bound and the other is memory bound. The cache is computed on the first group and sent over the network to the second, so moving the cache becomes a problem of its own.

Wrapping up

So the KV cache comes from a simple fact: because of the causal mask, a token’s keys and values never change. Storing them turns every decode step from “process the whole text again” into “process one token”, which was 100× faster at 1,000 tokens even on a laptop CPU, with exactly the same output. But the cache grows with every token and every user, and the memory it takes decides how many users a GPU can serve. That’s why serving engines page it, share it, shrink it and move it around.

In the next blogs, I will go deeper into these systems one at a time, starting with how a GPU actually works, so we can see exactly why decode is memory bound and why the cache matters so much. I am really enjoying writing these, and I hope they help someone:)

If you find anything wrong, please let me know!