Attention variants: MQA, GQA, and MLA

What it is and when you reach for it

The attention page derives scaled dot-product attention and multi-head attention (MHA), where every head owns its own query, key, and value projections. That design is symmetric, and the symmetry turns out to be wasteful in exactly one place: autoregressive decoding. When a language model generates token by token, it caches the keys and values of every past position so it never recomputes them, and each new token's query must be scored against that entire cache. The cache grows linearly with context, is duplicated per head and per layer, and must be streamed from GPU memory once per generated token. On modern hardware the matrix multiplies are cheap and the memory traffic is not, so decoding speed is set less by FLOPs than by how many bytes of KV cache each step has to read, and the variant family exists to shrink those bytes. Multi-query attention (MQA) shares a single K and V head across all query heads. Grouped-query attention (GQA) interpolates, sharing each KV head across a group of query heads. Multi-head latent attention (MLA) compresses keys and values into one small latent vector per token and caches only that. You reach for this family whenever you are building or serving a decoder model at real context lengths; every prominent open-weights LLM since 2023 ships one of these, not vanilla MHA. Sparse and linear-time attention methods, which shrink the other cost (the quadratic score matrix), live on the efficient attention page; the two families are orthogonal and are often combined.

The math

The KV cache bill, worked numerically

The cache stores two tensors per layer, K and V, each of shape (batch, kv_heads, seq, head_dim). Its size in bytes is

cache = 2 · n_layers · n_{kv_heads} · d_head · T · b,

where the leading 2 counts K and V, T is the sequence length, and b is bytes per element (2 for fp16 or bf16), all times the batch size. Take a Llama-2-7B-shaped model, which uses plain MHA: 32 layers, 32 heads, head dimension 128, context 4096, fp16. Per token, the cache costs 2 · 32 · 32 · 128 · 2 = 524,288 bytes, exactly 512 KiB. At the full 4096-token context that is 512 KiB · 4096 = 2 GiB for a single sequence. Now serve a batch: 32 concurrent 4k-context requests need 64 GiB of cache on top of the 14 GB of fp16 weights, which has already overflowed an 80 GB A100. And size is only half the problem. Every decode step must read the whole cache once; at 2 GiB per sequence and roughly 2 TB/s of HBM bandwidth, cache traffic alone caps a single sequence's contribution and, across a large batch, becomes the dominant cost per token. Shrinking n_{kv_heads} divides both the residency and the traffic by the same factor, which is why it is the highest-leverage knob in the formula: the other terms (layers, head_dim, context) are fixed by the model and the workload.

MQA: one KV head for everyone

Multi-query attention, proposed by Shazeer in 2019 in "Fast Transformer Decoding: One Write-Head is All You Need", keeps h query heads but computes a single key head and a single value head:

head_i = softmax(Q_iK^T/√d_head) V,   i = 1..h,

with one shared K = XW_K and V = XW_V of head dimension d_head, instead of per-head K_i, V_i. Each query head still asks its own question; all heads now consult the same directory. The cache shrinks by a factor of h (32× in the model above, from 512 KiB to 16 KiB per token), and the FLOP count barely moves because the score and value matmuls are unchanged, only the K and V projections got narrower. The cost is expressiveness on the memory side: the model loses the ability to organize different keys and values for different heads, and in practice MQA trained from scratch gives up a measurable slice of quality on some tasks and can be less stable in training. PaLM, Falcon-7B, and StarCoder shipped MQA; it is the aggressive end of the family.

GQA: interpolating with groups

Grouped-query attention (Ainslie et al., 2023) is the obvious middle point once you see MHA and MQA as endpoints. Choose g = n_{kv_heads} with 1 ≤ g ≤ h and h divisible by g; partition the h query heads into g groups of h/g; give each group one shared K and V head:

head_i = softmax(Q_iK_⌈i·g/h⌉^T/√d_head) V_⌈i·g/h⌉.

g = h recovers MHA exactly and g = 1 recovers MQA exactly, so a single implementation parameterized by g covers the whole family, which is what the code below does. The empirical finding that made GQA the default is that quality degrades very slowly as g drops until it gets close to 1: with g = 8 the cache shrinks 4 to 8× while benchmark quality stays essentially at MHA level, and the GQA paper further showed that an existing MHA checkpoint can be converted by mean-pooling each group's K and V heads and then uptrained for a small fraction of the original compute. The adoption list is the argument: Llama 2 70B and every Llama 3 model use GQA with 8 KV heads, and Mistral 7B uses 8 KV heads against 32 query heads; Falcon-40B, Qwen 2 and later, and most 2024-onward open models made the same choice. Eight groups appears to be a sweet spot for another reason: it matches 8-way tensor parallelism, so each GPU shard holds exactly one KV head and no KV replication is needed across shards.

MLA: cache a latent, not the keys and values

Multi-head latent attention, introduced with DeepSeek-V2 and carried into V3 and R1, takes a different route: instead of sharing full-width KV heads, compress. A learned down-projection squeezes each token's hidden state x_t into a small latent

c_t = x_t W_DKV,   W_DKV ∈ ℝ^{d_model × d_c},

with d_c much smaller than what MHA caches (DeepSeek-V2 uses d_c = 512 against 128 heads of dimension 128, so the latent is 64× smaller than the 2 · 128 · 128 = 32,768 numbers per token that full MHA would store). Per-head keys and values are reconstructed by up-projections k_t⁽ⁱ⁾ = c_t W_UK⁽ⁱ⁾ and v_t⁽ⁱ⁾ = c_t W_UV⁽ⁱ⁾. Written naively that looks like you must decompress the whole cache every step, but the up-projections are linear, and linear maps can be moved across the dot product. The score for head i is

q_t⁽ⁱ⁾ · k_s⁽ⁱ⁾ = q_t⁽ⁱ⁾ · (c_s W_UK⁽ⁱ⁾) = (q_t⁽ⁱ⁾ W_UK^(i)T) · c_s,

so W_UK is absorbed into the query path: each head projects its query once into latent space and scores directly against the cached latents. Symmetrically, the output Σ_s w_s v_s⁽ⁱ⁾ = (Σ_s w_s c_s) W_UV⁽ⁱ⁾, so attention averages latents and W_UV is applied once per query, where it can be folded into the output projection W_O. After absorption, decoding is structurally MQA with one shared "KV head" of width d_c whose keys and values are the same vector, yet each head still has its own learned view of that latent through its absorbed projections. That is the trick MQA cannot do: per-head diversity survives because it lives in the projections, not in the cache. Two real complications are worth naming honestly. RoPE positional rotations do not commute with the absorption (the rotation sits between q and k and depends on position), so DeepSeek adds a small decoupled rotary key of dimension d_r = 64 per token, cached alongside the latent; and queries get their own low-rank compression to save activation memory during training. With both, the cache per token per layer is d_c + d_r = 576 elements, and the paper reports a 93.3% KV cache reduction versus its MHA equivalent, comparable in size to GQA with about 2.25 groups, while reporting quality above the MHA baseline at their scale.

Implementation, twice

One module covers MHA, GQA, and MQA, because they are one algorithm: n_kv_heads = n_heads is MHA, 1 < n_kv_heads < n_heads is GQA, and n_kv_heads = 1 is MQA. The only new mechanics relative to the baseline module are narrower K and V projections and a broadcast of each KV head across its group of query heads, done here with repeat_interleave in PyTorch and jnp.repeat in JAX. This mirrors exactly how huggingface/transformers implements Llama's repeat_kv.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class KVEfficientAttention(nn.Module):
    """Causal self-attention with a configurable KV head count.

    n_kv_heads == n_heads          -> MHA (baseline)
    1 < n_kv_heads < n_heads       -> GQA (e.g. Llama 3: 32 q, 8 kv)
    n_kv_heads == 1                -> MQA
    """

    def __init__(self, n_embd, n_heads, n_kv_heads):
        super().__init__()
        assert n_embd % n_heads == 0
        assert n_heads % n_kv_heads == 0
        self.n_heads, self.n_kv_heads = n_heads, n_kv_heads
        self.head_dim = n_embd // n_heads
        # Only K and V get narrower; Q and O are full width, so the
        # model keeps h independent questions per position.
        self.q_proj = nn.Linear(n_embd, n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(n_embd, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(n_embd, n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(n_heads * self.head_dim, n_embd, bias=False)

    def forward(self, x):                                  # x: (B, T, C)
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        # Broadcast each KV head to its group of query heads.
        # Pedagogical: production kernels index into the small KV
        # instead of materializing the repeat (that would forfeit
        # the bandwidth win the variant exists to buy).
        g = self.n_heads // self.n_kv_heads
        k = k.repeat_interleave(g, dim=1)                  # (B, Hq, T, d)
        v = v.repeat_interleave(g, dim=1)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        future = torch.triu(torch.ones(T, T, dtype=torch.bool,
                                       device=x.device), diagonal=1)
        scores = scores.masked_fill(future, float('-inf'))
        w = F.softmax(scores, dim=-1)
        y = (w @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(y)

def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len,
                   batch=1, bytes_per_elem=2):
    """K and V, each (batch, n_kv_heads, seq_len, head_dim), per layer."""
    return 2 * n_layers * n_kv_heads * head_dim * seq_len * batch * bytes_per_elem

import jax
import jax.numpy as jnp

def init_kv_efficient_attention(key, n_embd, n_heads, n_kv_heads):
    assert n_embd % n_heads == 0 and n_heads % n_kv_heads == 0
    head_dim = n_embd // n_heads
    kq, kk, kv, ko = jax.random.split(key, 4)
    s = 1.0 / jnp.sqrt(n_embd)
    # Only K and V get narrower; Q and O stay full width.
    return {
        'w_q': jax.random.normal(kq, (n_embd, n_heads * head_dim)) * s,
        'w_k': jax.random.normal(kk, (n_embd, n_kv_heads * head_dim)) * s,
        'w_v': jax.random.normal(kv, (n_embd, n_kv_heads * head_dim)) * s,
        'w_o': jax.random.normal(ko, (n_heads * head_dim, n_embd)) * s,
    }

def kv_efficient_attention(params, x, n_heads, n_kv_heads):
    """n_kv_heads == n_heads is MHA, 1 is MQA, in between is GQA."""
    B, T, C = x.shape
    head_dim = C // n_heads
    q = (x @ params['w_q']).reshape(B, T, n_heads, head_dim).swapaxes(1, 2)
    k = (x @ params['w_k']).reshape(B, T, n_kv_heads, head_dim).swapaxes(1, 2)
    v = (x @ params['w_v']).reshape(B, T, n_kv_heads, head_dim).swapaxes(1, 2)
    # Broadcast each KV head across its query group. Pedagogical:
    # fused kernels index the small KV rather than materialize this.
    g = n_heads // n_kv_heads
    k = jnp.repeat(k, g, axis=1)                       # (B, Hq, T, d)
    v = jnp.repeat(v, g, axis=1)
    scores = jnp.einsum('bhqd,bhkd->bhqk', q, k) / jnp.sqrt(head_dim)
    visible = jnp.tril(jnp.ones((T, T), dtype=bool))
    scores = jnp.where(visible, scores, -jnp.inf)
    w = jax.nn.softmax(scores, axis=-1)
    y = jnp.einsum('bhqk,bhkd->bhqd', w, v)
    y = y.swapaxes(1, 2).reshape(B, T, C)
    return y @ params['w_o']

def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len,
                   batch=1, bytes_per_elem=2):
    """K and V, each (batch, n_kv_heads, seq_len, head_dim), per layer."""
    return 2 * n_layers * n_kv_heads * head_dim * seq_len * batch * bytes_per_elem

MLA deserves code too, but with its scope stated plainly. The version below shows the load-bearing mechanics, compressing to a latent, caching only the latent, and absorbing the K up-projection into the query path and the V up-projection into the output path, and it runs. What it omits is exactly what makes DeepSeek's production version subtle: the decoupled RoPE key that restores position information (rotations do not commute with the absorption), the query-side low-rank compression, and the fused kernels. Treat it as pseudocode that happens to execute, not as a faithful DeepSeek reimplementation, and do not benchmark quality conclusions off it.

class MLADecodeSimplified(nn.Module):
    """Multi-head latent attention, decode step, simplified.

    Omits DeepSeek's decoupled RoPE keys and query compression.
    The point being demonstrated: the cache holds ONE d_latent
    vector per token, shared by all heads, yet every head keeps
    its own learned view of it via absorbed projections.
    """

    def __init__(self, d_model, n_heads, head_dim, d_latent):
        super().__init__()
        self.n_heads, self.head_dim = n_heads, head_dim
        self.w_dkv = nn.Linear(d_model, d_latent, bias=False)   # compress
        self.w_q = nn.Linear(d_model, n_heads * head_dim, bias=False)
        # Per-head up-projections kept explicit so absorption is visible.
        self.w_uk = nn.Parameter(torch.randn(n_heads, d_latent, head_dim)
                                 / math.sqrt(d_latent))
        self.w_uv = nn.Parameter(torch.randn(n_heads, d_latent, head_dim)
                                 / math.sqrt(d_latent))
        self.w_o = nn.Linear(n_heads * head_dim, d_model, bias=False)

    def forward(self, x, latent_cache):
        # x: (B, 1, d_model) newest token
        # latent_cache: (B, T_past, d_latent), the ONLY cached state
        B = x.size(0)
        c_new = self.w_dkv(x)                          # (B, 1, d_latent)
        cache = torch.cat([latent_cache, c_new], dim=1)  # (B, T, d_latent)
        q = self.w_q(x).view(B, self.n_heads, self.head_dim)
        # Absorb W_uk into the query: score against latents directly.
        q_lat = torch.einsum('bhd,hcd->bhc', q, self.w_uk)   # (B, H, d_latent)
        s = torch.einsum('bhc,btc->bht', q_lat, cache)
        s = s / math.sqrt(self.head_dim)
        w = F.softmax(s, dim=-1)                       # (B, H, T)
        o_lat = torch.einsum('bht,btc->bhc', w, cache)  # attend over latents
        o = torch.einsum('bhc,hcd->bhd', o_lat, self.w_uv)  # absorb W_uv
        return self.w_o(o.reshape(B, 1, -1)), cache

def init_mla(key, d_model, n_heads, head_dim, d_latent):
    k1, k2, k3, k4, k5 = jax.random.split(key, 5)
    return {
        'w_dkv': jax.random.normal(k1, (d_model, d_latent)) / jnp.sqrt(d_model),
        'w_q': jax.random.normal(k2, (d_model, n_heads * head_dim))
               / jnp.sqrt(d_model),
        # Per-head up-projections kept explicit so absorption is visible.
        'w_uk': jax.random.normal(k3, (n_heads, d_latent, head_dim))
                / jnp.sqrt(d_latent),
        'w_uv': jax.random.normal(k4, (n_heads, d_latent, head_dim))
                / jnp.sqrt(d_latent),
        'w_o': jax.random.normal(k5, (n_heads * head_dim, d_model))
               / jnp.sqrt(n_heads * head_dim),
    }

def mla_decode_step(params, x, latent_cache, n_heads, head_dim):
    """Simplified MLA decode: no decoupled RoPE keys, no query
    compression. The cache holds ONE d_latent vector per token."""
    B = x.shape[0]                                 # x: (B, 1, d_model)
    c_new = x @ params['w_dkv']                    # (B, 1, d_latent)
    cache = jnp.concatenate([latent_cache, c_new], axis=1)  # (B, T, d_c)
    q = (x @ params['w_q']).reshape(B, n_heads, head_dim)
    # Absorb W_uk into the query: score against latents directly.
    q_lat = jnp.einsum('bhd,hcd->bhc', q, params['w_uk'])
    s = jnp.einsum('bhc,btc->bht', q_lat, cache) / jnp.sqrt(head_dim)
    w = jax.nn.softmax(s, axis=-1)                 # (B, H, T)
    o_lat = jnp.einsum('bht,btc->bhc', w, cache)   # attend over latents
    o = jnp.einsum('bhc,hcd->bhd', o_lat, params['w_uv'])
    return o.reshape(B, 1, -1) @ params['w_o'], cache

Using it on a real shape of problem

Two things are worth actually running: the module at the three settings, to see that the interface and output shape never change, and the cache calculator at Llama-2-7B dimensions, to see what the family buys. The parameter count also drops as n_kv_heads falls (the K and V projections narrow), which is a side benefit, not the point.

torch.manual_seed(0)
x = torch.randn(2, 256, 4096)          # (B, T, C) at 7B width

for n_kv in (32, 8, 1):                # MHA, GQA-8 (Llama 3), MQA
    attn = KVEfficientAttention(n_embd=4096, n_heads=32, n_kv_heads=n_kv)
    y = attn(x)
    n_params = sum(p.numel() for p in attn.parameters())
    print(n_kv, y.shape, f"{n_params/1e6:.1f}M params")
# 32 torch.Size([2, 256, 4096]) 67.1M params   (full MHA layer)
#  8 torch.Size([2, 256, 4096]) 41.9M params
#  1 torch.Size([2, 256, 4096]) 34.6M params

# The cache bill at Llama-2-7B shape: 32 layers, head_dim 128,
# context 4096, fp16, one sequence.
for n_kv in (32, 8, 1):
    gib = kv_cache_bytes(32, n_kv, 128, 4096) / 2**30
    print(f"n_kv_heads={n_kv:2d}: {gib:.3f} GiB")
# n_kv_heads=32: 2.000 GiB    (MHA: the 7B baseline)
# n_kv_heads= 8: 0.500 GiB    (GQA-8: 4x smaller)
# n_kv_heads= 1: 0.062 GiB    (MQA: 32x smaller)

key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (2, 256, 4096))   # (B, T, C) at 7B width

for n_kv in (32, 8, 1):                      # MHA, GQA-8, MQA
    params = init_kv_efficient_attention(key, 4096, 32, n_kv)
    apply = jax.jit(lambda p, x, n=n_kv: kv_efficient_attention(p, x, 32, n))
    y = apply(params, x)
    n_params = sum(p.size for p in jax.tree_util.tree_leaves(params))
    print(n_kv, y.shape, f"{n_params/1e6:.1f}M params")
# 32 (2, 256, 4096) 67.1M params   (full MHA layer)
#  8 (2, 256, 4096) 41.9M params
#  1 (2, 256, 4096) 34.6M params

# The cache bill at Llama-2-7B shape: 32 layers, head_dim 128,
# context 4096, fp16, one sequence.
for n_kv in (32, 8, 1):
    gib = kv_cache_bytes(32, n_kv, 128, 4096) / 2**30
    print(f"n_kv_heads={n_kv:2d}: {gib:.3f} GiB")
# n_kv_heads=32: 2.000 GiB    (MHA: the 7B baseline)
# n_kv_heads= 8: 0.500 GiB    (GQA-8: 4x smaller)
# n_kv_heads= 1: 0.062 GiB    (MQA: 32x smaller)

The numbers to internalize are the cache ones. GQA-8 turns the 2 GiB per-sequence cache into 512 MiB, which quadruples how many concurrent sequences fit in the same memory and cuts per-token cache reads by the same factor. Training-time forward FLOPs, by contrast, are nearly identical across the three settings, which is why you will not see the win in a training-shaped benchmark; the variant family pays off at decode, where the workload is one query token against a long cache. Exact speedups are machine- and batch-dependent, but the direction is robust: throughput at large batch scales roughly with how much cache traffic you removed.

Applications

This family is why long-context serving is economical at all. Concretely: Llama 2 70B, Llama 3 (8B and 70B), Mistral 7B, Mixtral, Falcon-40B, and Qwen 2 and later all ship GQA with 8 KV heads against 32 to 64 query heads; PaLM, Falcon-7B, and StarCoder shipped MQA; DeepSeek-V2, V3, and R1 ship MLA. The practical consequences show up in three places. First, batch size: a serving engine's throughput at fixed hardware is largely a function of how many requests' caches fit in HBM, so a 4× smaller cache is roughly 4× more concurrency, which is a direct cost-per-token reduction. Second, context length: 128k-token contexts at MHA cache rates would be tens of GiB per sequence; GQA and MLA are what make the long-context tiers of commercial APIs feasible. Third, on-device inference: llama.cpp-style local deployment on a laptop or phone lives inside a few GiB, and a GQA model leaves that budget to weights instead of cache. MLA's extra credit is that DeepSeek paired the 90%+ cache reduction with quality their paper reports as above the MHA baseline, which turned "compression as a necessary evil" into "compression as a free lunch" and is a large part of why DeepSeek-V3-class models serve cheaply at scale. When you read a model card, num_kv_heads is one of the first numbers to check: it tells you the serving economics before you run anything.

Against the real libraries

huggingface/transformers (my codebase notes are at /oss/transformers) exposes the whole family through one config field: num_key_value_heads. Set it equal to num_attention_heads and the model is MHA, set it to 1 for MQA, anything between for GQA; LlamaConfig, MistralConfig, and most modern architectures carry it. Inside the modeling code the broadcast is the same repeat_kv expand you saw above when running the eager path, but with attn_implementation="flash_attention_2" or "sdpa" the repeat never materializes. What the library adds over this page is everything around the module: RoPE application, cache management classes, quantized caches, and per-model wiring, all selected by config rather than code.

flash-attention (walkthrough at /oss/flash-attention) supports MQA and GQA natively: pass q with h heads and k, v with fewer heads, and as long as the query head count is divisible by the KV head count, the kernel indexes each query head into its group's KV tile inside SRAM, so the bandwidth saving is realized rather than simulated. PyTorch's F.scaled_dot_product_attention grew the same ability behind an enable_gqa=True flag (added around PyTorch 2.5; check your version), which lets you delete the repeat_interleave entirely.

Serving engines are where the variants cash out. vLLM (notes at /oss/vllm) manages the KV cache in fixed-size pages with PagedAttention, so cache memory is allocated like virtual memory and shared across sequences with common prefixes; a GQA model needs 4 to 8× fewer pages per request, which multiplies directly into batch capacity. vLLM and SGLang (notes at /oss/sglang) both added dedicated MLA paths for the DeepSeek models, including kernels in the lineage of DeepSeek's own open-sourced FlashMLA, because absorbed-MLA decode is a different kernel shape (one wide shared latent instead of many narrow heads) than GQA.

The from-scratch version is enough for architecture experiments, ablations at research scale, and any setting where you want the attention weights inspectable. The correctness check has two layers. First, internal consistency: with n_kv_heads = n_heads the module must match the baseline MHA module from the attention page given identical weights. Second, against the library kernel:

torch.manual_seed(0)
B, Hq, Hkv, T, d = 2, 8, 2, 64, 32
q = torch.randn(B, Hq, T, d, dtype=torch.float64)
k = torch.randn(B, Hkv, T, d, dtype=torch.float64)
v = torch.randn(B, Hkv, T, d, dtype=torch.float64)

# Ours: materialize the repeat, then plain causal attention.
kk = k.repeat_interleave(Hq // Hkv, dim=1)
vv = v.repeat_interleave(Hq // Hkv, dim=1)
scores = q @ kk.transpose(-2, -1) / math.sqrt(d)
scores = scores.masked_fill(
    torch.triu(torch.ones(T, T, dtype=torch.bool), 1), float('-inf'))
ours = F.softmax(scores, dim=-1) @ vv

# Reference: SDPA's native GQA path (PyTorch >= 2.5).
ref = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                     enable_gqa=True)
print((ref - ours).abs().max())        # ~1e-16 in float64
assert torch.allclose(ref, ours, atol=1e-12)

key = jax.random.PRNGKey(0)
kq, kk_, kv_ = jax.random.split(key, 3)
B, Hq, Hkv, T, d = 2, 8, 2, 64, 32
# jax.nn.dot_product_attention uses (B, T, H, d) layout and
# supports GQA when Hq is a multiple of Hkv.
q = jax.random.normal(kq, (B, T, Hq, d))
k = jax.random.normal(kk_, (B, T, Hkv, d))
v = jax.random.normal(kv_, (B, T, Hkv, d))

ref = jax.nn.dot_product_attention(q, k, v, is_causal=True)

# Ours: repeat KV heads, then plain causal attention in (B,H,T,d).
g = Hq // Hkv
qh = q.swapaxes(1, 2)
kh = jnp.repeat(k.swapaxes(1, 2), g, axis=1)
vh = jnp.repeat(v.swapaxes(1, 2), g, axis=1)
scores = jnp.einsum('bhqd,bhkd->bhqk', qh, kh) / jnp.sqrt(d)
visible = jnp.tril(jnp.ones((T, T), dtype=bool))
scores = jnp.where(visible, scores, -jnp.inf)
ours = jnp.einsum('bhqk,bhkd->bhqd',
                  jax.nn.softmax(scores, axis=-1), vh).swapaxes(1, 2)

print(jnp.abs(ref - ours).max())       # ~1e-6 in float32
assert jnp.allclose(ref, ours, atol=1e-4)

Traps and misconceptions

"GQA makes training faster." Mostly it does not, and measuring the wrong phase is the classic mistake. In a training step the sequence is processed in parallel and the dominant costs (score matmul, value matmul, MLP) are untouched by shrinking KV heads; only the K and V projections narrow. The win is at autoregressive decode, where each step is one query against a long cache and the cache read is the bottleneck. Benchmark decode throughput at realistic batch and context, or you will conclude the variant does nothing.

Materializing the repeat in production. repeat_interleave on K and V recreates full-size tensors and, if done on the cache itself, forfeits the entire bandwidth saving. It is the right teaching implementation and an acceptable eager fallback, but the deployed path must index heads into groups inside the kernel, which is what flash-attention's GQA support, SDPA's enable_gqa, and every serving engine actually do. Related bug: computing the cache, then repeating, then caching the repeated tensors, which silently restores MHA-size memory.

"MQA/GQA shrink the score matrix." No. Every query head still scores against every position, so the attention matrix is exactly as large and exactly as quadratic as MHA's; only the number of distinct K and V tensors changed. Quadratic-cost reduction is a different axis, covered on the efficient attention page, and models like Mistral combine both (GQA plus sliding-window attention) precisely because they are independent.

Treating head-count changes as checkpoint-compatible. You cannot flip num_key_value_heads on trained MHA weights and expect a working model: the K and V projection shapes change. The GQA paper's recipe is mean-pooling the heads within each group and then uptraining for a small fraction of the original tokens; skipping the uptraining step gives a visibly damaged model. Conversely MLA cannot be obtained by surgery on an MHA checkpoint at all in the general case; it is an architectural commitment made before pretraining, with recent research exploring conversions as an open problem.

"MLA is just GQA with extra steps." The cache sizes can look similar, but the mechanism differs where it matters: GQA forces groups of query heads to share literal keys and values, while MLA gives every head its own learned projection of a shared compressed latent, so per-head diversity survives compression. That is why DeepSeek could report quality above their MHA baseline at a cache size GQA only reaches by giving up head diversity. The price is real complexity: absorption changes the decode kernel, and RoPE requires a decoupled positional key path.