Grouped-Query Attention

Grouped-query attention does not remove any query heads, does not make the attention arithmetic smaller, and does not make the model meaningfully smaller. It removes copies of the keys and values. That is a memory decision, and it is worth four to eight times as many concurrent users on the same card.

Some vocabulary first, because this area is thick with three-letter names for the same dial. A transformer layer projects each token into a query, a key and a value, several times over in parallel — each parallel copy is a head. In multi-head attention (MHA), the original 2017 arrangement, every query head gets its own key head and its own value head. In multi-query attention (MQA), proposed by Shazeer in 2019, all the query heads share one key head and one value head. Grouped-query attention (GQA), from Ainslie et al. in 2023, is the setting in between: the query heads are divided into groups, and each group shares one key head and one value head. Every model config you will read states this as num_attention_heads and num_key_value_heads. They are the two ends and the middle of a single dial.

Only the second number is charged to memory, because only keys and values are kept between tokens — that stored tensor is the KV cache, and it costs:

bytes per token = 2 (K and V) × layers × kv_heads × head_dim × bytes_per_element

Query heads are absent from that formula. A query is used once, in the step that created it, and then thrown away. So halving the key/value heads halves the cache, exactly, with no effect on how many queries are computed.

Below is one card running one model. The dial marked kv heads sweeps from multi-head attention on the right to multi-query attention on the left. Move it and watch three numbers that most people expect to move together come apart: the cache, the parameters, and the arithmetic.

Start at the defaults — Qwen3-8B, H100, 8k tokens of context — and drag kv heads from 8 up to 32, which is what this model would have been before 2023. The cache goes from 144 KiB per token to 576 KiB, and the number of sequences the card can hold falls from 49 to 11. Throughput falls with it, 1,738 tokens per second to 413. Now look at what did not change: attention still does 4.83 GFLOP per sequence per step at every stop on the dial, because the shared key head is broadcast to all four of its query heads rather than skipped. The parameter count moves by about a tenth, 8.19B to 9.10B. A 10% change in one number and a 4× change in the other, from the same dial.

Three numbers that people expect to move together

The reason GQA is misread is that "fewer heads" sounds like "smaller model, less work, worse quality" — one lever with three consequences. It is not. The dial you just moved separates them.

The cache is linear in kv heads. Every kv head you remove removes its bytes from every token of every resident sequence, forever. Qwen3-8B at 8k context: 576 KiB per token at 32 kv heads, 144 KiB at 8, 18 KiB at 1. That term is multiplied by context length and by concurrency, so it is the term that decides how many users fit.

The parameters barely move. Only k_proj and v_proj get narrower. On Qwen3-8B that is 9.10B parameters at multi-head down to 8.19B at 8 kv heads — 10% — and only 8.19B to 7.93B for the whole rest of the journey to multi-query. Nobody adopts GQA to save parameters.

The arithmetic does not move at all. This is the part that surprises people. A shared key head is not computed once and used once; it is read once and broadcast to every query head in its group, and each of those query heads still does a full dot product against it. The readout proves it: 4.83 GFLOP per sequence per step on Qwen3-8B at 8k context, at every stop on the dial. Some kernels literally materialise the repeat — PyTorch's reference path calls repeat_interleave on the key and value tensors before the matmul; a GQA-aware kernel keeps the shared head in registers instead. Either way the FLOP count is the same and the byte count is not.

That is the whole mechanism, and it has a compact statement. The attention part of a decode step reads 2 × layers × kv_heads × head_dim × 2 bytes per token of context and does 4 × layers × q_heads × head_dim FLOPs on it. Divide:

attention FLOPs per byte of cache = q_heads ÷ kv_heads = the group size

The readout labelled "attention flops per byte" is that number, and it is the reason this lesson exists. Arithmetic intensity — FLOPs performed per byte moved from memory — is what decides whether a kernel waits on the memory system or on the tensor cores. An H100 can do about 295 FLOPs in the time it takes to fetch one byte from HBM (High Bandwidth Memory, the DRAM stacked next to the GPU die). Multi-head attention during decode achieves 1.0. It is not close, it is not marginal: the tensor cores are idle for 99.7% of that kernel. GQA with 8 groups gets 4.0, multi-query gets 32. Still bandwidth-bound, but each byte now does eight or thirty-two times as much useful work, and the cache it has to fetch is that many times smaller.

Why nobody ships one kv head

If the cache is linear in kv heads and the arithmetic is free, the dial should obviously go to 1. Almost no current model does. There are two reasons, and the simulation makes the second one visible.

Quality. Ainslie et al. uptrained a T5-XXL checkpoint into each variant using 5% of the original pre-training compute — the key and value projection matrices of each group are mean-pooled into one, which they found works better than picking one head or starting from random. Their measured result across summarisation, translation and question answering: multi-head scored 47.2 with 1.51 s per sample of inference, multi-query scored 46.6 at 0.24 s, and GQA with 8 groups scored 47.1 at 0.28 s. Read those two rows against each other. Almost the entire speed win of multi-query is already there at 8 groups, and almost the entire quality loss is in the last step from 8 groups to 1. Shazeer's original multi-query paper reported "only minor quality degradation"; the GQA paper adds that multi-query "can lead to training instability during fine-tuning, in particular combined with long input tasks". Those are measured results from the papers, not something this simulation models.

Tensor parallelism, which is the boundary you can reach here. Set tensor parallel size to 8 and then sweep the kv dial from 8 down to 1. Nothing happens. The cache per token stays at 18 KiB per rank, the resident sequence count stays put, throughput does not move. The head diagram shows why: the boxes turn magenta and start repeating themselves.

Tensor parallelism splits each weight matrix across GPUs by head, so a rank owns some query heads and the kv heads that go with them. A rank cannot own a third of a kv head. When there are fewer kv heads than ranks, the engine gives every rank its own full copy — literally, in vLLM's LlamaAttention:

if self.total_num_kv_heads >= tp_size:
    # Number of KV heads is greater than TP size, so we partition
    # the KV heads across multiple tensor parallel GPUs.
    assert self.total_num_kv_heads % tp_size == 0
else:
    # Number of KV heads is less than TP size, so we replicate
    # the KV heads across multiple tensor parallel GPUs.
    assert tp_size % self.total_num_kv_heads == 0
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

max(1, …) is the floor. Eight ranks each holding one replica of a single kv head store exactly as many kv bytes in total as eight distinct kv heads would have, so a multi-query model on eight GPUs has the aggregate cache footprint of a GQA-8 model — and the quality of a multi-query model. This is stated plainly in the GQA paper: "standard sharding for large models replicates the single key and value head by the number of model partitions; GQA removes the waste from such partitioning." Eight groups is not a coincidence, and neither is the fact that Qwen3-8B, Qwen3-32B, Qwen2.5-72B and Mistral-7B all ship exactly 8. It is the largest tensor-parallel degree people commonly deploy on a single node.

Real multi-query models do exist and their configs say so out loud: Falcon-7B sets multi_query: true with num_attention_heads: 71, so 71 query heads share one key and one value head. Seventy-one is prime, incidentally, which means that model has no grouped middle ground available at all — the only whole-number options are 1 and 71.

The other place the dial stops paying. Set context to 512 tokens and leave everything else alone. At 8 kv heads the pool would hold 791 sequences and the scheduler cap of 256 binds first: memory is no longer the constraint, so removing kv heads adds no concurrency. It still speeds up the step — 13.3 ms to 7.6 ms, because there is less cache to re-read — and at one kv head the decode step on this model finally becomes compute-bound, which is the only regime in this entire lesson where buying a GPU with more FLOPs would help. Short contexts and big batches are exactly where GQA matters least.

Checking it on a real system

Read it off the config before you benchmark anything. In config.json:

"num_attention_heads": 32,
"num_key_value_heads": 8,
"head_dim": 128,
"num_hidden_layers": 36

If num_key_value_heads is absent, the model is multi-head and the value equals num_attention_heads. If head_dim is absent, it is hidden_size ÷ num_attention_heads — but check rather than assume, because Qwen3-32B states 128 explicitly while that division gives 80. Then the cache per token is 2 × num_hidden_layers × num_key_value_heads × head_dim × dtype_bytes: for the config above, 2 × 36 × 8 × 128 × 2 = 147,456 bytes, 144 KiB per token, 1.13 GiB for an 8k-token request. Use kv heads. Using query heads is the single most common arithmetic error in capacity planning, and it overstates the cache by the group size — a factor of four here, eight on a 64-head model.

Then check what your serving engine actually built. vLLM prints its pool at boot as one line:

GPU KV cache size: 1,048,576 tokens, Maximum concurrency
for 32,768 tokens per request: 32.00x

Divide that token count by your real p95 context length to get your true concurrency ceiling — vLLM's own division uses the model's maximum context, which is usually far longer than what you serve. If the number is a factor of two to eight below what you expected, compare num_key_value_heads against your --tensor-parallel-size first. Any deployment where kv heads is less than the TP degree is paying for replicas: 8 kv heads on 16 ranks stores twice the cache the config implies. The fix is fewer ranks, or tensor-parallel 8 combined with pipeline or data parallelism for the rest, not a scheduler knob.

To confirm on the hardware that you are still bandwidth-bound after all this — you almost certainly are — sample DCGM's DCGM_FI_PROF_DRAM_ACTIVE against DCGM_FI_PROF_PIPE_TENSOR_ACTIVE during steady-state generation. High DRAM activity with tensor-pipe activity in the low single digits is the signature, and it is what "attention flops per byte = 4" looks like from the outside.

One last trap worth naming, because it costs people a week: you cannot convert an existing multi-head checkpoint to GQA by dropping heads at load time. The mean-pooled projections have to be uptrained, which is what the 5% of pre-training compute in the paper buys. GQA is an architecture decision made before training, not a serving flag. The serving-time levers on the same bytes are KV cache quantization, which halves them again, and paged attention, which stops you wasting the ones you have. What is not a lever on them is FlashAttention: it removes the N × N score matrix from HBM, not the cache.

Next: the general form of the bandwidth-versus-arithmetic argument this lesson keeps leaning on, arithmetic intensity; what the cache costs and why decode re-reads all of it, the KV cache; and the other thing people believe shrinks it, FlashAttention. The head-splitting arrangement GQA modifies is multi-head attention, and the step it performs is scaled dot-product attention.