DEV Community

Cover image for The KV Cache From Scratch
Sachin Kr. Rajput
Sachin Kr. Rajput

Posted on Edited on Originally published at blog.sachinrajput.dev

The KV Cache From Scratch

In every part of this series, one thing kept showing up: the notebook. In Part 1 it grew by 72 KB with every word. In Part 3 it was the part of the work that batching couldn't share. In Part 4 it stayed stuck on the memory side of the roof, however many users we added.

Engineers call it the KV cache. Today we build one by hand, check that it really gives the same answer, weigh it in bytes, and try the tricks real servers use to shrink it.

What goes in the notebook

Quick picture of how a word decides what comes next (it's called attention). On every floor of the Word Factory, each word makes three little cards:

  • a question card (Q): "what am I looking for?"
  • a label card (K, for key): "here's what I'm about"
  • a content card (V, for value): "here's what I can tell you"

To choose the next word, the newest word holds up its question and compares it with every earlier word's label. The better a label matches, the more of that word's content card gets mixed in.

Now the important bit: a word only looks backwards. Once "sky" has made its label and content cards, nothing that comes later can change them. So why make them again?

The notebook is just those two cards, K and V, saved for every earlier word, on every floor, for every head. When a new word arrives, the factory makes cards for that one word, adds them to the notebook, and reads the rest from the notebook instead of rebuilding them.

Without the notebook vs with it

I wrote a tiny factory in kv_cache.py, the same shape as before (GPT-2 small: 12 floors, width 768, 12 heads, random weights). One function, run, can work both ways:

def run(tokens, start, cache=None):
    for i, p in enumerate(L):                       # 12 floors
        q, k, v = (heads(m) for m in np.split(p["qkv"] @ norm(x), 3))
        if cache is not None:                       # write the new cards...
            cache["K"][i, :, start:start + t] = k
            cache["V"][i, :, start:start + t] = v
            k = cache["K"][i, :, :start + t]        # ...then read ALL saved cards
            v = cache["V"][i, :, :start + t]
        att = softmax(q @ k.transpose(0, 2, 1) / HD ** 0.5)
        ...
Enter fullscreen mode Exit fullscreen mode

(Shortened for reading. The full file is what I ran.)

  • Without a notebook: every time we want one more word, we push the whole text so far through the factory again.
  • With a notebook: we read the question once, then push only the newest word through, reading everything else from the notebook.

First, the honesty check. I gave both versions the same 128-token question and let each one write 32 words:

1) same words with and without the notebook: True
   biggest difference in any score: 0.0000021
Enter fullscreen mode Exit fullscreen mode

Same 32 words. The scores the factory gave every possible next word differed by at most two millionths, which is just rounding. The notebook doesn't change the answer. It only saves work.

How much work?

2) time per new word (ms)
   text length   no notebook   with notebook
           129         231.7            30.6
           136         249.7            20.4
           144         308.8            20.7
           159         328.2            21.4
   total for 32 words: 9.68 s without, 0.66 s with
Enter fullscreen mode Exit fullscreen mode
  • With the notebook, each word costs about 21 ms, and that barely moves as the text grows.
  • Without it, each word costs 230 to 330 ms, and it keeps climbing, because every step re-reads more text than the last.
  • In total, the notebook made typing about 15 times faster. And the longer the answer, the bigger that gap gets: without a notebook, writing word 1,000 means re-reading 1,000 words.

So the trade is simple: spend memory to save math. Now let's see how much memory.

How big is the notebook?

Every word saves two cards (K and V), on every floor, for every head, and each card is a list of numbers. So:

notebook bytes per word = 2 × floors × heads × numbers per head × bytes per number

For our toy factory that's 2 × 12 × 12 × 64 × 4 bytes. The script measures it straight from the arrays:

3) notebook per token: 73,728 bytes = 72 KB
    1,000 tokens:      74 MB   (the model itself is 340 MB)
    8,000 tokens:     590 MB   (the model itself is 340 MB)
   32,000 tokens:    2359 MB   (the model itself is 340 MB)
Enter fullscreen mode Exit fullscreen mode

After about 4,600 words, one user's notebook is bigger than the whole model. And every user has their own.

Real models are bigger, so I plugged their published settings into the same formula with kv_sizes.py:

model                                       per token  8K-token chat
GPT-2 small shape, float32 (our toy)             72 KB        0.60 GB
Llama 3.1 8B, if every head kept notes          512 KB        4.29 GB
Llama 3.1 8B, BF16 (8 shared notebooks)         128 KB        1.07 GB
Llama 3.1 8B, FP8 notebook                       64 KB        0.54 GB
Llama 3.1 70B, BF16                             320 KB        2.68 GB
DeepSeek-V3, if every head kept notes         4,880 KB       40.94 GB
DeepSeek-V3, BF16 (compressed notes)             69 KB        0.58 GB
Enter fullscreen mode Exit fullscreen mode

Notebook size for one 8K-word chat: Llama 3.1 8B drops from 4.29 GB to 1.07 GB by sharing notebooks between heads, and to 0.54 GB with 1-byte numbers. DeepSeek-V3 drops from 40.94 GB to 0.58 GB by compressing its notes.

Two things jump out. A single 8K-word chat on Llama 3.1 8B needs about 1 GB of notebook. And the "if every head kept notes" rows are much bigger than the real ones, because real models already use tricks. We'll get to those in a moment.

Why the notebook fills your GPU

Let's put Llama 3.1 8B on one 80 GB NVIDIA H100. The model's tables take 16 GB. Keep about 6 GB for everything else, and about 58 GB is left for notebooks:

how many users fit on one 80 GB H100 with Llama 3.1 8B in BF16?
  BF16 notebook: 58 GB free -> 442K tokens of notes ->    54 users at 8K,  3.4 users at 128K
  FP8 notebook: 58 GB free -> 884K tokens of notes ->   108 users at 8K,  6.7 users at 128K
Enter fullscreen mode Exit fullscreen mode
  • With 8K-word chats, about 54 people can be served at once.
  • With the long 128K-word chats this model supports, about 3.

The chip still has plenty of math left over (remember Part 4: one decode step barely uses it). What runs out first is room for notebooks. That's why "how long is the context?" is the first question anyone sizing a chatbot server should ask.

And it's not just room, it's time. Every typing step reads every user's whole notebook:

one decode step, Llama 3.1 8B, 48 users at 8K each:
  tables 16 GB + notebooks 52 GB -> the notebooks are 3.2x the model
Enter fullscreen mode Exit fullscreen mode

The forklift from Part 3 spends three quarters of each step carrying notebooks, not the model.

Tricks to make the notebook smaller

Since the notebook is just numbers, there are only a few ways to shrink it: keep fewer cards, use smaller numbers, squeeze the cards, or stop keeping copies of the same thing.

1. Let heads share a notebook

A floor has many heads, each asking its own kind of question. But the heads don't each need their own label and content cards. Several heads can share one set. Engineers call this grouped-query attention (GQA). Llama 3.1 8B has 32 question heads but keeps only 8 notebooks, so its notebook is 4 times smaller than it would be.

I measured what that does to the reading time, for one new word looking back over an 8,192-word notebook on all 12 floors:

4) reading a 8,192-token notebook for ONE new word, all 12 floors
   12 query heads, 12 notebook heads:   604 MB read,  43.5 ms
   12 query heads,  3 notebook heads:   151 MB read,  19.2 ms
   12 query heads,  1 notebook heads:    50 MB read,   9.8 ms
Enter fullscreen mode Exit fullscreen mode

Four times fewer bytes, more than twice as fast. Twelve times fewer bytes, more than four times as fast. It's the roofline again: reading notebooks is a forklift job, so fewer bytes means less waiting. (It isn't exactly 4x and 12x because some fixed costs don't shrink.)

2. Use smaller numbers

Store each number in 1 byte (FP8) instead of 2 and the notebook halves. On our H100 that's the difference between 54 and 108 users at 8K. The catch is a tiny loss of precision, which is why teams test answer quality before switching.

3. Compress the notes

DeepSeek-V3 goes further. Instead of saving full label and content cards for 128 heads, it saves one small compressed note per floor and unpacks it when needed (its authors call this multi-head latent attention). The result, from the table above: 4,880 KB per word becomes 69 KB, about 70 times smaller. A giant model ends up with a notebook about the size of our toy's.

4. Don't keep copies of the same thing

If a thousand users all start with the same 2,000-word system prompt, their first 2,000 words of notes are identical. Smart servers keep one copy and share it. Engineers call this prefix caching, and it's one of the cheapest wins in serving.

5. Don't waste the space you have

Notebooks grow one word at a time and you don't know how long an answer will be. Reserving the maximum for everyone wastes most of the memory. The fix is to hand out notebook space in small pages, like a library handing out shelves as needed. That's PagedAttention, and it gets its own part later in this series.

Cheat sheet

Idea One line
KV cache Saved label (K) and content (V) cards for every earlier word
Why it works Words only look backwards, so old cards never change
What it buys Same answer, about 15x faster typing in our test
What it costs 2 × floors × heads × head size × bytes, per word, per user
Llama 3.1 8B 128 KB per word, about 1 GB per 8K-word chat
On one H100 About 54 users at 8K, about 3 at 128K
Shrink it Share between heads, smaller numbers, compress, share prefixes, page it

What's next

We now know the two big things that live in GPU memory: the model's tables and everyone's notebooks. Part 6: How Much GPU Memory Do I Need? puts them together into one simple worksheet, so you can size a server for any model, context length and number of users.


Every number in this post comes from two scripts: kv_cache.py (measured on a 2-core cloud CPU with float32 NumPy) and kv_sizes.py (the notebook formula applied to each model's published config). The tiny model used random weights, so it shows the pattern, not the speed of a real GPU.


This is Part 5 of Inference Foundations, a 10-part series that explains how AI models answer you, one measured step at a time. Next up is Part 6: How Much GPU Memory Do I Need? Follow me here on DEV to get it when it's out, or read the whole series on blog.sachinrajput.dev.

Top comments (0)