Watching a KV cache grow (part 2)

In part 1 I watched a KV cache fill up one token at a time on my laptop, and worked out that SmolLM2-135M’s cache costs 23,040 bytes per token. The post ended with a question: every decode step has to read the weights and the whole cache, so as the cache grows, does each step get slower in proportion to the extra bytes?

The short answer is yes, in a straight line, but each step moves the cache through memory about 3 times rather than once. With a different cache class it’s closer to 7.

Same machine (M4 Pro, 24 GB), same model, same versions. The new script is kv_bandwidth.py in the same repo.

First, how fast is the memory? Link to heading

In part 1 I used Apple’s quoted 273 GB/s for a back-of-envelope estimate and said I didn’t measure it. This time I did: copy a 1 GiB tensor on the GPU 20 times and time it.

Measured copy bandwidth: 212 GB/s (read + write)
Apple's quoted figure:   273 GB/s

That’s about 78% of the quoted number. Everything below uses 212.

The experiment Link to heading

For each context length from 512 to 32,768 tokens:

1. Fill a fresh cache with one prefill of that many (random) tokens. Not timed.
2. Run 5 decode steps to warm up.
3. Time 50 decode steps one by one, with a GPU sync before and after each, and keep the median.

SmolLM2 was trained on up to 8,192 tokens, so the text it generates past that is garbage. The bytes each step moves are still real, and bytes are what I’m measuring.

I ran it twice, once with Transformers’ default DynamicCache and once with StaticCache. StaticCache preallocates the whole buffer up front, so on paper it avoids the copying DynamicCache does. It didn’t come out faster.

At the end the script repeats the first measurement to check the laptop didn’t throttle. Drift was +2.2%.

Drag the slider or hover the chart. The dashed line is where the model’s trained context ends.

Both lines are straight, so latency grows linearly with context. StaticCache also gets slower much faster than DynamicCache: at 32k it’s 32 ms per step vs 18 ms.

Reading the line Link to heading

A straight line has two numbers, and both mean something here.

Where it starts (the intercept) is the cost of a decode step with an empty cache. Both caches land at about 8 ms. That matches part 1’s ~8 ms, and it’s mostly fixed per-call overhead. Reading the 269 MB of weights at 212 GB/s is only about 1.3 ms of it.

How steep it is (the slope) tells you how many extra milliseconds each extra byte of cache costs. Turn that into GB/s and compare it with the 212 GB/s the memory can do:

dynamic:
  empty-cache cost (intercept): 8.01 ms
  effective bandwidth (slope):  77 GB/s of KV per extra ms
  passes over the cache/step:   2.7
  fit R^2:                      0.9970
static:
  empty-cache cost (intercept): 8.74 ms
  effective bandwidth (slope):  33 GB/s of KV per extra ms
  passes over the cache/step:   6.4
  fit R^2:                      0.9983

“Passes” is 212 divided by the effective number. If each decode step read the cache exactly once, it would be about 1. It’s 2.7 for DynamicCache and 6.4 for StaticCache. So each step is pushing the cache through memory several times.

Put another way: if a step read the cache once, going from 512 to 32k tokens would add about 3.5 ms. The measured increase is 9.7 ms for DynamicCache and 22.6 ms for StaticCache.

Where the extra passes come from Link to heading

DynamicCache appends the new token’s K and V with torch.cat. That allocates a new tensor and copies the entire old cache into it, for every layer, on every step. So one step reads the cache, writes a full copy, and then attention reads it again. That’s 3 passes, and the fit says 2.7.

StaticCache took more digging. I wrapped the attention function for one extra step and printed what it received:

dynamic  keys [1, 3, 1025, 64]  mask None          repeat_kv calls: 0
static   keys [1, 3, 4096, 64]  mask [1,1,1,4096]  repeat_kv calls: 60

First, attention runs over the whole preallocated buffer (4,096 slots for 1,025 real tokens), with a mask to hide the empty ones. In my sweep the buffer is only ~56 slots bigger than the data, so this part barely matters here, but it would if you preallocated generously.

Second, because a mask is passed, Transformers turns off its grouped-query attention shortcut and calls repeat_kv to expand K and V from 3 heads to 9 before attention, on every step (60 calls = K and V × 30 layers). So a step reads the 3-head cache, writes a 9-head copy (3× the bytes), and attention reads that. That’s about 7 passes, and the fit says 6.4.

In part 1 the whole point of GQA was that 3 KV heads instead of 9 makes the cache 3× smaller. With StaticCache on this setup, that saving exists in memory but not in traffic: every step rebuilds the 9-head version anyway.

On paper it’s worse than that: a model with 9 KV heads would have nothing to expand, so with StaticCache it would move less data per step than this 3-head one. I didn’t measure that.

As far as I know, StaticCache mainly exists so tools like torch.compile can work with fixed shapes, which I’m not using. So this doesn’t mean StaticCache is broken. On MPS with Transformers 5.17 and no compile, it’s the slower choice.

What that does to the crossover Link to heading

Part 1’s crossover (cache bigger than weights at ~11,700 tokens) was about size. What matters for speed is traffic. If every step moves the cache 2.7 times, cache traffic passes weight traffic at about 4,300 tokens. At 6.4 passes it’s about 1,800 tokens. Those are estimates from the line fit, but the point stands: the cache starts to dominate much earlier than its size alone suggests.

Takeaways for infra folks Link to heading

• Decode latency grew in a straight line with context. A straight line doesn’t prove memory is the bottleneck on its own, since the work per step also grows with context. What convinced me is that the slope lines up with counting the copies.
• But the bytes that matter are the bytes moved, not the bytes stored. Here the software moved the cache 2.7 to 6.4 times per step. If you’re sizing memory bandwidth for inference, the cache size is a floor, and the multiplier comes from the framework.
• The same model on the same laptop varied by almost 2× at 32k depending on one class name.

Caveats Link to heading

One small model, one laptop, one version of Transformers. The “passes” number is inferred from a line fit, not from a memory trace. The timestamps in the lab are there so I can line them up with an Instruments / Metal System Trace capture later and check it properly. Beyond 8k tokens the model is outside what it was trained for; the timings are real, but the output isn’t.

Update: removing the extra copies Link to heading

The explanation above rests on counting: the passes lined up with the copies I found in the code. But numbers lining up isn’t proof, so I removed the suspected causes and measured again.

• For DynamicCache, I replaced the copy-on-every-append with a buffer that’s allocated once and written in place.
• For StaticCache, I patched the attention call to use only the filled slots and skip the 3-to-9 head expansion.

Both still produce the same tokens as before, bit for bit. At 32k tokens, measured in the same run:

Both fixed versions come out at about 1.3 passes over the cache per step, down from about 3 and 6. So the extra copies were the cost, and without them decode reads the cache roughly once per step. I haven’t pinned down the last 0.3.

Both fixes are in the repo (--cache all). They patch library internals at runtime, which is fine for a lab but not something I’d ship.

Next, I moved to a datacenter GPU to ask a different question: is moving a cache cheaper than rebuilding it?

Part 3 is here.