E2.1 · LESSON 1 OF 6

What is cached

Calculate KV-cache memory and explain how layers, KV heads, context length and concurrent sequences affect it.

THE IDEA

Remembering keys and values

When the model generates the next token, the new token's query must be compared with the keys of all earlier tokens and then mixes their values. Those keys and values were already computed in earlier steps and do not change. So the server stores them: the KV cache. Each layer keeps a key vector and a value vector for every earlier token (per KV head). Each new step computes only the new token's key and value, appends them and reuses the rest.

Token 1 of 5

cached token	key (per layer, per KV head)	value
The ← new this step	k1	v1

step 1: compute k, v for "The" only, append them, then its query attends to k1…k1 and v1…v1 cache size = 1 tokens × 32 numbers = 32 numbers (tiny model: 2 layers × 2 KV heads × 4 × 2) without a cache this step would re-project 1 tokens; with it, 1

The cache holds the key and value vectors, not the attention probabilities, and not a copy of the answer. Each new query reads the stored keys and values.

Step through five tokens in a tiny model and watch the cache grow.

Predict first. Does the ordinary KV cache primarily store query-key probability matrices?

NumPy

cache["K"].append(x @ Wk); cache["V"].append(x @ Wv)   # new token only
out = softmax(K_cache @ q / sqrt(d)) @ V_cache            # q is the new token's query

Chapter context

Causal attention can reuse keys and values for earlier tokens rather than recomputing them at every generated step. This saves work but creates state that grows with the number of retained token positions.

TRY IT YOURSELF

Part 1. Do it hands-on

Count what a tiny cache holds, then implement a real cached attention step and check that it matches recomputing from scratch.

Token 1 of 5

cached token	key (per layer, per KV head)	value
The ← new this step	k1	v1

The cache holds the key and value vectors, not the attention probabilities, and not a copy of the answer. Each new query reads the stored keys and values.

numbers per token cache numbers after 5 tokens projections without a cache for 5 steps

To do: A tiny model: 2 layers, 2 KV heads per layer, head size 4. Each token stores one key and one value vector per KV head per layer. How many numbers per token?
To do: After 5 tokens, how many numbers does the cache hold?
To do: Without a cache, generating each new token would recompute keys and values for ALL tokens so far. For 5 tokens that is 1 + 2 + 3 + 4 + 5 token projections. With a cache only the new token is projected each step (5 in total). How many projections without the cache?

Part 2. Do it in code

attend_last(X, Wq, Wk, Wv): attention output for the LAST row of X (T, d) using all rows as keys and values, computed from scratch (no cache): q = X[-1] @ Wq, K = X @ Wk, V = X @ Wv, softmax(K @ q / sqrt(d_k)) @ V. Return a vector.
new_cache(): an empty cache: a dict with empty lists K and V.
step_with_cache(x, cache, Wq, Wk, Wv): process ONE new token vector x. Compute its k and v, append them to the cache, and attend with its q over all cached keys and values. Return the output vector.
cache_numbers(layers, kv_heads, head_dim, tokens): 2 · layers · kv_heads · head_dim · tokens.
projections_without_cache(T) and projections_with_cache(T): number of token key/value projections computed over T generation steps (1+2+…+T versus T).

import numpy as np

def attend_last(X, Wq, Wk, Wv):
    pass

def new_cache():
    pass

def step_with_cache(x, cache, Wq, Wk, Wv):
    pass

def cache_numbers(layers, kv_heads, head_dim, tokens):
    pass

def projections_without_cache(T):
    pass

def projections_with_cache(T):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Does the ordinary KV cache primarily store query-key probability matrices?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a uniform full-attention decoder:

M_{KV}=2L H_{KV}d_h T B b.

L is layers, H_KV KV heads, d_h head dimension, T retained tokens per sequence, B equal-length unshared sequences and b bytes per value. The factor 2 counts keys and values. For unequal lengths, replace TB by the sum of retained lengths.

Worked example

For L=32, H_KV=8, d_h=128 and b=2, each retained token needs 2×32×8×128×2=131,072 bytes, or 128 KiB. At T=4096, one sequence needs 512 MiB. Four equally long independent sequences need 2 GiB, before allocator overhead and other memory.

Reference code

The following example uses the Python standard library and does not call a remote service.

layers,heads,dim,bytes_per_value=32,8,128,2
per_token=2*layers*heads*dim*bytes_per_value
assert per_token == 131072
assert per_token/1024 == 128
assert per_token*4096/1024**2 == 512
assert per_token*4096*4/1024**3 == 2

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Use KV heads rather than query heads. Distinguish decimal GB from binary GiB. Hybrid, sliding-window and quantized-cache architectures require their own detailed accounting.