gqa: real grouped-query attention (repeat_kv op + both SDPA paths + wiring + tests)

- repeat_kv CUDA kernel: fwd head-block gather, bwd DETERMINISTIC group-sum (each
  kv head sums its group of query-head grads; no atomics) + Tensor/ops node.
- Config gains num_kv_heads (default = n_heads → MHA); wk/wv project to kv_dim;
  attention() repeat_kv-broadcasts K/V to nh heads before the UNCHANGED composed
  & flash SDPA → GQA on both paths. group=1 is identity → MHA bit-identical.
- --kv-heads flag on train/train_ddp/export_safetensors/greedy_sample; export
  writes real num_key_value_heads (xserv repeat_kv grouping aligned).
- Tests: repeat_kv grad-check (group>1 grad-sum + group=1 identity); model gqa.rs
  (GQA flash==composed fp32/bf16, group=1 bit-identical to MHA, kv-proj shape);
  parity_dump+parity.py GQA path (repeat_interleave) via XTRAIN_PARITY_KV_HEADS.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-model/tests/parity.py

@@ -52,6 +52,10 @@ cfg = read_cfg()
 DIM = int(cfg["dim"])
 NL = int(cfg["n_layers"])
 NH = int(cfg["n_heads"])
+# GQA (T15): num_kv_heads <= n_heads; each kv head shared by group query heads.
+# Default to NH (MHA) for fixtures dumped before the field existed.
+NKV = int(cfg.get("num_kv_heads", str(NH)))
+GROUP = NH // NKV
 HD = int(cfg["head_dim"])
 EPS = float(cfg["eps"])
 THETA = float(cfg["rope_theta"])
@@ -114,17 +118,23 @@ for L in layers:
     # Attention
     x = rms_norm(h, L["attn_norm"])
     q = (x @ L["wq"]).reshape(B * SEQ, NH, HD)
-    k = (x @ L["wk"]).reshape(B * SEQ, NH, HD)
-    v = (x @ L["wv"]).reshape(B * SEQ, NH, HD)
+    # GQA: K/V project to NKV heads, then repeat each kv head GROUP times to NH.
+    k = (x @ L["wk"]).reshape(B * SEQ, NKV, HD)
+    v = (x @ L["wv"]).reshape(B * SEQ, NKV, HD)
     # Per-head QK-norm (Qwen3-style), before RoPE.
     q = rms_norm(q, L["q_norm"])
     k = rms_norm(k, L["k_norm"])
     q = rope(q)  # [B*SEQ, nh, hd]
-    k = rope(k)
-    # Reshape to [B, NH, SEQ, HD] so attention runs within each sequence.
+    k = rope(k)  # [B*SEQ, nkv, hd]
+    # Reshape to [B, *, SEQ, HD]; broadcast kv heads to NH (repeat_interleave along
+    # the head axis: kv head kvh → query heads [kvh*GROUP, (kvh+1)*GROUP), matching
+    # xtrain repeat_kv + xserv repeat_kv).
     q = q.reshape(B, SEQ, NH, HD).transpose(1, 2)  # [B, nh, seq, hd]
-    k = k.reshape(B, SEQ, NH, HD).transpose(1, 2)
-    v = v.reshape(B, SEQ, NH, HD).transpose(1, 2)
+    k = k.reshape(B, SEQ, NKV, HD).transpose(1, 2)  # [B, nkv, seq, hd]
+    v = v.reshape(B, SEQ, NKV, HD).transpose(1, 2)
+    if GROUP > 1:
+        k = k.repeat_interleave(GROUP, dim=1)  # [B, nh, seq, hd]
+        v = v.repeat_interleave(GROUP, dim=1)
     scale = 1.0 / math.sqrt(HD)
     scores = (q @ k.transpose(-1, -2)) * scale + mask  # [B, nh, seq, seq]
     probs = torch.softmax(scores, dim=-1)