Cache Arithmetic
Cache arithmetic is the memory calculation for a transformer’s KV cache: key and value tensors saved for each processed token, layer, KV head, and head dimension. It tells you how much device memory long contexts and concurrent requests consume, using the formula 2 × layers × KV heads × head_dim × seq_len × element_bytes.
The need for cache arithmetic appears during autoregressive decoding. After prefill, each new token must attend to the previous tokens. Recomputing keys and values for the whole prefix at every step would waste work, so inference engines keep them. That saved work becomes saved state, and at long context lengths the state can rival or exceed the model weights. This is why a request that looks small by parameter count can still fail to fit in memory.
Mechanically, each layer stores two tensors per token: keys and values. For ordinary multi-head attention, multiply key plus value by layer count, KV head count, head dimension, sequence length, and bytes per element. For grouped-query or multi-query attention, use the number of KV heads, not query heads. That distinction is commonly missed, and it can make a memory estimate far too pessimistic for modern grouped-query models.
The trade-off is straightforward: the cache saves repeated projection work, but grows linearly with resident tokens and batch size. Weights are paid once per model replica, while KV memory is paid per active sequence and grows as output is generated. FlashAttention can reduce attention-matrix traffic, and paged cache systems can reduce allocation waste, but neither removes the persistent key and value bytes.
Engineers meet cache arithmetic when setting max context, batch limits, admission control, tensor parallel layouts, or serving policies. It is behind paged KV caches, prefix-sharing systems, prefill/decode separation, and schedulers that pack mixed-length requests. The honest answer to “will it fit?” depends on usable HBM after weights, runtime overhead, fragmentation, temporary buffers, and how many sequences are resident at once.
Common questions
- Why does the formula have a 2 in it?
- Each processed token stores both a key vector and a value vector for attention. The query for the next token is computed on the fly, but the old keys and values must remain available so the new token can attend over the prefix. That key-plus-value pair is the leading factor of 2.
- Should I use attention heads or KV heads in the calculation?
- Use KV heads. In full multi-head attention they may match the query-head count, but grouped-query and multi-query attention share keys and values across multiple query heads. Using query heads for those models overstates cache memory. This is one of the most common mistakes in long-context capacity planning.
- Why does KV cache dominate at long context?
- Because it grows with every resident token and every concurrent sequence, while weights are fixed once loaded. During decode, the model also repeatedly reads the existing keys and values for each new token. At long context, both memory capacity and memory bandwidth can become governed by KV rather than by parameter storage.
- Do paged attention or FlashAttention make the KV cache smaller?
- Not in the basic arithmetic. Paged attention reduces wasted or fragmented allocation by storing KV in manageable blocks. FlashAttention improves attention kernel IO and avoids materialising large intermediate attention matrices. The persistent cache is still keys plus values across layers, KV heads, head dimension, tokens, and element size.