Grouped-query attention (GQA)

Overview

Autoregressive decoding reads the whole key and value cache for every new token, so with many heads the cache is large and memory bandwidth limits speed 2305.13245 §1. Grouped-query attention (GQA) keeps all H query heads but gives them fewer key and value heads: the query heads are divided into G groups, and every query head in a group attends with that group's single key head and value head 2305.13245 §2.2. The per-token cache then holds G key and G value heads instead of H, a factor of H/G smaller (arithmetic from the definition; the paper states the factor for MQA).

The paper makes two separate contributions. The first is the attention variant itself. The second is a conversion recipe: take an existing multi-head checkpoint, mean-pool the key and value projections within each group, then continue pre-training for a fraction \alpha of the original steps 2305.13245 §2.1. All experiments use T5 encoder-decoder models; the paper reports no decoder-only results. Later decoder-only models use GQA too. The Gemma 2 card describes it as part of the architecture and does not describe a conversion step 2408.00118 §2. The Mistral 7B card lists 8 key/value heads for 32 query heads as part of its architecture 2310.06825 tab. 1.

Variants

Choosing the group count G

G=H is multi-head attention, G=1 is multi-query attention, and values in between trade cache size against quality 2305.13245 §2.2. The paper uses G=8 for its main T5-XXL results and calls it a favorable middle ground between the speed of MQA and the quality of MHA 2305.13245 §3.3 fig. 6. The GQA paper does not give a general rule for choosing G beyond the sweep in fig. 6. Note that G is a count of key/value heads, so the cache saving H/G depends on the head count: Mistral 7B also uses 8 key/value heads, but with 32 query heads its ratio is 4 2310.06825 tab. 1, while the paper's T5-XXL setting has 64 heads and a ratio of 8 fig. 6.

Gemma 2 does not sweep G. It fixes two query heads per key/value head at every size, which gives 8, 16 and 32 query heads with 4, 8 and 16 K/V heads for the 2B, 9B and 27B models 2408.00118 §2 tab. 1. The authors chose this from ablations that showed faster inference with no clear loss. On the 9B model, GQA averages 50.8 against 50.3 for MHA over four benchmarks 2408.00118 §2 tab. 8.

Converted and uptrained checkpoints (the paper's recipe)

Mean-pool the K and V projection matrices of the heads in each group, then uptrain with the original pre-training objective, data and optimizer settings for \alpha times the original number of steps 2305.13245 §2.1 §3.1. The paper's main setting is \alpha=0.05, about 600 TPUv3 chip-days for T5-XXL 2305.13245 §3.1. This variant needs access to the original pre-training data and recipe.

Decoder-only models with GQA

Gemma 2 is a family of decoder-only models at 2B, 9B and 27B parameters. Every size uses GQA in its attention layers 2408.00118 §2 tab. 1. This is the decoder-only setting the GQA paper did not test, but the card's evidence is limited to the 9B ablation tab. 8.

Mistral 7B is a decoder-only model with 32 query heads and 8 key/value heads, so each key/value head serves 4 query heads 2310.06825 tab. 1. Its card names GQA as one of its inference-cost choices and says it reduces decoding memory 2310.06825 §1. The card does not report a GQA-versus-MHA comparison, so its quality and speed effect is not established by that card. The same card pairs GQA with sliding-window attention and a rolling-buffer cache, which the GQA paper does not cover 2310.06825 §2.

Mixtral 8x7B reuses the Mistral 7B attention block, with 32 query heads and 8 key/value heads 2401.04088 tab. 1 §2. Its feed-forward blocks are sparse mixture-of-experts layers, which the GQA paper does not cover 2401.04088 §2.1. The card reports no GQA-only comparison, so the attention effect is not established by that card.

Training from scratch

The GQA paper never compares a GQA model trained from scratch with the converted one, so its quality and stability in that setting are unknown 2305.13245 Limitations. Gemma 2's 27B model is trained from scratch, while its 2B and 9B models are distilled from a larger teacher 2408.00118 §3.2. The card has no scratch-versus-converted comparison for GQA, and it does not say whether the ablation models were trained from scratch or converted 2408.00118 §5.

When to use which

Implementation sketch

Two steps: the one-off conversion of an MHA layer, and one decode step with the grouped cache. Batch, causal mask and position encoding are omitted.

# Conversion of one attention layer [arxiv:2305.13245 §2.1]
# W_K, W_V: [d_model, H*h]; head i uses columns i*h .. (i+1)*h - 1
# Head-to-group assignment is not stated in the paper; blocks are assumed here
for g in range(G):
    # mean-pool the K and V projections of the H/G heads in group g
    members = heads_of_group(g)                        # H/G head indices
    W_K_g = mean([W_K[:, i*h:(i+1)*h] for i in members], axis=0)
    W_V_g = mean([W_V[:, i*h:(i+1)*h] for i in members], axis=0)
# Keep the H query projections and the output projection unchanged (§2.1)
# Then uptrain for alpha * original_steps with the original recipe (§2.1)

# One decode step for a new token x: [d_model] [arxiv:2305.13245 §2.2]
q = reshape(x @ Wq, [H, h])               # H query heads
key = reshape(x @ Wk, [G, h])             # G key heads
val = reshape(x @ Wv, [G, h])             # G value heads
cache.K.append(key); cache.V.append(val)  # cache per token: 2*G*h values

# Query head i reads the K/V of its own group g(i)
for i in range(H):
    g = group_of(i)
    scores = q[i] @ cache.K[:, g, :].T / sqrt(h)   # [T]
    probs = softmax(scores)
    out[i] = probs @ cache.V[:, g, :]              # [h]
y = reshape(out, [H*h]) @ Wo

Setting G=H reproduces MHA and G=1 reproduces MQA, so one code path covers the whole design space. Only the K/V projections and the cache change, which is why the conversion touches only W_K and W_V. Gemma 2's 2B model has 8 query heads and 4 K/V heads, so its G is 4 and each group holds 2 query heads 2408.00118 tab. 1.

Contradictions

Two timing numbers for the same configurations come from different setups and do not match:

• MQA-XXL takes 0.24 s per sample and GQA-8-XXL 0.28 s in the main table, measured on averages over the tasks 2305.13245 tab. 1.
• In the group sweep with fixed input 2048 and output 512, 1 group (MQA) takes 0.489 s and 8 groups 0.514 s 2305.13245 fig. 6.

The paper does not reconcile these. The axis of fig. 3 is labelled "ms" while its values match the seconds in tab. 1, so the units there also need care 2305.13245 fig. 3. Read the two speed sources as separate setups, not as one measurement.

The word "groups" also has two counts in use:

• The GQA paper's G is the number of key/value heads, so GQA-8 has 8 K/V heads 2305.13245 §2.2. Mistral 7B's n_kv_heads = 8 uses the same count 2310.06825 tab. 1.
• Gemma 2 writes num_groups = 2 for two query heads per K/V head 2408.00118 §2. Under the GQA paper's convention that is G = H/2, not G = 2.

Copying the Gemma value into the GQA paper's G would give the 9B model 2 K/V heads where the model has 8 2408.00118 tab. 1. The two sources report no conflicting results; this is a naming conflict only.

Papers

• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints: introduces GQA, the mean-pooling conversion and uptraining recipe, and reports T5-XXL results against MHA and MQA.
• Mistral 7B: a decoder-only 7B model with GQA (32 query heads, 8 K/V heads) combined with sliding-window attention; the card reports no GQA-only ablation.
• Mixtral of Experts: a sparse mixture-of-experts decoder whose attention block is Mistral 7B's (32 query heads, 8 K/V heads); the card reports no GQA-only comparison.
• Gemma 2: Improving Open Language Models at a Practical Size: a decoder-only family that uses GQA with two query heads per K/V head at 2B, 9B and 27B, with a 9B ablation against MHA.