Grouped-Query Attention

An attention variant that reduces key-value cache memory by sharing key-value heads across query groups.

Grouped-query attention (GQA) is an attention variant that lets several query heads share fewer key-value heads, cutting key-value cache size without collapsing all heads into one shared pair.

At a glance

Released

May 2023

Authors

Joshua Ainslie, James Lee-Thorp, Seth R. Robertson, et al.

Optimizes

  • Kv Cache
  • Memory Bandwidth
  • Long Context Inference

What It Is

Grouped-query attention (GQA) is an attention variant derived from multi-head attention. It keeps multiple query heads but groups them so each group reads from the same key-value head pair.

Why It Exists

GQA targets key-value cache size, memory bandwidth during autoregressive decoding, and the cost of serving long contexts when head count would otherwise multiply cache size.

How It Works

Queries are partitioned into groups. Each group shares one key head and one value head. Attention scores are computed per query head against the shared key-value pair for that group, and the outputs are then projected as usual.
Grouped-query attention reduces key-value head count by letting several query heads share each key-value pair.

Math Or Compute Schema

With multiple query heads and several key-value groups, each group serves a subset of the query heads. The formulas below contrast how multi-head attention pairs every query head with its own key-value heads versus grouped-query attention, which routes query heads through shared key-value pairs per group.
Multi-head attention (MHA)
Attention(Qi,Ki,Vi)=softmax ⁣(QiKi⊤dk)Vi\text{Attention}(Q_i, K_i, V_i) = \mathrm{softmax}\!\left(\frac{Q_i K_i^{\top}}{\sqrt{d_k}}\right) V_i
QQ
Query vectors for head i.
KK
Key vectors for head i.
VV
Value vectors for head i.
HH
Number of query heads.
dkd_k
Key dimension per head.
ii
Query head index.
Grouped-query attention (GQA)
Attention(Qi,Kg(i),Vg(i))=softmax ⁣(QiKg(i)⊤dk)Vg(i)\text{Attention}(Q_i, K_{g(i)}, V_{g(i)}) = \mathrm{softmax}\!\left(\frac{Q_i K_{g(i)}^{\top}}{\sqrt{d_k}}\right) V_{g(i)}
QQ
Query vectors for head i.
KK
Key vectors for group g(i).
VV
Value vectors for group g(i).
HH
Number of query heads.
GG
Number of shared key-value groups.
dkd_k
Key dimension per head.
ii
Query head index.
g(i)g(i)
Key-value group index for query head i.

Compared To Nearby Modules

Compared with multi-head attention, GQA reduces the number of key-value tensors that must be stored. Compared with multi-query attention, GQA keeps more distinct query heads while still sharing key-value heads within each group.
Comparison dimensionGrouped-Query AttentionMulti-Head AttentionMulti-Query Attention
Key-value head countG key heads and G value headsH key heads and H value heads1 key head and 1 value head
Query-head flexibilityH distinct query heads grouped into G shared key-value pairsH independent query and key-value head pairsH query heads share one key-value pair
Cache footprint per token2G tensors (G keys plus G values)2H tensors (full multi-head attention cache)2 tensors (single shared key-value cache)

Example Architectures

GQA appears in large decoder-only language models that prioritize efficient long-context inference while retaining multi-head training dynamics.

Limitations And Tradeoffs

Sharing key-value heads can reduce representational flexibility versus full multi-head attention. The optimal group count depends on model scale and quality targets.

Why It Still Matters

As context windows grow, key-value cache efficiency remains a first-order serving constraint, so GQA stays a practical default between multi-head fidelity and multi-query compression.

Tags

References

  1. Ainslie, Joshua, et al. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints." arXiv, 2023, https://arxiv.org/abs/2305.13245.