Autoregressive transformer inference stores key and value vectors from earlier tokens so each new token does not have to recompute them. With standard multi-head attention, every attention head has its own key and value projections, so the KV cache grows with the number of key-value heads.
Grouped-query attention changes that head layout. It keeps multiple query heads but lets several query heads share one key head and one value head. The result reduces cached key-value state without collapsing all query heads into a single shared projection.
Head grouping changes the cache dimension
Consider an attention layer with Hq query heads and Hkv key-value heads, where Hq is divisible by Hkv. The number of query heads assigned to each key-value head is:
group_size = Hq / HkvStandard multi-head attention uses one key-value head per query head, so Hkv = Hq. Multi-query attention uses one shared key-value head, so Hkv = 1. Grouped-query attention occupies the range between those layouts when 1 < Hkv < Hq.
For a new token, each query head still produces its own query vector. Query heads in the same group attend against the same cached key and value vectors. They can still produce different attention weights because their query vectors differ.
The grouping therefore shares the representation stored on the key-value side; it does not make the grouped query heads identical.
KV cache size follows key-value heads
Ignoring implementation padding and metadata, a decoder layer stores one key vector and one value vector per cached token and per key-value head. If each head has dimension d_head, the element count for a sequence of length L is proportional to:
2 * L * Hkv * d_headThe factor of two accounts for keys and values. Batch size and layer count multiply this quantity across the full inference workload.
Reducing Hkv directly reduces this head-dependent part of the cache. For example, moving from 32 key-value heads to 8 while keeping the same head dimension stores one quarter as many key-value elements per token for that layer.
That ratio describes element count, not a guaranteed end-to-end memory reduction. Allocator behavior, cache paging, quantization, tensor alignment, other model state, and runtime workspaces also contribute to device memory use.
Query capacity and cache width become separate choices
In ordinary multi-head attention, the query-head count and key-value-head count are tied. Grouped-query attention separates them.
A model can retain many query projections while using fewer key and value projections. This gives different query heads room to form different query directions even when several of them read from the same key-value representation.
The sharing constraint still matters. Query heads within a group cannot select among independently projected key and value spaces because only one key projection and one value projection exist for that group. Changing the number of groups therefore changes model parameterization, not just a runtime storage option.
For that reason, an arbitrary multi-head checkpoint cannot generally be switched to fewer key-value heads at inference time merely by changing a configuration field. The projection weights and head layout must agree with the architecture represented by the checkpoint.
Group mapping must be consistent
A common implementation arranges query heads into contiguous groups. With eight query heads and two key-value heads, a conceptual mapping can look like:
query heads 0 1 2 3 -> key-value head 0
query heads 4 5 6 7 -> key-value head 1The exact tensor operation may avoid physically repeating key and value tensors. A kernel can map each query head to the relevant key-value head during attention computation.
Another implementation may expand shared keys and values to a query-head-shaped view before calling a more general attention operation. That can be functionally equivalent if the mapping is correct, but materializing repeated tensors can erase some memory advantages during the attention calculation even though the persistent KV cache remains smaller.
Head ordering is part of the checkpoint convention. Reassigning query heads to different key-value groups changes which projected keys and values each query consumes.
Cache bandwidth can matter during decoding
During single-token autoregressive decoding, each layer forms queries for the new token and reads cached keys and values for prior tokens. As the sequence grows, those cache reads cover more stored positions.
Fewer key-value heads mean fewer distinct key and value vectors need to be stored and fetched per cached position. This can reduce memory traffic associated with the KV cache. The effect on latency depends on the attention kernel, hardware, batch shape, sequence length, cache representation, and other work in the model.
It is therefore more precise to treat grouped-query attention as reducing the amount of distinct key-value state than to claim a fixed decoding speedup. A runtime that fails to exploit shared heads can realize less of the potential bandwidth advantage.
Prefill has a different workload shape. Processing many prompt tokens at once involves larger matrix operations and may be governed by different bottlenecks than token-by-token decoding. The same architectural sharing applies, but its latency effect need not match the decode phase.
Tensor shapes expose implementation errors
Grouped-query attention is easy to describe but shape assumptions can cause subtle failures. A typical projected layout before attention might use:
Q: [batch, Hq, tokens, d_head]
K: [batch, Hkv, tokens, d_head]
V: [batch, Hkv, tokens, d_head]An attention implementation must account for Hq being larger than Hkv. A generic routine that requires equal head counts cannot consume these tensors directly unless the shared key-value heads are expanded or the routine has native grouped-query support.
The divisibility condition also needs to match the model design. If query heads are assigned evenly, Hq % Hkv must be zero. Architecture definitions can impose additional constraints, so checkpoint metadata remains the source of truth for actual head counts.
Cache serialization needs the key-value shape rather than the query shape. Allocating cache storage as though every query head had a distinct key and value head wastes memory and defeats the central storage reduction.
Sharing does not remove attention’s sequence dependence
Grouped-query attention reduces the number of cached vectors per token, but the cache still grows with cached sequence length. Holding Hkv fixed, doubling the number of retained tokens still doubles the key-value element count.
It also does not change the basic need for each new query to attend over retained positions under dense causal attention. Techniques that limit the attended context, evict cache entries, compress cache precision, or use sparse attention address different dimensions of the inference cost.
Those methods can be combined with grouped-query attention, but their effects should be accounted for separately. Fewer key-value heads reduces state per position; a shorter retained context reduces the number of positions; lower cache precision reduces bytes per stored element.
Grouped-query attention is most useful to reason about as an architectural sharing rule with a direct storage consequence. The key implementation boundary is simple: query heads may remain numerous, but the persistent cache should contain only the distinct key and value heads defined by the checkpoint.