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/src/config.rs

@@ -10,8 +10,15 @@ pub struct Config {
     pub dim: usize,
     /// Number of decoder blocks.
     pub n_layers: usize,
-    /// Number of attention heads.
+    /// Number of attention (query) heads.
     pub n_heads: usize,
+    /// Number of key/value heads (Phase T15, GQA). Each KV head is shared by a
+    /// group of `n_heads / num_kv_heads` query heads (repeat_kv). Must divide
+    /// `n_heads`. `num_kv_heads == n_heads` (the default) = MHA, bit-identical to
+    /// the pre-T15 path; `num_kv_heads < n_heads` = real grouped-query attention,
+    /// shrinking the K/V projections to `num_kv_heads * head_dim` and exported as a
+    /// real `num_key_value_heads`.
+    pub num_kv_heads: usize,
     /// Per-head dimension (`dim / n_heads`).
     pub head_dim: usize,
     /// SwiGLU hidden width (gate/up project to this, down projects back).
@@ -37,6 +44,7 @@ impl Config {
             dim: n_heads * head_dim,
             n_layers: 2,
             n_heads,
+            num_kv_heads: n_heads, // default = MHA
             head_dim,
             ffn_hidden: 64,
             eps: 1e-5,
@@ -62,6 +70,7 @@ impl Config {
             dim: n_heads * head_dim,
             n_layers,
             n_heads,
+            num_kv_heads: n_heads, // default = MHA; set via with_kv_heads for GQA
             head_dim,
             ffn_hidden,
             eps: 1e-5,
@@ -70,6 +79,27 @@ impl Config {
         }
     }
 
+    /// Set the number of K/V heads (Phase T15, GQA). Builder-style so existing
+    /// `from_arch` call sites stay MHA unless they opt in. Asserts `num_kv_heads`
+    /// divides `n_heads`.
+    pub fn with_kv_heads(mut self, num_kv_heads: usize) -> Self {
+        assert!(num_kv_heads > 0, "num_kv_heads must be > 0");
+        assert_eq!(
+            self.n_heads % num_kv_heads,
+            0,
+            "n_heads {} not divisible by num_kv_heads {num_kv_heads}",
+            self.n_heads
+        );
+        self.num_kv_heads = num_kv_heads;
+        self
+    }
+
+    /// KV projection width (`num_kv_heads * head_dim`). For GQA this is smaller than
+    /// `dim`; for MHA it equals `dim`.
+    pub fn kv_dim(&self) -> usize {
+        self.num_kv_heads * self.head_dim
+    }
+
     /// Transformer-core parameter count: everything except the token embedding and
     /// the LM head (the two `vocab × dim` tables). This is the figure the scaling
     /// ladder is sized against — the 50257-vocab embed+lm_head adds a fixed ~25M on
@@ -82,7 +112,8 @@ impl Config {
     pub fn num_params(&self) -> usize {
         let per_layer = 2 * self.dim                       // 2 rmsnorm gammas
             + 2 * self.head_dim                             // q/k per-head norm gammas
-            + 3 * self.dim * self.dim                       // q/k/v proj
+            + self.dim * self.dim                           // q proj [dim,dim]
+            + 2 * self.dim * self.kv_dim()                  // k/v proj [dim,kv_dim] (GQA: smaller)
             + self.dim * self.dim                           // out proj
             + 2 * self.dim * self.ffn_hidden                // gate/up proj
             + self.ffn_hidden * self.dim; // down proj