During autoregressive token generation, Large Language Models are rarely bounded by raw floating-point compute; they are severely throttled by Memory Bandwidth. For every single generated token, the model must read all its weights and the entire growing history of Key-Value tensors from High Bandwidth Memory (HBM) into on-chip cache. Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) represent revolutionary structural innovations designed to slash the KV-cache bandwidth footprint by up to 8x without sacrificing representational capacity.
1. The Autoregressive Memory Wall & Arithmetic Intensity
The operational efficiency of neural inference is governed by Arithmetic Intensity: the ratio of floating-point operations (FLOPs) performed to bytes transferred from memory:
$mathcal{I} = frac{text{FLOPs}}{text{Bytes Transferred}}$
In standard Multi-Head Attention (MHA), given hidden dimension $d$, number of query heads $H_Q$, and sequence length $S$, each head maintains its own Key and Value projections:
$mathbf{K}_h = mathbf{X} mathbf{W}_K^{(h)}, quad mathbf{V}_h = mathbf{X} mathbf{W}_V^{(h)} quad (h = 1, dots, H_Q)$
During decoding of a single token ($N=1$), computing the attention output requires fetching the entire historical KV-cache of size $2 times H_Q times S times d_k$ bytes. Because the FLOPs scale as $O(S)$ while memory traffic also scales as $O(S)$, the arithmetic intensity drops below 1 FLOP/byte. Modern accelerators (NVIDIA H100) boast 2,000 TFLOPs of FP16 compute but only 3.35 TB/s of memory bandwidth. Under MHA decoding, tensor cores sit idle more than 90% of the time, starved for data.
2. Multi-Query Attention (MQA): Extreme KV Sharing
Introduced by Noam Shazeer in 2019, Multi-Query Attention (MQA) eliminates head-specific Key and Value projections entirely. While Query heads remain independent ($H_Q$ heads), all Query heads share a single common Key head and Value head ($H_{KV} = 1$):
$mathbf{Q} in mathbb{R}^{B times H_Q times 1 times d_k}, quad mathbf{K} in mathbb{R}^{B times 1 times S times d_k}, quad mathbf{V} in mathbb{R}^{B times 1 times S times d_k}$
The attention computation broadcasts the single Key and Value tensor across all $H_Q$ query heads:
$mathbf{A}_h = text{Softmax}left(frac{mathbf{Q}_h mathbf{K}^T}{sqrt{d_k}}right) mathbf{V}$
MQA slashes the memory transferred per token by a factor of $H_Q$ (typically 32x to 64x). However, forcing all attention heads to share identical historical representations can degrade performance in complex reasoning, syntactic parsing, and retrieval-heavy benchmarks.
3. Grouped-Query Attention (GQA): The Pareto-Optimal Compromise
To resolve the quality gap between MHA and MQA, Ainslie et al. (2023) formulated Grouped-Query Attention (GQA), the architecture adopted by LLaMA-2/3, Mistral, and modern open frontier models. GQA partitions $H_Q$ query heads into $G$ disjoint groups of size $H_Q / G$. Each group shares a single Key-Value head pair ($H_{KV} = G$):
- When $G = H_Q$: GQA reduces identically to standard Multi-Head Attention (MHA).
- When $G = 1$: GQA reduces identically to Multi-Query Attention (MQA).
- When $1 < G < H_Q$ (typically $G = 8$ for $H_Q = 32$ or $64$): GQA provides an 8x reduction in KV cache memory bandwidth while retaining 99% of MHA’s perplexity and reasoning capacity.
4. Mathematical Formulation of Grouped-Query Softmax
For query head $h in {1, dots, H_Q}$, its assigned group index is $g(h) = lfloor h / (H_Q / G) rfloor$. The grouped attention tensor contraction is defined as:
$mathbf{O}_h = text{Softmax}left(frac{mathbf{Q}_h mathbf{K}_{g(h)}^T}{sqrt{d_k}}right) mathbf{V}_{g(h)}$
The aggregated multi-head output $mathbf{O} = [mathbf{O}_1, dots, mathbf{O}_{H_Q}] mathbf{W}_O$ preserves multi-perspective expressive capacity while decoding throughput surges by 300% to 500% in production serving engines.
