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

@@ -1291,6 +1291,97 @@ impl Tensor {
         (dq.to_dtype(dt), dk.to_dtype(dt), dv.to_dtype(dt))
     }
 
+    // --- GQA repeat_kv head broadcast (the T15 op) ---
+
+    /// GQA repeat_kv (Phase T15). `self` is a K or V tensor `[batch·num_kv, seq,
+    /// head_dim]`; returns `[batch·nh, seq, head_dim]` with each KV head broadcast
+    /// to its `group = nh/num_kv` query heads (query head `qh` ← kv head `qh/group`,
+    /// query heads contiguous within a group — matching xserv's repeat_kv). When
+    /// `nh == num_kv` (group 1) this is a plain copy (identity → MHA bit-identical).
+    /// bf16 callers are upcast→f32→kernel→downcast (the attention stack's fp32 policy).
+    #[cfg(not(no_cuda))]
+    pub fn repeat_kv(&self, nh: usize, batch: usize) -> Self {
+        assert_eq!(
+            self.ndim(),
+            3,
+            "repeat_kv input must be [batch·num_kv,seq,hd]"
+        );
+        let (bh_kv, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
+        assert_eq!(
+            bh_kv % batch,
+            0,
+            "repeat_kv: bh_kv {bh_kv} not div by batch {batch}"
+        );
+        let num_kv = bh_kv / batch;
+        assert_eq!(
+            nh % num_kv,
+            0,
+            "repeat_kv: nh {nh} not divisible by num_kv {num_kv}"
+        );
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .repeat_kv(nh, batch)
+                .to_dtype(DType::BF16);
+        }
+        assert_eq!(self.dtype, DType::F32, "repeat_kv supports F32/BF16");
+        let out = Tensor::zeros(&[batch * nh, seq, hd], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_repeat_kv_fwd_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                batch as i32,
+                nh as i32,
+                num_kv as i32,
+                seq as i32,
+                hd as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        out
+    }
+
+    /// Backward of [`repeat_kv`](Self::repeat_kv). `dout`:[batch·nh,seq,hd] →
+    /// `din`:[batch·num_kv,seq,hd], summing the `group` query heads that share each
+    /// KV head (the multi-group grad accumulation GQA correctness hinges on).
+    /// Deterministic (no atomics: each kv element serially sums its group rows).
+    #[cfg(not(no_cuda))]
+    pub fn repeat_kv_backward(dout: &Tensor, num_kv: usize, batch: usize) -> Self {
+        assert_eq!(
+            dout.ndim(),
+            3,
+            "repeat_kv_backward dout must be [batch·nh,seq,hd]"
+        );
+        let (bh_q, seq, hd) = (dout.shape[0], dout.shape[1], dout.shape[2]);
+        assert_eq!(bh_q % batch, 0, "repeat_kv_backward: bh_q not div by batch");
+        let nh = bh_q / batch;
+        assert_eq!(
+            nh % num_kv,
+            0,
+            "repeat_kv_backward: nh not divisible by num_kv"
+        );
+        let dt = dout.dtype;
+        if dt == DType::BF16 {
+            return Tensor::repeat_kv_backward(&dout.to_dtype(DType::F32), num_kv, batch)
+                .to_dtype(DType::BF16);
+        }
+        assert_eq!(dt, DType::F32, "repeat_kv_backward supports F32/BF16");
+        let din = Tensor::zeros(&[batch * num_kv, seq, hd], DType::F32, dout.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_repeat_kv_bwd_f32(
+                dout.data_ptr() as *const f32,
+                din.data_ptr() as *mut f32,
+                batch as i32,
+                nh as i32,
+                num_kv as i32,
+                seq as i32,
+                hd as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        din
+    }
+
     /// 4D axis-(1,2) transpose: `self`:[a,b,c,d] → [a,c,b,d],
     /// `out[i,k,j,l]=self[i,j,k,l]`. Lays out batched multi-head attention
     /// (`[B,S,nh,hd] <-> [B,nh,S,hd]`). Its own backward is the same op (swap b,c).