I kept reading that a long context "needs a lot of memory" without anyone saying how much, so I sat down and did the arithmetic. It turns out to be one multiplication, and the only hard part is knowing which two numbers to put in it.
bytes = 2 × bytes_per_element
× num_hidden_layers
× num_key_value_heads
× head_dim
× context_length
× batch_size
The leading 2 is there because every layer caches a key tensor
and a value tensor. bytes_per_element is 2 for bf16 or
fp16, 1 for an fp8 cache, and 0.5 for int4. Everything else comes straight
out of the model's config.json.
Use num_key_value_heads, not
num_attention_heads. Grouped-query attention keeps a
number of key/value heads that is more than one and fewer than the number of
query heads, and it is the key/value count that decides the cache
(Ainslie et al., 2023).
Multi-query attention, the extreme case, was introduced precisely to shrink
the keys and values loaded during incremental decoding
(Shazeer, 2019). Get this term
wrong on an 8-key-value-head model with 32 query heads and your answer is
four times too big.
Read head_dim from the file if the file has it.
The habit of computing hidden_size / num_attention_heads is a
shortcut, and it is only right when the config happens to agree. Several
published configs do not.
Its config gives 32 layers, 8 key/value heads and a head dimension of 128. At bf16 that is:
2 × 2 × 32 × 8 × 128 = 131,072 bytes per token
= 128 KiB per token
So an 8,192-token context for one sequence needs exactly 1 GiB of cache, on top of the weights. Eight concurrent sequences at that context need 8 GiB, which is usually the number that decides whether a deployment fits.
| Head count used | Per token | 8k context |
|---|---|---|
| num_key_value_heads = 8 (correct) | 128 KiB | 1 GiB |
| num_attention_heads = 32 (wrong) | 512 KiB | 4 GiB |
It is not really about capacity. During decoding the whole cache is read for every single token generated, so its size sets how much memory bandwidth each token costs. That is the reason the cache, rather than the weights, is what limits how many sequences you can serve at once.
I checked my numbers against the KV cache calculator on ml0x.com, which runs the same formula client-side over the published configs of about nineteen models and, usefully, will also show you what the wrong head count would have told you. It agreed with the 128 KiB per token above, which is why I am confident enough to publish this.
Caveats worth keeping: this counts the cache only, not weights, activations or the allocator's overhead, and real serving stacks vary with paged attention, sharding and fused kernels. It is a floor, not a budget.