Grouped Query Attention (GQA)

Grouped-query attention (GQA) is a form of multi-head attention for transformers in which groups of query heads share one key head and one value head. Llama 3.1 8B, for example, pairs 32 query heads with 8 key-value heads in every layer. It sits between multi-head attention (MHA), where every query head has its own keys and values, and multi-query attention (MQA), where all query heads share a single key head and value head. Joshua Ainslie and colleagues at Google Research introduced it in GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints at EMNLP 2023, showing that GQA models reach quality close to MHA at decoding speed close to MQA. Open large language models (LLMs) such as Llama 2 70B, the Llama 3 family, Mistral 7B, and Gemma 2 use it, and vision-language models (VLMs) built on them, such as Google's PaliGemma 2 on Gemma 2, inherit it in their language decoders.

How GQA Works#

In self-attention, every token is projected into a query, a key, and a value for each head, and each head compares its queries against its keys. GQA keeps the full set of query heads but projects fewer key and value heads, and each group of query heads attends with its group's shared keys and values. With 8 query heads and 2 key-value heads, each key-value head serves 4 query heads:

The attention computation inside each head is unchanged scaled dot-product attention. With one group the layer becomes MQA, and with as many groups as query heads it becomes MHA. PyTorch supports the layout directly through the enable_gqa argument of scaled_dot_product_attention:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

It prints Output (1, 8, 16, 64), same as copied heads: True and 256 with GQA, 1024 with MHA: all 8 query heads produce outputs from a quarter of the keys and values that MHA would store.

Why GQA Shrinks the KV Cache#

During text generation, a model keeps the keys and values of every earlier token in its KV cache and reads the whole cache for each new token, so decoding speed is limited by memory bandwidth more than by arithmetic. Cache size scales with the number of key-value heads. Llama 3.1 8B has 32 layers, each with 8 key-value heads of 128 dimensions, so it stores 128 KiB per token in 16-bit precision, or 16 GiB at its full 131,072-token context window. With MHA, the same model would need 512 KiB per token and 64 GiB at full length. The freed memory goes to longer contexts, larger batches for continuous batching, and lower inference latency per token. The Ultralytics LLM interface accepts a base_url for any OpenAI-compatible server, such as a self-hosted Llama 3.1 8B, where the 8 key-value heads let the server hold four times as many tokens of conversation in a given cache memory as a 32-head MHA layout would.

Converting Models to GQA#

GQA does not require training from scratch. The authors converted an existing MHA checkpoint by mean-pooling the key and value heads within each group, then continued pretraining with about 5% of the original compute, a step they call uptraining. MQA, the extreme of one shared head, was introduced by Noam Shazeer in 2019 and is the fastest layout, but it degraded quality and could make training unstable, and an intermediate number of groups recovers most of that quality.

GQA vs. Related Techniques#

• Multi-head latent attention (MLA): Introduced with DeepSeek-V2, it compresses the keys and values of each token into a small shared latent vector.
• PagedAttention: Allocates the KV cache in fixed-size blocks to avoid fragmentation, independent of how many heads fill it.
• Sliding window attention: Caps the number of cached tokens per layer, and Mistral 7B combines it with GQA.
• Flash attention: A GPU kernel that computes exact attention faster and supports GQA layouts natively.