GQA Review: Grouped Query Attention for Faster LLM Inference

Why Multi-Head Attention Has a Memory Problem

Here’s a number that might surprise you: in a 65B parameter LLM with 80 attention heads, the key-value cache during inference can eat up 40GB of memory just to hold context. That’s before you’ve even loaded the model weights.

You can read the original paper here (Ainslie et al., 2023).

Multi-Head Attention (MHA) — the backbone of every Transformer since Vaswani et al. (2017) — stores separate key and value projections for each head. During autoregressive decoding, you cache these KV pairs for every token you’ve generated. With 80 heads and a context length of 8K tokens, you’re caching $80 times 8192 times d_{head} times 2$ floats per layer. Multiply by 80 layers and suddenly your GPU is sweating.

Multi-Query Attention (MQA), proposed by Shazeer (2019), took the nuclear option: share a single key-value head across all query heads. Memory footprint drops by 8-80x depending on your head count. Inference flies. But there’s a catch — quality degrades. The authors of GQA measured a non-trivial accuracy gap when MQA was applied to larger models. You save memory but lose capability.

GQA sits in the middle. Instead of N heads (MHA) or 1 head (MQA), you get G groups. Each group of query heads shares one KV head. It’s embarrassingly simple.

$$text{Memory}_{text{GQA}} = frac{G}{H} times text{Memory}_{text{MHA}}$$

Where $H$ is the number of query heads and $G$ is the number of KV groups. Llama 2 70B uses $H=64$ and $G=8$, giving an 8x reduction in KV cache size.

The Clever Part: Uptraining Instead of Training from Scratch

Training a 65B model from scratch costs millions of dollars. Nobody wants to throw that away just to try a new attention pattern.

GQA’s key practical insight is uptraining: take an existing MHA checkpoint, convert it to GQA by mean-pooling adjacent KV heads, then continue training for a small fraction of the original compute. The paper reports using only 5% of the original pretraining compute for uptraining.

The conversion process is dead simple:

import torch

def convert_mha_to_gqa(mha_kv_weights, num_groups):
    """
    Convert MHA key/value projection weights to GQA.
    mha_kv_weights: (num_heads, head_dim, model_dim)
    num_groups: number of KV groups (must divide num_heads evenly)
    """
    num_heads, head_dim, model_dim = mha_kv_weights.shape
    assert num_heads % num_groups == 0, "num_heads must be divisible by num_groups"

    heads_per_group = num_heads // num_groups

    # Mean pool adjacent heads into groups
    grouped = mha_kv_weights.view(num_groups, heads_per_group, head_dim, model_dim)
    gqa_weights = grouped.mean(dim=1)  # (num_groups, head_dim, model_dim)

    return gqa_weights

# Example: 64 heads -> 8 groups
mha_k = torch.randn(64, 128, 4096)  # 64 heads, d_head=128, d_model=4096
gqa_k = convert_mha_to_gqa(mha_k, num_groups=8)
print(f"MHA K shape: {mha_k.shape}")  # (64, 128, 4096)
print(f"GQA K shape: {gqa_k.shape}")  # (8, 128, 4096)
print(f"Memory reduction: {mha_k.numel() / gqa_k.numel()}x")  # 8.0x

I’m not entirely sure why mean-pooling works so well versus other aggregation methods (max, learned projection). The paper doesn’t ablate this deeply. My guess is that adjacent heads in trained models already learn somewhat similar patterns, so averaging doesn’t destroy too much information.

What Actually Happens at Inference Time

Here’s the attention computation with GQA. Each query head $q_i$ attends to its group’s shared keys $K_g$ and values $V_g$:

$$text{Attention}(Q, K_g, V_g) = text{softmax}left(frac{Q K_g^T}{sqrt{d_k}}right) V_g$$

The difference from MHA is that multiple query heads share the same $K_g$ and $V_g$. In practice, this means:

import torch
import torch.nn.functional as F

def gqa_attention(q, k, v, num_groups, scale=None):
    """
    Grouped Query Attention forward pass.
    q: (batch, seq_len, num_heads, head_dim)
    k: (batch, seq_len, num_kv_heads, head_dim)  # num_kv_heads = num_groups
    v: (batch, seq_len, num_kv_heads, head_dim)
    """
    batch, seq_len, num_heads, head_dim = q.shape
    num_kv_heads = k.shape[2]

    if scale is None:
        scale = head_dim ** -0.5

    # Expand KV heads to match Q heads
    # This is where the "grouping" happens
    heads_per_group = num_heads // num_kv_heads

    # Repeat each KV head for its group of Q heads
    k = k.unsqueeze(3).expand(-1, -1, -1, heads_per_group, -1)
    k = k.reshape(batch, seq_len, num_heads, head_dim)
    v = v.unsqueeze(3).expand(-1, -1, -1, heads_per_group, -1)
    v = v.reshape(batch, seq_len, num_heads, head_dim)

    # Standard attention from here
    # (batch, num_heads, seq_len, head_dim)
    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)

    attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale
    attn_weights = F.softmax(attn_weights, dim=-1)

    output = torch.matmul(attn_weights, v)
    return output.transpose(1, 2)  # (batch, seq_len, num_heads, head_dim)

# Test with 8 query heads, 2 KV groups
batch, seq, n_heads, d_head = 2, 128, 8, 64
n_kv_heads = 2

q = torch.randn(batch, seq, n_heads, d_head)
k = torch.randn(batch, seq, n_kv_heads, d_head)
v = torch.randn(batch, seq, n_kv_heads, d_head)

out = gqa_attention(q, k, v, num_groups=n_kv_heads)
print(f"Output shape: {out.shape}")  # (2, 128, 8, 64)

The expand + reshape pattern is how most implementations handle this. It’s not actually duplicating memory — PyTorch’s expand creates a view with stride=0 on the repeated dimension.

The Numbers That Actually Matter

The paper compares MHA, MQA, and GQA on T5 XXL (13B params) for summarization tasks. Here’s what stood out:

MQA loses about 0.5-0.6 ROUGE points. GQA-8 loses only 0.1 points while being nearly as fast as MQA.

But the really interesting numbers come from the uptraining ablation. The paper tries different amounts of uptraining compute:

• 0% (direct conversion): Quality tanks. ROUGE drops by ~1 point.
• 5% of original pretraining: Recovers almost all quality.
• 10%: Marginal further improvement, not worth the cost.

This 5% number is crucial for practitioners. It means if your original model took 1000 GPU-hours to train, you can convert to GQA with ~50 GPU-hours of uptraining.

The Ablation That Surprised Me

The paper includes an ablation on the number of groups. Intuition says more groups = better quality. But the curve isn’t linear.

Going from 1 group (MQA) to 8 groups recovers most of the quality gap to MHA. Going from 8 to 16 groups gives diminishing returns. And going all the way to 64 groups (approaching MHA) barely helps further.

$$text{Quality}(G) approx text{Quality}_{text{MHA}} – frac{c}{sqrt{G}}$$

This square-root relationship (roughly) means you get most of the benefit from just 4-8 groups. That’s why Llama 2 70B uses 8 groups despite having 64 heads — going higher isn’t worth the memory cost.

What I haven’t seen is analysis on whether the optimal number of groups changes with model scale. My best guess is that larger models can tolerate fewer groups (relatively), but this would need systematic study.

Where It Falls Short

GQA isn’t a free lunch. Some limitations the authors acknowledge:

1. Uptraining requires the original training setup. If you only have the checkpoint but not the training data/pipeline, you can’t uptrain. Direct conversion without uptraining hurts quality.

2. Batch size sensitivity. The speedups assume inference scenarios where you’re bottlenecked by memory bandwidth (small batches, long sequences). With large batch sizes where compute dominates, MHA vs GQA makes less difference.

3. No help for training. GQA reduces inference memory, not training memory. During training you need the full activations for backprop anyway.

One thing the paper doesn’t discuss much: what happens with very long contexts? At 32K or 128K context lengths (common in newer models), the KV cache dominates everything. GQA helps proportionally, but you might still OOM. Techniques like sliding window attention (used in Mistral) or chunked prefill might be needed alongside GQA.

Should You Use This in Production?

Absolutely — but you probably already are without knowing it.

Llama 2, Mistral, Mixtral, and most modern open-weight LLMs ship with GQA by default. If you’re running inference with these models via vLLM, TGI, or llama.cpp, you’re benefiting from GQA’s memory savings.

If you’re training your own model from scratch, just use GQA from the start. There’s no reason to use full MHA anymore unless you have a specific research question about attention patterns.

If you have an existing MHA checkpoint you want to convert:
1. Mean-pool KV heads into groups (code above)
2. Uptrain for ~5% of original compute
3. Validate on your eval set
4. Enjoy 8x smaller KV cache

One practical gotcha: some inference frameworks assume specific numbers of KV heads. I ran into issues with an older version of Hugging Face Transformers (pre-4.31) that didn’t handle GQA properly. Make sure your inference stack supports the num_key_value_heads config field.

How It Relates to Other Work

GQA fits into a broader trend of making attention more efficient. Compared to other approaches:

FlashAttention (Dao et al., 2022) optimizes the attention computation via IO-aware algorithms. GQA optimizes the memory footprint of KV caches. They’re complementary — you should use both.

Multi-Query Attention (Shazeer, 2019) is GQA with G=1. GQA generalizes it with a quality-efficiency knob.

Sliding Window Attention (Beltagy et al., 2020; used in Mistral) limits attention to recent tokens. This is orthogonal to GQA — you can combine them.

For a deeper look at Chinchilla’s scaling laws (which informed many of the training decisions in modern LLMs), I covered that in Paper Review: Chinchilla — Why Most LLMs Were Massively Overtrained on Too Few Tokens.

FAQ

Q: Can I convert any MHA model to GQA without retraining?
Yes, but quality will degrade noticeably — expect 0.5-1 point drops on benchmarks. For production use, you really need that 5% uptraining to recover quality. Direct conversion is fine for experimentation but not deployment.

Q: Why do Llama 2 and Mistral use different numbers of GQA groups?
Llama 2 70B uses 8 groups for 64 heads. Mistral 7B uses 8 groups for 32 heads (meaning 4 query heads per KV head). The “right” number depends on your quality/speed tradeoff requirements and hardware. 8 groups seems to be a sweet spot that most teams landed on empirically.

Q: Does GQA help during training or only inference?
Primarily inference. During training, you need full activations for backpropagation regardless of KV head sharing. The training compute is roughly the same for MHA vs GQA. The big wins are at inference time when you’re caching KV states for autoregressive generation.

References

• Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245
• Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150
• Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135

I’m still curious whether there’s a principled way to determine the optimal number of groups based on model architecture rather than empirical search. The square-root relationship I mentioned is just an observation from the ablations — there might be a theoretical grounding I’m missing. If anyone has insights on this, I’d love to hear them.