Efficient Attention: MHA → MQA → GQA → MLA

Module 08 ended on a waste: at each decode step the model recomputes K and V for every previous token, identically. The fix is obvious once seen — cache them.

# naive: recompute everything, every step
for step in range(max_new):
    K = W_k(all_tokens_so_far)     # recomputes positions 0..t-1 AGAIN
    V = W_v(all_tokens_so_far)

# cached: compute only the new token, append
for step in range(max_new):
    k_new = W_k(newest_token)      # ONE token
    v_new = W_v(newest_token)
    K = torch.cat([K_cache, k_new], dim=-2)
    V = torch.cat([V_cache, v_new], dim=-2)

Cost per step drops from O(T) to O(1) in projection work. Every serving stack does this.

But it converts a compute problem into a memory problem, and the memory problem is worse.

The cache-size formula

Every term except H_kv is fixed by the model or the workload. H_kv is the one you can design.

How bad is it?

Take a Llama-3-70B-shaped model — L = 80, H = 64 query heads, d_head = 128, bf16 — serving a single 32,768-token request:

80 GB for one user's context. An H100 has 80 GB total, and the model weights alone need ~140 GB. Full MHA at long context is simply not servable.

(Llama 3 70B actually uses GQA with 8 KV heads — the MHA row shows what it would cost without it.)

There is a second, subtler cost: decoding is memory-bandwidth-bound, not compute-bound. Generating one token requires reading the entire cache from HBM. Smaller cache means fewer bytes read means faster tokens/sec — the saving is speed as well as capacity.

Every query head owns its own K and V heads. H_kv = H.

Head 1:  Q1  K1  V1
Head 2:  Q2  K2  V2
Head 3:  Q3  K3  V3
Head 4:  Q4  K4  V4

Maximum expressiveness, maximum cache. Still used by OLMo 2 and Olmo 3 7B — Raschka notes OLMo 2 "still uses traditional Multi-Head Attention (MHA) instead of MLA or GQA," though the later 32B variant switched to GQA.

(Shazeer, 2019)

All query heads share a single K and a single V head. H_kv = 1.

Head 1:  Q1 ─┐
Head 2:  Q2 ─┼─  K_shared  V_shared
Head 3:  Q3 ─┤
Head 4:  Q4 ─┘

The cache shrinks by a factor of H — 64× for a 64-head model. Enormous.

The problem is equally large: every head now attends using the same key space. Module 04's entire argument was that different heads should capture different relational views; MQA forces them to share the lens they look through. Quality degrades measurably, and training can become unstable.

MQA is rarely used alone today, but it matters as the endpoint of a spectrum — and, as §9.5 notes, MLA effectively reaches it at inference.

(Ainslie et al., 2023)

The interpolation, and the reason it won: query heads are partitioned into groups, and each group shares one K/V pair.

Raschka's description: "if there are 2 key-value groups and 4 attention heads, then heads 1 and 2 might share one set of keys and values, while heads 3 and 4 share another."

Group 1:  Q1, Q2  ──►  K1  V1
Group 2:  Q3, Q4  ──►  K2  V2

H_kv is now a dial: H_kv = H is MHA, H_kv = 1 is MQA, anything between is GQA. Typical production values are 4 or 8.

The quality claim: ablations in the original GQA paper and the Llama 2 paper show GQA "performs comparably to standard MHA in terms of LLM modeling performance" (Raschka). An 8× cache reduction for approximately no quality loss is why GQA became the default.

Implementation

The mechanism is: project fewer K/V heads, then repeat them to match the query heads.

import torch, torch.nn as nn, torch.nn.functional as F

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, d_head=None):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.H, self.H_kv = n_heads, n_kv_heads
        self.group_size = n_heads // n_kv_heads
        self.d_head = d_head or d_model // n_heads

        self.W_q = nn.Linear(d_model, self.H    * self.d_head, bias=False)
        self.W_k = nn.Linear(d_model, self.H_kv * self.d_head, bias=False)  # SMALLER
        self.W_v = nn.Linear(d_model, self.H_kv * self.d_head, bias=False)  # SMALLER
        self.W_O = nn.Linear(self.H * self.d_head, d_model, bias=False)

    def forward(self, x, cos, sin, kv_cache=None):
        B, T, _ = x.shape
        q = self.W_q(x).view(B, T, self.H,    self.d_head).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.H_kv, self.d_head).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.H_kv, self.d_head).transpose(1, 2)

        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        if kv_cache is not None:                    # cache the SMALL tensors
            k, v = kv_cache.update(k, v)

        # expand K,V to match query heads — this is a VIEW, not a copy
        k = k.repeat_interleave(self.group_size, dim=1)   # (B, H, T, d_head)
        v = v.repeat_interleave(self.group_size, dim=1)

        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).reshape(B, T, -1)
        return self.W_O(out)

The saving is real because only the pre-expansion tensors are cached. The repeat_interleave happens on the fly, per step, and good kernels avoid materialising it at all.

Who uses GQA

(DeepSeek-V2, 2024)

DeepSeek took a different route: instead of sharing K and V, compress them.

Raschka: "Instead of sharing key and value heads like GQA, MLA compresses the key and value tensors into a lower-dimensional space before storing them in the KV cache. At inference time, these compressed tensors are projected back to their original size before being used... This adds an extra matrix multiplication but reduces memory usage."

The mechanism

standard:   x ──W_k──► K  (H · d_head)     cache K and V directly
                                            
MLA:        x ──W_dkv──► c   (d_c ≈ 512)   cache ONLY c
                          │
                          ├──W_uk──► K     reconstruct at use time
                          └──W_uv──► V

c is a low-rank latent: a single compressed vector per token per layer standing in for all heads' K and V.

The RoPE complication

There is a genuine wrinkle. RoPE is position-dependent, so it cannot be applied after decompression without breaking the low-rank structure — the up-projection and the rotation do not commute cleanly.

DeepSeek's solution: split each head's dimensions in two.

• a compressed part, reconstructed from the latent, carrying no RoPE
• a small RoPE-carrying part (64 dims), cached separately in full

So the cache holds the latent c_KV (512) plus a decoupled RoPE key (64) — 576 values per token per layer, instead of 2 × 128 heads × 128 dims = 32,768.

Also worth noting, per Raschka: "the queries are also compressed, but only during training, not inference."

The numbers

DeepSeek V3 — L = 61, 128 heads, d_head = 128, bf16, 32k context:

~57× smaller. That is a different regime, not an improvement.

The surprising part: it is also better

Compression usually costs quality. Raschka reports the opposite, from the DeepSeek-V2 ablations: "GQA appears to perform worse than MHA, whereas MLA offers better modeling performance than MHA, which is likely why the DeepSeek team chose MLA over GQA."

He adds a fair caveat — the paper does not show a "KV Cache per Token" comparison between MLA and GQA directly — and elsewhere notes that "according to the ablation studies in the DeepSeek-V2 paper, MLA is on par with regular full attention when the hyperparameters are carefully chosen." So: MLA is somewhere between "on par with" and "slightly better than" MHA, while using a fraction of the cache.

Why might compression help? The low-rank bottleneck is a regulariser. Forcing K/V through 512 dimensions may discard noise that full-rank projections preserve.

The cost is complexity. MLA is substantially harder to implement, and kernel support lagged for a long time.

Adoption

MLA has spread beyond DeepSeek:

Raschka's note on Mistral 3 Large is a good summary of where the field landed: "why change what ain't broke? A lot of the secret sauce these days is in the training pipeline as well as the inference scaling strategies."

Choosing: GQA unless you have a reason not to. It captures most of the benefit at a fraction of the implementation cost, which is why it is the 2026 default. MLA is worth it at frontier scale where cache dominates the serving budget and you can afford the engineering.

All three shrink the cache by a constant factor. None change that it grows linearly with T, or that attention compute is still O(T²).

Halving the constant buys you 2× context. It does not buy you 100×.

For that you need to change the attention pattern (module 10) or the kernel (module 11). Those are orthogonal to this module — Gemma 3 combines GQA with sliding windows; DeepSeek V3.2 combines MLA with sparse attention; everything runs on FlashAttention.