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-train/src/bin/export_safetensors.rs

@@ -174,7 +174,7 @@ fn config_json(cfg: &Config) -> String {
         ffn = cfg.ffn_hidden,
         layers = cfg.n_layers,
         heads = cfg.n_heads,
-        kv_heads = cfg.n_heads, // xtrain is MHA → kv heads == query heads
+        kv_heads = cfg.num_kv_heads, // GQA (T15): real num_key_value_heads (= n_heads for MHA)
         head_dim = cfg.head_dim,
         eps = cfg.eps,
         theta = cfg.rope_theta,
@@ -206,6 +206,8 @@ fn main() {
     let head_dim = flag(&args, "--head-dim", 16usize);
     let n_layers = flag(&args, "--layers", 4usize);
     let ffn = flag(&args, "--ffn", 64usize);
+    // GQA (Phase T15): num K/V heads (must match the trained ckpt; default = --heads).
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
 
     assert!(device::device_count().unwrap() > 0, "no CUDA device");
     device::set_device(0).unwrap();
@@ -214,15 +216,16 @@ fn main() {
     // Size the model from the arch flags + gpt2 vocab; must match the checkpoint.
     let tok = Tokenizer::from_file(&tok_path);
     let vocab = tok.vocab_size();
-    let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
+    let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
     println!(
-        "export: ckpt {} → {} (vocab {}, dim {}, layers {}, heads {}, head_dim {})",
+        "export: ckpt {} → {} (vocab {}, dim {}, layers {}, heads {}, kv_heads {}, head_dim {})",
         ckpt.display(),
         out_dir.display(),
         cfg.vocab,
         cfg.dim,
         cfg.n_layers,
         cfg.n_heads,
+        cfg.num_kv_heads,
         cfg.head_dim,
     );
 

crates/xtrain-train/src/bin/greedy_sample.rs

@@ -72,6 +72,8 @@ fn main() {
     let head_dim = flag(&args, "--head-dim", 16usize);
     let n_layers = flag(&args, "--layers", 4usize);
     let ffn = flag(&args, "--ffn", 64usize);
+    // GQA (Phase T15): num K/V heads (must match the ckpt; default = --heads).
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
     let max_new = flag(&args, "--max-tokens", 40usize);
 
     assert!(device::device_count().unwrap() > 0, "no CUDA device");
@@ -79,14 +81,16 @@ fn main() {
     let device = Device::Cuda(0);
 
     let tok = Tokenizer::from_file(&tok_path);
-    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn);
+    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn)
+        .with_kv_heads(kv_heads);
     println!(
-        "greedy_sample: ckpt {} (vocab {}, dim {}, layers {}, heads {}, head_dim {})",
+        "greedy_sample: ckpt {} (vocab {}, dim {}, layers {}, heads {}, kv_heads {}, head_dim {})",
         ckpt.display(),
         cfg.vocab,
         cfg.dim,
         cfg.n_layers,
         cfg.n_heads,
+        cfg.num_kv_heads,
         cfg.head_dim,
     );
 

crates/xtrain-train/src/bin/train.rs

@@ -88,6 +88,8 @@ fn main() {
     let head_dim = flag(&args, "--head-dim", 16usize);
     let n_layers = flag(&args, "--layers", 4usize);
     let ffn = flag(&args, "--ffn", 64usize);
+    // GQA (Phase T15): num K/V heads (must divide --heads). Default = --heads (MHA).
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
     // `--dim` is informational; dim is always n_heads*head_dim. Warn on mismatch.
     let dim_flag = flag(&args, "--dim", 0usize);
     if dim_flag != 0 && dim_flag != n_heads * head_dim {
@@ -160,14 +162,16 @@ fn main() {
         (corpus, None)
     };
 
-    let mut cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
+    let mut cfg =
+        Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
     cfg.dropout = dropout;
     println!(
-        "model: dim {} layers {} heads {} head_dim {} ffn {} → core {:.3}M params \
+        "model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
          (+ embed/lm {:.2}M = {:.2}M total)",
         cfg.dim,
         cfg.n_layers,
         cfg.n_heads,
+        cfg.num_kv_heads,
         cfg.head_dim,
         cfg.ffn_hidden,
         cfg.core_params() as f32 / 1e6,