KV cache arithmetic for a hybrid-attention model, checked against two real machines
The textbook KV formula overstates Gemma 4 31B's cache by about fifteen times at 32K context, because fifty of its sixty layers only ever look back 1,024 tokens. I work the real number for my Mac Studio, then compare the first-principles figure with what vLLM actually allocated on a rented B300, and explain the gap.
Written October 2026. The B300 and A100 numbers are from a batch run on 2026-06-07. The Mac Studio figures are from the box as configured on 2026-08-31. The formula and the dense-model table are from my inference field guide, Part 3 and Part 6.
The formula everyone knows
For a dense transformer with grouped-query attention, the key-value cache costs, per token:
KV bytes per token = 2 × layers × KV heads × head dim × bytes per value
The 2 counts keys and values. Bytes per value is 2 for a BF16 cache, 1 for FP8. Multiply by context length for one sequence. Divide the pool you have left after weights by that figure and you get how many tokens you can hold, and from that, how many users fit at a given context.
For Llama 3 70B that is 320 KB per token in BF16. At 32K context one sequence needs 10.5 GB. At 128K it needs 42 GB, which does not fit beside the weights on one 80 GB card at all. The sentence from the guide that matters is this one: if your product promises long context to many users, the KV pool decides your GPU count, not the weights.
That formula is right for Llama. It is wrong by a large factor for the model I actually run.
Gemma 4 31B is not a dense-attention model
Reading the config file for Gemma 4 31B instead of assuming:
| Field | Value |
|---|---|
| Layers | 60 |
| Layer pattern | 5 sliding, 1 full, repeating |
| Sliding layers | 50, window 1,024 tokens, 16 KV heads, head dim 256 |
| Full layers | 10, 4 KV heads, head dim 512 |
| K equals V on full layers | true, so one tensor stored, not two |
| Max context | 262,144 |
Fifty of the sixty layers only attend to the last 1,024 tokens. Their cache is bounded. It does not grow with context. Only ten layers carry a cache that grows, and those ten store a single tensor per token because the model shares keys and values on its global layers.
Worked in BF16:
A sliding layer holds keys and values for at most 1,024 tokens. Per token that is 2 × 16 × 256 × 2 bytes, which is 16 KB. Times 1,024 tokens, times 50 layers, is 800 MB per sequence, fixed, no matter how long the context gets.
A full layer stores one tensor of 4 × 512 × 2 bytes, which is 4 KB per token. Ten layers is 40 KB per token, growing with context.
| Context | Growing part | Fixed part | Per sequence | Naive formula |
|---|---|---|---|---|
| 8K | 0.33 GB | 0.8 GB | 1.1 GB | 7.9 GB |
| 32K | 1.3 GB | 0.8 GB | 2.1 GB | 31 GB |
| 128K | 5.2 GB | 0.8 GB | 6.0 GB | 126 GB |
| 256K | 10.5 GB | 0.8 GB | 11.3 GB | 252 GB |
The naive column applies the dense formula to all 60 layers at the sliding-layer shape. At 32K it is wrong by fifteen times. At the model’s full 256K context it says one sequence needs a quarter terabyte. The real figure fits on a laptop GPU. A capacity plan built on the naive number would buy the wrong hardware.
On the Mac Studio
The box is an M3 Ultra with 96 GB of unified memory at 819 GB/s. It serves Gemma 4 31B through an MLX gateway from an 8-bit bundle, and the resident weights measure about 29 GB. It also runs a text-to-speech model and the gateway’s own session cache, which I capped at 12 GB after it grew toward 62 GB on its own. Call it 45 GB spoken for, which leaves roughly 40 GB for KV.
At 32K per session that is about 19 concurrent sessions. At 128K it is 6. At the full 256K it is 3. For a household box with one to three users, long context is not a constraint on this model. On a dense 70B at 4-bit the same 40 GB holds one session at 128K.
Decode speed follows from the roofline the same way. Each token reads the weights plus the current cache. At 29 GB of weights and a 2 GB cache at 32K, 819 GB/s gives a ceiling of about 26 tokens per second, and the real number is usually 60 to 80 percent of that. The bounded sliding cache also means the per-token read barely grows with context, so decode on this model does not slow down the way a dense model does at long context.
On the rented GPUs
In June I ran this model over 104,532 articles with vLLM on Verda spot capacity. Two machines, two outcomes.
A100 80 GB, BF16 weights. The weights alone are 62 GB. What was left after workspace was a few gigabytes of cache. The run was cache-bound at about 23 concurrent requests and roughly one article per second. Ampere has no native FP8, so quantizing the weights to free pool space barely helped. The compute was never the limit. The pool was.
B300 288 GB, FP8 weights and FP8 cache. Weights drop to 31 GB. vLLM logged a KV pool of about 720,000 tokens. Throughput went to about seven articles per second with 256 to 384 concurrent requests, and the whole corpus finished in about four hours. Seven times the throughput came almost entirely from widening the pool, not from faster math.
Now the check. First principles says the growing part of this model’s cache costs 20 KB per token in FP8. A pool of roughly 220 GB after weights and workspace should hold about 11 million tokens of context. vLLM said 720,000. That is fifteen times less.
The gap is the engine’s accounting, not the model’s. vLLM’s hybrid KV cache manager groups layers by attention type and allocates per group. Its own design doc notes that the grouping heuristic pads some groups with unused layers, that this happens on the Gemma family, and that the result wastes memory. At 720,000 tokens over 220 GB the engine was effectively charging about 300 KB per token, which is close to treating the sliding layers as if they grew with context. Whether that version was ignoring the window or padding the groups, the effect was the same: the admission number I had to plan around was the engine’s, not mine.
What I take from it
Three rules, in order.
Read the architecture before you do the arithmetic. The textbook formula assumes every layer attends to everything. Hybrid sliding-window models are now common, and the formula overstates their cache by an order of magnitude.
Do the arithmetic anyway. It tells you what the hardware can hold, so you know when the engine is leaving capacity on the table. Without the first-principles number, 720,000 tokens looks like a fact about the model. It is a fact about one engine version.
Then read the engine’s number and plan on that one. vLLM prints the KV pool size in tokens and the maximum concurrency at your chosen max model length at startup. The second line is the admission limit. When the two numbers disagree by fifteen times, you have found either a flag you have not set or an engine limitation worth tracking, and either way you have found it before your users did.
Sources for the architecture are the model’s own config file, the NVIDIA Megatron Bridge documentation for Gemma 4, and the vLLM hybrid KV cache manager design document.