KV caches, grouped heads, and compute budgets

When a request starts, the model processes the prompt to construct hidden states and attention keys and values. This prompt-processing stage is often called prefill. During generation, the model repeatedly processes a new token and attends to the previous context. A key/value cache retains earlier representations so they need not be recomputed from scratch at every generation step.

Caching works for a causal decoder because an earlier position's representation does not change when later tokens are appended. The new token still needs to compare its query with relevant cached keys and combine values. Caching removes repeated prefix computation; it does not make long-context attention free. Memory traffic, cache organization, and batch scheduling can strongly influence token-generation latency.

For a conventional dense KV layout, cache bytes are approximately 2 times layers times batch size times sequence length times KV heads times head dimension times bytes per element. The factor two accounts for keys and values. This estimate excludes metadata, padding, allocator overhead, temporary activations, and model weights. State those exclusions when using the number for capacity planning.

For 24 layers, batch size 1, 4,096 cached tokens, 8 KV heads, head dimension 64, and two bytes per element, the cache is 201,326,592 bytes, or 192 MiB. Doubling sequence length doubles this cache estimate. Increasing batch size to four quadruples it. A service that fits one request can therefore fail when several long conversations decode concurrently, even though the model weights themselves are unchanged.

In ordinary multi-head attention, each query head has its own key and value head. Multi-query attention shares a single key/value head across query heads. Grouped-query attention lies between these cases: several query heads share each key/value head. The number of query heads can remain unchanged while the cache stores fewer key/value heads.

In the previous example, reducing KV heads from eight to two reduces the approximate cache from 192 MiB to 48 MiB. It does not imply every compute component becomes four times faster, because query projections, feed-forward layers, and other operations remain. Sharing representations is an architectural tradeoff that must be trained or adapted appropriately. You cannot simply discard heads from an arbitrary checkpoint and assume quality is preserved.

Full attention over a prompt has a quadratic number of query-key relationships in sequence length, while each incremental query compares against a growing prefix. Specialized kernels can reduce materialized intermediate memory and improve execution without changing the underlying attention definition. Cache quantization can reduce bytes per stored element, but introduces its own implementation and accuracy considerations.

Model weight memory is another term: a billion parameters at two bytes each require about two billion bytes before extra serving state. Training adds gradients, optimizer states, and activations, so an inference memory estimate does not describe training memory. Keep all units explicit, especially decimal GB versus binary GiB. A clean resource worksheet separates weights, KV cache, temporary memory, and operational headroom.

Empirical scaling research studies how loss changes with model size, training data, and compute within measured regimes. Compute-optimal training results show that allocating a fixed budget only to larger models can undertrain them; model size and token count must be considered together. These are fitted relationships under specific experimental choices, not universal guarantees about every architecture, dataset, or downstream agent task.

The code compares cache estimates for two head layouts and two context lengths. The project extends this into a budget inspector that records assumptions rather than advertising benchmark results. Use the same discipline when reading a scaling paper: identify the objective, data distribution, compute definition, experimental range, and extrapolation. A better language-model loss may help an agent, but end-to-end tool use and task completion still need their own evaluations.

The worksheet uses explicit positive integer dimensions and binary MiB. Reducing KV heads divides the cache term, while doubling tokens doubles it. This is an architectural estimate, not a measurement of a serving engine. Add model-weight and headroom terms before using it for concurrency planning.