Serving large language models at scale presents a fundamental hardware bottleneck: autoregressive decoding is bound by memory bandwidth rather than compute. While the initial prompt processing phase (prefill) operates as compute-bound matrix multiplications, token-by-token generation requires loading billions of cached attention states from GPU High-Bandwidth Memory (HBM) to on-chip SRAM for every single generated token.
Multi-Head Attention (MHA), introduced in the foundational Vaswani et al. (2017) transformer architecture, allocates an independent key and value head for every query head. As context windows expand to tens of thousands of tokens and concurrent user batches grow, the resulting Key-Value (KV) cache consumes hundreds of gigabytes of VRAM.
Grouped-Query Attention (GQA), proposed by Ainslie et al. (2023), solved this serving crisis. By grouping query heads to share a smaller number of key and value heads, GQA achieves an order-of-magnitude reduction in KV cache memory footprint and memory traffic while matching the modeling capacity and accuracy of standard Multi-Head Attention.
The Memory Bandwidth Bottleneck in Autoregressive Decoding
To understand why attention heads dominate inference costs, consider the operational difference between prefill and decode phases in an LLM inference engine.
During the prefill phase, all prompt tokens are available simultaneously. The self-attention operation computes interactions across the entire sequence via dense matrix-matrix multiplications (GEMM). Modern GPU tensor cores operate near peak throughput because arithmetic intensity (the ratio of floating-point operations to memory bytes transferred) is high.
During autoregressive generation, the model predicts one token at a time. To generate token t + 1, the attention mechanism must compute the dot product between the single new query vector q_t and the key vectors of all previous tokens k_1, ..., k_t, followed by multiplying the resulting attention weights with all previous value vectors v_1, ..., v_t.
Because previous keys and values are static, inference engines cache them in VRAM to avoid redundant computation. However, at each decoding step, the entire accumulated KV cache for the sequence must be fetched from GPU HBM into local SRAM. This operation is a matrix-vector multiplication (GEMV) with an arithmetic intensity of approximately 1 to 2 FLOPs per byte. Modern accelerators like the NVIDIA H100 provide up to 3.35 TB/s of memory bandwidth compared to nearly 2,000 TFLOPs of 16-bit tensor compute. Under low arithmetic intensity, tensor cores remain idle while waiting for memory transfers over the HBM bus.
The KV Cache Footprint in Multi-Head Attention
In standard Multi-Head Attention, an attention layer with H query heads has H corresponding key heads and H value heads (H_Q = H_{KV}).
The memory required to store the KV cache for a model across a generation request scales linearly with sequence length, batch size, layer count, and the number of KV heads:
Memory = 2 * L * H_{KV} * d_{head} * P * B * S
Where:
Lis the number of transformer layers.H_{KV}is the number of Key-Value heads per layer.d_{head}is the dimension of each attention head.Pis the precision in bytes (2 bytes for FP16 or BF16).Bis the concurrent batch size.Sis the total sequence length (prompt plus generated tokens).- The factor of 2 accounts for storing both keys and values.
Consider a 70-billion parameter model such as Llama 2 70B or Llama 3 70B, which uses 80 layers, a hidden dimension of 8,192, and 64 heads with d_{head} = 128. Under standard Multi-Head Attention (H_{KV} = 64):
Each token consumes: 2 * 80 * 64 * 128 * 2 bytes = 2,621,440 bytes ≈ 2.50 MB per token
For a moderate batch size of 32 requests with a 4,096-token context window: 32 * 4,096 * 2.50 MB = 327,680 MB ≈ 320 GB
Storing the KV cache alone requires four 80 GB GPUs, entirely separate from the 140 GB needed just to hold the 70B model weights in 16-bit precision. At long context lengths (such as 32,768 or 131,072 tokens), MHA KV caches quickly reach several terabytes, rendering high-throughput multi-user serving computationally and economically infeasible.
Multi-Query Attention: The Extreme Compression Approach
In 2019, Noam Shazeer introduced Multi-Query Attention (MQA) in the paper "Fast Transformer Decoding: One Write-Head is All You Need".
Shazeer proposed collapsing all key and value heads into a single shared key head and single shared value head per layer (H_{KV} = 1), while retaining H_Q distinct query heads.
Under MQA, the KV cache size is reduced by a factor equal to the number of query heads (a 32x to 64x reduction). During autoregressive decoding, the memory bandwidth overhead drops proportionally, allowing inference engines to increase batch sizes significantly and achieve near-peak compute utilization.
However, Multi-Query Attention introduced distinct architectural drawbacks:
- Capacity Loss: Collapsing all keys and values into a single projection restricts the representational capacity of the attention layer. While small models or general language modeling tasks tolerated the compression with minor perplexity increases, larger models and tasks requiring precise associative recall, multi-document extraction, or code syntax modeling experienced noticeable quality degradation.
- Training and Checkpoint Incompatibility: Converting existing pretrained MHA checkpoints to MQA required expensive retraining. Training large MQA models from scratch also exhibited optimization instability across various scaling regimes.
Grouped-Query Attention: The Generalized Balance
To eliminate the quality gap of MQA while preserving its inference speed, Ainslie et al. (2023) formulated Grouped-Query Attention (GQA).
Instead of using one KV head per query head (MHA) or one KV head for the entire layer (MQA), GQA partitions the H_Q query heads into G distinct groups. All query heads within a given group share a single Key head and a single Value head. The number of KV heads per layer is therefore equal to G (H_{KV} = G), where 1 < G < H_Q.
GQA acts as a continuous generalization across attention topologies:
- When
G = H_Q, the architecture is standard Multi-Head Attention. - When
G = 1, the architecture collapses to Multi-Query Attention. - When
1 < G < H_Q(such asG = 8for a 64-query-head model, yielding an 8:1 query-to-KV ratio), GQA cuts the KV cache size and memory traffic by an exact factor of 8.
Empirical evaluations in the original Google Research paper demonstrated that an 8-group GQA configuration (GQA-8) on large language models recovers virtually 100% of MHA performance across summarization (CNN/DailyMail, Multi-News), question answering (TriviaQA, SQuAD), translation (WMT benchmarks), and reasoning tasks, while achieving decoding latencies nearly identical to MQA.
Uptraining: Converting MHA Models to GQA at 5% Compute
A major practical innovation introduced in the GQA paper was the "uptraining" method, which allows converting existing pretrained Multi-Head Attention checkpoints into Grouped-Query Attention models without training from scratch.
The conversion workflow operates in two phases:
- Mean-Pooling Initialization: For each of the
Ggroups, the originalH_Q / Gseparate key projection weight matrices are averaged together via mean-pooling to create a single consolidated key projection matrix. The same mean-pooling operation is applied to the corresponding value projection weight matrices. The query projection matrices and output projection matrices are left unmodified. - Continued Pretraining (Uptraining): The resulting model is trained on the original pretraining dataset for approximately 5% of the original training token budget using the original learning rate schedule.
Ainslie et al. demonstrated that 5% uptraining compute is sufficient for the model to adapt to the shared key-value representations, restoring validation perplexity and benchmark performance to the original MHA baseline.
Industry Adoption Across Modern Architectures
Following its publication, Grouped-Query Attention rapidly replaced Multi-Head Attention as the industry standard architecture for open and proprietary LLMs:
- Meta Llama Series: Meta transitioned to GQA starting with the 70B variant in Llama 2 (Touvron et al., 2023). In Llama 3 (Dubey et al., 2024), Meta adopted GQA across all model sizes (8B, 70B, and 405B), utilizing 8 KV heads (
G = 8) to support native 128k context windows. - Mistral AI: The architecture was adopted in Mistral 7B (Jiang et al., 2023) with 32 query heads and 8 KV heads (4:1 ratio), and carried forward into the Mixtral 8x7B mixture-of-experts model.
- Alibaba Qwen: The Qwen 2 and Qwen 2.5 families utilize GQA across their dense and MoE model variants to facilitate long-context processing up to 128k tokens.
- Google Gemma 2: Google integrated GQA into Gemma 2 (9B and 27B) along with alternating sliding-window local and global attention layers.
Architectural Successors: Multi-Head Latent Attention
Building on the principles of KV cache reduction established by GQA, newer architectures have explored low-rank projection techniques to compress KV cache representations even further.
DeepSeek introduced Multi-head Latent Attention (MLA) in DeepSeek-V2 and DeepSeek-V3. Rather than grouping discrete heads, MLA projects the key and value states into a compressed low-rank latent vector (typically 512 dimensions) during generation. Only the compressed latent vector is stored in the KV cache, which is dynamically projected back into full multi-head keys and values inside SRAM during attention computation. MLA reduces KV cache memory consumption to roughly 15% of standard MHA while allowing every query head to interact with a distinct subspace.
Grouped-Query Attention remains the foundational mechanism that made high-throughput serving, larger batch sizes, and 128k+ context windows viable on standard datacenter hardware.
Sources
- Vaswani et al. (2017) Attention Is All You Need: https://arxiv.org/abs/1706.03762
- Shazeer (2019) Fast Transformer Decoding: One Write-Head is All You Need: https://arxiv.org/abs/1911.02150
- Ainslie et al. (2023) GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints: https://arxiv.org/abs/2305.13245
- Touvron et al. (2023) Llama 2: Open Foundation and Fine-Tuned Chat Models: https://arxiv.org/abs/2307.09288
- Dubey et al. (2024) The Llama 3 Herd of Models: https://arxiv.org/abs/2407.21783
- Jiang et al. (2023) Mistral 7B: https://arxiv.org/abs/2310.06825
