Multi-Query and Grouped-Query Attention

Generating text one token at a time means a transformer must remember, for every previous token, the key and value vectors every head computed for it — the The KV Cache and Incremental Decoding — since recomputing them at every step would be wasteful. The problem: standard multi-head attention gives every head its own keys and values, so the cache grows with the number of heads as well as sequence length. For a large model generating long outputs, that cache can dwarf the model's own weights and become the real bottleneck on how fast text streams out. Multi-query and grouped-query attention are two related fixes: stop giving every head its own keys and values, and share them instead.

The analogy: one shared reference desk, many readers

Picture a library reading room with 32 researchers (the attention heads), each keeping a personal, photocopied set of every reference book they might need (the keys and values) under standard multi-head attention. Photocopying 32 full sets costs shelf space. Multi-query attention keeps just one shared master set at a central desk, and lets all 32 researchers consult it — each still asks their own personal questions (queries stay separate), but all read from the same shelf. Grouped-query attention is the middle ground: split the 32 researchers into, say, 8 groups of 4, each group with its own shared shelf — less duplication than 32 personal copies, more flexibility than one shared shelf.

The mechanics

In standard multi-head attention with $h$ heads, each head $i$ has its own projection matrices $W_i^K, W_i^V$, producing distinct keys $K_i$ and values $V_i$ per head — the cache stores $h$ separate $K, V$ pairs.

Multi-query attention (MQA) keeps $h$ separate query projections $W_i^Q$ (each head still asks its own questions) but uses a single shared $W^K, W^V$ for all heads:

$$\text{head}_i = \text{softmax}\!\left(\frac{Q_i K^\top}{\sqrt{d_k}}\right) V$$

In words: every head computes its own query $Q_i$ and its own attention weights, but $K$ and $V$ — the thing that has to be cached across every generation step — are the same for all $h$ heads. That shrinks the KV cache by a factor of $h$.

Grouped-query attention (GQA) splits the $h$ heads into $g$ groups (with $1 < g < h$), and each group shares one $K, V$ pair. MQA is the special case $g=1$; ordinary multi-head attention is the special case $g=h$.

Worked example 1: cache size for a real-sized model

Take a model with $h = 32$ heads, head dimension $d_k = 128$, sequence length $n = 4{,}096$, in 16-bit precision, across $L = 32$ layers, caching both $K$ and $V$.

Standard multi-head cache: $2 \times L \times h \times n \times d_k \times 2\text{ bytes} = 2 \times 32 \times 32 \times 4096 \times 128 \times 2 = 2{,}147{,}483{,}648$ bytes $\approx 2.0$ GB.

Multi-query cache ($h=1$ shared instead of 32): divide by 32 → $\approx 67$ MB.

Grouped-query with $g = 8$ groups: divide by $32/8 = 4$ → $\approx 512$ MB.

The KV cache — the part of memory that scales with how many tokens you have generated so far — shrinks by roughly 32× under MQA and 4× under GQA with 8 groups, for the exact same sequence length.

Worked example 2: how many tokens fit in a fixed memory budget

Suppose a GPU has 8 GB free for KV cache after loading weights, with the cache costs from above. MHA lets you cache roughly $8\text{ GB} / (2.0\text{ GB} / 4096\text{ tokens}) \approx 16{,}384$ tokens before running out. Under MQA, the same 8 GB stretches to about $8 / (0.067/4096) \approx 490{,}000$ tokens — far beyond what's needed, meaning the real win is serving far more simultaneous requests with the same memory, not just longer single contexts.

Drag the exponent above — cache size as a function of the number of shared groups $g$ shrinks in roughly this shape, falling fastest as $g$ first drops from $h$ and flattening once $g$ is already small.

What this means in practice

Almost every large language model deployed for fast serving today (Llama 2 70B, Mistral, and others) uses GQA rather than full multi-head attention, because inference speed at long context is bottlenecked by moving the KV cache through memory, not raw compute — GQA directly shrinks the data that has to move. The choice of $g$ is a knob teams tune explicitly: more groups preserves quality closer to full MHA, fewer groups saves more memory and bandwidth.

Practice

1. With $h=16$ heads and $g=4$ groups, by what factor does the KV cache shrink relative to full multi-head attention?
2. Why does MQA save cache memory but not reduce the number of query projections computed per head?