Skip to content
LLMs from Scratch

Chapter 4 · Bonus material

The KV Cache

The single most important inference optimization, in one idea: never compute the same key and value twice.

Adapted from the original notebook by Sebastian Raschka · the Python shown is the book's; the demos run a faithful JavaScript port in your browser

Bonus · The KV cache

Time
KV
flies
KV

two tokens in — each computes a key and a value vector

01 / 05

Attention turns each token into a query, a key, and a value. Start generating from a prompt — “Time flies” — and every token has its K and V.

The lab

The waste is quadratic; the fix is linear. Set a prompt length and a token count, toggle the cache, and watch the gap open up.

KV cache · compute counter

Count the key/value vectors a plain loop recomputes vs. what a cache reuses. Amber = computed now, teal = reused from cache.

each token computed once
K/V vectors computed (with cache)
15
102
no cache
15
with cache
6.8 ×
less work

The cache skips 87 redundant K/V computations here — and the gap widens with every token.

The real code

The change is small. Inside MultiHeadAttention.forward, a use_cache flag appends the new token's keys and values to a stored buffer instead of recomputing everything:

gpt_with_kv_cache.py
def forward(self, x, use_cache=False):
    keys_new = self.W_key(x)      # only for the NEW tokens
    values_new = self.W_value(x)
    queries = self.W_query(x)
    # ...
    if use_cache:
        if self.cache_k is None:
            self.cache_k, self.cache_v = keys_new, values_new
        else:                                            # append, don't recompute
            self.cache_k = torch.cat([self.cache_k, keys_new], dim=1)
            self.cache_v = torch.cat([self.cache_v, values_new], dim=1)
        keys, values = self.cache_k, self.cache_v
    else:
        keys, values = keys_new, values_new

The generation loop then feeds the model only the single new token each step, instead of the whole sequence:

gpt_with_kv_cache.py
def generate_text_simple_cached(model, idx, max_new_tokens, use_cache=True):
    model.eval()
    model.reset_kv_cache()
    logits = model(idx, use_cache=True)          # prime the cache with the prompt
 
    for _ in range(max_new_tokens):
        next_idx = logits[:, -1].argmax(dim=-1, keepdim=True)   # greedy
        idx = torch.cat([idx, next_idx], dim=1)
        logits = model(next_idx, use_cache=True)  # feed ONLY the new token
    return idx

Two things to remember

  1. A KV cache trades memory for time. Keys and values are computed once per token and reused, turning per-step attention work from O(n²) to O(n) — about 5× faster on a 124M model. The cache grows with the sequence, so it costs memory and is inference-only.
  2. The output is unchanged. Caching is pure bookkeeping — same math, same tokens. If cached and uncached generation ever disagree, it's an indexing bug, not a feature.

Adapted from Build a Large Language Model (From Scratch) by Sebastian Raschka — original bonus material (Apache 2.0). The tokens/sec figures are its published benchmark; the compute counts are exact. The TypeScript KV-cache port is verified to match the uncached path token-for-token.