GQA
How it works
Instead of one KV pair per head (MHA) or one for the whole model (MQA), GQA groups query heads — each group shares one KV pair. E.g., 8 heads in 2 groups of 4 → only 2 KV pairs instead of 8.
Problem solved
Multi-Head Attention requires a separate key-value pair per head, which is memory-intensive. GQA reduces KV cache memory usage by grouping query heads.
Components
H independent query projections, as in Multi-Head Attention.
G key-value pairs (1 < G < H); each pair is shared by all query heads in its group.
Uptraining procedure: the K/V heads of the original MHA checkpoint are mean-pooled within each group, then the model is further trained for about 5% of the pre-training budget.
Official
Implementation
G=1 is MQA (maximum savings, quality loss), G=H is MHA (no savings). Optimal G depends on task and model size — no universal rule exists.
Models pre-trained with MHA cannot be directly fine-tuned as GQA without converting KV head weights (e.g. by averaging or pruning). Requires a dedicated conversion step.
Evolution
Noam Shazeer introduces MQA: a single K/V pair for all heads, speeding up decoding at the cost of quality.
GQA introduced as a generalization of MQA, together with a recipe for uptraining MHA checkpoints at ~5% of pre-training compute.
GQA enters production open-weight LLMs, including Llama 2 70B and Mistral 7B.
GQA becomes the default attention mechanism across most new LLM families.
Hyperparameters (configurable axes)
Number of K/V pairs. G = num_attention_heads → MHA; G = 1 → MQA; 1 < G < H → GQA. The main knob controlling the quality/memory trade-off.
Number of query heads; must be divisible by the number of K/V groups.
Dimension of a single attention head (d_head), typically hidden_size / num_attention_heads.
Computational complexity
Time complexity: O(n² · d). Space complexity: O(n · d_head · G).
Compute bottleneck
Autoregressive decoding is memory-bound: the dominant cost is repeatedly loading the KV cache from memory. GQA shrinks the KV cache, reducing this memory traffic.
Execution paradigm
GQA is a dense attention mechanism — all heads are active for every token; there is no routing (unlike Mixture of Experts).
Parallelism
Attention computation remains parallel across heads and positions; sharing K/V introduces no sequentiality. Autoregressive decoding is sequential across tokens — as in any decoder — independently of GQA.
Hardware requirements
GQA reduces KV cache size — particularly valuable on GPUs with limited VRAM for long contexts (128k+ tokens). Natively supported by FlashAttention-2 and vLLM, among others.
GQA was developed at Google with efficient inference in mind; the KV memory and bandwidth reduction is beneficial on TPU as well.