post-train: M2c — device-side KV cache (cat_seq), profile-first bottleneck shift

Device-resident KV cache: keep K/V on the GPU as [bh,T,hd], grow by one token
per step via a new cat_seq kernel (concat along seq) — removes the M2a/M2b
per-layer host round-trip (to_cpu/from_slice/re-upload) AND the transpose_3d01.
Both single-seq and batched decode refactored to it; cache is Option<Tensor>
per layer (cleaner than the host Vec + rebuild).

Gates all hold: cat_seq == host concat; decode_kv single-seq + decode_batch
G-way both still TOKEN-IDENTICAL; GQA training path unaffected.

Honest measurement (the point): removing the host round-trip buys ~10% on pure
single-seq decode (133 → 147 tok/s @128) but does NOT move the GRPO step
(~8.5 s/step unchanged) — because after M2b batching the rollout is no longer
the step's bottleneck; the per-sample per_token_logp captures + the PG-update
forwards/backwards (model.forward, full-seq) now dominate. Measure-first lesson
(cf. T11/T17/M2a): the long pole shifted to the training-side forwards; the next
decode lever (ragged batched prefill) targets those, not the cache.

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

crates/xtrain-tensor/src/tensor.rs

@@ -856,6 +856,41 @@ impl Tensor {
         out
     }
 
+    /// Concatenate along the sequence (middle) dim: `self`:[bh,ta,hd] ++
+    /// `other`:[bh,tb,hd] → `[bh,ta+tb,hd]`. The device-side KV-cache append (M2c):
+    /// the cache stays on the GPU and grows by one token per decode step, removing
+    /// the M2a/M2b host round-trip. Mirrors the bf16 cast handling of the other
+    /// structural kernels.
+    #[cfg(not(no_cuda))]
+    pub fn cat_seq(&self, other: &Tensor) -> Self {
+        assert_eq!(self.ndim(), 3, "cat_seq requires [bh,t,hd]");
+        assert_eq!(other.ndim(), 3, "cat_seq requires [bh,t,hd]");
+        assert_eq!(self.dtype, other.dtype, "cat_seq dtype mismatch");
+        let (bh, ta, hd) = (self.shape[0], self.shape[1], self.shape[2]);
+        let (bh2, tb, hd2) = (other.shape[0], other.shape[1], other.shape[2]);
+        assert_eq!(bh, bh2, "cat_seq bh mismatch");
+        assert_eq!(hd, hd2, "cat_seq head_dim mismatch");
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .cat_seq(&other.to_dtype(DType::F32))
+                .to_dtype(DType::BF16);
+        }
+        let out = Tensor::zeros(&[bh, ta + tb, hd], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_cat_seq_f32(
+                self.data_ptr() as *const f32,
+                other.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                bh as i32,
+                (ta * hd) as i32,
+                (tb * hd) as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        out
+    }
+
     /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
     /// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
     #[cfg(not(no_cuda))]