Two transformer models of the same size can have very different memory requirements during inference, i.e., when they generate responses. According to the source article, the main reason lies in the structure of their attention layers: the number of key-value heads, the stored representation for each token, and the number of previous tokens retained by each layer determine the amount of KV cache, i.e., the state memory used by a query.
When deploying an existing checkpoint, one does not choose the attention variant it uses, whether it’s multi-head attention (MHA), grouped-query attention (GQA), multi-query attention (MQA), sliding window attention, or multi-head latent attention (MLA). These choices were made during training. The task is to inspect the architecture, calculate memory requirements, and use a compatible cache layout and attention kernel. Replacing one variant with another does not guarantee identical behavior, as the model was trained with its original architecture.
The KV Cache, a Persistent Memory
Attention generates two types of memory. Persistent memory remains between two decoding steps, and the KV cache is the main example. Temporary memory, on the other hand, serves as workspace during a forward pass, for attention score tiles or computation buffers. Since persistent memory grows with the number of active queries and cached tokens, it often determines the maximum number of simultaneous queries.
For a batch of queries of the same length, the cache memory is expressed by the formula 2 × B × S × L × T × HKV × D, where B is the batch size, S the number of bytes per element, L the number of layers, T the cache length, HKV the number of key-value heads, and D the dimension of each head. The factor 2 corresponds to keys and values. In classic multi-head attention, the number of key-value heads equals the number of heads, and the cache grows linearly with batch size and sequence length.
Variants That Reduce the Cache
Multi-query attention (MQA) gives all query heads a single shared key-value head. Grouped-query attention (GQA) is a compromise: several query heads share each key-value head. In the article’s example, 32 query heads and 8 key-value heads form four query heads per group. Compared to MHA, GQA reduces cache storage and the number of key and value projection parameters by a factor of four, and MQA by a factor of 32, factors that correspond to the reduction in the number of key-value heads in this example.
Sliding window attention limits each token to the last W tokens. The cache can then be implemented as a circular buffer of capacity W, which bounds the memory of the layer. Some models alternate local layers and global layers. In this case, the cache can no longer be bounded in its entirety: each token beyond the local window only adds cache for the full-attention layers. As for multi-head latent attention (MLA), it compresses the key-value representation into a smaller latent vector, which reduces cache reads at the cost of additional projection computations, a trade-off the article deems favorable when decoding is limited by memory bandwidth.
What the Runtime Can Change
The checkpoint fixes the logical rules of attention, while the inference engine controls execution and memory management. FlashAttention computes exact attention by tiles to reduce temporary memory, without changing the number of key-value heads. Paged attention stores the cache in fixed-size blocks, which are not necessarily contiguous in physical memory. Cache precision modifies the number of bytes per element, offloading changes the cache location, and prefix sharing avoids duplicated blocks. These techniques change the consumption or placement of physical memory, but do not replace the attention variant learned by the model.
To estimate service capacity, the article proposes converting cache bytes per token into a deployment limit, based on the cache budget remaining after reserving space for model weights, workspaces, and safety margins. This result remains a theoretical bound, as blocking, fragmentation, and uneven query lengths consume additional capacity. The article adds that the next installment in the series will study these implementation techniques in detail.
