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>

csrc/ops/nn.cu

@@ -296,6 +296,25 @@ void launch_rope_pos_f32(const float* x, const int* positions, float* y,
     rope_pos_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, positions, y, heads, head_dim, theta);
 }
 
+// Concatenate along the sequence (middle) dim: a:[bh,ta,hd], b:[bh,tb,hd] →
+// out:[bh,ta+tb,hd] with out[:, :ta]=a, out[:, ta:]=b. The device-side KV-cache
+// append (M2c): keeps K/V on the GPU and grows by one token per step, removing the
+// host round-trip the M2a/M2b host cache paid. One block per bh row.
+__global__ void cat_seq_k(const float* a, const float* b, float* out,
+                          int ta_hd, int tb_hd) {
+    int i = blockIdx.x; // bh row
+    int o_hd = ta_hd + tb_hd;
+    const float* ar = a + (long)i * ta_hd;
+    const float* br = b + (long)i * tb_hd;
+    float* outr = out + (long)i * o_hd;
+    for (int j = threadIdx.x; j < ta_hd; j += blockDim.x) outr[j] = ar[j];
+    for (int j = threadIdx.x; j < tb_hd; j += blockDim.x) outr[ta_hd + j] = br[j];
+}
+void launch_cat_seq_f32(const float* a, const float* b, float* out,
+                        int bh, int ta_hd, int tb_hd, void* s) {
+    cat_seq_k<<<bh, 256, 0, (cudaStream_t)s>>>(a, b, out, ta_hd, tb_hd);
+}
+
 // Per-row scale: y[r,c] = x[r,c] * s[r]. One block per row. Used by the GRPO
 // (M4) policy-gradient backward, where each completion token's row of
 // (probs − onehot) is scaled by its own per-token coefficient.