Attention Variants for Inference

Bookkeeping of Variants of attention

$n_h$ = query heads, $n_{kv}$ = key/value heads, $g=n_h/n_{kv}$ = group size.

Variant	$n_{kv}$	Per-token KV	$L=32$, $n_h=32$, $d_{\text{head}}=128$, $b=2$
MHA	$n_h$	$2L{n}_h{d}_{\text{head}}b$	512 KiB/token
GQA-$g$	$n_h/g$	$2L(n_h/g)d_{\text{head}}b$	128 KiB ($g=4$), 64 KiB ($g=8$)
MQA	1	$2L{d}_{\text{head}}b$	16 KiB/token
MLA	latent	$L(d_c+d_r)b$	~36 KiB/token ($d_c=512$, $d_r=64$)

MLA caches a per-token low-rank latent $c_t\in {\mathbb{R}}^{d_c}$ shared by all heads, plus a decoupled RoPE vector $d_r$. $K$ and $V$ are reconstructed by up-projections that can be absorbed into $W_Q$ and $W_O$.

Decode attention reads the whole KV cache to produce one token. Arithmetic intensity of that read: $I_{\text{attn}}=\frac{4n_h{d}_{\text{head}}S}{2n_{kv}{d}_{\text{head}}Sb}=\frac{2}{b}\cdot \frac{n_h}{n_{kv}}=\frac{2g}{b}$ The $S$ cancels: intensity is independent of context length and equals the group size (times $2/b$). MHA gives $I=1$ FLOP/byte in FP16, hopelessly memory-bound. GQA-8 gives 8; MQA gives 32. Each query group reuses one KV read, which is the whole point.

Same effect via the batch axis: $B$ sequences each read their own KV, so unlike weights, KV traffic does not amortize over the batch. That is why $k_{\text{tok}}$ appears in the throughput ceiling ${\text{BW}}_{\text{eff}}/(S{k}_{\text{tok}})$ in Batch Size and Latency Tradeoff.

Quality tradeoff

• MQA is the aggressive end: measurable quality loss and training instability, largely abandoned for large models.
• GQA with $g\in \{4,8\}$ recovers MHA quality within noise while cutting the cache 4–8×. The current default.
• MLA matches or beats MHA quality at an MQA-sized cache, because the latent is a learned low-rank bottleneck rather than a hard head-sharing constraint. Cost: a more complex kernel and more compute in the projections.

Also note: GQA changes the KV tensor layout, so tensor parallelism can only shard KV heads $n_{kv}$ ways. Beyond $TP>n_{kv}$ you must replicate the KV cache across ranks, an easy and expensive mistake at high TP. MLA has the same issue in sharper form (one shared latent), which is one reason MLA models lean on data/expert parallelism instead.

Whiteboard version

Write $k_{\text{tok}}$, divide free HBM by $k_{\text{tok}}\cdot S$ for concurrency (KV Cache Sizing), divide ${\text{BW}}_{\text{eff}}$ by $S{k}_{\text{tok}}$ for the throughput ceiling. Everything else about the variant is quality and kernel availability.

Rule of thumb

GQA group size $g$ is simultaneously the KV memory divisor and the decode-attention arithmetic-intensity multiplier. Pick $g=8$ unless you have a reason not to.