Grouped-Query Attention (GQA)

GQA reduces the size of the KV-Cache by a factor equal to the number of query heads per key-value head. In standard Multi-Head Attention, each query head has a dedicated key and value head, causing the cache to grow linearly with the number of heads. By sharing a single key-value pair across a group of query heads, GQA slashes the memory bandwidth required during autoregressive decoding. For a model with 32 query heads and 8 key-value groups, the KV-Cache size is reduced by 4x, directly lowering the hardware cost of serving long sequences.

GQA acts as a tunable trade-off between Multi-Head Attention (MHA) and Multi-Query Attention (MQA). MHA provides the highest model quality but the largest KV-Cache. MQA uses a single key-value head for all queries, minimizing memory but often degrading quality on tasks requiring fine-grained attention. GQA generalizes both: setting the number of key-value groups equal to the number of query heads recovers MHA; setting it to one recovers MQA. Intermediate group counts achieve near-MHA quality with near-MQA efficiency.

GQA models can be efficiently created from pre-trained MHA checkpoints through a process called uptraining. The original key and value head weight matrices are mean-pooled into the target number of groups, initializing a GQA model. This converted model is then fine-tuned for a small fraction (typically 5-10%) of the original training steps. This avoids the prohibitive cost of training a GQA model from scratch and preserves the knowledge already learned by the base model.

GQA has been widely adopted in state-of-the-art production models due to its favorable memory-quality trade-off:

• Llama 2 70B uses GQA with 8 key-value heads
• Llama 3 and Llama 3.1 use GQA across all model sizes
• Mistral models leverage GQA for efficient inference
• Gemma and other open-weight models incorporate GQA The technique is now a standard architectural choice for large-scale autoregressive decoders.

The primary bottleneck in LLM inference for long sequences is memory bandwidth, not compute. The KV-Cache must be read from GPU HBM for every new token generated. By shrinking the KV-Cache, GQA directly increases the number of tokens the system can generate per second. This improvement is most pronounced in batch inference scenarios where multiple sequences are processed simultaneously, as the aggregate KV-Cache size across the batch is the dominant memory consumer.

GQA is complementary to other attention efficiency techniques:

• FlashAttention: Reduces HBM reads/writes via tiling; GQA reduces the total data that must be stored
• Multi-Query Attention (MQA): The extreme case of GQA with one key-value group
• Multi-Head Latent Attention (MLA): Used in DeepSeek-V2, compresses KV-Cache into a low-rank latent vector, achieving even greater compression than GQA
• Sliding Window Attention: Limits attention to a local context; can be combined with GQA for further memory savings

The foundational attention mechanism in the original Transformer where each attention head has its own independent set of query, key, and value weight matrices. This allows each head to learn distinct relational patterns—one head might track syntactic dependencies while another captures long-range semantic links. The outputs of all heads are concatenated and linearly projected. While expressive, MHA's memory footprint scales linearly with the number of heads, making the KV-Cache a significant bottleneck for long-sequence inference.

An extreme memory-saving variant where all query heads share a single key-value head. This eliminates the multi-head KV-Cache overhead entirely, dramatically reducing memory bandwidth requirements during autoregressive decoding.

• Trade-off: Significant reduction in model quality and training stability
• Use case: Deployments where memory is the absolute binding constraint
• Adoption: Used in PaLM and early efficient inference models before GQA emerged as a superior middle ground

The KV-Cache stores pre-computed key and value tensors to avoid redundant computation during autoregressive generation. Its size scales as:

• MHA: 2 × num_layers × num_heads × d_head × seq_len
• MQA: 2 × num_layers × 1 × d_head × seq_len
• GQA: 2 × num_layers × num_groups × d_head × seq_len

For a 70B model at 100k tokens, MHA requires ~140 GB just for the KV-Cache, while GQA with 8 groups reduces this to ~17.5 GB—an 8× reduction that makes long-context deployment feasible on a single GPU.

FlashAttention is an exact-attention algorithm that minimizes high-bandwidth memory (HBM) reads and writes by fusing attention operations in SRAM using tiling and recomputation. When combined with GQA:

• The already-reduced KV-Cache from GQA means fewer bytes to transfer from HBM to SRAM
• FlashAttention's tiling further optimizes how those key-value tensors are streamed
• Together they enable 2-4× faster attention computation for long sequences compared to MHA with standard attention implementations
• This combination is standard in modern inference engines like vLLM and TensorRT-LLM

GQA introduces a structural constraint that requires careful handling during training:

• Uptraining from MHA: Mean-pooling the key and value heads of a pre-trained MHA checkpoint into GQA groups provides a strong initialization, minimizing quality loss
• From-scratch training: Requires tuning the number of groups; too few groups (approaching MQA) can cause training instability and loss spikes
• Group count sweet spot: Empirical evidence from Llama 2 and Mistral shows 8 groups for a 32-query-head model preserves >99% of MHA quality while achieving near-MQA memory efficiency
• Gradient flow: The shared KV heads receive gradients from multiple query heads, effectively increasing their learning signal, which can accelerate convergence for the key-value projections

GQA has become the de facto standard for large-scale production LLMs due to its quality-efficiency Pareto frontier:

• Llama 2 70B: 8 KV groups across 64 query heads
• Llama 3 70B & 405B: Continued GQA adoption with optimized group ratios
• Mistral Large: GQA for efficient long-context serving
• Gemma 2: Google's open models use GQA for inference efficiency
• Command R+: Cohere's RAG-optimized model leverages GQA for 128k context windows

The pattern is clear: any model designed for real-world deployment with context windows exceeding 32k tokens overwhelmingly adopts GQA.