Multi-Query and Grouped-Query Attention

Standard multi-head attention gives every head its own query, key, and value projections. At inference, that means every head needs its own cached key and value vectors for every past token (see kv-cache), and that cache, not the model's weights, is often what actually runs a GPU out of memory during long-context serving. Shazeer, 2019 proposed a specific, surgical fix: keep multiple query heads, but share a single key and value head across all of them.

The asymmetric cut

Multi-query attention (MQA) keeps h separate query heads but collapses the key and value projections down to one shared head each:

# standard multi-head: h separate K, V pairs
K_i, V_i for i in 1..h

# multi-query: one shared K, V for every query head
K, V   (shared)
Q_i for i in 1..h

Every query head still computes its own attention pattern against the same shared keys, so the model retains its ability to run multiple independent matching criteria per position (see multi-head-attention-deep). What it loses is the ability to have differently-projected keys and values per head. The cache that must be stored per token shrinks by a factor of h for keys and h for values, because there is only one key vector and one value vector per token instead of h of each. For long sequences and large batch sizes, where the KV cache dominates memory far more than the model weights do, this is the single largest lever available for fitting more context or more concurrent requests into the same GPU memory.

Why keys and values, not queries

The asymmetry is deliberate. Queries are used once per generated token and then discarded; they never need to be cached, because a new query is computed fresh at every decoding step from the current token only. Keys and values for every past token, in contrast, must be kept around for the entire generation, because every future token's attention needs to compare against all of them. Cutting the query heads down would save nothing at inference (queries were never the memory cost); cutting keys and values down attacks the actual bottleneck directly.

Grouped-query attention: the compromise that shipped

MQA's aggressive sharing (one key/value head for the entire model) sometimes cost noticeable quality relative to full multi-head attention, since every query head is now forced to share the exact same keys and values, with no room for even a few independently projected views. Ainslie et al., 2023 proposed grouped-query attention (GQA): partition the h query heads into g groups, and give each group its own shared key/value head, with g somewhere between 1 (MQA) and h (full multi-head).

g groups, each of size h/g query heads
each group shares one K, V pair
g = 1        -> multi-query attention
g = h        -> standard multi-head attention
1 < g < h    -> grouped-query attention