Build a KV cache for LLM generation

Problem

When you need this. The generator recalculates the opening sentence every time it adds a token. Reuse the earlier attention state and check that the generated result stays the same.

The generator in the previous build runs the whole model on every token so far, once for every token it adds. A story of about 28 generated tokens therefore processes the opening sentence again and again. This build removes that repeated work with a key-value (KV) cache and checks that nothing else changes.

The contract matches the previous build exactly: generate_cached() takes the same arguments as generate() and must return the same text, the same chosen tokens with the same probabilities and the same stop reason. Underneath, forward_step(params, cache, token_id, position) runs the model on one new token and returns its next-token scores.

Success has three checks, all recorded on the openings of the 60 test stories. Nothing is chosen on them: this build compares two ways to compute the same output, not settings. The cached scores equal the full forward pass’s scores exactly, not approximately. Cached and uncached generation give identical results with greedy choice and with sampling. And the trace counts how many model rows and attention scores each version computes.

Concepts

Build causal self-attention and inspect intermediate outputs tested that changing a later token never changes an earlier position’s output. That property is what makes a cache possible:

KV cache: The key and value vectors of every position processed so far, kept for one sequence. Because of the causal mask, a position’s key and value never change once computed: later tokens cannot affect them.
Prefill: Processing the prompt’s tokens to fill the cache before the first new token is chosen. This build feeds the prompt one token at a time through the same step function.
Decode step: Running the model on one new token: compute its query, key and value, append the key and value to the cache, score the query against every cached key and mix the cached values. Everything else in the block, layer normalization, the MLP and the skip paths, works on one position at a time and needs nothing from earlier positions.

Architecture

One step with and without a KV cache

The uncached generator recomputes all prefix positions. The cached path processes each prompt token and then one new token at a time, appending its key and value and reading retained earlier keys and values. Both paths use the same token selection function. The cache belongs to one sequence and parameter set; invalid identity, positions or token IDs are refused.

The figure compares one generation step on two lanes. On the without a cache lane, generate() passes all n token IDs so far to model.forward(), which computes all n rows; only the last row of scores is used. On the with a cache lane, generate_cached() passes one ID at a time to forward_step(): first each of the prompt IDs, then each chosen token, which returns on the path labeled chosen token. forward_step() computes one row, appends its key and value to KVCache and reads the keys and values of every earlier position. It raises ValueError for a cache built with other parameters, a position out of order, a position past the model’s last one, or a token ID outside the vocabulary. Both lanes feed the same next_token_probs(). The comparison below the run command reports the work saved, measured time and output equivalence.

kv_cache.py holds the KVCache class, forward_step() and generate_cached(). forward_step() repeats each operation of model.forward() for one row, in the same order, including the order of additions; floating-point addition is not associative, and the tests show that even one reordered sum changes the last bit. generate_cached() reuses next_token_probs() and check_settings() from generate.py, so the selection logic is shared code, not a copy. trace_kv.py runs both generators on every opening and writes kv-trace.json and kv-time.json.

Working code

Runtime: Python 3.10 or newer. Dependencies: Python standard library only.

Download and extract the complete lab using Download the complete lab above. Open build-a-kv-cache-for-llm-generation/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.

kv_cache.pySource · pythonDownload

"""Generation with a key-value (KV) cache for the one-block model, standard library only.

Without a cache, every new token runs the whole model on every token so far.
With a cache, each step runs the model on the new token only: its query is
compared with the keys and values that earlier steps saved. Every other part of
the block works on one position at a time, so nothing else needs saving.
"""
import math
import random
from operator import mul

import generate as g
import model as m


class KVCache:
    """Keys and values of every position processed so far, for one sequence and one set of weights."""

    def __init__(self, params):
        self.params_id = id(params)
        self.keys, self.values = [], []

    def __len__(self):
        return len(self.keys)


def _row(x, w):
    cols = list(zip(*w))
    return [math.fsum(map(mul, x, col)) for col in cols]


def forward_step(p, cache, token_id, position):
    """Scores for the next token after appending one token at `position`, updating the cache."""
    if cache.params_id != id(p):
        raise ValueError("this cache was built with different parameters; start a new cache")
    if position != len(cache):
        raise ValueError(f"expected position {len(cache)}, got {position}: tokens must arrive in order")
    if position >= len(p["pos_emb"]):
        raise ValueError(f"position {position} is beyond the {len(p['pos_emb'])} positions the model has")
    if isinstance(token_id, bool) or not 0 <= token_id < len(p["tok_emb"]):
        raise ValueError(f"token ID {token_id!r} is outside the vocabulary")
    x0 = [a + b for a, b in zip(p["tok_emb"][token_id], p["pos_emb"][position])]
    h1 = m.layernorm_forward([x0], p["ln1_g"], p["ln1_b"])[0][0]
    q, k, v = _row(h1, p["Wq"]), _row(h1, p["Wk"]), _row(h1, p["Wv"])
    cache.keys.append(k)
    cache.values.append(v)
    scale = 1 / math.sqrt(len(q))
    weights = m.softmax_row([math.fsum(map(mul, q, kj)) * scale for kj in cache.keys])
    mixed = [math.fsum(map(mul, weights, col)) for col in zip(*cache.values)]
    x1 = [a + b for a, b in zip(x0, _row(mixed, p["Wo"]))]
    h2 = m.layernorm_forward([x1], p["ln2_g"], p["ln2_b"])[0][0]
    act = [v if v > 0 else 0.0 for v in (a + b for a, b in zip(_row(h2, p["W1"]), p["b1"]))]
    x2 = [a + (b + c) for a, b, c in zip(x1, _row(act, p["W2"]), p["b2"])]  # same order as model.forward
    hf = m.layernorm_forward([x2], p["lnf_g"], p["lnf_b"])[0][0]
    return [a + b for a, b in zip(_row(hf, p["W_out"]), p["b_out"])]


def generate_cached(p, tok, prompt, max_new_tokens=60, temperature=0.0, top_k=None, seed=0):
    """Same contract and output as generate.generate(), with one model row per step."""
    g.check_settings(temperature, top_k, len(p["tok_emb"]))
    ids = tok.encode(prompt)
    if not ids:
        raise ValueError("empty prompt: the model needs at least one token to continue")
    max_len = len(p["pos_emb"])
    if len(ids) > max_len:
        raise ValueError(f"prompt of {len(ids)} tokens is longer than the {max_len} positions the model has")
    cache = KVCache(p)
    for pos, token_id in enumerate(ids):  # prefill: one step per prompt token
        logits = forward_step(p, cache, token_id, pos)
    rng = random.Random(seed)
    steps, reason = [], "max_new_tokens"
    newline = tok.ids["\n"]
    for _ in range(max_new_tokens):
        if len(ids) == max_len:
            reason = "context full"
            break
        probs = g.next_token_probs(logits, temperature, top_k)
        choice = rng.choices(range(len(probs)), weights=probs)[0] if temperature > 0 else probs.index(1.0)
        ids.append(choice)
        steps.append({"token": tok.vocab[choice], "probability": round(probs[choice], 4)})
        if choice == newline:
            reason = "end of story"
            break
        if len(ids) < max_len:
            logits = forward_step(p, cache, choice, len(ids) - 1)
    return {"text": tok.decode(ids), "steps": steps, "stopped": reason}

forward_step() checks four things before it computes anything: the cache belongs to these parameters, the position is the next one in order, the position exists in the table, and the token ID is in the vocabulary. A wrong cache would otherwise give plausible-looking but wrong scores with no error.

Read trace_kv.py

trace_kv.pySource · pythonDownload

"""Compare cached and uncached generation on the openings of the 60 test stories and record the results.

Run: python3 trace_kv.py          (print the trace; under a minute on a laptop CPU)
     python3 trace_kv.py --write  (rewrite kv-trace.json and kv-time.json)
     python3 trace_kv.py --check  (exit 1 if kv-trace.json differs)
kv-time.json holds wall-clock times from the --write run and is not compared,
because timings change from run to run and machine to machine.
"""
import json
import platform
import sys
import time
from pathlib import Path

import checkpoint as ck
import generate as g
import kv_cache as kv
import model as m
import train as t
from bpe import Tokenizer

HERE = Path(__file__).parent


def work(prompt_tokens, generated_tokens):
    """Rows the model processes and attention scores it computes, with and without a cache.

    Without a cache, the step that produces token number n runs the model on all n-1
    tokens before it. With a cache, every token is processed once.
    """
    lengths = range(prompt_tokens, prompt_tokens + generated_tokens)
    return {
        "uncached_rows": sum(lengths),
        "cached_rows": prompt_tokens + generated_tokens - 1 if generated_tokens else prompt_tokens,
        "uncached_scores": sum(n * (n + 1) // 2 for n in lengths),
        "cached_scores": sum(range(1, prompt_tokens + generated_tokens)),
    }


def run():
    tok = Tokenizer.load(HERE / "tokenizer.json")
    p = ck.load(HERE / "best.json", t.CONFIG, ck.input_signature(tok, (HERE / "corpus.txt").read_text()))["params"]
    _, _, test = t.split_stories((HERE / "corpus.txt").read_text())
    openings = [s.split(". ")[0] + "." for s in test]

    # 1. Every logit of a cached forward pass equals the matching row of the full forward pass.
    ids = tok.encode(test[0])[:-1]
    full, _ = m.forward(p, ids)
    cache = kv.KVCache(p)
    largest_gap = max(max(abs(a - b) for a, b in zip(kv.forward_step(p, cache, i, pos), full[pos]))
                      for pos, i in enumerate(ids))

    # 2. Greedy and sampled generation give the same text with and without the cache.
    same, totals, seconds = 0, {}, {}
    start = time.perf_counter()
    uncached = [g.generate(p, tok, o) for o in openings]
    seconds["uncached"] = time.perf_counter() - start
    start = time.perf_counter()
    cached = [kv.generate_cached(p, tok, o) for o in openings]
    seconds["cached"] = time.perf_counter() - start
    for o, a, b in zip(openings, uncached, cached):
        same += a == b
        for key, value in work(len(tok.encode(o)), len(a["steps"])).items():
            totals[key] = totals.get(key, 0) + value
    sampled_same = sum(g.generate(p, tok, o, temperature=1.0, seed=i) == kv.generate_cached(p, tok, o, temperature=1.0, seed=i)
                       for i, o in enumerate(openings))

    record = {
        "stories": len(openings),
        "largest_logit_gap_first_story": largest_gap,
        "positions_checked": len(ids),
        "greedy_identical": same,
        "sampled_identical": sampled_same,
        "generated_tokens": sum(len(a["steps"]) for a in uncached),
        "work": totals,
        "cache_floats_per_position": 2 * len(p["Wk"][0]),
        "cache_floats_full_context": 2 * len(p["Wk"][0]) * len(p["pos_emb"]),
    }
    timing = {"uncached_seconds": round(seconds["uncached"], 2), "cached_seconds": round(seconds["cached"], 2),
              "stories": len(openings), "python": platform.python_version(), "machine": platform.machine(),
              "system": platform.system()}
    return record, timing


if __name__ == "__main__":
    record, timing = run()
    text = json.dumps(record, indent=1) + "\n"
    if "--write" in sys.argv:
        (HERE / "kv-trace.json").write_text(text)
        (HERE / "kv-time.json").write_text(json.dumps(timing, indent=1) + "\n")
    elif "--check" in sys.argv:
        sys.exit(0 if (HERE / "kv-trace.json").read_text() == text else "kv-trace.json is out of date")
    else:
        print(text, end="")

Run the comparison:

Terminal commandsTerminal · bash

python3 trace_kv.py

python3 trace_kv.py --check exits with an error if your output differs from the recorded kv-trace.json.

Check that nothing changed

The trace runs forward_step() over the 30 positions of the first test story and compares each row of scores with the full forward pass. The largest difference is 0.0: every score is the same number, bit for bit. Over all 60 openings, cached generation gives the identical result in 60 of 60 stories with greedy choice and in 60 of 60 with sampling at temperature 1.0.

Count the work removed

Together the 60 greedy stories generate 1,673 tokens:

Work	Without a cache	With a cache	Ratio
Rows run through the model	35,805	2,073	17.3
Attention scores computed	464,956	37,419	12.4

These counts follow from the algorithm and are the same on every machine. The recorded run, measured on one machine, took 5.98 seconds without the cache and 0.38 seconds with it. That is one measurement, not a benchmark, and timings vary between runs.

The cache costs memory instead: 32 numbers per position, a key and a value of the model’s width, or 2,048 for all positions of one sequence. Each sequence being generated needs its own cache.

Check your understanding

You generate one story, then start a second story with a new prompt but reuse the same KVCache object. Before you open the comparison, predict what happens, and why the code is written so.

Compare your prediction

forward_step() raises ValueError: the cache already holds positions, so the new prompt’s first token at position 0 is out of order. Without that check, the second story’s tokens would attend to the first story’s keys and values, and the model would produce scores for a text that does not exist. This implementation starts a new cache for every sequence; a library that reuses one does so deliberately, for text that really continues.

Read the explanation

Causal attention makes each position's key and value depend only on tokens up to that position. A new token cannot change them, so a cache can keep them and the result stays identical.

Inspect the recorded trace

kv-trace.jsonRecorded output · jsonDownload

{
 "stories": 60,
 "largest_logit_gap_first_story": 0.0,
 "positions_checked": 30,
 "greedy_identical": 60,
 "sampled_identical": 60,
 "generated_tokens": 1673,
 "work": {
  "uncached_rows": 35805,
  "cached_rows": 2073,
  "uncached_scores": 464956,
  "cached_scores": 37419
 },
 "cache_floats_per_position": 32,
 "cache_floats_full_context": 2048
}

Failure cases

Each case has a named test in test_kv_cache.py.

Case	Result with this implementation	What it tells you
Reusing a cache for a new prompt	`forward_step()` raises `ValueError` about the expected position	A cache belongs to one sequence
Tokens out of order	`forward_step()` raises `ValueError`	Each key and value belongs to one position
A cache built with other parameters	`forward_step()` raises `ValueError`	Keys from one checkpoint are meaningless to another
A position beyond the table	`forward_step()` raises `ValueError`	The cache cannot extend the model’s positions
A token outside the vocabulary	`forward_step()` raises `ValueError`	The same boundary check as the full model

The complete lab includes these tests and their fixtures. Run the 11 standard-library tests:

Terminal commandsTerminal · bash

python3 -m unittest -v test_kv_cache.py

Read the tests

test_kv_cache.pySource · pythonDownload

"""Tests for kv_cache.py. Run: python3 -m unittest -v test_kv_cache.py"""
import copy
import subprocess
import sys
import unittest
from pathlib import Path

import checkpoint as ck
import generate as g
import kv_cache as kv
import model as m
import train as t
import trace_kv as tk
from bpe import Tokenizer

HERE = Path(__file__).parent
TOK = Tokenizer.load(HERE / "tokenizer.json")
P = ck.load(HERE / "best.json", t.CONFIG, ck.input_signature(TOK, (HERE / "corpus.txt").read_text()))["params"]
STORIES = t.split_stories((HERE / "corpus.txt").read_text())[2]


class Recorded(unittest.TestCase):
    def test_recorded_trace(self):
        result = subprocess.run([sys.executable, "trace_kv.py", "--check"], cwd=HERE)
        self.assertEqual(result.returncode, 0)


class Equivalence(unittest.TestCase):
    def test_every_cached_row_equals_the_full_forward_pass(self):
        ids = TOK.encode(STORIES[1])[:-1]
        full, _ = m.forward(P, ids)
        cache = kv.KVCache(P)
        for pos, token_id in enumerate(ids):
            self.assertEqual(kv.forward_step(P, cache, token_id, pos), full[pos])

    def test_cached_generation_matches_uncached(self):
        for i, story in enumerate(STORIES[:5]):
            opening = story.split(". ")[0] + "."
            for settings in ({}, {"temperature": 1.0, "seed": i}, {"temperature": 1.0, "top_k": 3, "seed": i}):
                with self.subTest(story=i, settings=settings):
                    self.assertEqual(kv.generate_cached(P, TOK, opening, **settings),
                                     g.generate(P, TOK, opening, **settings))

    def test_context_full_stops_at_the_same_point(self):
        prompt = TOK.decode(TOK.encode((HERE / "corpus.txt").read_text())[:63])
        self.assertEqual(kv.generate_cached(P, TOK, prompt), g.generate(P, TOK, prompt))

    def test_work_counts_match_the_rows_the_model_processes(self):
        rows = []
        original = m.forward
        m.forward = lambda p, ids, *a: (rows.append(len(ids)), original(p, ids, *a))[1]
        try:
            opening = STORIES[0].split(". ")[0] + "."
            result = g.generate(P, TOK, opening)
        finally:
            m.forward = original
        counts = tk.work(len(TOK.encode(opening)), len(result["steps"]))
        self.assertEqual(sum(rows), counts["uncached_rows"])
        self.assertEqual(sum(n * (n + 1) // 2 for n in rows), counts["uncached_scores"])

    def test_cached_work_counts_match_the_forward_step_calls(self):
        positions = []
        original = kv.forward_step
        kv.forward_step = lambda p, cache, token_id, position: (positions.append(position), original(p, cache, token_id, position))[1]
        try:
            opening = STORIES[0].split(". ")[0] + "."
            result = kv.generate_cached(P, TOK, opening)
        finally:
            kv.forward_step = original
        counts = tk.work(len(TOK.encode(opening)), len(result["steps"]))
        self.assertEqual(len(positions), counts["cached_rows"])
        self.assertEqual(sum(pos + 1 for pos in positions), counts["cached_scores"])  # one score per cached key


class FailureCases(unittest.TestCase):
    def test_reusing_a_cache_for_a_new_prompt(self):
        cache = kv.KVCache(P)
        kv.forward_step(P, cache, TOK.ids["mia"], 0)
        with self.assertRaisesRegex(ValueError, "expected position 1, got 0"):
            kv.forward_step(P, cache, TOK.ids["mia"], 0)

    def test_tokens_out_of_order(self):
        with self.assertRaisesRegex(ValueError, "expected position 0, got 2"):
            kv.forward_step(P, kv.KVCache(P), TOK.ids["mia"], 2)

    def test_cache_from_other_parameters(self):
        other = copy.deepcopy(P)
        cache = kv.KVCache(other)
        with self.assertRaisesRegex(ValueError, "different parameters"):
            kv.forward_step(P, cache, TOK.ids["mia"], 0)

    def test_position_beyond_the_table(self):
        cache = kv.KVCache(P)
        cache.keys = cache.values = [[0.0]] * 64
        with self.assertRaisesRegex(ValueError, "beyond the 64 positions"):
            kv.forward_step(P, cache, TOK.ids["mia"], 64)

    def test_token_outside_the_vocabulary(self):
        for bad in (-1, 96, True):
            with self.subTest(token=bad), self.assertRaisesRegex(ValueError, "outside the vocabulary"):
                kv.forward_step(P, kv.KVCache(P), bad, 0)


if __name__ == "__main__":
    unittest.main()

test_work_counts_match_the_rows_the_model_processes counts the rows the uncached generator actually passes to model.forward() and compares them with the formula that produced the table.

Deployment

The next build, Build a server for your trained LLM, serves the model. A server generating for several requests at once needs one cache per request, created at the start and dropped at the end, and must never let one request’s cache reach another.

Scaling

The row and score counts above come from formulas in trace_kv.py, and two tests check those formulas against the rows model.forward() and forward_step() actually process; the timing is one run. For a fixed prompt and growing output, the uncached model-row count is quadratic in the number of generated tokens. Its attention-score count is cubic: each step computes a triangular number of scores over the whole growing prefix. With a cache, rows grow linearly and attention scores grow quadratically, because each new query still reads all earlier keys. These statements count operations at fixed model width, before the position limit; they are not latency predictions. The price is memory that grows with sequence length, model width, number of layers and the number of sequences in flight.

Key takeaways

The causal mask means earlier keys and values never change, so a generation step can reuse them instead of recomputing them.
A correct cache changes how much work is done, not the result: cached scores and texts are identical to uncached ones.
In this implementation a cache belongs to one sequence and one set of weights. Libraries that reuse a cache, across the turns of a conversation or for a shared prompt prefix, do so explicitly by copying or cropping it; reuse by accident gives wrong scores without an error unless the code checks.

Best practices

Compare cached and uncached output exactly, row by row and story by story, before trusting a cache.
Count work in rows and scores, which do not vary, and report timings as single measurements with the machine and Python version.
Create a fresh cache per sequence and check the position of every token added to it.

References

Sources checked October 20, 2026.

Hugging Face Transformers v5.20: how caching works: keys and values of processed tokens are stored and reused, only the newest token is processed each step, and the cache grows with sequence length.
Hugging Face Transformers v5.20: cache strategies: cache types, reuse across conversation turns, and copying a prefilled prompt cache for new prompts.
Pope et al., Efficiently Scaling Transformer Inference (MLSys 2023), arXiv 2211.05102, section 2: prefill and decode as separate phases, and a KV cache unique to each sequence whose memory can dominate at long contexts and large batches.
Vaswani et al., Attention Is All You Need (NIPS 2017), sections 3 and 3.2.3: the decoder is auto-regressive, and masking keeps each position from using later ones. The paper does not describe caching.