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/tests/integration.rs

@@ -197,3 +197,29 @@ fn rope_pos_matches_rope_and_rope_at() {
     }
     println!("rope_pos OK: == full rope for [0..n] and == rope_at(P) per row for uniform P");
 }
+
+/// (f) `cat_seq` (device-side KV-cache append, M2c): concatenating [bh,ta,hd] ++
+/// [bh,tb,hd] along the seq dim equals the host-side interleaved concat (per bh row,
+/// a's block then b's block). This is the device append that removes the M2a/M2b
+/// host round-trip.
+#[test]
+fn cat_seq_matches_host_concat() {
+    assert!(device::device_count().expect("device count") > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let (bh, ta, tb, hd) = (4usize, 3usize, 2usize, 5usize);
+    let ah: Vec<f32> = (0..bh * ta * hd).map(|i| i as f32 * 0.1).collect();
+    let bhost: Vec<f32> = (0..bh * tb * hd).map(|i| -(i as f32) - 1.0).collect();
+    let a = Tensor::from_slice(&ah, &[bh, ta, hd]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&bhost, &[bh, tb, hd]).to_device(Device::Cuda(0));
+
+    let got = a.cat_seq(&b).to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    // Host reference: per bh row, a's ta*hd then b's tb*hd.
+    let mut want = vec![0f32; bh * (ta + tb) * hd];
+    for r in 0..bh {
+        let (oa, ob, oo) = (r * ta * hd, r * tb * hd, r * (ta + tb) * hd);
+        want[oo..oo + ta * hd].copy_from_slice(&ah[oa..oa + ta * hd]);
+        want[oo + ta * hd..oo + (ta + tb) * hd].copy_from_slice(&bhost[ob..ob + tb * hd]);
+    }
+    assert_eq!(got, want, "cat_seq != host interleaved concat");
+    println!("cat_seq OK: [bh={bh},{ta}+{tb},{hd}] == host concat");
+}