Glenn Lockwood’s What are KV caches, really? explains prefill, decode, and why a KV cache exists, and then takes apart some vendor speedup claims for KV offload. It’s a good read and it got me curious.
I’m an infra person, not an ML person, so I wanted to see a KV cache for real, the way I’d look at any other buffer: what’s in it, how it grows, and how big it gets.
So I built a small lab on my laptop. It runs SmolLM2-135M-Instruct with PyTorch on an M4 Pro (MPS backend, bf16), hooks the Q/K/V projections of the first transformer layer, and does generation by hand instead of calling model.generate(), so every piece of state is visible. Code is here.
This is part 1. It covers the mechanics and a few things I noticed along the way. Part 2 will grow the context and look at bandwidth.
Setup Link to heading
Machine: MacBook Pro (Mac16,8)
Chip: Apple M4 Pro, 14-core CPU (10 performance + 4 efficiency), 20-core GPU
Memory: 24 GB unified (CPU and GPU share it)
OS: macOS 26.6.2
Software: Python 3.11.15, PyTorch 2.14.0 (MPS backend), Transformers 5.17.0
Apple quotes 273 GB/s of memory bandwidth for the M4 Pro. I use that number once, later, for a back-of-envelope estimate. I didn’t measure it.
The model Link to heading
Everything below is read from the model config at runtime:
Layers: 30
Hidden size: 576
Attention heads: 9
KV heads: 3
Head dimension: 64
Model dtype: torch.bfloat16
Attention type: GQA (grouped-query attention), 3 query heads share each KV head
The 9 vs 3 matters later. The prompt is "The capital of France is", fed as raw text without the chat template (the template would wrap 5 tokens in ~30 tokens of system prompt). It tokenizes to 5 tokens, no BOS:
idx id text
0 504 'The'
1 3575 ' capital'
2 282 ' of'
3 4649 ' France'
4 314 ' is'
Prefill Link to heading
Prefill is one call with the whole prompt:
outputs = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=True)
cache = outputs.past_key_values
Forward hooks on q_proj, k_proj, v_proj in layer 0 record shapes. They only append metadata to a list, and I print it after the call, so the terminal I/O stays out of the timing:
q_proj out [1, 5, 576] in [1, 5, 576] = [1, 9 heads, 5 tok, 64]
k_proj out [1, 5, 192] in [1, 5, 576] = [1, 3 heads, 5 tok, 64]
v_proj out [1, 5, 192] in [1, 5, 576] = [1, 3 heads, 5 tok, 64]
All 5 tokens go through at once. K and V come out 192 wide instead of 576 because there are only 3 KV heads. The cache after prefill, for layer 0:
K shape: [1, 3, 5, 64]
V shape: [1, 3, 5, 64]
[0] batch = 1
[1] kv_heads = 3
[2] sequence_length = 5
[3] head_dim = 64
The logits at the last position pick the first generated token, ' Paris'.
Decode Link to heading
Each decode call feeds exactly one token plus the cache from the previous call:
for step in range(1, decode_steps + 1):
token_id = generated[-1]
kv_before = cache.get_seq_length() # DynamicCache mutates in place, read it first
attention_mask = torch.cat([attention_mask, attention_mask.new_ones(1, 1)], dim=1)
outputs = model(input_ids=torch.tensor([[token_id]], device=device),
attention_mask=attention_mask, past_key_values=cache, use_cache=True)
cache = outputs.past_key_values
generated.append(outputs.logits[0, -1].argmax().item())
Same hooks, decode step 1:
Input token text: ' Paris'
KV length before: 5
q_proj out [1, 1, 576] in [1, 1, 576] = [1, 9 heads, 1 tok, 64]
k_proj out [1, 1, 192] in [1, 1, 576] = [1, 3 heads, 1 tok, 64]
v_proj out [1, 1, 192] in [1, 1, 576] = [1, 3 heads, 1 tok, 64]
Output token text: '.'
KV length after: 6
Sequence dimension 5 during prefill, 1 during every decode step. The K/V for the first five tokens are never recomputed; the new token computes its own K and V, they get appended, and its query attends over all 6 cached keys. After 10 steps:
KV progression:
0 -> 5 prefill
5 -> 6 decode 1
6 -> 7 decode 2
...
14 -> 15 decode 10
Full text: 'The capital of France is Paris. Paris is a major city in France, known'
Here’s the same run as a slider:
Drag the slider or use the arrows. Real tokens and timings from the M4 Pro run. Orange is work done in that call, green came out of the cache.
One small thing that tripped me up while writing the spec for this: 10 decode calls produce 11 tokens, not 10. Prefill already picks ' Paris', and each decode call picks one more. The last one (' known') is never fed back, so it has no cache entry and the cache ends at 5 + 10 = 15, not 16.
To make sure the hand-rolled loop is correct, I compared it with model.generate() using greedy decoding. Same tokens on MPS and CPU.
How big is it Link to heading
I summed the bytes of all 60 K/V tensors (30 layers × K and V) and divided by the token count:
Total KV bytes: 115200 (0.110 MiB)
Cached tokens: 5
bytes_per_cached_token: 23040
Formula check: 30 layers * 2 (K,V) * 3 kv_heads * 64 head_dim * 2 bytes = 23040
So each token costs:
bytes per token = layers × 2 (K and V) × kv_heads × head_dim × bytes per element
This is the number that made the rest click for me.
Architecture matters a lot here. With plain multi-head attention (9 KV heads instead of 3), SmolLM2 would need 69,120 bytes per token, 3× more. Fewer KV heads means a smaller cache, which as far as I can tell is the main reason GQA exists.
The number also adds up fast. SmolLM2’s weights are 269 MB, and at 23,040 bytes per token its cache passes that at about 11,700 tokens. That’s past the 8,192 tokens SmolLM2 was trained on, so for this model it’s a number on paper.
For a bigger model, plug in the published config. Llama 3 70B has 80 layers, 8 KV heads, head_dim 128, which in bf16 is 327,680 bytes (320 KiB) per token. At 112,000 tokens, the context size in one of the vendor examples in Glenn’s post, that’s about 34 GiB of cache for one conversation. That’s what the KV offload systems in that post are moving to and from SSDs and remote storage.
Recomputing the cache Link to heading
Glenn’s post points out that “keys and values for each output token never change after they’re generated”, which is what makes them cacheable. I got curious what that looks like in practice. So after the decode loop, the lab takes the same 15 tokens, runs them through the model in one uncached call, and compares the resulting K/V against the cache built one token at a time:
fed_ids = prompt_ids + generated[:-1]
fresh = model(input_ids=torch.tensor([fed_ids], device=device), use_cache=True)
for cached, recomputed in zip(cache.layers, fresh.past_key_values.layers):
diff = (cached.keys.float() - recomputed.keys.float()).abs()
...
bf16 on MPS:
Bit-identical elements: 87812 / 172800 (50.8%)
Max |diff|, decode positions: 1.25e-01 (largest |K| in cache: 19.2)
Greedy tokens from the full pass match decode: False
generated token 5: decode chose ' a', full pass chose ' the' (full-pass top-2 logit gap 0.1250)
generated token 6: decode chose ' major', full pass chose ' city' (full-pass top-2 logit gap 0.1250)
Only about half of the cached K/V elements match the recomputed ones bit for bit. Layers 0 and 1 are identical, and the error grows as you go deeper.
Two of the eleven generated tokens also come out different (the lab counts generated tokens from 1, so token 5 is ' a'). The full pass scores decode’s own token sequence, so each position is an independent check: after “Paris is” it picks ' the' where decode picked ' a', and after “Paris is a” it picks ' city' where decode picked ' major'.
The top two scores were 22.375 vs 22.25 in the first case and 20.875 vs 20.75 in the second. In both cases that’s as close as bf16 can store two numbers without them being equal, so rounding picked the winner.
Same thing on CPU (bf16): 45.5% bit-identical, one token flipped. The results repeat exactly run to run, so the same calls always give the same bits.
What differs is that a 15-token call and a 1-token call run different kernels, and my guess is they add the same numbers up in a different order. With floats, a different order can give a slightly different sum. I haven’t tracked down which op diverges first.
In fp32 (--dtype float32) the differences shrink to around 0.00001 and every token matches.
So recomputing gets you approximately the same cache, and in bf16 “approximately” is enough to flip a token that was that close. Neither answer is wrong.
One practical takeaway for infra folks: if you’re testing a KV offload system, diff it against a normal cached run, not a recompute. Offload is a byte copy, so a correct system should match the cached run bit for bit. A recompute baseline will show mismatches like the ones above that aren’t bugs.
The timings don’t show much (yet) Link to heading
Prefill of 5 tokens took 9.91 ms. Each decode step took about 8.3 ms (7.92 to 8.63 across the ten steps). Five times the tokens for 20% more time.
Glenn’s post describes prefill as compute-bound and decode as memory-bandwidth-bound. At this size neither shows up in the timings. Back of the envelope: 269 MB of weights at 273 GB/s is about 1 ms, so most of the 8 ms is going somewhere else. My guess is fixed overhead per call, since 30 layers means a lot of small GPU kernel launches driven from Python.
That’s fine for part 1, where the point was mechanics. It does mean that anyone quoting tokens/sec from a toy setup like this one is mostly measuring framework overhead.
Next Link to heading
Every decode step has to read all the weights and all of the cache. For SmolLM2 that really is all the weights: it shares one matrix between the input embeddings and the output layer, so the whole vocabulary table gets read on every step to score the next token. The weights are fixed and the cache grows by 23,040 bytes per token, so if each step reads the cache once, somewhere around 11,700 tokens it becomes the bigger read.
Part 2 will hold the model fixed, grow the context, and see whether time-per-token on Apple Silicon’s unified memory actually follows (weight bytes + KV bytes) / bandwidth, and where the launch overhead stops hiding it.