Continuous Function

Foundation reference

Grouped-Query Attention (GQA)

Generating each new token means reading the whole key-value cache from memory, so a smaller cache means less memory traffic and room for longer contexts or larger batches; MQA was introduced to cut exactly this cost (Shazeer, 2019).

Concept 80 of 100EfficiencyStudy phase 2
Open the demoAll concepts
Key equation
KV cache=gh×MHA cache\text{KV cache} = \frac{g}{h} \times \text{MHA cache}
Reading map and next steps

Study guide

Where to start

Concept 80 of 100GQAEfficiency · Study phase 2: Architecture fundamentals
Question

Why can several query heads share one key head and one value head with little loss of quality, when each query head still computes its own attention pattern?

Open the demoGo to the mathAll conceptsNext concept
How to study this pagePredict · Check · Generalize · Continue
PredictSay what you expect.

Before you move a control in the demo, say what you expect to change.

CheckCheck against the math and the papers.

Compare what the demo shows with the key equation and with the papers cited below.

GeneralizeName what stays the same.

Identify the part of the idea that survives a change of notation or model, so you can recognize it in another paper or system.

ContinueFollow the connections.

Use the prerequisites and related concepts at the end of the page to fill the next gap.

Why It Matters for Modern Models

  • Generating each new token means reading the whole key-value cache from memory, so a smaller cache means less memory traffic and room for longer contexts or larger batches; MQA was introduced to cut exactly this cost (Shazeer, 2019).
  • Ainslie et al. (2023) converted existing multi-head checkpoints to GQA using about 5% of the original pre-training compute and found quality close to MHA at a speed close to MQA.
  • Llama 2 70B, the largest released Llama 2 model, uses GQA; the paper also reports a 34B model with GQA that was trained but not released (Touvron et al., 2023).

Questions to Think About

  • GQA leaves the attention arithmetic per query head unchanged; why does token-by-token generation still get faster, and why is the gain smaller when processing the prompt?
  • How should the number of key-value heads g be chosen, and what is lost in going from MHA to GQA and from GQA to MQA (Ainslie et al., 2023)?

Interactive Demo

Core Math

If you prefer intuition first, try the demo above, then come back for the math.

Key Equation
KV cache=gh×MHA cache\text{KV cache} = \frac{g}{h} \times \text{MHA cache}

In multi-head attention (MHA) each of the hh query heads has its own key head and value head. During generation the keys and values of all earlier tokens are cached, so for one sequence the cache holds

2⋅L⋅g⋅dh⋅T2 \cdot L \cdot g \cdot d_h \cdot T

numbers, with LL layers, gg key-value heads, head dimension dhd_h and TT tokens; the factor 2 counts keys and values. For MHA, g=hg = h. The demo on this page calls these two numbers Hq=hH_q = h and Hkv=gH_{kv} = g.

Multi-query attention (MQA) keeps all hh query heads but gives them a single shared key-value head, g=1g = 1 (Shazeer, 2019).

Grouped-query attention (GQA) sits in between: the hh query heads are divided into gg equal groups, and the heads in a group share one key head and one value head (Ainslie et al., 2023):

headi=Attention(XWiQ, XWj(i)K, XWj(i)V),j(i)=⌈i gh⌉\text{head}_i = \text{Attention}\big(XW^Q_i,\ XW^K_{j(i)},\ XW^V_{j(i)}\big), \qquad j(i) = \left\lceil \frac{i\, g}{h} \right\rceil

Memory is reduced by a factor of hg\frac{h}{g}; equivalently, GQA retains the fraction gh\frac{g}{h} of the MHA KV cache. g=hg = h recovers MHA and g=1g = 1 recovers MQA.

Key Papers

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Ainslie et al.2023EMNLP
Read paper →

Fast Transformer Decoding: One Write-Head is All You Need

Shazeer2019arXiv
Read paper →

Llama 2: Open Foundation and Fine-Tuned Chat Models

Touvron et al.2023arXiv
Read paper →

Connections

Prerequisites

Efficient AttentionLong Context