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