Grouped-Query Attention (GQA)

GQA exists on a spectrum between two extremes:

• MHA (G=H): Every query head has its own unique key-value head. Highest quality, highest memory bandwidth.
• MQA (G=1): All query heads share a single key-value head. Fastest inference, potential quality degradation.
• GQA (1 < G < H): Query heads are partitioned into G groups, with each group sharing one key-value head. This balances the trade-off.

The primary motivation for GQA is reducing the size of the KV cache during autoregressive decoding. The memory footprint scales inversely with the number of groups:

• MHA stores H key-value heads per layer.
• MQA stores 1 key-value head per layer.
• GQA stores G key-value heads per layer, reducing memory by a factor of H/G compared to MHA. This allows for larger batch sizes and longer context lengths under the same memory budget.

Empirical results from the original paper demonstrate that GQA achieves quality nearly indistinguishable from MHA while approaching the speed of MQA. Key findings include:

• Using G=8 for a model with H=32 (4 query heads per KV head) performs comparably to the MHA baseline on summarization and translation tasks.
• The performance gap between MQA and MHA widens with model size, but GQA closes this gap effectively.
• The shared KV heads learn to generalize, acting as a form of regularization.

GQA can be efficiently bootstrapped from existing MHA models through a process called uptraining:

• Mean Pooling: Initialize the single key-value head for a group by mean-pooling the key-value heads of all query heads in that group.
• Continued Training: Fine-tune the converted model on a small fraction (typically <10%) of the original pre-training data.
• This allows rapid conversion of existing MHA models to the GQA architecture without training from scratch, preserving the original knowledge.

GQA has become the standard attention mechanism for production-grade open-weight models due to its inference efficiency:

• Llama 2 70B and Llama 3 use GQA with 8 key-value heads.
• Mistral models leverage GQA alongside sliding window attention.
• Gemma and other Google-derived architectures implement GQA.
• The mechanism is a critical enabler for serving large models on commodity hardware with constrained memory bandwidth.

While GQA reduces memory reads, the total FLOPs remain largely unchanged from MHA. The speedup comes from memory-bound operations:

• During autoregressive generation, loading the KV cache from HBM is the bottleneck.
• GQA reduces the bytes transferred per token, shifting the operation from memory-bound toward compute-bound.
• The effective throughput gain is most pronounced on hardware with limited memory bandwidth relative to compute, such as consumer GPUs and edge devices.

The foundational attention mechanism from the original Transformer paper. MHA uses separate key, query, and value projections for each attention head, providing maximum representational capacity at the cost of significant memory bandwidth during autoregressive decoding. The full KV cache must store unique key-value pairs for every head, making MHA the slowest option for inference but the highest quality baseline.

An extreme optimization where all query heads share a single key-value head. MQA dramatically reduces the size of the KV cache and memory bandwidth requirements, leading to the fastest possible decoding speed. However, this aggressive sharing can degrade model quality and training stability, particularly on large-scale models. GQA was developed as a direct compromise between MQA's speed and MHA's quality.

The KV cache stores pre-computed key and value tensors to avoid redundant computation during autoregressive generation. Its size grows linearly with batch size, sequence length, and the number of key-value heads. GQA directly reduces the memory footprint of this cache by a factor equal to the ratio of query heads to KV heads, enabling larger batch sizes and longer context windows on memory-constrained hardware.

A complementary memory management technique that partitions the KV cache into non-contiguous blocks, analogous to virtual memory paging in operating systems. When combined with GQA, PagedAttention eliminates internal fragmentation and enables near-optimal memory utilization. This pairing is a key architectural insight behind the high throughput of inference engines like vLLM.

An IO-aware exact attention algorithm that fuses operations and tiles computation to minimize reads and writes to high-bandwidth memory (HBM). FlashAttention is orthogonal to GQA—it optimizes how attention is computed, while GQA optimizes what is stored. Modern implementations like FlashAttention-2 and FlashAttention-3 support GQA natively, providing compounding speedups for long-context inference.

GQA has been adopted by several prominent open-weight models to balance quality and efficiency:

• Llama 2 70B: Used MHA, requiring substantial KV cache memory
• Llama 3 70B: Adopted GQA with 8 key-value heads, significantly improving inference throughput
• Mistral 7B: Uses GQA alongside sliding window attention for efficient long-context processing
• Gemma 7B: Implements GQA to enable competitive performance on consumer hardware