7.4. Attention Variants

Attention Variants

Standard multi-head attention (MHA) uses $H$ independent heads, each with its own query, key, and value projections ${\mathbf{W}}_h^Q,{\mathbf{W}}_h^K,{\mathbf{W}}_h^V$. This means the KV cache during autoregressive inference scales as $O(H\cdot T\cdot d_k)$ per layer, which becomes a memory bottleneck for long sequences and large batch sizes. Several variants reduce this cost by sharing key-value projections across heads while keeping separate queries.

Multi-Query Attention (MQA) uses a single shared key-value pair across all heads:
${\mathbf{q}}_i^{(h)}={\mathbf{x}}_i{\mathbf{W}}_h^Q,{\mathbf{k}}_j={\mathbf{x}}_j{\mathbf{W}}^K,{\mathbf{v}}_j={\mathbf{x}}_j{\mathbf{W}}^V,$
where each head $h$ has its own query projection but all heads share the same ${\mathbf{W}}^K$ and ${\mathbf{W}}^V$. This reduces the KV cache by a factor of $H$ and speeds up inference significantly, at the cost of slightly reduced model quality.

Grouped-Query Attention (GQA) interpolates between MHA and MQA by dividing the $H$ query heads into $G$ groups, each sharing one set of key-value projections:
${\mathbf{q}}_i^{(h)}={\mathbf{x}}_i{\mathbf{W}}_h^Q,{\mathbf{k}}_j^{(g)}={\mathbf{x}}_j{\mathbf{W}}_g^K,{\mathbf{v}}_j^{(g)}={\mathbf{x}}_j{\mathbf{W}}_g^V,$
where head $h$ belongs to group $g=\lfloor hG/H\rfloor$. Setting $G=1$ recovers MQA; setting $G=H$ recovers full MHA. GQA achieves quality close to MHA with memory and speed close to MQA, and is used in LLaMA 2/3, Mistral, and Gemma.

KV Cache and Inference Cost $2\cdot L\cdot H\cdot T\cdot d_k$ values (factor 2 for K and V, $L$ layers, $H$ heads, $T$ sequence length). For a 70B model with $T=8192$, this can exceed 10 GB per sequence. GQA with $G=8$ instead of $H=64$ reduces this by $8\times$.

During autoregressive generation, each new token requires attending to all previous tokens. In MHA, the KV cache stores

Intuition

The key insight is that different query heads often need similar contextual information from keys and values—they don’t all need completely independent KV representations. By sharing keys and values, the model can maintain most of its representational power (via diverse queries) while dramatically reducing the memory footprint. This is especially important at inference time, where the KV cache is the dominant memory cost and the main bottleneck for serving long contexts or large batches.

Properties

Multi-Head Attention (MHA)

(Standard MHA) Each head $h\in \{1,\dots ,H\}$ computes independent projections $Q_h=X{\mathbf{W}}_h^Q$, $K_h=X{\mathbf{W}}_h^K$, $V_h=X{\mathbf{W}}_h^V$ with ${\mathbf{W}}_h^Q,{\mathbf{W}}_h^K\in {\mathbb{R}}^{d\times d_k}$ and ${\mathbf{W}}_h^V\in {\mathbb{R}}^{d\times d_v}$, where typically $d_k=d_v=d/H$. The outputs are concatenated and projected: $\mathrm{MHA}(X)=[{\mathrm{head}}_1\Vert \cdots \Vert {\mathrm{head}}_H]{\mathbf{W}}^O$.

(Parameter Count) MHA has $3Hd\cdot d_k+H{d}_v\cdot d=4d^2$ parameters per layer (when $d_k=d_v=d/H$), plus $d^2$ for ${\mathbf{W}}^O$.

(Implementation) In practice, instead of computing $H$ separate matrix multiplications, a single large projection $X{\mathbf{W}}^{QKV}\in {\mathbb{R}}^{T\times 3d}$ is computed and then split/reshaped into $H$ heads. This is more efficient on GPUs.

Multi-Query Attention (MQA)

(MQA) (Shazeer, 2019) All $H$ heads share a single key and value projection. Only the query projection is per-head. The KV cache per layer drops from $2H{d}_k$ to $2d_k$ per token.

(Quality Trade-off) MQA typically incurs a small quality degradation (0.5–1% on benchmarks) compared to full MHA, but inference throughput can increase by $2$–$5\times$ for memory-bound workloads.

Grouped-Query Attention (GQA)

(GQA) (Ainslie et al., 2023) The $H$ query heads are partitioned into $G$ groups of $H/G$ heads each. Each group shares one set of key-value projections. The KV cache per layer is $2G{d}_k$ per token—a factor $H/G$ reduction over MHA.

(Uptraining) Existing MHA models can be converted to GQA by mean-pooling the key and value projection matrices within each group, then continuing pre-training for a small fraction of the original compute. This is how LLaMA 2 70B was adapted to GQA.

Examples

KV Cache Size Comparison

For a model with $d=4096$, $H=32$ heads, $d_k=128$, $L=32$ layers, and sequence length $T=4096$:

Cache sizes

• MHA: $2\times 32\times 32\times 4096\times 128=1.07$ GB (in FP16)
• GQA ($G=8$): $2\times 32\times 8\times 4096\times 128=0.27$ GB — a $4\times$ reduction
• MQA ($G=1$): $2\times 32\times 1\times 4096\times 128=0.034$ GB — a $32\times$ reduction

For batch inference serving many sequences simultaneously, this reduction is the difference between fitting in GPU memory or not.