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_dump.rs

@@ -58,8 +58,20 @@ fn dump_for_parity() {
     // sequence-major to [B*S]=8 ids. Per-sequence RoPE position (resets at the
     // sequence boundary) + per-sequence causal masking (no cross-sequence
     // attention) are both checked against PyTorch.
+    // Default: tiny MHA (2 heads). With XTRAIN_PARITY_KV_HEADS=k set, dump a real
+    // GQA config (8 query heads / k kv heads) so parity.py checks GQA at B>1 — the
+    // kv-projection shapes + the repeat_kv group-sum backward against PyTorch.
     let mut cfg = Config::tiny();
     cfg.vocab = 12;
+    if let Ok(kv) = std::env::var("XTRAIN_PARITY_KV_HEADS") {
+        let kv: usize = kv.parse().expect("XTRAIN_PARITY_KV_HEADS");
+        cfg = Config::from_arch(cfg.vocab, 8, cfg.head_dim, cfg.n_layers, cfg.ffn_hidden)
+            .with_kv_heads(kv);
+        println!(
+            "parity: GQA config (n_heads {} kv_heads {})",
+            cfg.n_heads, cfg.num_kv_heads
+        );
+    }
     let batch = 2usize;
     let seq = 4usize;
     let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major
@@ -92,6 +104,7 @@ fn dump_for_parity() {
         writeln!(f, "dim {}", cfg.dim).unwrap();
         writeln!(f, "n_layers {}", cfg.n_layers).unwrap();
         writeln!(f, "n_heads {}", cfg.n_heads).unwrap();
+        writeln!(f, "num_kv_heads {}", cfg.num_kv_heads).unwrap();
         writeln!(f, "head_dim {}", cfg.head_dim).unwrap();
         writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
         writeln!(f, "eps {:e}", cfg.eps).unwrap();