post-train: M2b — batched KV-cache decode (G-way, token-identical)

The rollout long-pole fix deferred from M2a: decode the G samples of one prompt
in lockstep (one forward per step over the group → G× fewer kernel launches).

- rope_pos(x, positions[]): RoPE with a per-row absolute position (new forward-
  only kernel) — G rows share one decode position. Gate: == full rope for
  [0..n], == rope_at(P) per row for uniform P (bit-identical).
- generate_cached_batch: BatchKVCache [T, G·num_kv, hd] + batched decode_step.
  decode_attention is already batch-agnostic (bh = G·nh); repeat_kv(nh, batch=G)
  broadcasts per group. No finished-mask / ragged prompts yet (perf-only / next).
- Gate (tests/decode_batch.rs): all G greedy rows token-identical to the single-
  sequence decode (8 query / 2 kv heads → exercises repeat_kv batching).

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

crates/xtrain-tensor/tests/integration.rs

@@ -159,3 +159,41 @@ fn decode_attention_matches_full_attention_last_row() {
     );
     println!("decode_attention OK: matches full causal last row (bh={bh}, t={t}, max|Δ|={max_abs:.2e})");
 }
+
+/// (e) `rope_pos` (per-row positions, M2b batched decode): with positions
+/// [0,1,…,n-1] it is bit-identical to the full-sequence `rope` (period=n); with a
+/// uniform position P every row matches `rope_at(·, P)` of that single row. This is
+/// the primitive the batched decode uses (G rows sharing one decode position).
+#[test]
+fn rope_pos_matches_rope_and_rope_at() {
+    assert!(device::device_count().expect("device count") > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let (n, heads, hd) = (7usize, 3usize, 8usize);
+    let theta = 10000.0f32;
+    let host: Vec<f32> = (0..n * heads * hd).map(|i| ((i * 37 % 101) as f32 / 50.0) - 1.0).collect();
+    let x = Tensor::from_slice(&host, &[n, heads, hd]).to_device(Device::Cuda(0));
+
+    // positions [0,1,…,n-1] ⇒ identical to the full-sequence rope.
+    let seq_pos: Vec<i32> = (0..n as i32).collect();
+    let pos_t = Tensor::from_slice(&seq_pos, &[n]).to_device(Device::Cuda(0));
+    let got = x.rope_pos(&pos_t, theta).to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    let want = x.rope(theta, n).to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    assert_eq!(got, want, "rope_pos [0..n] != full rope");
+
+    // uniform position P ⇒ each row matches rope_at(single row, P).
+    let p = 5i32;
+    let uni = Tensor::from_slice(&vec![p; n], &[n]).to_device(Device::Cuda(0));
+    let got_u = x.rope_pos(&uni, theta).to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    let row_len = heads * hd;
+    for t in 0..n {
+        let row = &host[t * row_len..(t + 1) * row_len];
+        let want_row = Tensor::from_slice(row, &[1, heads, hd])
+            .to_device(Device::Cuda(0))
+            .rope_at(theta, p as usize)
+            .to_device(Device::Cpu)
+            .as_slice::<f32>()
+            .to_vec();
+        assert_eq!(&got_u[t * row_len..(t + 1) * row_len], want_row.as_slice(), "uniform pos row {t}");
+    }
+    println!("rope_pos OK: == full rope for [0..n] and == rope_at(P) per row for uniform P");
+}