From a603d953c1c1456d0ce009f7cd86c4fd04b64c69 Mon Sep 17 00:00:00 2001 From: jzzliu Date: Sat, 7 Nov 2026 20:25:52 -0700 Subject: [PATCH] Add grouped-query attention with parity tests against multi-head GQA/MQA via GPTConfig.n_kv_head; the cache stores only the KV heads. Attention also accepts per-row positions and a key padding mask, used by batched decoding. --- scripts/bench_gqa.py | 46 ++++++++++++++++++ src/microgpt/model.py | 86 +++++++++++++++++++++++++++------- src/microgpt/train.py | 2 + tests/helpers.py | 19 ++++++++ tests/test_gqa.py | 105 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 241 insertions(+), 17 deletions(-) create mode 100644 scripts/bench_gqa.py create mode 100644 tests/helpers.py create mode 100644 tests/test_gqa.py diff --git a/scripts/bench_gqa.py b/scripts/bench_gqa.py new file mode 100644 index 0000000..99d7605 --- /dev/null +++ b/scripts/bench_gqa.py @@ -0,0 +1,46 @@ +"""KV-cache memory and decode speed for multi-head vs grouped-query vs multi-query attention. + + python scripts/bench_gqa.py +""" + +from __future__ import annotations + +import time + +import torch + +from microgpt.model import GPT, GPTConfig, kv_cache_bytes + + +def main() -> None: + torch.manual_seed(0) + n_head, seq, new = 8, 448, 48 + print(f"{n_head} query heads, 6 layers, 256-dim, cache holds {seq + new} tokens\n") + print("| KV heads | variant | params | KV cache | vs MHA | decode (tok/s) |") + print("|---:|:--|---:|---:|---:|---:|") + base = None + for n_kv in (8, 4, 2, 1): + cfg = GPTConfig( + vocab_size=96, block_size=512, n_layer=6, n_head=n_head, n_embd=256, n_kv_head=n_kv + ) + model = GPT(cfg).eval() + prompt = torch.randint(0, cfg.vocab_size, (1, seq)) + with torch.no_grad(): + _, _, cache = model(prompt) # prefill + tok = prompt[:, -1:] + t0 = time.perf_counter() + for _ in range(new): + logits, _, cache = model(tok, cache=cache) + tok = logits[:, -1].argmax(-1, keepdim=True) + rate = new / (time.perf_counter() - t0) + nbytes = kv_cache_bytes(cache) + base = base or nbytes + name = {8: "MHA", 1: "MQA"}.get(n_kv, "GQA") + print( + f"| {n_kv} | {name} | {model.num_params() / 1e6:.2f}M | {nbytes / 2**20:.1f} MiB " + f"| {nbytes / base:.2f}x | {rate:.0f} |" + ) + + +if __name__ == "__main__": + main() diff --git a/src/microgpt/model.py b/src/microgpt/model.py index 4041ff8..dc333b5 100644 --- a/src/microgpt/model.py +++ b/src/microgpt/model.py @@ -3,7 +3,8 @@ Architecture (the Llama recipe, at toy scale): * RMSNorm pre-normalisation * Rotary position embeddings (RoPE) applied to queries and keys - * Causal multi-head self-attention, written out explicitly (no nn.Transformer) + * Causal self-attention, written out explicitly (no nn.Transformer), with optional + grouped-query attention (GQA/MQA) to shrink the KV cache * SwiGLU feed-forward network * Tied input/output embeddings * Per-layer key/value cache for incremental decoding @@ -32,6 +33,16 @@ class GPTConfig: n_head: int = 4 n_embd: int = 128 rope_base: float = 10000.0 + n_kv_head: Optional[int] = None # None -> multi-head; 1 -> multi-query; else grouped-query + + @property + def kv_heads(self) -> int: + return self.n_head if self.n_kv_head is None else self.n_kv_head + + +def kv_cache_bytes(cache: KVCache) -> int: + """Total bytes held by a key/value cache.""" + return sum(t.numel() * t.element_size() for kv in cache for t in kv) class RMSNorm(nn.Module): @@ -67,40 +78,71 @@ class CausalSelfAttention(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() assert cfg.n_embd % cfg.n_head == 0, "n_embd must be divisible by n_head" + if cfg.n_head % cfg.kv_heads != 0: + raise ValueError(f"n_head {cfg.n_head} not divisible by kv heads {cfg.kv_heads}") self.n_head = cfg.n_head + self.n_kv = cfg.kv_heads self.head_dim = cfg.n_embd // cfg.n_head assert self.head_dim % 2 == 0, "head_dim must be even for RoPE" - self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False) + # One projection produces n_head query heads but only n_kv key/value heads. + self.qkv = nn.Linear(cfg.n_embd, (cfg.n_head + 2 * self.n_kv) * self.head_dim, bias=False) self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False) cos, sin = rope_tables(self.head_dim, cfg.block_size, cfg.rope_base) self.register_buffer("cos", cos, persistent=False) self.register_buffer("sin", sin, persistent=False) def forward( - self, x: torch.Tensor, cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + self, + x: torch.Tensor, + cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + positions: Optional[torch.Tensor] = None, + key_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """x: (B, T, C). Optional batching arguments, used for variable-length batches: + + positions: (B, T) true position of each new token (so left-padding does not shift RoPE). + key_mask: (B, past + T) bool, False for padding columns that must never be attended to. + """ B, T, C = x.shape + hd = self.head_dim past = 0 if cache is None else cache[0].size(2) - q, k, v = self.qkv(x).split(C, dim=2) - q, k, v = (t.view(B, T, self.n_head, self.head_dim).transpose(1, 2) for t in (q, k, v)) + q, k, v = self.qkv(x).split([self.n_head * hd, self.n_kv * hd, self.n_kv * hd], dim=2) + q = q.view(B, T, self.n_head, hd).transpose(1, 2) # (B, n_head, T, hd) + k = k.view(B, T, self.n_kv, hd).transpose(1, 2) # (B, n_kv, T, hd) + v = v.view(B, T, self.n_kv, hd).transpose(1, 2) # Positions are absolute, so cached keys keep the rotation they were stored with. - cos, sin = self.cos[past : past + T], self.sin[past : past + T] + if positions is None: + cos, sin = self.cos[past : past + T], self.sin[past : past + T] + else: + cos, sin = self.cos[positions].unsqueeze(1), self.sin[positions].unsqueeze(1) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) if cache is not None: k = torch.cat((cache[0], k), dim=2) v = torch.cat((cache[1], v), dim=2) - new_cache = (k, v) + new_cache = (k, v) # cached at n_kv heads: this is where GQA saves memory + # Each group of g query heads shares one key/value head. Instead of copying K/V g times + # we view q as (B, n_kv, g, T, hd) and let matmul broadcast over the group axis. S = k.size(2) # total keys = past + T - scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) - # Query i sits at absolute position past + i and may see keys j <= past + i. + g = self.n_head // self.n_kv + q = q.reshape(B, self.n_kv, g, T, hd) + scores = (q @ k.unsqueeze(2).transpose(-2, -1)) / math.sqrt(hd) # (B, n_kv, g, T, S) + + # Query i sits at absolute column past + i and may see columns j <= past + i. allowed = torch.ones(T, S, dtype=torch.bool, device=x.device).tril(diagonal=S - T) + if key_mask is not None: + # A padded query would otherwise have no visible key (softmax over all -inf = NaN, + # which would then poison real tokens through v). Let every query see itself. + own = torch.arange(T, device=x.device)[:, None] + (S - T) == torch.arange( + S, device=x.device + ) + allowed = ((allowed[None] & key_mask[:, None, :]) | own[None])[:, None, None] scores = scores.masked_fill(~allowed, float("-inf")) - out = F.softmax(scores, dim=-1) @ v - out = out.transpose(1, 2).contiguous().view(B, T, C) + out = F.softmax(scores, dim=-1) @ v.unsqueeze(2) # (B, n_kv, g, T, hd) + out = out.reshape(B, self.n_head, T, hd).transpose(1, 2).reshape(B, T, C) return self.proj(out), new_cache @@ -125,8 +167,8 @@ class Block(nn.Module): self.norm2 = RMSNorm(cfg.n_embd) self.mlp = SwiGLU(cfg) - def forward(self, x, cache=None): - a, new_cache = self.attn(self.norm1(x), cache) + def forward(self, x, cache=None, positions=None, key_mask=None): + a, new_cache = self.attn(self.norm1(x), cache, positions, key_mask) x = x + a x = x + self.mlp(self.norm2(x)) return x, new_cache @@ -160,17 +202,27 @@ class GPT(nn.Module): idx: torch.Tensor, targets: Optional[torch.Tensor] = None, cache: Optional[KVCache] = None, + positions: Optional[torch.Tensor] = None, + key_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], KVCache]: - """Return (logits, loss, new_cache). Pass `cache` to decode incrementally.""" + """Return (logits, loss, new_cache). Pass `cache` to decode incrementally. + + `positions` / `key_mask` enable variable-length batches (see microgpt.batching). + """ _, T = idx.shape past = 0 if cache is None else cache[0][0].size(2) - if past + T > self.cfg.block_size: - raise ValueError(f"sequence length {past + T} exceeds block_size {self.cfg.block_size}") + if positions is None: + if past + T > self.cfg.block_size: + raise ValueError( + f"sequence length {past + T} exceeds block_size {self.cfg.block_size}" + ) + elif int(positions.max()) >= self.cfg.block_size: + raise ValueError(f"position exceeds block_size {self.cfg.block_size}") x = self.tok_emb(idx) new_cache: KVCache = [] for i, block in enumerate(self.blocks): - x, layer_cache = block(x, None if cache is None else cache[i]) + x, layer_cache = block(x, None if cache is None else cache[i], positions, key_mask) new_cache.append(layer_cache) logits = self.lm_head(self.norm_f(x)) diff --git a/src/microgpt/train.py b/src/microgpt/train.py index c667f51..2f3b2a6 100644 --- a/src/microgpt/train.py +++ b/src/microgpt/train.py @@ -43,6 +43,7 @@ def main() -> None: p.add_argument("--n-layer", type=int, default=4) p.add_argument("--n-head", type=int, default=4) p.add_argument("--n-embd", type=int, default=128) + p.add_argument("--n-kv-head", type=int, default=None, help="KV heads (GQA); default = n-head") p.add_argument("--lr", type=float, default=3e-3) p.add_argument("--warmup", type=int, default=50) p.add_argument("--eval-every", type=int, default=100) @@ -62,6 +63,7 @@ def main() -> None: n_layer=args.n_layer, n_head=args.n_head, n_embd=args.n_embd, + n_kv_head=args.n_kv_head, ) model = GPT(cfg) print(f"vocab={tok.vocab_size} params={model.num_params():,} train_chars={len(train_data):,}") diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..b95f214 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,19 @@ +import torch +import torch.nn as nn + +from microgpt.model import GPT, GPTConfig + + +def make_model(cfg: GPTConfig, seed: int = 0, std: float = 0.3) -> GPT: + """Random model with *confident* outputs. + + The default init (std 0.02) gives near-uniform logits, where greedy argmax can flip on + float noise and tests would compare ties. A wider init makes distinct, stable argmaxes. + """ + torch.manual_seed(seed) + model = GPT(cfg) + with torch.no_grad(): + for mod in model.modules(): + if isinstance(mod, (nn.Linear, nn.Embedding)): + mod.weight.normal_(0.0, std) + return model.eval() diff --git a/tests/test_gqa.py b/tests/test_gqa.py new file mode 100644 index 0000000..850e7ac --- /dev/null +++ b/tests/test_gqa.py @@ -0,0 +1,105 @@ +import pytest +import torch +from helpers import make_model + +from microgpt.model import GPT, GPTConfig, kv_cache_bytes + +BASE = dict(vocab_size=50, block_size=32, n_layer=2, n_head=4, n_embd=32) + + +def _cfg(n_kv_head): + return GPTConfig(**BASE, n_kv_head=n_kv_head) + + +def test_cache_shrinks_in_proportion_to_kv_heads(): + idx = torch.randint(0, 50, (2, 16)) + sizes = {} + for n_kv in (4, 2, 1): + _, _, cache = make_model(_cfg(n_kv))(idx) + assert cache[0][0].shape == (2, n_kv, 16, 8) + sizes[n_kv] = kv_cache_bytes(cache) + assert sizes[2] * 2 == sizes[4] + assert sizes[1] * 4 == sizes[4] + + +def test_fewer_kv_heads_means_fewer_parameters(): + assert GPT(_cfg(2)).num_params() < GPT(_cfg(None)).num_params() + + +@pytest.mark.parametrize("n_kv", [2, 1]) +def test_gqa_equals_mha_with_duplicated_kv_heads(n_kv): + """GQA must be exactly multi-head attention whose K/V heads are repeated per group. + + Build an MHA model, copy the GQA weights into it with each K/V head duplicated across its + query group, and require identical logits (full forward and cached decoding). + """ + gqa = make_model(_cfg(n_kv), seed=1) + mha = make_model(_cfg(None), seed=2) + g, hd, C = 4 // n_kv, 8, 32 + + state = {k: v.clone() for k, v in gqa.state_dict().items()} + for name in list(state): + if name.endswith("attn.qkv.weight"): + w = state[name] + q, k, v = w.split([4 * hd, n_kv * hd, n_kv * hd], dim=0) + dup = lambda t: t.view(n_kv, hd, C).repeat_interleave(g, dim=0).reshape(4 * hd, C) # noqa: E731 + state[name] = torch.cat([q, dup(k), dup(v)], dim=0) + mha.load_state_dict(state) + + idx = torch.randint(0, 50, (2, 20)) + with torch.no_grad(): + full_g, full_m = gqa(idx)[0], mha(idx)[0] + assert torch.allclose(full_g, full_m, atol=1e-5) + + _, _, cg = gqa(idx[:, :9]) + _, _, cm = mha(idx[:, :9]) + for t in range(9, 20): + lg, _, cg = gqa(idx[:, t : t + 1], cache=cg) + lm, _, cm = mha(idx[:, t : t + 1], cache=cm) + assert torch.allclose(lg, lm, atol=1e-5) + assert torch.allclose(lg[:, 0], full_g[:, t], atol=1e-5) + + +@pytest.mark.parametrize("n_kv", [None, 2, 1]) +def test_cached_decoding_matches_full_forward(n_kv): + model = make_model(_cfg(n_kv)) + idx = torch.randint(0, 50, (2, 18)) + with torch.no_grad(): + full = model(idx)[0] + logits, _, cache = model(idx[:, :5]) + steps = [logits] + for t in range(5, 18): + step, _, cache = model(idx[:, t : t + 1], cache=cache) + steps.append(step) + assert torch.allclose(full, torch.cat(steps, dim=1), atol=1e-5) + + +def test_no_future_leakage_with_gqa(): + model = make_model(_cfg(2)) + a = torch.randint(0, 50, (1, 12)) + b = a.clone() + b[0, 8:] = (b[0, 8:] + 1) % 50 + with torch.no_grad(): + la, lb = model(a)[0], model(b)[0] + assert torch.allclose(la[:, :8], lb[:, :8], atol=1e-6) + + +def test_invalid_group_size_is_rejected(): + with pytest.raises(ValueError): + GPT(_cfg(3)) # 4 query heads cannot be split into 3 groups + + +def test_gqa_model_can_learn(): + torch.manual_seed(0) + m = GPT(_cfg(2)) + x = torch.randint(0, 50, (4, 16)) + y = torch.roll(x, -1, dims=1) + opt = torch.optim.AdamW(m.parameters(), lr=3e-3) + first = None + for _ in range(80): + loss = m(x, y)[1] + first = first if first is not None else loss.item() + opt.zero_grad() + loss.backward() + opt.step() + assert loss.item() < 0.3 * first