Attention Variants and Memory Efficiency

Two transformer models with the same parameter count can have very different memory requirements during inference. A major reason is the structure of their attention layers. The number of key-value heads, the representation stored for each token, and the number of earlier tokens retained by each layer determine how much KV cache a request needs. When serving an existing checkpoint, you do not choose whether its layers use multi-head attention, grouped-query attention, multi-query attention, sliding-window attention, or multi-head latent attention. The trained model has already made those choices. Your task during inference is to inspect the architecture, calculate its memory requirements, and use a compatible cache layout and attention kernel. You can compare checkpoints before deployment, but you cannot replace one attention variant with another and expect identical behavior. In this post, you will learn about: How attention architecture determines persistent KV cache size How grouped heads, sliding windows, and latent attention change memory needs How to estimate serving capacity from a checkpoint and runtime settings Let’s get started. Attention Variants and Memory Efficiency.Photo by Chris Boland. Some rights reserved. Overview This post is divided into eight parts; they are: The Checkpoint Fixes the Attention Architecture A Closer Look at Attention Memory Multi-Head Attention Multi-Query and Grouped-Query Attention Sliding-Window and Mixed Attention Multi-Head Latent Attention What the Inference Runtime Can Change From Cache Size to Serving Capacity The Checkpoint Fixes the Attention Architecture The checkpoint fixes the model configuration; the inference engine does not change it. The configuration and weights define the attention projections, numbers of query and key-value heads, local-attention pattern, and latent dimensions. Replacing MHA with GQA after training modifies the model and may require additional training to recover quality. Replacing full attention with a sliding window would discard context that the trained model expects to see. The attention variant is still relevant to deployment in two ways: It determines the memory and bandwidth cost of serving a selected checkpoint. It can be one criterion when choosing among trained models with acceptable quality. Use the PyTorch and transformers environment set up in Decoding Strategies and Output Control. Models in the Hugging Face transformers library have corresponding configuration objects. You can inspect common fields as follows: from transformers import AutoConfigconfig = AutoConfig.from_pretrained("Qwen/Qwen2.5-0.5B")head_dim = getattr( config, "head_dim", config.hidden_size // config.num_attention_heads,)print("layers:", config.num_hidden_layers)print("attention heads:", config.num_attention_heads)print("key-value heads:", getattr(config, "num_key_value_heads", None))print("hidden size:", config.hidden_size)print("head dimension:", head_dim)print("uses sliding window:", getattr(config, "use_sliding_window", None))print("sliding window:", getattr(config, "sliding_window", None)) Field names and layer patterns vary across model families. A configuration may also contain a value for a feature that another field disables, so inspect both the value and its enable flag. Some models use one attention type in every layer, while others alternate local and global layers. The configuration and model implementation are the source of truth. A Closer Look at Attention Memory Attention creates two different kinds of inference memory: Persistent memory remains between decode steps. The KV cache is the main example. Temporary memory is workspace used during a forward pass, such as attention-score tiles, softmax statistics, and kernel scratch buffers. Training retains many forward-pass intermediates for the backward pass, whereas inference can discard temporary values immediately after use. Persistent memory grows with the number of live requests and cached tokens, so it often determines maximum concurrency. A memory-efficient attention algorithm can reduce temporary workspace without changing the logical KV cache required by the model. For a standard key-value representation, cache memory for one request is: $$ M^{\text{KV}} = 2S\sum_{i=1}^{L} T_i H_i^{\text{KV}}D_i $$ Here, $T_i$ is the number of tokens retained by layer $i$, $H_i^{\text{KV}}$ is its number of key-value heads, $D_i$ is its head dimension, and $S$ is the number of bytes per cache element. The factor of two represents keys and values. This layer-wise form also handles models with different attention patterns across layers. When every layer uses the same configuration and retains the full context, the expression simplifies to: $$ M^{\text{KV}}=2SLTH^{\text{KV}}D $$ For a batch of equal-length requests, multiply this memory by the batch size. An implementation that represents the batch as padded sequences may also allocate cache for padding positions. A ragged implementation can instead allocate only the cache required by each request. Multi-Head Attention Multi-head attention projects each token into queries, keys, and values. If a layer has $H$ heads and each head has dimension $D$, the tensors have shapes: Q: [batch, H, query_tokens, D]K: [batch, H, key_tokens, D]V: [batch, H, key_tokens, D] Each query head attends to the corresponding key and value head. During autoregressive decode, the keys and values from earlier tokens are stored in the KV cache. Queries are used immediately and do not need to remain in the cache. Ignoring allocator metadata, KV cache memory for a batch of equal-length requests is: $$ M^{\text{KV}} = 2BSLTH^{\text{KV}}D $$ Here, $B$ is the batch size, $L$ is the number of layers, $T$ is the cached sequence length, $H^{\text{KV}}$ is the number of key-value heads, $D$ is the head dimension, and $S$ is the number of bytes per element. The factor of two accounts for keys and values. This simplified formula applies when all layers retain the same number of tokens and have the same KV shape. For ordinary multi-head attention, $H^{\text{KV}}=H$. This cache grows linearly with batch size and sequence length. During decode, each layer must read cached keys and values to produce the next token. KV cache size therefore affects both capacity and memory bandwidth. Multi-Query and Grouped-Query Attention Multi-query attention, or MQA, gives all query heads a single shared key-value head. The query representation remains wide, but the KV cache becomes much smaller. Grouped-query attention, or GQA, is an intermediate design between MHA and MQA. Several query heads share each key-value head. For example, 32 query heads and eight key-value heads create four query heads per key-value group. The following example illustrates this grouping: MHA: 32 query heads, 32 key-value headsGQA: 32 query heads, 8 key-value headsMQA: 32 query heads, 1 key-value head Relative to MHA, GQA reduces KV cache storage and the parameter count of the key and value projections by a factor of four; MQA reduces them by a factor of 32. These factors match the reduction in key-value heads in this example. During attention computation, each shared key-value head must serve multiple query heads. A simple implementation repeats them: import torchdef repeat_kv(x, num_query_heads): # x: [batch, num_kv_heads, tokens, head_dim] num_kv_heads = x.size(1) assert num_query_heads % num_kv_heads == 0 repeats = num_query_heads // num_kv_heads return x.repeat_interleave(repeats, dim=1)k = torch.randn(2, 8, 128, 64)k_for_queries = repeat_kv(k, num_query_heads=32)print(k_for_queries.shape) This code is useful for understanding the shape, but an efficient attention kernel should reuse the keys and values internally without physically copying them. Sliding-Window and Mixed Attention Full causal attention allows a new token to attend to every earlier token. Sliding-window attention limits attention to the most recent $W$ tokens. In a pure sliding-window layer, a token’s key and value are no longer needed once that token falls outside the window. For query position $t$, the allowed key positions are approximately $\max(0,t-W+1), \ldots, t$. The attention work per decode token stops growing once the context exceeds $W$. The layer’s cache can also be implemented as a ring buffer with capacity $W$. This bounds memory for that layer even when the total conversation is longer. A circular cache can overwrite its oldest slot: 1234567891011121314151617181920212223 class SlidingKVCache: def __init__(self, window_size): self.window_size = window_size self.items = [None] * window_size self.length = 0 def append(self, key, value): slot = self.length % self.window_size self.items[slot] = (key, value) self.length += 1 def ordered(self): count = min(self.length, self.window_size) start = max(0, self.length - count) return [ self.items[index % self.window_size] for index in range(start, self.length) ]cache = SlidingKVCache(window_size=3)for position in range(5): cache.append(key=position, value=position * 10)print(cache.ordered()) # [(2, 20), (3, 30), (4, 40)] Production caches store tensors in contiguous blocks and avoid Python objects, but the indexing rule is the same. Position information still uses the token’s logical position, not its slot in the ring buffer. The trade-off is the loss of direct access to distant tokens within a local-attention layer. The model must be trained or adapted for this pattern. Some models alternate local and global-attention layers, preserving a route for long-range information while reducing total cost. Their cache memory must be calculated layer by layer. Assume that these layers use the same number of KV heads, head dimension, and cache precision. For a model with $L_1$ local layers using a window $W$ and $L_2$ full-attention layers retaining a context of length $T$, the standard KV representation requires approximately: $$ M^{\text{KV}}=2H^{\text{KV}}DS\left(L_1\min(T,W)+L_2T\right) $$ This mixed design does not bound total cache memory when full-attention layers remain. Instead, it reduces the rate of growth: each token beyond the local window adds cache only to the full-attention layers. The runtime must follow the model’s exact eviction rule. A cache allocator cannot discard old KV entries from a full-attention layer merely because other layers use a sliding window. Likewise, allocating full-context cache for every local layer is correct but may waste a large amount of memory. Multi-Head Latent Attention Multi-head latent attention, or MLA, compresses the key-value representation into a smaller latent vector. Instead of storing a full key and value for every head, the model stores a compressed state from which head-specific information is derived during attention. The cache reduction depends on the latent dimension and on which positional components must be cached separately. The decode step trades smaller cache reads for additional projection work. This can be favorable because decode is often limited by memory bandwidth: spending more arithmetic to move fewer bytes is useful when arithmetic units would otherwise be idle. MLA is not a generic compression wrapper around an arbitrary MHA checkpoint. Its projections and positional treatment are part of the trained architecture. Efficient inference also needs specialized kernels that operate on the compressed representation without expanding a large intermediate tensor in memory. When estimating an MLA cache, use the exact tensors stored by the implementation. Substituting num_key_value_heads into the MHA formula can be misleading because the cached representation is structurally different. What the Inference Runtime Can Change It is useful to separate four questions: What tokens may attend to what tokens? Full attention and sliding-window attention answer this question. How many key-value representations are learned? MHA, GQA, and MQA answer this question. What representation is cached? Standard KV tensors and MLA latent state answer this question. How are operations executed and memory managed? FlashAttention, paged attention, fused kernels, and cache allocators answer this question. The checkpoint answers the first three questions at the logical level. The inference runtime controls the fourth and may alter the cache’s physical storage. Changing the logical attention behavior, however, requires a model conversion that may affect quality. FlashAttention computes exact attention using tiles so that large intermediate score matrices do not have to be written to accelerator memory. It reduces temporary attention memory and memory traffic, but it does not change the persistent logical cache or number of KV heads. Paged attention stores a request’s KV cache in fixed-size blocks that need not be contiguous in physical memory. It improves allocation and sharing, but it does not by itself reduce the logical number of cached values. Cache precision changes the number of bytes used by each stored element. Offload changes where the cache resides. Prefix sharing prevents duplicate blocks for matching prefixes. These techniques change physical memory consumption or placement while preserving the attention architecture. Kernel fusion combines neighboring operations to reduce memory traffic and launch overhead. It may be used with MHA, GQA, or other variants when an appropriate kernel exists. The kernel must still implement the head mapping, attention window, positional treatment, and cached representation expected by the checkpoint. The next post studies these implementation techniques in detail. From Cache Size to Serving Capacity Cache bytes per token are more useful when converted into a deployment limit. The following code handles a uniform standard KV cache: 12345678910111213141516171819202122232425262728 def kv_bytes_per_token( num_layers, num_kv_heads, head_dim, bytes_per_element,): return 2 * num_layers * num_kv_heads * head_dim * bytes_per_elementdef maximum_equal_length_requests(cache_budget_bytes, context_length, **model): bytes_per_request = context_length * kv_bytes_per_token(**model) return cache_budget_bytes // bytes_per_requestmodel_config = { "num_layers": 32, "num_kv_heads": 8, "head_dim": 128, "bytes_per_element": 2,}per_token = kv_bytes_per_token(**model_config)requests = maximum_equal_length_requests( cache_budget_bytes=24 * 1024**3, context_length=8192, **model_config,)print("KV bytes per token:", per_token)print("maximum equal-length requests:", requests) The cache budget is the accelerator memory remaining after reserving space for model weights, runtime workspaces, communication buffers, and safety headroom. The result is a theoretical upper bound because block rounding, fragmentation, allocator metadata, and unequal request lengths consume additional capacity. When comparing candidate checkpoints that meet the quality requirement, record: KV bytes per token and weight memory Local and global layer patterns Maximum supported and maximum useful context length Maximum concurrency under a fixed memory budget Prefill and decode performance at representative lengths Availability of kernels that support the exact attention configuration Realistic traffic gives a more useful capacity estimate than the advertised maximum context alone. GQA’s cache advantage may be small for one short request but decisive for many long conversations. A sliding-window model may have bounded cache in local layers while its global layers continue to grow. An MLA checkpoint may save cache bytes but require a specialized implementation to avoid expanding its latent representation. Further Reading Below are some resources you may find useful: Ashish Vaswani et al. (2017). Attention Is All You Need. This paper introduces the transformer and multi-head attention. Noam Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need. This paper introduces multi-query attention for faster autoregressive decoding. Joshua Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. This paper describes grouped-query attention and methods for converting MHA checkpoints through additional training. Albert Q. Jiang et al. (2023). Mistral 7B. This paper describes a model using grouped-query and sliding-window attention. DeepSeek-AI et al. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. This paper describes multi-head latent attention and its compressed KV cache. Summary In this post, you learned that the trained checkpoint fixes the attention architecture used during inference. Your task is to inspect that architecture and account for the state it requires. MHA stores one key and value per query head, GQA and MQA reduce the number of KV heads, sliding-window layers limit retained positions, and MLA stores a compressed latent representation. You also learned to separate persistent KV cache from temporary attention workspace and to convert cache bytes per token into an estimate of serving capacity. FlashAttention, paged attention, cache quantization, offload, and fusion can change how memory is used, but they do not replace the model’s learned attention variant. No comments yet.