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

@@ -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()

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 = self.qkv(x).split([self.n_head * hd, self.n_kv * hd, self.n_kv * hd], 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 = 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.
+        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)
+        g = self.n_head // self.n_kv
-        # Query i sits at absolute position past + i and may see keys j <= past + i.
+        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 = F.softmax(scores, dim=-1) @ v.unsqueeze(2)  # (B, n_kv, g, T, hd)
-        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        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):
+    def forward(self, x, cache=None, positions=None, key_mask=None):
-        a, new_cache = self.attn(self.norm1(x), cache)
+        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 positions is None:
             if past + T > self.cfg.block_size:
-            raise ValueError(f"sequence length {past + T} exceeds block_size {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))
 

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):,}")

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()

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