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

Holding other model dimensions and the workload fixed, H_kv is the design variable studied here. Depth, head width and precision are also design choices.

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 GiB for one user's hypothetical MHA context, before roughly 140 decimal GB of BF16 model weights. This does not fit on one 80-GB-class accelerator; serving requires sharding, offload, reduced precision or a smaller memory requirement.

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

There is a second cost: small-batch decoding is often memory-bandwidth-bound. Full attention reads the relevant cache, though locality and implementation affect where those bytes are served from. 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)

        # Explicit teaching expansion allocates copies; native GQA may avoid it.
        k = k.repeat_interleave(self.group_size, dim=1)   # (B, H, T, d_head)
        v = v.repeat_interleave(self.group_size, dim=1)

        past = k.shape[2] - T
        allowed = (torch.arange(k.shape[2], device=x.device)[None, :]
                   <= past + torch.arange(T, device=x.device)[:, None])
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed)
        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 explicit repeat_interleave above allocates temporary expanded tensors each step. Native grouped kernels can avoid this expansion; this Python reference does not. The offset mask lets query i see keys through past+i.

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

One can decompress keys and then apply RoPE correctly. The difficulty is efficient weight absorption: position-dependent rotations generally cannot be absorbed into one fixed query/up-projection product across all key positions. Decoupling the positional branch avoids rebuilding the large content keys.

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.

Queries may also use a low-rank projection. Whether its factors are explicitly evaluated or algebraically fused is an implementation choice; query compression is not inherently a training-only operation.

The numbers

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

~57× smaller against this chosen equal-width MHA baseline. DeepSeek's content and positional key widths differ; this is not an exact count of every possible expanded-key implementation or a universal MLA compression ratio.

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.

Use row-vector notation with cached C of shape (T,r). For one head let U_K:(r,d_k), U_V:(r,d_v), query q:(1,d_k), and output slice O_h:(d_v,d_model). Expanded content scores are q (C U_K).T; absorbed scores are (q U_K.T) C.T. If a is the resulting attention row, the head output (a C U_V) O_h equals (a C)(U_V O_h). The latter aggregates in latent space. Head-specific softmax remains necessary; absorption does not collapse different heads into one attention distribution.

import torch
torch.manual_seed(9)
dtype = torch.float64
c, uk, uv = torch.randn(5, 3, dtype=dtype), torch.randn(3, 4, dtype=dtype), torch.randn(3, 2, dtype=dtype)
q, out_proj = torch.randn(1, 4, dtype=dtype), torch.randn(2, 6, dtype=dtype)
expanded_scores = q @ (c @ uk).T
absorbed_scores = (q @ uk.T) @ c.T
torch.testing.assert_close(expanded_scores, absorbed_scores)
a = (expanded_scores / 4**0.5).softmax(-1)
torch.testing.assert_close((a @ (c @ uv)) @ out_proj, (a @ c) @ (uv @ out_proj))
past, new = 3, 2
allowed = torch.arange(past + new)[None, :] <= past + torch.arange(new)[:, None]
assert allowed.tolist() == [[True, True, True, True, False], [True] * 5]
print("Expanded and absorbed content branches agree; cached masks include offsets.")

RoPE adds a positional score branch. Rotating expanded content keys is valid but obstructs the single position-independent absorption above. Decoupled positional keys retain that efficient content path. The latent algebra here is a shape-level derivation, not the entire DeepSeek projection stack.

For GQA, persistent cache bytes are B*T*L*H_kv*(d_k+d_v)*s. The familiar factor two assumes equal key/value widths. A latent-plus-positional cache uses B*T*L*(r+d_rope)*s, before metadata/alignment. Divide by 2**30 for GiB, 10**9 for GB. Sharding requires per-rank accounting: replicated KV heads or latent states do not automatically divide by tensor-parallel degree. Failure to fit on one GPU does not imply failure to serve across devices.