Grouped-Query Attention Shares Key-Value Heads Across Query Heads

Grouped-query attention changes a specific part of multi-head attention: several query heads use the same key and value head. The query projections remain separate, so those query heads can produce different attention weights, but they read keys and values from a shared projected representation.

That distinction matters during inference. A decoder cache stores past keys and values, not past queries. Reducing the number of key-value heads can therefore reduce KV cache storage without reducing the number of query heads by the same factor.

Head counts define the sharing pattern

Let an attention layer have Hq query heads and Hkv key-value heads, with Hq divisible by Hkv. A common grouped arrangement assigns:

G = Hq / Hkv

query heads to each key-value head. With Hq = 32 and Hkv = 8, each key-value head serves four query heads.

For a token position, query head i is associated with a key-value head determined by its group. Its attention computation still has the familiar form:

softmax(q_i K_g^T / sqrt(d)) V_g

where g identifies the shared key-value head for query head i. Two query heads in the same group can produce different score distributions because their q_i vectors differ, even though K_g and V_g are shared.

This is not equivalent to collapsing those query heads into one head. Their query projections and resulting attention distributions remain distinct.

KV cache size follows key-value heads

During autoregressive decoding, implementations commonly cache the key and value tensors produced for earlier positions. Ignoring metadata, padding, allocator effects, and implementation-specific layouts, the element count for one layer scales as:

2 * sequence_length * Hkv * head_dim

The factor of two represents keys and values. If the cached dtype and head_dim stay fixed, changing Hkv changes this storage approximately in direct proportion.

A layer with 32 query heads does not require 32 cached key heads when its architecture uses only 8 key-value heads. Reconstructing or duplicating shared keys into 32 physical cache heads would discard the storage advantage of the grouped representation, although an implementation may temporarily expand views or tensors for a kernel.

Cache memory is only one part of serving memory. Model weights, activations, temporary workspaces, allocator fragmentation, batching policy, and framework-specific buffers remain separate. The head-count relation should therefore be treated as a KV cache property rather than a complete memory estimate.

GQA sits between two endpoint arrangements

Standard multi-head attention commonly uses one key-value head per query head, giving Hkv = Hq. Multi-query attention uses a single key head and a single value head shared across all query heads, giving Hkv = 1.

Grouped-query attention occupies the parameterization between those endpoints:

1 < Hkv < Hq

for the non-degenerate grouped case. The architectural distinction is the sharing ratio, not a separate attention scoring rule.

This also means the labels describe structure, not a universal runtime guarantee. Kernel implementations can use different tensor layouts, fused operations, cache formats, or replication strategies. An architecture with grouped key-value heads can still be served poorly by a kernel that fails to exploit that structure.

Sharing changes representational capacity

Reducing Hkv is not merely a cache compression switch applied after the model has produced ordinary multi-head keys and values. In a model parameterized for GQA, the key and value projections themselves have fewer output heads.

As a result, query heads within a group cannot each receive an independently projected key-value representation. They can still ask different questions of the shared keys through distinct query vectors, but the key and value subspaces available to those heads are coupled by construction.

This boundary is relevant when converting an existing model. Simply reshaping a checkpoint with Hq independent key-value heads into fewer heads is not generally semantics-preserving. A conversion needs a defined method for producing the grouped key-value parameters, and any resulting model quality is an empirical property of that conversion or subsequent training procedure.

Cache layout must preserve the architectural mapping

Serving code needs to know both query-head and key-value-head counts. Treating them as interchangeable can create shape errors, unnecessary replication, or incorrect head-to-group mapping.

A cache tensor may conceptually use a layout such as:

[batch, Hkv, sequence, head_dim]

while a query tensor uses:

[batch, Hq, current_tokens, head_dim]

Exact dimension order is implementation-specific. The invariant is that cached keys and values represent Hkv heads, while the attention operation maps Hq query heads onto those shared heads according to the architecture.

This is especially relevant for model-serving adapters and custom kernels that infer cache shapes from attention configuration. num_attention_heads alone is insufficient when the model exposes a separate key-value-head count.

Grouped-query attention therefore sets a precise implementation boundary: query-head parallelism and KV cache cardinality are different quantities. Preserving that distinction keeps cache sizing, tensor shapes, and head mapping aligned with the model’s actual attention parameterization.