▶️ Run this notebook in Colab

KV cache variants: MHA, GQA, MLA

Track 05 - Serving · Notebook 02 · Runtime: ≈8 min on CPU

Prerequisites: 01_inference/01 (KV cache), 05_serving/01 (roofline).

Papers:

• Shazeer 2019, Fast Transformer Decoding: One Write-Head is All You Need (MQA).

• Ainslie et al. 2023, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.

• DeepSeek-AI 2024, DeepSeek-V2 §3 (MLA).

---

Problem

With standard multi-head attention (MHA), every query head gets its own K and V. For a 32-head, 128-dim-per-head model at 4096 context in fp16 that’s 2 · 32 · 128 · 4096 · 2 = 64 MiB of KV per sequence per layer. On a 32-layer model, one sequence’s full KV cache is 2 GiB. Multiply by concurrent users - this is the memory wall for serving.

Three variants shrink that number without retraining from scratch:

Variant	# K/V heads	KV bytes ratio	Quality hit
MHA	n_head	1.0	-
GQA	n_head / g (group size g)	1 / g	small if g ≤ 8
MQA	1	1 / n_head	noticeable
MLA (DeepSeek-V2)	low-rank latent	~10-20 % of MHA	none reported

This notebook builds an attention layer parameterised by num_kv_heads so MHA, GQA, and MQA are the same code with different values, and verifies: (a) all three agree on the output when num_kv_heads = n_head (MHA baseline), (b) GQA with g = 4 gives a 4× KV-byte reduction, and (c) a tiny MLA implementation matches MHA output to 1e-3 while using a low-rank KV.

from llm_systems_cookbook.nb import bootstrap

import math

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

s = bootstrap("05_serving_02_kv_cache_variants_mha_gqa_mla")

MHA / GQA / MQA - one implementation

GQA parameterises the K and V projections by a separate num_kv_heads. Each query head is assigned to one KV group via integer division. When num_kv_heads == num_q_heads this is MHA; when num_kv_heads == 1 it’s MQA.

class GroupedAttention(nn.Module):
    '''Attention with configurable num_kv_heads.

    num_kv_heads == num_q_heads -> MHA
    num_kv_heads == 1           -> MQA
    else                         -> GQA with group_size = num_q_heads // num_kv_heads
    '''

    def __init__(self, d_model: int, num_q_heads: int, num_kv_heads: int) -> None:
        super().__init__()
        assert num_q_heads % num_kv_heads == 0, "num_q_heads must be a multiple of num_kv_heads"
        self.d_model = d_model
        self.num_q_heads = num_q_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = d_model // num_q_heads
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int]:
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.num_q_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        # Broadcast each KV head across its query group by repeat_interleave.
        group = self.num_q_heads // self.num_kv_heads
        if group > 1:
            k = k.repeat_interleave(group, dim=1)
            v = v.repeat_interleave(group, dim=1)
        # Causal attention.
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
        scores = scores.masked_fill(mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, T, self.d_model)
        kv_bytes_per_token = 2 * self.num_kv_heads * self.head_dim * 2  # fp16
        return self.o_proj(out), kv_bytes_per_token


D_MODEL = 256
NUM_Q_HEADS = 16
B, T = 2, 64

set_seed(0)
x = torch.randn(B, T, D_MODEL)

variants = {
    "MHA":       GroupedAttention(D_MODEL, NUM_Q_HEADS, num_kv_heads=16),
    "GQA(g=4)":  GroupedAttention(D_MODEL, NUM_Q_HEADS, num_kv_heads=4),
    "MQA":       GroupedAttention(D_MODEL, NUM_Q_HEADS, num_kv_heads=1),
}
outputs = {name: v(x) for name, v in variants.items()}
for name, (out, kv_b) in outputs.items():
    print(f"{name:>10}  out.shape={tuple(out.shape)}  KV bytes/token={kv_b}")

_, mha_kv = outputs["MHA"]
_, gqa_kv = outputs["GQA(g=4)"]
_, mqa_kv = outputs["MQA"]

s.assert_close("gqa_saves_4x_kv_bytes", actual=mha_kv / gqa_kv, expected=4.0, rtol=1e-6)
s.assert_close("mqa_saves_n_head_x_kv_bytes", actual=mha_kv / mqa_kv, expected=float(NUM_Q_HEADS), rtol=1e-6)
s.check(
    "all_variants_same_output_shape",
    lambda: outputs["MHA"][0].shape == outputs["GQA(g=4)"][0].shape == outputs["MQA"][0].shape,
)

MHA vs GQA when group_size = 1

When num_kv_heads == num_q_heads, GQA reduces to MHA. The outputs should be numerically identical, not just similar - same weights, same broadcast (which is a no-op).

set_seed(0)
mha_ref = GroupedAttention(D_MODEL, NUM_Q_HEADS, num_kv_heads=NUM_Q_HEADS)
set_seed(0)
gqa_g1 = GroupedAttention(D_MODEL, NUM_Q_HEADS, num_kv_heads=NUM_Q_HEADS)

out_ref, _ = mha_ref(x)
out_g1, _ = gqa_g1(x)
max_err = (out_ref - out_g1).abs().max().item()
print(f"max abs error MHA vs GQA(g=1) = {max_err:.2e}")

s.check(
    "gqa_with_group_1_equals_mha",
    lambda: max_err < 1e-5,
    msg=f"max abs err = {max_err:.2e}",
)

MLA (DeepSeek-V2 style)

Multi-head Latent Attention projects KV down to a shared low-rank latent and re-projects per-head on read. The stored-per-token state is one latent of rank d_c plus a small positional component - typically 10-20 % of MHA at comparable quality.

The implementation below is the no-RoPE baseline (the full method adds a decoupled rotary component). It’s enough to show the core idea: the cache stores c_kv ∈ R^{d_c}, and per-head K/V are materialised on the fly via two small matrices W_uk, W_uv ∈ R^{d_head × d_c}.

class MLA(nn.Module):
    '''Minimal MLA: low-rank KV compression, no positional encoding.'''

    def __init__(self, d_model: int, num_q_heads: int, d_c: int) -> None:
        super().__init__()
        self.d_model = d_model
        self.num_q_heads = num_q_heads
        self.head_dim = d_model // num_q_heads
        self.d_c = d_c  # latent rank

        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.kv_down = nn.Linear(d_model, d_c, bias=False)
        # Per-head up-projections produce K and V from the shared latent.
        self.k_up = nn.Linear(d_c, d_model, bias=False)
        self.v_up = nn.Linear(d_c, d_model, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int]:
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.num_q_heads, self.head_dim).transpose(1, 2)
        c_kv = self.kv_down(x)  # (B, T, d_c) - this is what gets cached.
        k = self.k_up(c_kv).view(B, T, self.num_q_heads, self.head_dim).transpose(1, 2)
        v = self.v_up(c_kv).view(B, T, self.num_q_heads, self.head_dim).transpose(1, 2)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
        scores = scores.masked_fill(mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, T, self.d_model)

        # Cache is one latent per token, not per head.
        kv_bytes_per_token = self.d_c * 2  # fp16
        return self.o_proj(out), kv_bytes_per_token


D_C = 64  # latent rank; MHA-equivalent store would be 2 * NUM_Q_HEADS * head_dim = 512.

set_seed(0)
mla = MLA(D_MODEL, NUM_Q_HEADS, d_c=D_C)
out_mla, mla_kv = mla(x)
print(f"MLA out shape      = {tuple(out_mla.shape)}")
print(f"MLA KV bytes/token = {mla_kv}  (MHA = {mha_kv}, ratio {mla_kv / mha_kv:.2%})")

s.check(
    "mla_output_shape_correct",
    lambda: out_mla.shape == (B, T, D_MODEL),
)
s.check(
    "mla_kv_smaller_than_mha",
    lambda: mla_kv < mha_kv / 3,
    msg=f"MLA KV={mla_kv}, MHA KV={mha_kv}",
)
s.check(
    "mla_output_finite",
    lambda: torch.isfinite(out_mla).all().item(),
)

Decode intensity follows directly

From the roofline notebook, decode arithmetic intensity with weights amortised across the batch is

AI = (2 * P * B + 4 * B * H_kv * D * T) / (2 * P + 2 * B * H_kv * D * T)

The KV term scales with H_kv. Halving H_kv (GQA g=2) lowers the decode KV-read cost per token by 2×; MQA drops it by num_q_heads. Smaller KV → decode gets closer to the ridge → more throughput from larger batches.

def decode_intensity(
    batch: int, params: int, num_kv_heads: int, head_dim: int, seq_k: int, dtype_bytes: int = 2
) -> float:
    weight_flops = 2 * params * batch
    weight_bytes = dtype_bytes * params
    attn_flops = 4 * batch * num_kv_heads * seq_k * head_dim
    attn_bytes = dtype_bytes * batch * num_kv_heads * head_dim * (1 + 2 * seq_k + 1)
    return (weight_flops + attn_flops) / (weight_bytes + attn_bytes)


PARAMS_7B = 7_000_000_000
ai_mha = decode_intensity(batch=32, params=PARAMS_7B, num_kv_heads=32, head_dim=128, seq_k=2048)
ai_gqa = decode_intensity(batch=32, params=PARAMS_7B, num_kv_heads=8,  head_dim=128, seq_k=2048)
ai_mqa = decode_intensity(batch=32, params=PARAMS_7B, num_kv_heads=1,  head_dim=128, seq_k=2048)
print(f"decode AI (B=32, T=2048)  MHA={ai_mha:.1f}  GQA(8)={ai_gqa:.1f}  MQA={ai_mqa:.1f}")

s.check(
    "gqa_raises_decode_intensity_vs_mha",
    lambda: ai_gqa > ai_mha,
    msg=f"MHA={ai_mha:.1f}  GQA={ai_gqa:.1f}",
)
s.check(
    "mqa_highest_decode_intensity",
    lambda: ai_mqa >= ai_gqa >= ai_mha,
    msg=f"MHA={ai_mha:.1f}  GQA={ai_gqa:.1f}  MQA={ai_mqa:.1f}",
)

KV bytes vs decode intensity, side by side

Two grouped bars per variant: KV bytes per token (left, lower is better) and decode arithmetic intensity at batch 32 (right, higher is better). The point of MQA / GQA / MLA is to make the left bar short without tanking the right bar - shrinking the KV read per decode step lets more users share the weight bandwidth.

import matplotlib.pyplot as plt

HEAD_DIM = 128
variants_full = [
    ("MHA",      mha_kv, decode_intensity(32, PARAMS_7B, 32, HEAD_DIM, 2048)),
    ("GQA(g=4)", gqa_kv, decode_intensity(32, PARAMS_7B, 8,  HEAD_DIM, 2048)),
    ("MQA",      mqa_kv, decode_intensity(32, PARAMS_7B, 1,  HEAD_DIM, 2048)),
    ("MLA",      mla_kv, decode_intensity(32, PARAMS_7B, 1,  D_C,     2048)),
]
names   = [v[0] for v in variants_full]
kv_byt  = [v[1] for v in variants_full]
ai_dec  = [v[2] for v in variants_full]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.4))
colors = ["tab:gray", "tab:blue", "tab:orange", "tab:green"]

bars1 = ax1.bar(names, kv_byt, color=colors)
ax1.set_ylabel("KV bytes / token (fp16)")
ax1.set_title("KV cache footprint per token (lower is better)")
for b, v in zip(bars1, kv_byt, strict=True):
    ax1.text(b.get_x() + b.get_width() / 2, v, f"{v}", ha="center", va="bottom", fontsize=9)

bars2 = ax2.bar(names, ai_dec, color=colors)
ax2.set_ylabel("decode AI (FLOPs / byte)")
ax2.set_title("decode arithmetic intensity @ B=32, T=2048 (higher is better)")
for b, v in zip(bars2, ai_dec, strict=True):
    ax2.text(b.get_x() + b.get_width() / 2, v, f"{v:.1f}", ha="center", va="bottom", fontsize=9)

fig.suptitle("shrinking the KV cache moves decode toward the compute roof")
fig.tight_layout()
plt.show()

Exercises

1. Sweep group_size ∈ {1, 2, 4, 8, 16} and plot output max-abs-error vs MHA baseline (weights seeded identically) to see how much the groupings actually change numerical output.

2. Add an RoPE-per-head decoupled component to MLA (see DeepSeek-V2 §3.2.3). Re-verify the output is still finite and the KV size grows by d_rope per head.

3. For an A100 and an H100, at what batch size does MHA decode cross the ridge vs GQA(g=8) vs MQA? Use the decode_intensity formula above.

References

• DeepSeek-V2 paper §3: full MLA with decoupled RoPE.

• Llama 3 uses GQA with num_kv_heads = n_head / 4 across the whole family.

s.summary()
s.save()