Attention variants architecture — MHA, MQA, GQA, sliding window, linear

One mechanism, three knobs

Every variant computes the same core

Attn(Q, K, V) = softmax( Q K^T / sqrt(d_k) + M ) V

Q: [N, h_q * d_h]   K, V: [N, h_kv * d_h]   M: [N, N] mask

and differs only in how it fills in three slots.

Knob 1 — how many KV heads. h_kv = h_q is multi-head attention (MHA); h_kv = 1 is multi-query (MQA); 1 < h_kv < h_q is grouped-query (GQA); replacing K and V with a low-rank latent you decompress on the fly is multi-head latent attention (MLA).

Knob 2 — which entries of M are finite. Full causal attention allows every past token; sliding window allows the last W; sparse patterns allow a block-local or strided subset plus a few global tokens.

Knob 3 — whether softmax survives. Replace it with a kernel feature map φ and the product reassociates into a recurrent state: linear attention and its state-space cousins.

The knobs are orthogonal; production models turn all three at once.

The three numbers to keep score with

Compare variants by cost, not by name. All three are per layer, for a model of width d with h_q query heads, h_kv KV heads, head dim d_h = d / h_q, and b bytes per cache element.

params   = d*d (W_Q) + 2 * d * h_kv*d_h (W_K, W_V) + d*d (W_O)
KV/token = 2 * h_kv * d_h * b   bytes
prefill  = 4 * N^2 * d          FLOPs for QK^T and AV
decode   = 4 * N   * d          FLOPs per new token, over an N-long cache

Read the last two lines twice, because they hold the most misunderstood fact in this space: the attention FLOP count does not contain h_kv at all. MQA and GQA broadcast their few KV heads back out to all h_q query heads, so the same multiply-adds happen as in MHA. Head sharing buys parameters and memory traffic, not arithmetic; only knob 2 and knob 3 delete FLOPs. Keeping those effects in separate columns keeps the comparison honest.

Knob 1: sharing KV heads, from MHA to MQA

This axis is a clean spectrum indexed by one integer, h_kv. With h_q = 32, d_h = 128, d = 4096:

The cache column moves 32× while the parameter column moves less than 2×, which tells you what this knob is for. Quality does not fall linearly with h_kv: the drop from MHA to GQA-8 is small, while the drop from GQA to MQA is disproportionately large, and MQA is the least stable to train. That asymmetry — nearly all the memory win for nearly none of the quality loss at h_kv of 4 to 8 — is why GQA became the default in the Llama-2/3, Mistral and Qwen families.

Multi-head latent attention: compress instead of share

MLA (DeepSeek-V2) attacks the same cache but refuses the same trade. Instead of deleting KV heads it projects each token to a shared latent c_t = W_DKV x_t of width d_c, caches only c_t, and reconstructs per-head K and V at use time with up-projections that fold into W_Q and W_O.

MHA cache/layer/token = 2 * h_q * d_h            = 8192 elements
MLA cache/layer/token = d_c + d_h_rope ≈ 512 + 64 = 576 elements

That is roughly a 14× reduction — MQA territory — while every query head still sees a distinct reconstructed key and value, so the representational collapse that costs MQA its quality does not happen. The price is different in kind: extra up-projection FLOPs every step, a rotary term that must ride in a separate un-compressed slice (RoPE does not commute with the absorbed projections), and a serving stack that understands the layout. Strongest point on the cache-versus-quality frontier, and the most implementation-heavy.

Knob 2: restricting the mask

Sparsity attacks a different term. Sliding-window attention lets each token see only the previous W, so the score matrix has N·W finite entries instead of N^2/2, and a rolling buffer caps the cache at W. Prefill falls from quadratic to linear in N; the cache becomes constant in N. Depth buys back reach: with L layers the receptive field is about L·W tokens, so 32 layers at W = 4096 can propagate across 128k — indirectly, through many hops, with attenuation.

Block-sparse and strided patterns (Longformer, BigBird) generalize this: local blocks plus a few global tokens plus optional random blocks, giving O(N√N) cost while the attention graph still connects every pair within a couple of hops. Attention sinks — keeping the first few tokens permanently resident — are the cheap fix for the softmax-mass collapse that otherwise wrecks windowed models on long streams.

Knob 3: dropping softmax

Linear attention removes the quadratic term outright. Write sim(q, k) = φ(q) · φ(k) and associativity lets you compute φ(Q) (φ(K)^T V) instead of (φ(Q) φ(K)^T) V. The inner product is a [d_φ, d_h] matrix independent of N, so cost becomes O(N · d^2) and decoding becomes an RNN over a fixed state.

S_t = S_{t-1} + φ(k_t) v_t^T      z_t = z_{t-1} + φ(k_t)
y_t = φ(q_t)^T S_t / (φ(q_t)^T z_t)   →  O(1) memory per step

This is the only family whose decode memory is genuinely constant — no cache, just a state. The catch is capacity: a fixed S cannot store an unbounded set of key-value associations, so exact recall degrades in a way perplexity hides. Gated variants (RetNet, GLA, Mamba-2, DeltaNet) add decay and better update rules, narrowing but not closing the gap on retrieval.

The comparative map

All three knobs on one grid. N is context length, W the window, d the model width.

Nothing dominates. Head sharing wins on traffic and loses nothing on FLOPs; sparsity wins on both but forfeits exactness; linearization wins on everything except the one thing transformers are prized for — exact recall.

A worked example: 7B at 32k

L = 32, d = 4096, h_q = 32, d_h = 128, fp16 cache, N = 32768.

MHA:            2*32*128*2 B = 16 KiB/layer/token → 512 KiB/token → 16.0 GiB
GQA-8:                          4 KiB/layer/token → 128 KiB/token →  4.0 GiB
MQA:                          0.5 KiB/layer/token →  16 KiB/token →  0.5 GiB
MLA (d_c=512):               1.13 KiB/layer/token →  36 KiB/token →  1.1 GiB
SW W=4096 + GQA-8:  128 KiB * 4096 tokens        (flat) →  0.5 GiB
Linear:             one [d_φ, d_h] state per head       →  ~0 GiB

A 7B model in 4-bit weights is about 3.5 GB. The MHA cache here is four and a half times the model itself; GQA-8 is comparable to it; the windowed and latent caches are a rounding error. That ratio is why long-context serving pushed the industry off plain MHA.

Why cache size is speed, especially on CPU

Decode is memory-bound. Generating one token streams every weight and the entire KV cache through the memory bus while doing a few FLOPs per byte — arithmetic intensity near 1, far below any machine’s balance point. Bytes moved per token, divided by bandwidth, is a hard floor on latency. Take a CPU box at roughly 50 GB/s of usable DDR5 bandwidth, 8k context, the 7B model above:

Add the ~3.5 GB of weights — a fixed ~70 ms per token — and the crossover appears: with MHA the cache overtakes the weights at roughly 7k tokens, with GQA-8 not until roughly 28k. Choosing h_kv is choosing where your throughput cliff sits.

Hybrid stacks: the variant is a per-layer decision

Because the knobs are orthogonal and cheap to mix, the strongest current designs do not pick one variant — they interleave. Gemma 2 alternates sliding-window and full-attention layers in a 5:1 ratio, keeping a flat cache in most layers while a minority of global layers preserve exact long-range lookup. Jamba interleaves Mamba blocks with a sprinkling of full attention. Nearly every recent dense model runs GQA underneath whatever else it does.

The intuition holds up empirically: local structure is cheap and abundant, exact retrieval is expensive and rare. A few full-attention layers serve the retrieval need, so paying quadratic cost in all thirty-two layers is waste. So ‘which attention variant does this model use’ is usually the wrong question; the right one is the per-layer schedule, and the cache arithmetic is the weighted sum over it.

Choosing, and the comparisons that mislead

A short decision path. Training from scratch for normal contexts: GQA with h_kv of 4 to 8, no justification needed. If the cache still dominates, add a sliding window to most layers and keep a few global ones. If you own the serving stack and the cache is still the ceiling, MLA is the strongest frontier point. Reach for linear or SSM layers when sequences are very long and the task is streaming rather than exact lookup — and interleave rather than commit. Uptraining an existing MHA checkpoint to GQA costs a low single-digit percentage of the original compute.

Three pitfalls void a comparison. Comparing at equal parameters instead of equal cache flatters MHA, since the saved parameters should have been spent elsewhere. Comparing on perplexity alone hides exactly the failure these variants cause — run needle-in-a-haystack and multi-hop retrieval instead. And comparing theoretical FLOPs ignores whether a fused kernel exists: a variant with 4× fewer FLOPs and no FlashAttention-class implementation is usually slower in wall-clock than the dense baseline it replaced.