test: T21-for-proc — clear ENV_DROPOUT across tests to sever ordering coupling

libtest with --test-threads=1 (the documented invariant for this file's DDP
tests) runs tests alphabetically. The new
proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout ('d') runs BEFORE
proc_per_gpu_matches_single_gpu_and_thread_path ('m'). It sets ENV_DROPOUT=0.2
via std::env::set_var; if left in place, the correctness test's spawned workers
would inherit it (Command inherits parent env by default) and build with
cfg.dropout=0.2 while its single-GPU baseline (run_single_gpu → test_config →
dropout=0) stays at 0 — GATE (a) `max_rel_single < 1e-3` would blow up by
orders of magnitude.

Two defenses:
- correctness test remove_var(ENV_DROPOUT) before spawn (belt): even if the
  dropout test forgot to clean up, this test starts from a clean env.
- dropout test remove_var(ENV_DROPOUT, ENV_DUMP_DIR) at exit (suspenders):
  keep the invariant "each test leaves the env as it found it" so any future
  test added after these two starts clean too.

Same --test-threads=1 SAFETY comment applies (no concurrent env access).

crates/xtrain-autodiff/src/ops.rs

@@ -439,3 +439,245 @@ pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
         }),
     )
 }
+
+/// Per-sequence log-probability: `Σ log πθ(target)` over the non-ignored
+/// (`target ≥ 0`) positions — the quantity DPO (M3) compares between policy and
+/// reference. `target` is `[rows]` I32 carrying `-100` (ignore) at masked positions
+/// (e.g. the prompt) and the gold token id elsewhere; ignored positions contribute
+/// 0, exactly like the SFT cross-entropy masking. Returns a scalar `[1]` Var.
+///
+/// Reuses the CE forward (per-row `−log p(target)`) and backward, so no new kernel:
+/// `seq_logprob = −Σ per_row`, and `d(seq_logprob)/d(logits) = −(probs − onehot)`
+/// = `cross_entropy_backward(probs, target, −upstream)` (a SUM, so no mean
+/// division — contrast [`cross_entropy`], which divides by `valid_rows`).
+pub fn seq_logprob(x: &Var, target: &Tensor) -> Var {
+    let logit_dtype = x.value().dtype();
+    let (probs, per_row) = x.value().cross_entropy(target);
+    // per_row[r] = −log p(target_r), and is 0 for ignored rows (target < 0), so the
+    // sum already counts only the supervised (completion) positions.
+    let sum_neg_lp: f32 = per_row
+        .to_device(xtrain_tensor::Device::Cpu)
+        .as_slice::<f32>()
+        .iter()
+        .sum();
+    let out = Tensor::from_slice(&[-sum_neg_lp], &[1]).to_device(x.value().device());
+
+    let target = target.clone();
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
+            // d(Σ log p)/d(logits) = −(probs − onehot); SUM, so no /valid_rows.
+            let dx = Tensor::cross_entropy_backward(&probs, &target, -upstream);
+            Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
+        }),
+    )
+}
+
+/// DPO loss (Rafailov et al., M3) for one preference pair, as a scalar `[1]` Var
+/// whose two parents are the POLICY sequence-logprobs of the chosen and rejected
+/// completions (from [`seq_logprob`]); the REFERENCE logprobs are constants
+/// (precomputed once from the frozen SFT model). With
+///   `Δ = β·[(lpθ_chosen − lpref_chosen) − (lpθ_rejected − lpref_rejected)]`
+/// the loss is `L = −log σ(Δ) = softplus(−Δ)`. Only the policy terms carry gradient:
+///   `∂L/∂lpθ_chosen = −β·(1−σ(Δ))`,  `∂L/∂lpθ_rejected = +β·(1−σ(Δ))`.
+/// Degenerate points the M3 gate pins: `πθ == πref` ⇒ `Δ = 0`, `L = log 2`, implicit
+/// reward 0; `β → 0` ⇒ gradient → 0. Same formula as TRL
+/// (`-logsigmoid(β·(pol_c − pol_r − (ref_c − ref_r)))`).
+pub fn dpo_loss(
+    lp_pol_chosen: &Var,
+    lp_pol_rejected: &Var,
+    lp_ref_chosen: f32,
+    lp_ref_rejected: f32,
+    beta: f32,
+) -> Var {
+    use xtrain_tensor::Device;
+    let scalar = |v: &Var| v.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let pc = scalar(lp_pol_chosen);
+    let pr = scalar(lp_pol_rejected);
+    let delta = beta * ((pc - lp_ref_chosen) - (pr - lp_ref_rejected));
+    // L = softplus(−Δ) = log(1 + e^{−Δ}) (numerically stable).
+    let nd = -delta;
+    let l = nd.max(0.0) + (-(nd.abs())).exp().ln_1p();
+    let dev = lp_pol_chosen.value().device();
+    let out = Tensor::from_slice(&[l], &[1]).to_device(dev);
+
+    Var::from_op(
+        out,
+        vec![lp_pol_chosen.clone(), lp_pol_rejected.clone()],
+        Box::new(move |d, parents| {
+            let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
+            // s = σ(−Δ) = 1 − σ(Δ); ∂L/∂Δ = −s, and ∂Δ/∂pc = β, ∂Δ/∂pr = −β.
+            let s = 1.0 / (1.0 + delta.exp());
+            let g = up * beta * s;
+            let dev = parents[0].value().device();
+            Var::push_grad(&parents[0], Tensor::from_slice(&[-g], &[1]).to_device(dev));
+            Var::push_grad(&parents[1], Tensor::from_slice(&[g], &[1]).to_device(dev));
+        }),
+    )
+}
+
+/// GRPO clipped policy-gradient loss (M4) for ONE completion, a scalar `[1]` Var
+/// with the policy logits as the single parent. Per non-ignored (completion) token
+/// `t` (`target[t] ≥ 0`):
+///   `logπθ_t = log softmax(logits[t])[target_t]` (`= −per_row[t]` of cross_entropy)
+///   `ρ_t = exp(logπθ_t − logp_old[t])`
+///   `pg_t = min(ρ_t·A, clip(ρ_t, 1−ε, 1+ε)·A)`
+///   `kl_t = exp(logp_ref[t] − logπθ_t) − (logp_ref[t] − logπθ_t) − 1`  (k3 estimator)
+/// `L = −mean_t pg_t + β·mean_t kl_t` over the `N` completion tokens.
+///
+/// `advantage` `A` is the group-relative advantage (constant per completion in
+/// GRPO); `logp_old`/`logp_ref` are per-position constants (old policy at rollout
+/// time / frozen reference). Backward reuses the CE machinery + the per-row
+/// `scale_rows`: `dL/dlogits[t,:] = g_t·(onehot − probs)[t,:]` with
+///   `g_t = −(1/N)A·ρ_t·[unclipped active] + (β/N)(1 − exp(logp_ref_t − logπθ_t))`.
+/// Degenerate points the gate pins: `A=0` ⇒ only the KL term; `ε→∞` ⇒ vanilla PG
+/// (no clip); `β=0` ⇒ no KL term.
+#[allow(clippy::too_many_arguments)]
+pub fn clipped_pg_loss(
+    logits: &Var,
+    target: &Tensor,
+    logp_old: &[f32],
+    logp_ref: &[f32],
+    advantage: f32,
+    eps: f32,
+    beta: f32,
+) -> Var {
+    use xtrain_tensor::Device;
+    let logit_dtype = logits.value().dtype();
+    let (probs, per_row) = logits.value().cross_entropy(target);
+    let rows = per_row.shape()[0];
+    let per_row_h = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    let target_h = target.to_device(Device::Cpu).as_slice::<i32>().to_vec();
+    assert_eq!(logp_old.len(), rows, "logp_old must have one entry per position");
+    assert_eq!(logp_ref.len(), rows, "logp_ref must have one entry per position");
+
+    let mut s = vec![0f32; rows]; // per-row scale for cross_entropy_backward(·,·,1.0)
+    let (mut pg_sum, mut kl_sum, mut n) = (0f32, 0f32, 0f32);
+    for t in 0..rows {
+        if target_h[t] < 0 {
+            continue; // masked (prompt) position — no contribution, no gradient
+        }
+        n += 1.0;
+        let lp = -per_row_h[t]; // logπθ_t
+        let ratio = (lp - logp_old[t]).exp();
+        let clipped = ratio.clamp(1.0 - eps, 1.0 + eps);
+        let (unclipped_term, clipped_term) = (ratio * advantage, clipped * advantage);
+        pg_sum += unclipped_term.min(clipped_term);
+        let active = unclipped_term <= clipped_term; // min picks unclipped ⇒ grad flows
+        let d = logp_ref[t] - lp;
+        kl_sum += d.exp() - d - 1.0;
+        let pg_grad = if active { -advantage * ratio } else { 0.0 };
+        let kl_grad = beta * (1.0 - d.exp());
+        s[t] = -(pg_grad + kl_grad); // dL/dlogits = g·(onehot−probs) = −g·(probs−onehot)
+    }
+    let inv_n = if n > 0.0 { 1.0 / n } else { 1.0 };
+    for v in &mut s {
+        *v *= inv_n;
+    }
+    let loss_val = -pg_sum * inv_n + beta * kl_sum * inv_n;
+    let dev = logits.value().device();
+    let out = Tensor::from_slice(&[loss_val], &[1]).to_device(dev);
+    let s_dev = Tensor::from_slice(&s, &[rows]).to_device(dev);
+
+    let target = target.clone();
+    Var::from_op(
+        out,
+        vec![logits.clone()],
+        Box::new(move |d, parents| {
+            let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
+            // (probs − onehot), masked rows already 0; per-row scale by s; × upstream.
+            let ce = Tensor::cross_entropy_backward(&probs, &target, 1.0);
+            let mut dx = ce.scale_rows(&s_dev);
+            if up != 1.0 {
+                dx = dx.scale(up);
+            }
+            Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
+        }),
+    )
+}
+
+/// Batched GRPO clipped-PG loss over `N` ragged completions packed into ONE
+/// `forward_batched` (M2d): `logits` is `[R, vocab]` with `R = N·Lmax` rows in
+/// sequence-major order (sample 0's `Lmax` rows, then sample 1's, …), each ragged
+/// completion right-padded to the batch's `Lmax`. Prompt AND pad rows are masked
+/// (`target < 0`), so they contribute nothing and carry no gradient — the
+/// **right-pad-is-free-under-causal-attention** property (a real completion row
+/// never attends to the trailing pad rows, so its logits equal the unpadded
+/// single-sequence forward's).
+///
+/// Unlike the per-sample [`clipped_pg_loss`] (which folds a single scalar
+/// `advantage` and a global `1/N_tokens` normaliser), this op takes **per-row**
+/// `advantage[t]` (the owning sample's group-relative `A`) and **per-row**
+/// `weight[t]` (the full normaliser, e.g. `1/(N_samples · n_s)` for sample `s`'s
+/// completion rows, `0` at masked rows). It does NOT compute its own `inv_n`. With
+/// `weight[t] = 1/(N_samples·n_s)` and `advantage[t] = A_s` this is **bit-equivalent
+/// to the looped path** `Σ_s scale·(1/n_s)·clipped_pg_loss_s` (`scale = 1/N_samples`):
+/// the per-row backward is local (`cross_entropy_backward` is row-wise), so the
+/// batched row-`t` gradient equals the looped sample-`s` row-`t` gradient, and the
+/// scalar loss equals the looped weighted sum. (`tests/autograd.rs`:
+/// `clipped_pg_loss_batched_matches_looped`.) Degenerate points match
+/// [`clipped_pg_loss`] (`A=0` ⇒ KL only; `ε→∞` ⇒ vanilla PG; `β=0` ⇒ no KL).
+#[allow(clippy::too_many_arguments)]
+pub fn clipped_pg_loss_batched(
+    logits: &Var,
+    target: &Tensor,
+    logp_old: &[f32],
+    logp_ref: &[f32],
+    advantage: &[f32],
+    weight: &[f32],
+    eps: f32,
+    beta: f32,
+) -> Var {
+    use xtrain_tensor::Device;
+    let logit_dtype = logits.value().dtype();
+    let (probs, per_row) = logits.value().cross_entropy(target);
+    let rows = per_row.shape()[0];
+    let per_row_h = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    let target_h = target.to_device(Device::Cpu).as_slice::<i32>().to_vec();
+    assert_eq!(logp_old.len(), rows, "logp_old must have one entry per row");
+    assert_eq!(logp_ref.len(), rows, "logp_ref must have one entry per row");
+    assert_eq!(advantage.len(), rows, "advantage must have one entry per row");
+    assert_eq!(weight.len(), rows, "weight must have one entry per row");
+
+    let mut s = vec![0f32; rows]; // per-row scale for cross_entropy_backward(·,·,1.0)
+    let mut loss_val = 0f32;
+    for t in 0..rows {
+        if target_h[t] < 0 {
+            continue; // masked (prompt or pad) row — no contribution, no gradient
+        }
+        let (a, w) = (advantage[t], weight[t]);
+        let lp = -per_row_h[t]; // logπθ_t
+        let ratio = (lp - logp_old[t]).exp();
+        let clipped = ratio.clamp(1.0 - eps, 1.0 + eps);
+        let (unclipped_term, clipped_term) = (ratio * a, clipped * a);
+        let pg_t = unclipped_term.min(clipped_term);
+        let active = unclipped_term <= clipped_term; // min picks unclipped ⇒ grad flows
+        let d = logp_ref[t] - lp;
+        let kl_t = d.exp() - d - 1.0;
+        let pg_grad = if active { -a * ratio } else { 0.0 };
+        let kl_grad = beta * (1.0 - d.exp());
+        // The full per-row normaliser is folded into s (no global inv_n here).
+        s[t] = -(pg_grad + kl_grad) * w;
+        loss_val += (-pg_t + beta * kl_t) * w;
+    }
+    let dev = logits.value().device();
+    let out = Tensor::from_slice(&[loss_val], &[1]).to_device(dev);
+    let s_dev = Tensor::from_slice(&s, &[rows]).to_device(dev);
+
+    let target = target.clone();
+    Var::from_op(
+        out,
+        vec![logits.clone()],
+        Box::new(move |d, parents| {
+            let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
+            let ce = Tensor::cross_entropy_backward(&probs, &target, 1.0);
+            let mut dx = ce.scale_rows(&s_dev);
+            if up != 1.0 {
+                dx = dx.scale(up);
+            }
+            Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
+        }),
+    )
+}

crates/xtrain-autodiff/tests/autograd.rs

@@ -1005,3 +1005,266 @@ fn transpose_var(x: &Var) -> Var {
         }),
     )
 }
+
+// seq_logprob (M3 DPO): Σ log p(target) over non-ignored rows. Grad-check with a
+// completion mask — rows 0,1 are -100 (prompt, contribute 0), rows 2..6 supervised.
+#[test]
+fn seq_logprob_bwd() {
+    require_gpu();
+    let (rows, cols) = (6usize, 9usize);
+    let x_h = fill(rows * cols, 202);
+    let targets: Vec<i32> = (0..rows)
+        .map(|r| if r < 2 { -100 } else { (r * 2 % cols) as i32 })
+        .collect();
+    let target = Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
+
+    let x = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let lp = ops::seq_logprob(&x, &target);
+    lp.backward();
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+
+    // Numeric scalar = seq_logprob = −Σ per_row (per_row is 0 for ignored rows).
+    let tgt = targets.clone();
+    let lx = move |v: &[f32], s: &[usize]| {
+        let t = Tensor::from_slice(&tgt, &[rows]).to_device(Device::Cuda(0));
+        let (_, per_row) = cuda(v, s).cross_entropy(&t);
+        -per_row
+            .to_device(Device::Cpu)
+            .as_slice::<f32>()
+            .iter()
+            .sum::<f32>()
+    };
+    report(
+        "seq_logprob dX",
+        &grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_nonlinear()),
+    );
+}
+
+// dpo_loss (M3): scalar DPO loss with the two policy logprobs as parents. Grad-check
+// each parent (finite diff of softplus(−Δ)) + the degenerate points the gate pins:
+// policy==reference ⇒ Δ=0, L=log2, grads ∓β/2; β=0 ⇒ grads 0.
+#[test]
+fn dpo_loss_bwd_and_degenerate() {
+    require_gpu();
+    let (ref_c, ref_r, beta) = (0.5f32, 0.9f32, 0.1f32);
+    let (pc0, pr0) = (1.2f32, 0.7f32);
+    let softplus = |z: f32| z.max(0.0) + (-(z.abs())).exp().ln_1p();
+
+    let pc = Var::leaf(cuda(&[pc0], &[1]));
+    let pr = Var::leaf(cuda(&[pr0], &[1]));
+    let l = ops::dpo_loss(&pc, &pr, ref_c, ref_r, beta);
+    l.backward();
+    let dpc = pc.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let dpr = pr.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+
+    let l_of_pc = move |v: &[f32], _s: &[usize]| softplus(-(beta * ((v[0] - ref_c) - (pr0 - ref_r))));
+    report("dpo_loss dpc", &grad_check(&[pc0], &[1], &l_of_pc, &[dpc], cfg_nonlinear()));
+    let l_of_pr = move |v: &[f32], _s: &[usize]| softplus(-(beta * ((pc0 - ref_c) - (v[0] - ref_r))));
+    report("dpo_loss dpr", &grad_check(&[pr0], &[1], &l_of_pr, &[dpr], cfg_nonlinear()));
+
+    // Degenerate 1: policy == reference ⇒ Δ=0 ⇒ L=log2, grads = (∓β/2).
+    let pc2 = Var::leaf(cuda(&[ref_c], &[1]));
+    let pr2 = Var::leaf(cuda(&[ref_r], &[1]));
+    let l2 = ops::dpo_loss(&pc2, &pr2, ref_c, ref_r, beta);
+    let lval = l2.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+    l2.backward();
+    let d2c = pc2.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let d2r = pr2.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    assert!((lval - 2f32.ln()).abs() < 1e-5, "L at Δ=0 must be log2, got {lval}");
+    assert!(
+        (d2c + beta * 0.5).abs() < 1e-5 && (d2r - beta * 0.5).abs() < 1e-5,
+        "grads at Δ=0 must be ∓β/2, got ({d2c},{d2r})"
+    );
+
+    // Degenerate 2: β=0 ⇒ grads 0.
+    let pc3 = Var::leaf(cuda(&[pc0], &[1]));
+    let pr3 = Var::leaf(cuda(&[pr0], &[1]));
+    let l3 = ops::dpo_loss(&pc3, &pr3, ref_c, ref_r, 0.0);
+    l3.backward();
+    let d3c = pc3.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    assert!(d3c.abs() < 1e-9, "β=0 ⇒ grad 0, got {d3c}");
+    println!("dpo_loss OK: grad-check (dpc,dpr) + degenerate (Δ=0→log2 & ∓β/2, β=0→0)");
+}
+
+// clipped_pg_loss (M4 GRPO): per-token clipped PG + k3 KL, one completion. Grad-check
+// the active (in-trust-region) path + the A=0 (KL-only) path, plus value-level
+// degenerate checks (ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL).
+#[test]
+fn clipped_pg_loss_bwd_and_degenerate() {
+    require_gpu();
+    let (rows, cols) = (6usize, 10usize);
+    let x_h = fill(rows * cols, 303);
+    // rows 0,1 masked (prompt); 2..6 supervised (completion).
+    let targets: Vec<i32> = (0..rows)
+        .map(|r| if r < 2 { -100 } else { (r * 2 % cols) as i32 })
+        .collect();
+    let mk_target = || Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
+
+    // logp_old = logπθ at the base logits ⇒ ρ≈1 (in trust region → active path).
+    let (_, per_row0) = cuda(&x_h, &[rows, cols]).cross_entropy(&mk_target());
+    let logp_old: Vec<f32> = per_row0
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .iter()
+        .map(|p| -p)
+        .collect();
+    let logp_ref: Vec<f32> = logp_old.iter().map(|l| l - 0.3).collect(); // exercise KL
+    let (eps, beta) = (0.2f32, 0.1f32);
+
+    // Host replica of the forward loss as a function of per-row CE values.
+    let host_loss = {
+        let (tg, lo, lr) = (targets.clone(), logp_old.clone(), logp_ref.clone());
+        move |per_row_h: &[f32], a: f32, e: f32, b: f32| -> f32 {
+            let (mut pg, mut kl, mut n) = (0f32, 0f32, 0f32);
+            for t in 0..per_row_h.len() {
+                if tg[t] < 0 {
+                    continue;
+                }
+                n += 1.0;
+                let lp = -per_row_h[t];
+                let ratio = (lp - lo[t]).exp();
+                let clipped = ratio.clamp(1.0 - e, 1.0 + e);
+                pg += (ratio * a).min(clipped * a);
+                let d = lr[t] - lp;
+                kl += d.exp() - d - 1.0;
+            }
+            let inv = if n > 0.0 { 1.0 / n } else { 1.0 };
+            -pg * inv + b * kl * inv
+        }
+    };
+    let per_row_of = |v: &[f32], s: &[usize]| {
+        let (_, pr) = cuda(v, s).cross_entropy(&mk_target());
+        pr.to_device(Device::Cpu).as_slice::<f32>().to_vec()
+    };
+
+    // (1) grad-check the active PG path (A>0, ρ≈1).
+    let adv = 0.7f32;
+    let x = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let loss = ops::clipped_pg_loss(&x, &mk_target(), &logp_old, &logp_ref, adv, eps, beta);
+    loss.backward();
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let hl = host_loss.clone();
+    let lx = move |v: &[f32], s: &[usize]| hl(&per_row_of(v, s), adv, eps, beta);
+    report(
+        "clipped_pg dX (active)",
+        &grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_nonlinear()),
+    );
+
+    // (2) grad-check the A=0 path (loss = β·mean KL; PG gradient must vanish).
+    let x0 = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let loss0 = ops::clipped_pg_loss(&x0, &mk_target(), &logp_old, &logp_ref, 0.0, eps, beta);
+    loss0.backward();
+    let dx0 = x0.grad().unwrap().to_device(Device::Cpu);
+    let hl0 = host_loss.clone();
+    let lx0 = move |v: &[f32], s: &[usize]| hl0(&per_row_of(v, s), 0.0, eps, beta);
+    report(
+        "clipped_pg dX (A=0, KL only)",
+        &grad_check(&x_h, &[rows, cols], &lx0, dx0.as_slice::<f32>(), cfg_nonlinear()),
+    );
+
+    // (3) ε→∞ ⇒ vanilla PG (no clip): loss value == −mean(ρA) + β·mean KL.
+    let big = 1e9f32;
+    let lv = ops::clipped_pg_loss(&Var::leaf(cuda(&x_h, &[rows, cols])), &mk_target(), &logp_old, &logp_ref, adv, big, beta);
+    let got = lv.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let pr0 = per_row_of(&x_h, &[rows, cols]);
+    let want = host_loss(&pr0, adv, big, beta);
+    assert!((got - want).abs() < 1e-4, "ε→∞ vanilla loss mismatch: {got} vs {want}");
+
+    // (4) β=0 ⇒ no KL term (loss == −mean pg only).
+    let lvb = ops::clipped_pg_loss(&Var::leaf(cuda(&x_h, &[rows, cols])), &mk_target(), &logp_old, &logp_ref, adv, eps, 0.0);
+    let gotb = lvb.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let wantb = host_loss(&pr0, adv, eps, 0.0);
+    assert!((gotb - wantb).abs() < 1e-5, "β=0 loss mismatch: {gotb} vs {wantb}");
+    println!("clipped_pg_loss OK: grad-check (active + A=0) + degenerate (ε→∞ vanilla, β=0 no KL)");
+}
+
+// clipped_pg_loss_batched (M2d): N ragged completions packed + right-padded into ONE
+// forward must equal the looped per-sample path Σ_s (1/N)·clipped_pg_loss_s. The
+// per-row CE backward is row-local, so folding weight = 1/(N·n_s) into the batched
+// op reproduces the looped gradient and weighted-sum loss bit-for-bit (f32 path).
+#[test]
+fn clipped_pg_loss_batched_matches_looped() {
+    require_gpu();
+    let (n, lmax, cols) = (3usize, 5usize, 10usize);
+    let rows = n * lmax;
+    let x_h = fill(rows * cols, 909);
+    // Per sample: row 0 = prompt (-100); rows 1..real_len = completion; rest = pad
+    // (-100). Different real_len ⇒ n_s = {2, 3, 1} completion rows.
+    let real_len = [3usize, 4, 2];
+    let adv_s = [0.7f32, -0.5, 0.3];
+    let mut targets = vec![-100i32; rows];
+    for s in 0..n {
+        for r in 1..real_len[s] {
+            let t = s * lmax + r;
+            targets[t] = ((t * 3) % cols) as i32;
+        }
+    }
+    let mk_target = || Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
+
+    // logp_old ≈ logπθ at base logits (ρ≈1), logp_ref offset to exercise the KL term.
+    let (_, per_row0) = cuda(&x_h, &[rows, cols]).cross_entropy(&mk_target());
+    let logp_old: Vec<f32> = per_row0
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .iter()
+        .map(|p| -p)
+        .collect();
+    let logp_ref: Vec<f32> = logp_old.iter().map(|l| l - 0.3).collect();
+    let (eps, beta) = (0.2f32, 0.1f32);
+
+    // Per-row advantage (sample's A) + per-row weight 1/(N·n_s) (full normaliser).
+    let n_of = |s: usize| (0..lmax).filter(|&r| targets[s * lmax + r] >= 0).count() as f32;
+    let mut advantage = vec![0f32; rows];
+    let mut weight = vec![0f32; rows];
+    for s in 0..n {
+        let w = (1.0 / n as f32) * (1.0 / n_of(s));
+        for r in 0..lmax {
+            advantage[s * lmax + r] = adv_s[s];
+            weight[s * lmax + r] = w;
+        }
+    }
+
+    // Batched: one packed [R, vocab] forward + one backward.
+    let xb = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let lb = ops::clipped_pg_loss_batched(
+        &xb, &mk_target(), &logp_old, &logp_ref, &advantage, &weight, eps, beta,
+    );
+    lb.backward();
+    let gb = xb.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    let lb_val = lb.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+
+    // Looped reference: per-sample slice → clipped_pg_loss → scale(1/N) → backward.
+    let mut g_ref = vec![0f32; rows * cols];
+    let mut loss_ref = 0f32;
+    for s in 0..n {
+        let r0 = s * lmax;
+        let xs_h = x_h[r0 * cols..(r0 + lmax) * cols].to_vec();
+        let tgt_s: Vec<i32> = targets[r0..r0 + lmax].to_vec();
+        let lo_s = logp_old[r0..r0 + lmax].to_vec();
+        let lr_s = logp_ref[r0..r0 + lmax].to_vec();
+        let xs = Var::leaf(cuda(&xs_h, &[lmax, cols]));
+        let tgt = Tensor::from_slice(&tgt_s, &[lmax]).to_device(Device::Cuda(0));
+        let ls = ops::clipped_pg_loss(&xs, &tgt, &lo_s, &lr_s, adv_s[s], eps, beta);
+        let scaled = ops::scale(&ls, 1.0 / n as f32);
+        scaled.backward();
+        let gs = xs.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>().to_vec();
+        g_ref[r0 * cols..(r0 + lmax) * cols].copy_from_slice(&gs);
+        loss_ref += scaled.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+    }
+
+    let max_g = gb
+        .iter()
+        .zip(&g_ref)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0.0f32, f32::max);
+    assert!(
+        (lb_val - loss_ref).abs() < 1e-5,
+        "batched loss {lb_val} vs looped {loss_ref}"
+    );
+    assert!(max_g < 1e-5, "batched grad vs looped: max|Δ| = {max_g}");
+    println!(
+        "clipped_pg_loss_batched OK: loss Δ={:.2e}, grad max|Δ|={:.2e} (== looped Σ_s 1/N·pg_s)",
+        (lb_val - loss_ref).abs(),
+        max_g
+    );
+}

crates/xtrain-cuda/src/ffi.rs

@@ -139,6 +139,51 @@ unsafe extern "C" {
         period: i32,
         s: CudaStream,
     );
+    // RoPE at an absolute position offset (KV-cache decode, forward only): row
+    // `tok`'s position is `pos0 + tok` (no modulo). For a single decode token
+    // (tokens == 1) the one row sits at absolute position `pos0`.
+    pub fn launch_rope_at_f32(
+        x: *const f32,
+        y: *mut f32,
+        tokens: i32,
+        heads: i32,
+        head_dim: i32,
+        theta: f32,
+        pos0: i32,
+        s: CudaStream,
+    );
+    // RoPE with a per-row absolute position (batched KV-cache decode, M2b): row
+    // `tok`'s position is `positions[tok]`. Forward only.
+    pub fn launch_rope_pos_f32(
+        x: *const f32,
+        positions: *const i32,
+        y: *mut f32,
+        tokens: i32,
+        heads: i32,
+        head_dim: i32,
+        theta: f32,
+        s: CudaStream,
+    );
+    // Concatenate along the sequence dim: a:[bh,ta,hd], b:[bh,tb,hd] →
+    // out:[bh,ta+tb,hd] (device-side KV-cache append, M2c).
+    pub fn launch_cat_seq_f32(
+        a: *const f32,
+        b: *const f32,
+        out: *mut f32,
+        bh: i32,
+        ta_hd: i32,
+        tb_hd: i32,
+        s: CudaStream,
+    );
+    // Per-row scale: y[r,c] = x[r,c] * s[r] (GRPO policy-gradient backward).
+    pub fn launch_scale_rows_f32(
+        x: *const f32,
+        s: *const f32,
+        y: *mut f32,
+        rows: i32,
+        cols: i32,
+        stream: CudaStream,
+    );
     pub fn launch_rope_dx_f32(
         dy: *const f32,
         dx: *mut f32,

crates/xtrain-distributed/src/bin/train_ddp_mp.rs

@@ -10,7 +10,9 @@
 //!
 //! Versus `train_ddp` (thread-per-GPU, kept as the regression baseline) the ONLY
 //! difference is the launch model + cross-process UniqueId bootstrap. CLI flags
-//! are identical, so it doubles as the before→after throughput driver.
+//! mirror `train_ddp` (incl. `--dropout` — same T21 wiring: `cfg.dropout` set here
+//! and `train_rank` re-asserts `model.train()` each step), so it doubles as the
+//! before→after throughput driver.
 //!
 //! Run on dash5 (pick idle GPUs — dash5 is shared):
 //!   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
@@ -108,6 +110,11 @@ fn main() {
     let val_tokens: usize = flag(&args, "--val-tokens", 0);
     let eval_every: usize = flag(&args, "--eval-every", 0);
     let eval_batches: usize = flag(&args, "--eval-batches", 64);
+    // Dropout (Phase T18/T21): residual-path dropout prob, active at training time
+    // only (inverted scaling), identity at eval/sampling/export. Default 0 = off
+    // (bit-identical to the no-dropout path). Mirrors bin/train_ddp; propagates into
+    // cfg.dropout (below) and relies on T21's per-step model.train() in train_rank.
+    let dropout: f32 = flag(&args, "--dropout", 0.0f32);
     let opts = ModelOpts {
         bf16: args.iter().any(|a| a == "--bf16"),
         recompute: args.iter().any(|a| a == "--recompute"),
@@ -136,7 +143,9 @@ fn main() {
         (corpus, None)
     };
 
-    let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
+    let mut cfg =
+        Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
+    cfg.dropout = dropout;
 
     if env.rank == 0 {
         println!(
@@ -162,6 +171,9 @@ fn main() {
         if opts.flash {
             println!("flash-attention: ON (fused SDPA kernel, no materialized scores)");
         }
+        if dropout > 0.0 {
+            println!("dropout: ON (p={dropout}, residual-path, train-only inverted scaling)");
+        }
     }
 
     let dcfg = DdpConfig {

crates/xtrain-distributed/tests/ddp_proc.rs

@@ -10,6 +10,14 @@
 //!   (a) multi-process loss matches single-GPU within `<1e-3`,
 //!   (b) cross-rank params agree within `<1e-6` (KI-5 ULP tolerance),
 //!   (c) multi-process loss matches the thread-per-GPU `launch` path within `<1e-3`.
+//!
+//! T21-for-proc regression `proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout`
+//! (below) additionally proves that `--dropout` propagates through the process-per-
+//! GPU launcher — the analogue of the thread-per-GPU T21 fix. Pre-fix
+//! `train_ddp_mp` had no `--dropout` flag, so `cfg.dropout` stayed 0 regardless of
+//! what the user passed, silently disabling dropout under process-per-GPU. The
+//! GATE B loss-trace signal (>1e-3 gap between p=0 and p=0.2) sits orders of
+//! magnitude above the KI-5 cross-rank noise floor and catches that gap directly.
 
 #![cfg(not(no_cuda))]
 
@@ -74,8 +82,20 @@ fn dcfg(batch_size: usize) -> DdpConfig {
 // The dump dir is passed launcher→worker via this env key (separate from the
 // XTRAIN_* keys the launcher sets); workers write `rank{N}.dump` there.
 const ENV_DUMP_DIR: &str = "XTRAIN_TEST_DUMP_DIR";
+// Optional launcher→worker channel for `cfg.dropout`. Absent = 0.0 = the existing
+// correctness test's contract (no perturbation). The T21-for-proc regression test
+// below sets it before each `launch_processes` call to prove the process-per-GPU
+// path actually plumbs `--dropout` into every worker's model.
+const ENV_DROPOUT: &str = "XTRAIN_TEST_DROPOUT";
 const GLOBAL_BATCH: usize = 8;
 
+fn worker_dropout() -> f32 {
+    std::env::var(ENV_DROPOUT)
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(0.0)
+}
+
 // ── Worker entry: runs when this test binary is re-execed by launch_processes ─
 
 fn run_as_worker_if_needed() {
@@ -87,7 +107,13 @@ fn run_as_worker_if_needed() {
     // production `run_worker` wrapper is exercised by `bin/train_ddp_mp` on dash5.
     let ctx = DdpContext::init(env.rank, env.world, env.id, env.local_rank);
     let device = Device::Cuda(env.local_rank);
-    let model = build_model(test_config(), device);
+    // Mirrors bin/train_ddp_mp's `cfg.dropout = dropout` wiring — the T21-for-proc
+    // regression: if this line were missing (the pre-fix launcher's exact gap),
+    // `cfg.dropout` would stay 0 and the GATE B test below would find a bit-
+    // identical p=0 / p=0.2 loss trace and FAIL.
+    let mut cfg = test_config();
+    cfg.dropout = worker_dropout();
+    let model = build_model(cfg, device);
     let res = train_rank(
         &ctx,
         &model,
@@ -203,8 +229,16 @@ fn proc_per_gpu_matches_single_gpu_and_thread_path() {
     let dump_dir = std::env::temp_dir().join(format!("xtrain_t17_{}", std::process::id()));
     std::fs::create_dir_all(&dump_dir).unwrap();
     // SAFETY: single-threaded test (forced by --test-threads=1) sets this env
-    // before spawning workers; no concurrent env access.
+    // before spawning workers; no concurrent env access. ENV_DROPOUT is cleared
+    // defensively — libtest orders `--test-threads=1` runs alphabetically, so the
+    // sibling `proc_per_gpu_dropout_is_live_...` test (starts with 'd') runs BEFORE
+    // this one (starts with 'm'). If it happened to leak `ENV_DROPOUT=0.2` in this
+    // process's env, the workers here would inherit it (Command inherits parent
+    // env by default) and build with dropout=0.2 while the single-GPU baseline
+    // (run_single_gpu → test_config → dropout=0) stays at 0 — GATE (a) would blow up.
+    // Explicit remove here severs that ordering coupling.
     unsafe {
+        std::env::remove_var(ENV_DROPOUT);
         std::env::set_var(ENV_DUMP_DIR, &dump_dir);
     }
     // Re-exec the test binary but run ONLY this test, single-threaded, so the
@@ -273,6 +307,100 @@ fn proc_per_gpu_matches_single_gpu_and_thread_path() {
     let _ = std::fs::remove_dir_all(&dump_dir);
 }
 
+/// T21-for-proc regression: prove that `--dropout` actually reaches the model
+/// under process-per-GPU. The pre-fix `bin/train_ddp_mp` had no `--dropout` flag
+/// and never set `cfg.dropout`, so the launcher's worker built its model with
+/// dropout stuck at 0 — silent identity, regardless of what the user passed. The
+/// thread-per-GPU T21 fix caught the analogous gap; this test caps the same gap
+/// on the proc-per-GPU path with the same GATE-B pattern (loss trajectory of a
+/// p=0.2 run differs from p=0 by a large margin, well above the NCCL noise floor).
+///
+/// Both runs share the corpus, the initial params (via `build_model`'s deterministic
+/// LCG), and every other config knob; the ONLY difference is `cfg.dropout`. If the
+/// worker didn't plumb the env-provided dropout into `cfg.dropout` (the exact pre-
+/// fix regression), both traces would be bit-identical and this test would FAIL.
+/// The `>1e-3` threshold sits orders of magnitude above the KI-5 cross-rank ULP
+/// noise floor (~1e-7 on this PCIe box), so it's a hard signal for "dropout is
+/// active" rather than a noise measurement. Mirrors
+/// `ddp_dropout_is_live_and_p0_bit_identical` in ddp_correctness.rs for T21's
+/// thread-per-GPU fix.
+#[test]
+fn proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout() {
+    run_as_worker_if_needed();
+
+    let world = 2usize;
+    if device::device_count().unwrap_or(0) < world as i32 {
+        eprintln!("skip: need >= {world} GPUs");
+        return;
+    }
+
+    let base_dump_dir = std::env::temp_dir().join(format!("xtrain_t21mp_{}", std::process::id()));
+    std::fs::create_dir_all(&base_dump_dir).unwrap();
+    let worker_args = [
+        "--exact".to_string(),
+        "proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout".to_string(),
+        "--test-threads=1".to_string(),
+        "--nocapture".to_string(),
+    ];
+
+    // Helper: launch `world` workers with a specific dropout prob (via env), read
+    // rank 0's loss trace, clean up. Uses a subdir per run so the two invocations
+    // do not clobber each other's dumps.
+    let mut launch_with_dropout = |p: f32, tag: &str| -> Vec<f32> {
+        let dump_dir = base_dump_dir.join(tag);
+        std::fs::create_dir_all(&dump_dir).unwrap();
+        // SAFETY: single-threaded test (forced by --test-threads=1); no concurrent env access.
+        unsafe {
+            std::env::set_var(ENV_DUMP_DIR, &dump_dir);
+            std::env::set_var(ENV_DROPOUT, format!("{p}"));
+        }
+        launch_processes(world, &worker_args).expect("worker processes failed");
+        let (losses, _) = read_dump(dump_dir.to_str().unwrap(), 0);
+        losses
+    };
+
+    let loss_p0 = launch_with_dropout(0.0, "p0");
+    let loss_p1 = launch_with_dropout(0.2, "p02");
+
+    // GATE B — dropout is LIVE under process-per-GPU with p>0. If the worker
+    // didn't set `cfg.dropout` (the pre-fix gap), the two traces would match to
+    // the ~1e-7 NCCL noise floor. Anything above ~1e-3 is unambiguous evidence
+    // that dropout masks are actually applied in every worker's forward.
+    let max_live_diff = loss_p0
+        .iter()
+        .zip(&loss_p1)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0.0f32, f32::max);
+    println!(
+        "T21-proc GATE B (dropout live under proc-per-GPU): p0[last]={:.6} p0.2[last]={:.6} max |loss diff| = {max_live_diff:.3e}",
+        loss_p0.last().unwrap(),
+        loss_p1.last().unwrap()
+    );
+    assert!(
+        max_live_diff > 1e-3,
+        "p=0.2 proc-per-GPU loss matches p=0 — dropout NOT plumbed through the \
+         process-per-GPU launcher (cfg.dropout stayed 0 in the worker): max |loss diff| {max_live_diff:.3e}"
+    );
+
+    // No NaN/Inf in the p>0 run.
+    assert!(
+        loss_p1.iter().all(|l| l.is_finite()),
+        "p=0.2 proc-per-GPU loss has non-finite values"
+    );
+
+    // Clear the launcher→worker env keys so we don't leak state to anything that
+    // runs later in this process. `proc_per_gpu_matches_single_gpu_and_thread_path`
+    // clears ENV_DROPOUT defensively too, but keeping the invariant "each test
+    // leaves the env as it found it" costs nothing.
+    // SAFETY: single-threaded test (forced by --test-threads=1); no concurrent env access.
+    unsafe {
+        std::env::remove_var(ENV_DROPOUT);
+        std::env::remove_var(ENV_DUMP_DIR);
+    }
+
+    let _ = std::fs::remove_dir_all(&base_dump_dir);
+}
+
 fn max_rel(a: &[f32], b: &[f32]) -> f32 {
     a.iter()
         .zip(b)

crates/xtrain-model/src/decode.rs

@@ -0,0 +1,436 @@
+//! KV-cache incremental-decode engine (post-training M2a, single sequence).
+//!
+//! The naive sampler ([`crate::TinyTransformer`] via `train::sample::generate`)
+//! re-runs the full forward over the whole growing prefix every step — O(t²) and
+//! a fresh autograd graph per token. This is the inference engine that replaces it:
+//! a per-layer **K/V cache** + a **single-token incremental forward** that processes
+//! one new token at a time, attending to the cached keys/values.
+//!
+//! Built on three primitives, all gated by their own correctness tests:
+//! - [`Tensor::rope_at`](xtrain_tensor::Tensor::rope_at): RoPE at the token's
+//!   absolute position (not row-in-tile), so cached post-RoPE K matches the full
+//!   forward (bit-identical, `integration::rope_at_matches_full_rope_row`).
+//! - [`Tensor::decode_attention`](xtrain_tensor::Tensor::decode_attention): the
+//!   single-query × cached-K/V SDPA, equal to the full causal attention's last row
+//!   (`integration::decode_attention_matches_full_attention_last_row`).
+//! - this module's per-token block forward, mirroring `model::block_forward` at the
+//!   raw-Tensor level (no autograd tape — inference needs no gradients).
+//!
+//! Correctness gate (the M2 centerpiece): KV-cache greedy decode is **token-
+//! identical** to the naive full-recompute greedy (`tests/decode_kv.rs`).
+//!
+//! Prefill is just the first `prompt.len()` decode steps (one token at a time) —
+//! one code path, at the cost of a non-batched prefill (M2b adds batched prefill +
+//! ragged batch decode). The cache is host-accumulated (token-major f32) and the
+//! K/V tensor is rebuilt per step; the host round-trip is small (`num_kv·head_dim`
+//! floats/token/layer) and is the honest M2a baseline — M2b moves it device-side.
+
+#![cfg(not(no_cuda))]
+
+use crate::TinyTransformer;
+use xtrain_tensor::{DType, Device, Tensor};
+
+/// Per-layer K/V cache: token-major host accumulation. For each layer, `k[li]` and
+/// `v[li]` hold `[T, num_kv, head_dim]` (f32, flattened), grown by one token's
+/// `num_kv·head_dim` values per decode step. Stored f32 (an exact upcast of the
+/// bf16 projection output); rebuilt to the compute dtype when forming the K/V
+/// tensor, so bf16 values round-trip bit-for-bit.
+struct KVCache {
+    k: Vec<Option<Tensor>>,
+    v: Vec<Option<Tensor>>,
+}
+
+impl KVCache {
+    fn new(n_layers: usize) -> Self {
+        Self {
+            k: (0..n_layers).map(|_| None).collect(),
+            v: (0..n_layers).map(|_| None).collect(),
+        }
+    }
+
+    /// Append one token's K/V (`[bh,1,hd]`, compute dtype) to layer `li`, growing the
+    /// device-resident `[bh,T,hd]` cache via `cat_seq` (no host round-trip, M2c).
+    fn append(&mut self, li: usize, k_bh: Tensor, v_bh: Tensor) {
+        self.k[li] = Some(match self.k[li].take() {
+            Some(c) => c.cat_seq(&k_bh),
+            None => k_bh,
+        });
+        self.v[li] = Some(match self.v[li].take() {
+            Some(c) => c.cat_seq(&v_bh),
+            None => v_bh,
+        });
+    }
+}
+
+/// Linear `x @ W` in the compute dtype — mirrors `model::linear` (bf16 casts the
+/// fp32-master weight to bf16 on the fly; the activation stream is already bf16).
+fn linear_t(cdt: DType, x: &Tensor, w: &Tensor) -> Tensor {
+    match cdt {
+        DType::F32 => x.matmul(w),
+        DType::BF16 => x.matmul(&w.to_dtype(DType::BF16)),
+        _ => unreachable!("compute dtype must be F32/BF16"),
+    }
+}
+
+/// A norm/QK-norm gamma in the compute dtype — mirrors `model::norm_gamma`.
+fn gamma_t(cdt: DType, g: &Tensor) -> Tensor {
+    match cdt {
+        DType::F32 => g.clone(),
+        DType::BF16 => g.to_dtype(DType::BF16),
+        _ => unreachable!("compute dtype must be F32/BF16"),
+    }
+}
+
+/// Greedy KV-cache decode: continue `prompt` by `max_new` tokens, argmax each step.
+/// Returns the full token sequence (prompt + generated), matching the naive
+/// `sample::generate` interface for `temperature == 0`. Token-identical to the
+/// naive full-recompute greedy (gated by `tests/decode_kv.rs`).
+pub fn generate_greedy_cached(
+    model: &TinyTransformer,
+    device: Device,
+    prompt: &[i32],
+    max_new: usize,
+) -> Vec<i32> {
+    let mut rng = 0u64;
+    generate_cached(model, device, prompt, max_new, 0.0, &mut rng)
+}
+
+/// KV-cache decode with temperature sampling (`temperature == 0` → greedy argmax,
+/// matching [`generate_greedy_cached`]; otherwise sample from `softmax(logits/T)`).
+/// The KV-cache rollout the GRPO loop uses: each step allocates only a single-row
+/// `[1, vocab]` logits buffer (vs the naive sampler's `[seq, vocab]`), so it is far
+/// lighter on memory + the allocator — the naive sampler fragments the caching
+/// allocator over a long rollout, which is the M4 "rollout is the long pole" wall.
+pub fn generate_cached(
+    model: &TinyTransformer,
+    device: Device,
+    prompt: &[i32],
+    max_new: usize,
+    temperature: f32,
+    rng_state: &mut u64,
+) -> Vec<i32> {
+    assert!(!prompt.is_empty(), "prompt must be non-empty");
+    let cfg = model.config();
+    let cdt = model.compute_dtype();
+    let n_layers = cfg.n_layers;
+
+    // params() is a stable, documented order (see TinyTransformer::params):
+    //   [0]                         = embed              [vocab, dim]
+    //   [1 + li*11 .. +11]          = layer li's 11 leaves, in block_params order:
+    //       attn_norm, wq, wk, wv, q_norm, k_norm, wo, ffn_norm, w_gate, w_up, w_down
+    //   [1 + n_layers*11]           = final_norm         [dim]
+    //   [1 + n_layers*11 + 1]       = lm_head            [dim, vocab]
+    let params: Vec<Tensor> = model.params().iter().map(|p| p.value()).collect();
+    assert_eq!(
+        params.len(),
+        1 + n_layers * 11 + 2,
+        "unexpected param layout for decode"
+    );
+    let embed = &params[0];
+    let final_norm = &params[1 + n_layers * 11];
+    let lm_head = &params[1 + n_layers * 11 + 1];
+
+    let mut cache = KVCache::new(n_layers);
+    let mut tokens = prompt.to_vec();
+
+    // Prefill: feed each prompt token in order; the last step's logits are the
+    // distribution for the first generated token.
+    let mut logits = Vec::new();
+    for (pos, &tok) in prompt.iter().enumerate() {
+        logits = decode_step(&params, cfg, cdt, device, &mut cache, tok, pos, embed, final_norm, lm_head);
+    }
+
+    for _ in 0..max_new {
+        let next = if temperature <= 0.0 {
+            argmax(&logits) as i32
+        } else {
+            sample_temperature(&logits, temperature, rng_state) as i32
+        };
+        tokens.push(next);
+        let pos = tokens.len() - 1; // absolute position of the token just appended
+        logits = decode_step(&params, cfg, cdt, device, &mut cache, next, pos, embed, final_norm, lm_head);
+    }
+    tokens
+}
+
+/// Sample a token from `softmax(logits / temperature)` (numerically stable). Same
+/// LCG + inverse-CDF scheme as the naive `sample::sample_temperature`.
+fn sample_temperature(row: &[f32], temperature: f32, rng_state: &mut u64) -> usize {
+    let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+    let exps: Vec<f32> = row.iter().map(|&x| ((x - max) / temperature).exp()).collect();
+    let sum: f32 = exps.iter().sum();
+    *rng_state = rng_state
+        .wrapping_mul(6364136223846793005)
+        .wrapping_add(1442695040888963407);
+    let r = ((*rng_state >> 32) as f32 / u32::MAX as f32) * sum;
+    let mut acc = 0.0;
+    for (i, &e) in exps.iter().enumerate() {
+        acc += e;
+        if acc >= r {
+            return i;
+        }
+    }
+    exps.len() - 1
+}
+
+/// One incremental decode step for token `tok` at absolute position `pos`: append
+/// its K/V to the cache and return the next-token logits as host f32 `[vocab]`.
+#[allow(clippy::too_many_arguments)]
+fn decode_step(
+    params: &[Tensor],
+    cfg: &crate::Config,
+    cdt: DType,
+    device: Device,
+    cache: &mut KVCache,
+    tok: i32,
+    pos: usize,
+    embed: &Tensor,
+    final_norm: &Tensor,
+    lm_head: &Tensor,
+) -> Vec<f32> {
+    let (nh, hd, num_kv) = (cfg.n_heads, cfg.head_dim, cfg.num_kv_heads);
+    let dim = cfg.dim;
+    let scale = 1.0 / (hd as f32).sqrt();
+    let (theta, eps) = (cfg.rope_theta, cfg.eps);
+    let n_layers = cfg.n_layers;
+
+    // Embedding (fp32 table) → activation stream in the compute dtype.
+    let ids = Tensor::from_slice(&[tok], &[1]).to_device(device);
+    let mut h = embed.embedding(&ids); // [1, dim] f32
+    if cdt == DType::BF16 {
+        h = h.to_dtype(DType::BF16);
+    }
+
+    for li in 0..n_layers {
+        let base = 1 + li * 11;
+        let (attn_norm, wq, wk, wv) =
+            (&params[base], &params[base + 1], &params[base + 2], &params[base + 3]);
+        let (q_norm, k_norm, wo) = (&params[base + 4], &params[base + 5], &params[base + 6]);
+        let (ffn_norm, w_gate, w_up, w_down) =
+            (&params[base + 7], &params[base + 8], &params[base + 9], &params[base + 10]);
+
+        // --- Attention sub-block (pre-norm + cached-KV attention + residual) ---
+        let normed = h.rms_norm(&gamma_t(cdt, attn_norm), eps).0; // [1, dim]
+
+        // Q: project → per-head QK-norm → RoPE at absolute position `pos`.
+        let q = linear_t(cdt, &normed, wq).reshape(&[1, nh, hd]); // [1, nh, hd]
+        let q = q.reshape(&[nh, hd]).rms_norm(&gamma_t(cdt, q_norm), eps).0;
+        let q = q.reshape(&[1, nh, hd]).rope_at(theta, pos);
+        let q_bh = q.reshape(&[nh, 1, hd]); // seq=1 ⇒ the head-transpose is a no-op on data
+
+        // K: same as Q (QK-norm + RoPE). V: project only. Append each as [num_kv,1,hd]
+        // (bh-major) into the device cache; no host round-trip, no transpose (M2c).
+        let k = linear_t(cdt, &normed, wk).reshape(&[1, num_kv, hd]);
+        let k = k.reshape(&[num_kv, hd]).rms_norm(&gamma_t(cdt, k_norm), eps).0;
+        let k_bh = k.reshape(&[1, num_kv, hd]).rope_at(theta, pos).reshape(&[num_kv, 1, hd]);
+        let v_bh = linear_t(cdt, &normed, wv).reshape(&[num_kv, 1, hd]);
+        cache.append(li, k_bh, v_bh);
+
+        // repeat_kv the cached [num_kv,T,hd] to [nh,T,hd] for the SDPA.
+        let expand = |c: &Tensor| if num_kv == nh { c.clone() } else { c.repeat_kv(nh, 1) };
+        let k_full = expand(cache.k[li].as_ref().unwrap());
+        let v_full = expand(cache.v[li].as_ref().unwrap());
+
+        let attn = q_bh.decode_attention(&k_full, &v_full, scale); // [nh, hd]
+        let attn = attn.reshape(&[1, dim]); // concat heads (nh·hd == dim)
+        let attn_out = linear_t(cdt, &attn, wo); // [1, dim]
+        h = h.add(&attn_out);
+
+        // --- MLP sub-block (pre-norm + SwiGLU + residual) ---
+        let normed = h.rms_norm(&gamma_t(cdt, ffn_norm), eps).0;
+        let gate = linear_t(cdt, &normed, w_gate);
+        let up = linear_t(cdt, &normed, w_up);
+        let act = gate.silu().mul(&up); // swiglu = silu(gate) ∘ up
+        let down = linear_t(cdt, &act, w_down);
+        h = h.add(&down);
+    }
+
+    let h = h.rms_norm(&gamma_t(cdt, final_norm), eps).0;
+    let logits = linear_t(cdt, &h, lm_head); // [1, vocab]
+    logits
+        .to_dtype(DType::F32)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec()
+}
+
+fn argmax(row: &[f32]) -> usize {
+    row.iter()
+        .enumerate()
+        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+        .unwrap()
+        .0
+}
+
+// ===================================================================
+// M2b — batched KV-cache decode (G samples of one prompt, in lockstep)
+// ===================================================================
+
+/// Batched K/V cache: `G` sequences advancing together. Per layer, a device-resident
+/// `[G·num_kv, T, head_dim]` grown one token per step via `cat_seq` (M2c — no host
+/// round-trip). Same as M2a's device cache with a G dimension in `bh`.
+struct BatchKVCache {
+    k: Vec<Option<Tensor>>,
+    v: Vec<Option<Tensor>>,
+}
+
+impl BatchKVCache {
+    fn new(n_layers: usize) -> Self {
+        Self {
+            k: (0..n_layers).map(|_| None).collect(),
+            v: (0..n_layers).map(|_| None).collect(),
+        }
+    }
+    fn append(&mut self, li: usize, k_bh: Tensor, v_bh: Tensor) {
+        self.k[li] = Some(match self.k[li].take() {
+            Some(c) => c.cat_seq(&k_bh),
+            None => k_bh,
+        });
+        self.v[li] = Some(match self.v[li].take() {
+            Some(c) => c.cat_seq(&v_bh),
+            None => v_bh,
+        });
+    }
+}
+
+/// Batched KV-cache decode: roll out `n_samples` (G) completions of the SAME
+/// `prompt` in lockstep — all G share the prompt, so they advance at one common
+/// decode position each step (uniform RoPE via `rope_pos`). Returns G full token
+/// sequences (prompt + sampled continuation). The G-way batching amortises the
+/// per-step kernel launches across G (the rollout long-pole). Token-identical per
+/// row to G independent single-sequence decodes (gated by `tests/decode_batch.rs`).
+///
+/// `temperature == 0` ⇒ greedy (all G identical); `> 0` ⇒ independent samples
+/// (per-row draw from one shared `rng_state`). No finished-mask: all G generate
+/// `max_new` tokens; the caller cuts each at `<|endoftext|>` (a perf-only early
+/// stop is the M2b+ follow-up). Ragged (different-length prompts) is also deferred.
+pub fn generate_cached_batch(
+    model: &TinyTransformer,
+    device: Device,
+    prompt: &[i32],
+    n_samples: usize,
+    max_new: usize,
+    temperature: f32,
+    rng_state: &mut u64,
+) -> Vec<Vec<i32>> {
+    assert!(!prompt.is_empty(), "prompt must be non-empty");
+    assert!(n_samples > 0, "n_samples must be > 0");
+    let cfg = model.config();
+    let cdt = model.compute_dtype();
+    let n_layers = cfg.n_layers;
+    let params: Vec<Tensor> = model.params().iter().map(|p| p.value()).collect();
+    let embed = &params[0];
+    let final_norm = &params[1 + n_layers * 11];
+    let lm_head = &params[1 + n_layers * 11 + 1];
+
+    let g = n_samples;
+    let mut cache = BatchKVCache::new(n_layers);
+    let mut seqs: Vec<Vec<i32>> = vec![prompt.to_vec(); g];
+
+    // Prefill: feed each prompt token (identical across G) at its position.
+    let mut logits = Vec::new(); // [G, vocab] flattened
+    for (pos, &tok) in prompt.iter().enumerate() {
+        let toks = vec![tok; g];
+        logits = decode_step_batch(&params, cfg, cdt, device, &mut cache, &toks, pos, embed, final_norm, lm_head);
+    }
+
+    let vocab = cfg.vocab;
+    for _ in 0..max_new {
+        let mut next = Vec::with_capacity(g);
+        for row in 0..g {
+            let lg = &logits[row * vocab..(row + 1) * vocab];
+            let t = if temperature <= 0.0 {
+                argmax(lg) as i32
+            } else {
+                sample_temperature(lg, temperature, rng_state) as i32
+            };
+            next.push(t);
+            seqs[row].push(t);
+        }
+        let pos = seqs[0].len() - 1; // all G are at the same position
+        logits = decode_step_batch(&params, cfg, cdt, device, &mut cache, &next, pos, embed, final_norm, lm_head);
+    }
+    seqs
+}
+
+/// One batched decode step: `toks` is one current token per sequence (`[G]`), all at
+/// absolute position `pos`. Appends each sequence's K/V and returns logits `[G·vocab]`.
+#[allow(clippy::too_many_arguments)]
+fn decode_step_batch(
+    params: &[Tensor],
+    cfg: &crate::Config,
+    cdt: DType,
+    device: Device,
+    cache: &mut BatchKVCache,
+    toks: &[i32],
+    pos: usize,
+    embed: &Tensor,
+    final_norm: &Tensor,
+    lm_head: &Tensor,
+) -> Vec<f32> {
+    let (nh, hd, num_kv) = (cfg.n_heads, cfg.head_dim, cfg.num_kv_heads);
+    let dim = cfg.dim;
+    let g = toks.len();
+    let scale = 1.0 / (hd as f32).sqrt();
+    let (theta, eps) = (cfg.rope_theta, cfg.eps);
+    let n_layers = cfg.n_layers;
+    // Uniform per-row position (all G at the same decode step).
+    let positions = Tensor::from_slice(&vec![pos as i32; g], &[g]).to_device(device);
+
+    let ids = Tensor::from_slice(toks, &[g]).to_device(device);
+    let mut h = embed.embedding(&ids); // [G, dim] f32
+    if cdt == DType::BF16 {
+        h = h.to_dtype(DType::BF16);
+    }
+
+    for li in 0..n_layers {
+        let base = 1 + li * 11;
+        let (attn_norm, wq, wk, wv) =
+            (&params[base], &params[base + 1], &params[base + 2], &params[base + 3]);
+        let (q_norm, k_norm, wo) = (&params[base + 4], &params[base + 5], &params[base + 6]);
+        let (ffn_norm, w_gate, w_up, w_down) =
+            (&params[base + 7], &params[base + 8], &params[base + 9], &params[base + 10]);
+
+        let normed = h.rms_norm(&gamma_t(cdt, attn_norm), eps).0; // [G, dim]
+
+        // Q: project → per-head QK-norm → RoPE at `pos` for every row.
+        let q = linear_t(cdt, &normed, wq).reshape(&[g, nh, hd]);
+        let q = q.reshape(&[g * nh, hd]).rms_norm(&gamma_t(cdt, q_norm), eps).0;
+        let q = q.reshape(&[g, nh, hd]).rope_pos(&positions, theta);
+        let q_bh = q.reshape(&[g * nh, 1, hd]); // bh = G·nh
+
+        // K/V appended as [G·num_kv,1,hd] (bh-major) into the device cache (M2c).
+        let k = linear_t(cdt, &normed, wk).reshape(&[g, num_kv, hd]);
+        let k = k.reshape(&[g * num_kv, hd]).rms_norm(&gamma_t(cdt, k_norm), eps).0;
+        let k_bh = k
+            .reshape(&[g, num_kv, hd])
+            .rope_pos(&positions, theta)
+            .reshape(&[g * num_kv, 1, hd]);
+        let v_bh = linear_t(cdt, &normed, wv).reshape(&[g * num_kv, 1, hd]);
+        cache.append(li, k_bh, v_bh);
+
+        // repeat_kv the cached [G·num_kv,T,hd] to [G·nh,T,hd] for the SDPA.
+        let expand = |c: &Tensor| if num_kv == nh { c.clone() } else { c.repeat_kv(nh, g) };
+        let k_full = expand(cache.k[li].as_ref().unwrap());
+        let v_full = expand(cache.v[li].as_ref().unwrap());
+
+        let attn = q_bh.decode_attention(&k_full, &v_full, scale); // [G·nh, hd]
+        let attn = attn.reshape(&[g, dim]); // concat heads per sequence
+        let attn_out = linear_t(cdt, &attn, wo);
+        h = h.add(&attn_out);
+
+        let normed = h.rms_norm(&gamma_t(cdt, ffn_norm), eps).0;
+        let gate = linear_t(cdt, &normed, w_gate);
+        let up = linear_t(cdt, &normed, w_up);
+        let act = gate.silu().mul(&up);
+        let down = linear_t(cdt, &act, w_down);
+        h = h.add(&down);
+    }
+
+    let h = h.rms_norm(&gamma_t(cdt, final_norm), eps).0;
+    linear_t(cdt, &h, lm_head)
+        .to_dtype(DType::F32)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec()
+}

crates/xtrain-model/src/lib.rs

@@ -25,3 +25,8 @@ pub use config::Config;
 mod model;
 #[cfg(not(no_cuda))]
 pub use model::{TinyTransformer, batched_ids_tensor, ids_tensor, param_to_host};
+
+#[cfg(not(no_cuda))]
+pub mod decode;
+#[cfg(not(no_cuda))]
+pub use decode::{generate_cached, generate_cached_batch, generate_greedy_cached};

crates/xtrain-model/tests/ragged_batch.rs

@@ -0,0 +1,97 @@
+// M2d gate: does forward_batched on RIGHT-PADDED ragged sequences reproduce the
+// per-sequence single-seq forward on the real (non-pad) rows? The batched GRPO
+// training-side forwards depend on this "right-pad is free under causal attention"
+// property — a real completion row is at an earlier position than the trailing pad,
+// and causal masking forbids attending forward, so its logits should be unchanged.
+//
+// Tested in fp32 (exact) over both SDPA cores (composed + fused flash), since the
+// bench uses flash and a kernel could in principle leak the pad keys into the online
+// softmax.
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, ids_tensor};
+use xtrain_tensor::{DType, Device, Tensor};
+
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed.wrapping_mul(2862933555777941757).wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+fn build(cfg: Config, device: Device, dtype: DType, flash: bool) -> TinyTransformer {
+    let mut seed = 1u64;
+    let m = TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.08)
+        }
+    });
+    m.with_compute_dtype(dtype).with_flash(flash)
+}
+
+fn host(t: &Tensor) -> Vec<f32> {
+    t.to_dtype(DType::F32).to_device(Device::Cpu).as_slice::<f32>().to_vec()
+}
+
+#[test]
+fn forward_batched_ragged_matches_looped() {
+    if device::device_count().unwrap_or(0) == 0 {
+        eprintln!("no CUDA device; skipping");
+        return;
+    }
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let mut cfg = Config::tiny();
+    cfg.vocab = 32;
+    cfg.n_layers = 2;
+    let vocab = cfg.vocab;
+
+    // Ragged lengths incl. one crossing the flash tile (>32) and short ones.
+    let lens = [6usize, 40, 9, 4];
+    let lmax = *lens.iter().max().unwrap();
+    let n = lens.len();
+    let seqs: Vec<Vec<i32>> = lens
+        .iter()
+        .enumerate()
+        .map(|(b, &l)| (0..l).map(|i| ((b * 7 + i * 3 + 1) % vocab) as i32).collect())
+        .collect();
+
+    for (dtype, tol) in [(DType::F32, 2e-3f32), (DType::BF16, 3e-1f32)] {
+      for flash in [false, true] {
+        let m = build(cfg, device, dtype, flash);
+        // Looped: each sequence on its own (the ground truth).
+        let looped: Vec<Vec<f32>> = seqs.iter().map(|s| host(&m.forward(&ids_tensor(s, device)).value())).collect();
+
+        // Batched: right-pad each to lmax (pad id 0), one forward_batched(batch = n).
+        let mut flat = vec![0i32; n * lmax];
+        for (i, s) in seqs.iter().enumerate() {
+            flat[i * lmax..i * lmax + s.len()].copy_from_slice(s);
+        }
+        let ids = Tensor::from_slice(&flat, &[n * lmax]).to_device(device);
+        let batched = host(&m.forward_batched(&ids, n).value()); // [n*lmax, vocab]
+
+        let mut dmax = 0f32;
+        for (i, s) in seqs.iter().enumerate() {
+            for r in 0..s.len() {
+                for c in 0..vocab {
+                    let a = looped[i][r * vocab + c];
+                    let b = batched[(i * lmax + r) * vocab + c];
+                    dmax = dmax.max((a - b).abs());
+                }
+            }
+        }
+        println!("dtype={dtype:?} flash={flash}: ragged right-pad vs looped, max|Δlogit| (real rows) = {dmax:.3e}");
+        assert!(dmax < tol, "dtype={dtype:?} flash={flash}: right-pad NOT free under causal — max|Δ| = {dmax}");
+      }
+    }
+    println!("forward_batched_ragged_matches_looped OK: right-pad is free under causal (fp32+bf16, composed + flash)");
+}

crates/xtrain-tensor/src/tensor.rs

@@ -790,6 +790,107 @@ impl Tensor {
         out
     }
 
+    /// RoPE at an absolute position offset (KV-cache decode, forward only).
+    /// `self`:[tokens,heads,head_dim]; row `r`'s position is `pos0 + r` (no
+    /// modulo). For a single new decode token pass `tokens == 1` → the one row is
+    /// rotated at absolute position `pos0`. Mirrors [`rope`](Self::rope)'s dtype
+    /// handling (bf16 → f32 → bf16); no backward (inference path).
+    #[cfg(not(no_cuda))]
+    pub fn rope_at(&self, theta: f32, pos0: usize) -> Self {
+        assert_eq!(self.ndim(), 3, "rope_at requires [tokens,heads,head_dim]");
+        let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
+        assert_eq!(head_dim % 2, 0, "head_dim must be even");
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .rope_at(theta, pos0)
+                .to_dtype(DType::BF16);
+        }
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_rope_at_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                tokens as i32,
+                heads as i32,
+                head_dim as i32,
+                theta,
+                pos0 as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        out
+    }
+
+    /// RoPE with a PER-ROW absolute position (batched KV-cache decode, M2b).
+    /// `self`:[tokens,heads,head_dim]; row `t`'s position is `positions[t]` (an
+    /// I32 `[tokens]` tensor). For G-way batched decode all G rows share one decode
+    /// position; for ragged batches each row carries its own. Mirrors `rope_at`'s
+    /// dtype handling; forward only.
+    #[cfg(not(no_cuda))]
+    pub fn rope_pos(&self, positions: &Tensor, theta: f32) -> Self {
+        assert_eq!(self.ndim(), 3, "rope_pos requires [tokens,heads,head_dim]");
+        let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
+        assert_eq!(head_dim % 2, 0, "head_dim must be even");
+        assert_eq!(positions.dtype, DType::I32, "positions must be I32");
+        assert_eq!(positions.numel(), tokens, "one position per token");
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .rope_pos(positions, theta)
+                .to_dtype(DType::BF16);
+        }
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_rope_pos_f32(
+                self.data_ptr() as *const f32,
+                positions.data_ptr() as *const i32,
+                out.data_ptr() as *mut f32,
+                tokens as i32,
+                heads as i32,
+                head_dim as i32,
+                theta,
+                std::ptr::null_mut(),
+            );
+        }
+        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))]
@@ -909,6 +1010,31 @@ impl Tensor {
         dx
     }
 
+    /// Per-row scale: `out[r,c] = self[r,c] * s[r]`. `self`:[rows,cols] F32,
+    /// `s`:[rows] F32. 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 (the per-token clipped-PG + KL gradient). Forward-only.
+    #[cfg(not(no_cuda))]
+    pub fn scale_rows(&self, s: &Tensor) -> Self {
+        assert_eq!(self.ndim(), 2, "scale_rows requires a 2D tensor");
+        assert_eq!(self.dtype, DType::F32, "scale_rows is F32");
+        assert_eq!(s.dtype, DType::F32, "scale vector is F32");
+        let (rows, cols) = (self.shape[0], self.shape[1]);
+        assert_eq!(s.numel(), rows, "scale vector must have one entry per row");
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_scale_rows_f32(
+                self.data_ptr() as *const f32,
+                s.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        out
+    }
+
     // --- Structural / model ops (the T5 kernels) ---
 
     /// Reshape to `new_shape` (must keep `numel`). Pure metadata change on a
@@ -1076,6 +1202,76 @@ impl Tensor {
         (out, probs)
     }
 
+    /// Decode-time (incremental) attention: a SINGLE query position against a
+    /// cached K/V of length `t` (KV-cache decode, forward only). `self` = Q
+    /// `[bh,1,head_dim]`; `k`,`v` = `[bh,t,head_dim]`, already repeat_kv-expanded
+    /// to `bh` heads. Returns out `[bh,head_dim]` (= `[bh,1,head_dim]` flattened).
+    ///
+    /// No causal mask is needed — the one query sits at the end, so every cached
+    /// key (positions `0..t`) is visible. This is exactly the LAST query row of the
+    /// full causal [`attention`](Self::attention), so KV-cache greedy decode is
+    /// token-identical to full recompute. Softmax is computed in f32 (matching the
+    /// causal path) with `scale` folded in before the exponentials.
+    #[cfg(not(no_cuda))]
+    pub fn decode_attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> Self {
+        assert_eq!(self.ndim(), 3, "decode_attention Q must be [bh,1,head_dim]");
+        assert_eq!(self.shape[1], 1, "decode_attention Q seq must be 1");
+        assert_eq!(k.ndim(), 3, "decode_attention K must be [bh,t,head_dim]");
+        assert_eq!(k.shape(), v.shape(), "K/V shape mismatch");
+        assert_eq!(self.dtype, k.dtype, "Q/K dtype mismatch");
+        assert_eq!(self.dtype, v.dtype, "Q/V dtype mismatch");
+        let (bh, hd) = (self.shape[0], self.shape[2]);
+        assert_eq!(k.shape[0], bh, "Q/K batch-head mismatch");
+        assert_eq!(k.shape[2], hd, "Q/K head_dim mismatch");
+        let t = k.shape[1]; // cached length
+        let dt = self.dtype;
+        let dev = self.device();
+
+        // scores[bh,1,t] = Q[bh,1,hd] · Kᵀ[bh,hd,t]  (per-head batched GEMM).
+        // [bh,1,t] is stored identically to [bh,t]; allocate 2D so the rowwise
+        // softmax can run without a reshape.
+        let scores = Tensor::zeros(&[bh, t], dt, dev);
+        strided_batched_gemm(
+            dt,
+            false,
+            true,
+            1,
+            t,
+            hd,
+            self.data_ptr(),
+            hd,
+            k.data_ptr(),
+            t * hd,
+            scores.data_ptr(),
+            t,
+            bh,
+        );
+        // probs = softmax(scale · scores) over the t keys (f32, like the causal path).
+        let probs = scores
+            .to_dtype(DType::F32)
+            .scale(scale)
+            .softmax()
+            .to_dtype(dt);
+        // out[bh,1,hd] = probs[bh,1,t] · V[bh,t,hd].
+        let out = Tensor::zeros(&[bh, hd], dt, dev);
+        strided_batched_gemm(
+            dt,
+            false,
+            false,
+            1,
+            hd,
+            t,
+            probs.data_ptr(),
+            t,
+            v.data_ptr(),
+            t * hd,
+            out.data_ptr(),
+            hd,
+            bh,
+        );
+        out
+    }
+
     /// Backward of [`attention`](Self::attention). Inputs: forward `q`,`k`,`v`,
     /// the cached `probs`, the upstream `dout` (all batched `[bh,seq,*]`), and the
     /// same `scale`. Returns `(dq, dk, dv)`.

crates/xtrain-tensor/tests/integration.rs

@@ -56,3 +56,170 @@ fn elementwise_scale_kernel() {
         r.len()
     );
 }
+
+/// (c) `rope_at` (KV-cache decode RoPE at an absolute position) is bit-identical
+/// to the full-sequence `rope`'s corresponding row. This is the invariant the
+/// decode KV-cache relies on: a single new token RoPE'd at position `t` must equal
+/// what the full-sequence forward would have produced at row `t` (so cached
+/// post-RoPE K matches the full-recompute path → token-identical decode).
+#[test]
+fn rope_at_matches_full_rope_row() {
+    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;
+    // Deterministic pseudo-random fill in [-1, 1).
+    let host: Vec<f32> = (0..n * heads * hd)
+        .map(|i| ((i * 37 % 101) as f32 / 50.0) - 1.0)
+        .collect();
+
+    // Full-sequence rope (period = n → row r gets position r).
+    let full = Tensor::from_slice(&host, &[n, heads, hd]).to_device(Device::Cuda(0));
+    let roped_full = full
+        .rope(theta, n)
+        .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 roped_row = Tensor::from_slice(row, &[1, heads, hd])
+            .to_device(Device::Cuda(0))
+            .rope_at(theta, t)
+            .to_device(Device::Cpu)
+            .as_slice::<f32>()
+            .to_vec();
+        let expect = &roped_full[t * row_len..(t + 1) * row_len];
+        assert_eq!(
+            roped_row.as_slice(),
+            expect,
+            "rope_at(pos0={t}) != full rope row {t}"
+        );
+    }
+    println!("rope_at OK: bit-identical to full rope across {n} positions");
+}
+
+/// (d) `decode_attention` (single query vs cached K/V, no mask) equals the LAST
+/// query row of the full causal `attention`. This is the core decode-engine
+/// invariant: the incremental path must reproduce what the full-recompute forward
+/// computes for the final position, so KV-cache greedy decode is token-identical.
+/// Tolerance is fp rounding (different softmax kernel + reduction order), not bits.
+#[test]
+fn decode_attention_matches_full_attention_last_row() {
+    assert!(
+        device::device_count().expect("device count") > 0,
+        "no CUDA device"
+    );
+    device::set_device(0).unwrap();
+
+    let (bh, t, hd) = (6usize, 5usize, 8usize);
+    let scale = 1.0 / (hd as f32).sqrt();
+    let n = bh * t * hd;
+    let qh: Vec<f32> = (0..n).map(|i| ((i * 31 % 97) as f32 / 48.0) - 1.0).collect();
+    let kh: Vec<f32> = (0..n).map(|i| ((i * 53 % 89) as f32 / 44.0) - 1.0).collect();
+    let vh: Vec<f32> = (0..n).map(|i| ((i * 17 % 83) as f32 / 41.0) - 1.0).collect();
+    let q = Tensor::from_slice(&qh, &[bh, t, hd]).to_device(Device::Cuda(0));
+    let k = Tensor::from_slice(&kh, &[bh, t, hd]).to_device(Device::Cuda(0));
+    let v = Tensor::from_slice(&vh, &[bh, t, hd]).to_device(Device::Cuda(0));
+
+    // Reference: full causal attention, take each head's last query row.
+    let (full, _) = q.attention(&k, &v, scale);
+    let full_h = full.to_device(Device::Cpu).as_slice::<f32>().to_vec();
+
+    // Decode: build Q_last [bh,1,hd] from each head's last row, attend to all K/V.
+    let mut ql = vec![0f32; bh * hd];
+    for b in 0..bh {
+        let src = (b * t + (t - 1)) * hd;
+        ql[b * hd..(b + 1) * hd].copy_from_slice(&qh[src..src + hd]);
+    }
+    let q_last = Tensor::from_slice(&ql, &[bh, 1, hd]).to_device(Device::Cuda(0));
+    let dec = q_last
+        .decode_attention(&k, &v, scale)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec();
+    assert_eq!(dec.len(), bh * hd, "decode out shape");
+
+    let mut max_abs = 0f32;
+    for b in 0..bh {
+        for d in 0..hd {
+            let got = dec[b * hd + d];
+            let exp = full_h[(b * t + (t - 1)) * hd + d];
+            max_abs = max_abs.max((got - exp).abs());
+        }
+    }
+    assert!(
+        max_abs < 1e-4,
+        "decode_attention vs full last-row max abs diff {max_abs} exceeds 1e-4"
+    );
+    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");
+}
+
+/// (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");
+}

crates/xtrain-train/src/bin/bench_grpo_batch.rs

@@ -0,0 +1,268 @@
+//! Micro-benchmark + closeness gate for the M2d batched GRPO training-side forwards.
+//!
+//! After M2b/M2c the GRPO *step* is no longer rollout-bound — it is the `N = B·G`
+//! per-sample full-sequence forwards (the `per_token_logp` captures + the inner
+//! clipped-PG forward/backwards). This bin isolates exactly that, weight-independently
+//! (step wall-clock depends on shapes + launch counts, not on what the weights are), by
+//! synthesising `N` realistic ragged samples and A/B-timing the looped vs batched path
+//! for BOTH phases — plus asserting they agree numerically (the looped-vs-batched
+//! closeness gate; per-row bit-equivalence of the loss op is pinned by the autograd
+//! test `clipped_pg_loss_batched_matches_looped`).
+//!
+//!   bench_grpo_batch <tokenizer.json> --init-ckpt <base.ckpt> <arch flags> \
+//!       --n 48 --plen 12 --clen 24 --micro 16 --reps 3
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("bench_grpo_batch: built without CUDA (no_cuda); run on a GPU host.");
+}
+
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::{DType, Device, Tensor};
+#[cfg(not(no_cuda))]
+use xtrain_train::grpo_batch::{PgSample, inner_pg_step_batched, inner_pg_step_looped, per_token_logp, per_token_logp_batched};
+
+#[cfg(not(no_cuda))]
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed.wrapping_mul(2862933555777941757).wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+#[cfg(not(no_cuda))]
+fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
+    args.iter().position(|a| a == name).and_then(|i| args.get(i + 1)).and_then(|s| s.parse().ok()).unwrap_or(default)
+}
+
+#[cfg(not(no_cuda))]
+fn flag_value(args: &[String], name: &str) -> Option<String> {
+    args.iter().position(|a| a == name).and_then(|i| args.get(i + 1)).cloned()
+}
+
+#[cfg(not(no_cuda))]
+fn load_model(cfg: Config, device: Device, ckpt: &str) -> TinyTransformer {
+    let mut seed = 1u64;
+    let m = TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.04)
+        }
+    })
+    .with_compute_dtype(DType::BF16)
+    .with_flash(true);
+    xtrain_train::checkpoint::load_into(std::path::Path::new(ckpt), &m.params()).expect("load ckpt");
+    m.eval();
+    m
+}
+
+#[cfg(not(no_cuda))]
+fn elapsed_ms<F: FnMut()>(reps: usize, mut f: F) -> f32 {
+    let start = std::time::Instant::now();
+    for _ in 0..reps {
+        f();
+    }
+    start.elapsed().as_secs_f32() * 1e3 / reps as f32
+}
+
+/// Per-position argmax of the model over each ragged `input` (one `forward_batched`
+/// per `micro`-chunk). Used to teacher-force WELL-CONDITIONED targets (the top-1 token,
+/// high prob) so the closeness gate's logp isn't the ~−20 of a random token — where
+/// `−log p` amplifies bf16 noise. This matches real GRPO (targets are model samples).
+#[cfg(not(no_cuda))]
+fn model_argmax(model: &TinyTransformer, device: Device, inputs: &[Vec<i32>], vocab: usize, micro: usize) -> Vec<Vec<i32>> {
+    let mut out = Vec::with_capacity(inputs.len());
+    for chunk in inputs.chunks(micro.max(1)) {
+        let m = chunk.len();
+        let lmax = chunk.iter().map(|s| s.len()).max().unwrap();
+        let mut flat = vec![0i32; m * lmax];
+        for (i, s) in chunk.iter().enumerate() {
+            flat[i * lmax..i * lmax + s.len()].copy_from_slice(s);
+        }
+        let ids = Tensor::from_slice(&flat, &[m * lmax]).to_device(device);
+        let logits = model.forward_batched(&ids, m).value().to_dtype(DType::F32).to_device(Device::Cpu);
+        let v = logits.as_slice::<f32>();
+        for (i, s) in chunk.iter().enumerate() {
+            let mut row = Vec::with_capacity(s.len());
+            for r in 0..s.len() {
+                let base = (i * lmax + r) * vocab;
+                let mut best = 0usize;
+                for c in 1..vocab {
+                    if v[base + c] > v[base + best] {
+                        best = c;
+                    }
+                }
+                row.push(best as i32);
+            }
+            out.push(row);
+        }
+    }
+    out
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use xserv_tokenizer::Tokenizer;
+
+    let args: Vec<String> = std::env::args().collect();
+    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
+    let tok_path = positionals.first().expect("usage: bench_grpo_batch <tokenizer.json> [flags]");
+
+    let n_heads = flag(&args, "--heads", 52usize);
+    let head_dim = flag(&args, "--head-dim", 32usize);
+    let n_layers = flag(&args, "--layers", 22usize);
+    let ffn = flag(&args, "--ffn", 6656usize);
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
+    let n: usize = flag(&args, "--n", 48); // B·G samples per step
+    let plen: usize = flag(&args, "--plen", 12); // prompt tokens
+    let clen: usize = flag(&args, "--clen", 24); // max completion tokens
+    let micro: usize = flag(&args, "--micro", 16);
+    let reps: usize = flag(&args, "--reps", 3);
+    let (eps, beta) = (flag(&args, "--eps", 0.2f32), flag(&args, "--beta", 0.0f32));
+    let init_ckpt = flag_value(&args, "--init-ckpt").expect("--init-ckpt <base.ckpt> required");
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+    let tok = Tokenizer::from_file(std::path::Path::new(tok_path.as_str()));
+    let vocab = tok.vocab_size();
+    let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
+    let policy = load_model(cfg, device, &init_ckpt);
+    let params = policy.params();
+
+    // --- Synthesise N ragged samples (frame-shaped: prompt masked, ragged completion).
+    // Token IDs are random-but-valid; only the SHAPES drive the forward cost.
+    let mut rng = 0xC0FFEEu64;
+    let mut next = || {
+        rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
+        (rng >> 33) as usize
+    };
+    let mut io: Vec<(Vec<i32>, Vec<i32>)> = Vec::with_capacity(n);
+    let mut advs: Vec<f32> = Vec::with_capacity(n);
+    for _ in 0..n {
+        let pl = plen.saturating_sub(2) + next() % 5; // jitter prompt length a little
+        let cl = 4 + next() % clen.max(1); // completion 4..=clen
+        let total = pl + cl;
+        let toks: Vec<i32> = (0..total).map(|_| (next() % vocab) as i32).collect();
+        let mut labels = vec![-100i32; pl]; // prompt masked
+        labels.extend_from_slice(&toks[pl..]);
+        let l = toks.len();
+        io.push((toks[..l - 1].to_vec(), labels[1..l].to_vec())); // target masked at [..pl-1]
+        advs.push(if next() % 2 == 0 { 0.7 } else { -0.7 });
+    }
+    let toklens: Vec<usize> = io.iter().map(|(i, _)| i.len()).collect();
+    let (lmin, lmax) = (*toklens.iter().min().unwrap(), *toklens.iter().max().unwrap());
+    println!("samples N={n}, seq len {lmin}..{lmax} (ragged), micro={micro}, β={beta}\n");
+
+    // Replace random completion targets with the model's own argmax (teacher forcing):
+    // well-conditioned logp (top-1, not the ~−20 of a random token where bf16 noise
+    // blows up via −log p). The completion target positions are where the skeleton is
+    // ≥0; prompt positions stay masked (−100).
+    let inputs: Vec<Vec<i32>> = io.iter().map(|(i, _)| i.clone()).collect();
+    let preds = model_argmax(&policy, device, &inputs, vocab, micro);
+    for (s, (_, target)) in io.iter_mut().enumerate() {
+        for j in 0..target.len() {
+            if target[j] >= 0 {
+                target[j] = preds[s][j];
+            }
+        }
+    }
+
+    // ---------------- Phase 1: capture (per_token_logp) ----------------
+    let logp_loop: Vec<Vec<f32>> = io.iter().map(|(i, t)| per_token_logp(&policy, device, i, t)).collect();
+    let logp_batch = per_token_logp_batched(&policy, device, &io, micro);
+    let cap_dmax = logp_loop
+        .iter()
+        .zip(&logp_batch)
+        .flat_map(|(a, b)| a.iter().zip(b).map(|(x, y)| (x - y).abs()))
+        .fold(0.0f32, f32::max);
+    let t_cap_loop = elapsed_ms(reps, || {
+        let _: Vec<Vec<f32>> = io.iter().map(|(i, t)| per_token_logp(&policy, device, i, t)).collect();
+    });
+    let t_cap_batch = elapsed_ms(reps, || {
+        let _ = per_token_logp_batched(&policy, device, &io, micro);
+    });
+
+    // Build PgSamples from the (matching) capture; ref = old − 0.3 to exercise KL.
+    let batch: Vec<PgSample> = io
+        .iter()
+        .zip(&advs)
+        .zip(&logp_batch)
+        .map(|(((input, target), &adv), lp)| PgSample {
+            input: input.clone(),
+            target: target.clone(),
+            adv,
+            logp_old: lp.clone(),
+            logp_ref: lp.iter().map(|v| v - 0.3).collect(),
+        })
+        .collect();
+
+    // ---------------- Phase 2: inner clipped-PG (forward + backward) ----------------
+    // Representative grad snapshots: layer-0 wq (params[2]) + final_norm.
+    let wq0 = &params[2];
+    let fnorm = &params[1 + n_layers * 11];
+    let snap = |v: &xtrain_autodiff::Var| -> Vec<f32> {
+        v.grad().map(|g| g.to_device(Device::Cpu).as_slice::<f32>().to_vec()).unwrap_or_default()
+    };
+    let zero = |ps: &[xtrain_autodiff::Var]| ps.iter().for_each(|p| p.zero_grad());
+
+    zero(&params);
+    inner_pg_step_looped(&policy, device, &batch, eps, beta);
+    let (gq_loop, gn_loop) = (snap(wq0), snap(fnorm));
+    zero(&params);
+    inner_pg_step_batched(&policy, device, &batch, eps, beta, micro);
+    let (gq_batch, gn_batch) = (snap(wq0), snap(fnorm));
+    zero(&params);
+
+    let reldiff = |a: &[f32], b: &[f32]| -> f32 {
+        let num = a.iter().zip(b).map(|(x, y)| (x - y).abs()).fold(0.0f32, f32::max);
+        let den = a.iter().map(|x| x.abs()).fold(0.0f32, f32::max).max(1e-12);
+        num / den
+    };
+    let gq_rel = reldiff(&gq_loop, &gq_batch);
+    let gn_rel = reldiff(&gn_loop, &gn_batch);
+
+    // Time only forward+backward — the lever. opt.step + grad-clip are identical in
+    // both paths (one call over `params` after the per-sample loop), so they would
+    // only add a constant; excluding them also dodges the unrelated 1B-Adam-state
+    // memory wall (the M4 finding) that this diagnostic doesn't need to reproduce.
+    let t_inner_loop = elapsed_ms(reps, || {
+        inner_pg_step_looped(&policy, device, &batch, eps, beta);
+        zero(&params);
+    });
+    let t_inner_batch = elapsed_ms(reps, || {
+        inner_pg_step_batched(&policy, device, &batch, eps, beta, micro);
+        zero(&params);
+    });
+
+    // ---------------- Report ----------------
+    let spd = |a: f32, b: f32| if b > 0.0 { a / b } else { 0.0 };
+    println!("=== closeness gate (looped vs batched) ===");
+    println!("  capture per_token_logp : max|Δ| = {cap_dmax:.3e}");
+    println!("  inner grad wq[0]       : rel|Δ| = {gq_rel:.3e}");
+    println!("  inner grad final_norm  : rel|Δ| = {gn_rel:.3e}");
+    println!("\n=== timing (mean of {reps} reps, ms/phase) ===");
+    println!("  capture : looped {t_cap_loop:8.1}  batched {t_cap_batch:8.1}  ({:.2}× )", spd(t_cap_loop, t_cap_batch));
+    println!("  inner   : looped {t_inner_loop:8.1}  batched {t_inner_batch:8.1}  ({:.2}× )", spd(t_inner_loop, t_inner_batch));
+    let (step_loop, step_batch) = (t_cap_loop + t_inner_loop, t_cap_batch + t_inner_batch);
+    println!("  STEP    : looped {step_loop:8.1}  batched {step_batch:8.1}  ({:.2}× )", spd(step_loop, step_batch));
+
+    // The RIGOROUS correctness gates live in the test suite (exact, not bf16-noisy):
+    //   - xtrain-model  forward_batched_ragged_matches_looped  (forward+pad == looped)
+    //   - xtrain-autodiff clipped_pg_loss_batched_matches_looped (op == looped, f32)
+    // This is a smoke check at the 1B/bf16 scale: single-seq vs batched GEMM differ in
+    // batch-reduction order, so a loose band, with well-conditioned (argmax) targets.
+    assert!(cap_dmax < 0.2, "capture closeness smoke FAILED: max|Δlogp| = {cap_dmax}");
+    assert!(gq_rel < 0.2 && gn_rel < 0.2, "inner grad closeness smoke FAILED: wq {gq_rel}, fn {gn_rel}");
+    println!("\nSMOKE PASS (bf16 band): batched ≈ looped; rigorous gates are the two tests above.");
+}

crates/xtrain-train/src/bin/eval_arith.rs

@@ -103,6 +103,10 @@ fn main() {
     let n_show = flag(&args, "--show", 8usize);
     let prompts_file = flag_value(&args, "--prompts-file").expect("--prompts-file is required");
     let gold_file = flag_value(&args, "--gold-file").expect("--gold-file is required");
+    // M2: decode through the KV-cache incremental engine instead of the naive
+    // full-recompute sampler. Token-identical to the naive path (gated by
+    // tests/decode_kv.rs); this flag also lets us A/B the two for the speedup.
+    let use_cached = args.iter().any(|a| a == "--cached");
 
     // Prompts: skip the `#` header / blank lines and decode escaped newlines so the
     // count and order line up with the gold file.
@@ -148,18 +152,26 @@ fn main() {
     xtrain_train::checkpoint::load_into(&ckpt, &model.params()).expect("load checkpoint");
 
     println!(
-        "eval_arith: ckpt {} | {} prompts | max_new {}",
+        "eval_arith: ckpt {} | {} prompts | max_new {} | decode={}",
         ckpt.display(),
         prompts.len(),
-        max_new
+        max_new,
+        if use_cached { "kv-cache" } else { "naive" }
     );
 
     let (mut n_boxed, mut n_correct) = (0usize, 0usize);
     let mut shown = 0usize;
+    let mut gen_tokens = 0usize;
+    let t0 = std::time::Instant::now();
     for (prompt, &gold) in prompts.iter().zip(&golds) {
         let ids: Vec<i32> = tok.encode(prompt).into_iter().map(|t| t as i32).collect();
-        let mut rng = 7u64;
-        let out = generate(&model, device, &ids, max_new, 0.0, &mut rng);
+        let out = if use_cached {
+            xtrain_model::generate_greedy_cached(&model, device, &ids, max_new)
+        } else {
+            let mut rng = 7u64;
+            generate(&model, device, &ids, max_new, 0.0, &mut rng)
+        };
+        gen_tokens += out.len() - ids.len();
         let cont = tok.decode(&out[ids.len()..].iter().map(|&t| t as u32).collect::<Vec<_>>());
         let seg = first_answer_segment(&cont);
         if parse_boxed_answer(seg).is_some() {
@@ -176,6 +188,7 @@ fn main() {
         }
     }
 
+    let elapsed = t0.elapsed().as_secs_f64();
     let n = prompts.len() as f64;
     println!(
         "RESULT format(boxed)={}/{} ({:.1}%) | correct={}/{} ({:.1}%)",
@@ -186,4 +199,11 @@ fn main() {
         prompts.len(),
         100.0 * n_correct as f64 / n,
     );
+    println!(
+        "TIMING decode={} | {:.2}s | {} gen tokens | {:.1} tok/s",
+        if use_cached { "kv-cache" } else { "naive" },
+        elapsed,
+        gen_tokens,
+        gen_tokens as f64 / elapsed,
+    );
 }

crates/xtrain-train/src/bin/gen_dpo_pairs.rs

@@ -0,0 +1,157 @@
+//! Generate DPO preference pairs for the verifiable arithmetic task (M3).
+//!
+//! Per the aligned decision: **chosen = the gold answer** (`sft_answer`, always
+//! correct), **rejected = a sampled-incorrect completion from the SFT model** — a
+//! format-valid but wrong boxed answer, i.e. a hard negative drawn from the model's
+//! own distribution. Since the SFT model is only ~8% correct (M1), a single GREEDY
+//! decode is wrong ~92% of the time, so we use the KV-cache greedy engine (M2a) and
+//! simply skip the ~8% of prompts where greedy happens to be correct (no usable
+//! negative). Fast (cached), deterministic, and one clean hard negative per prompt.
+//!
+//! Writes `<out>` as `question<TAB>chosen<TAB>rejected` (bare text, like the SFT
+//! TSV — `train_dpo` adds the `User:/Assistant:` frame). Problems are deduped.
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("gen_dpo_pairs: built without CUDA (no_cuda); run on a GPU host.");
+}
+
+#[cfg(not(no_cuda))]
+use std::collections::HashSet;
+#[cfg(not(no_cuda))]
+use std::io::Write;
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer, generate_greedy_cached};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::Device;
+#[cfg(not(no_cuda))]
+use xtrain_train::task::{Op, GenConfig, check_answer, gen_problem, parse_boxed_answer};
+
+#[cfg(not(no_cuda))]
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+#[cfg(not(no_cuda))]
+fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}
+
+#[cfg(not(no_cuda))]
+fn flag_value(args: &[String], name: &str) -> Option<String> {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .cloned()
+}
+
+/// Keep only the first answer "turn": cut at the first `<|endoftext|>` then the
+/// first newline (mirrors eval_arith).
+#[cfg(not(no_cuda))]
+fn first_answer_segment(continuation: &str) -> &str {
+    let s = continuation
+        .split("<|endoftext|>")
+        .next()
+        .unwrap_or(continuation);
+    s.split('\n').next().unwrap_or(s)
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use xserv_tokenizer::Tokenizer;
+
+    let args: Vec<String> = std::env::args().collect();
+    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
+    let ckpt = positionals.first().expect("usage: gen_dpo_pairs <sft_ckpt> <tokenizer.json> [flags]");
+    let tok_path = positionals
+        .get(1)
+        .map(|s| s.as_str())
+        .unwrap_or("/opt/wjh/models/gpt2/tokenizer.json");
+
+    let n_heads = flag(&args, "--heads", 52usize);
+    let head_dim = flag(&args, "--head-dim", 32usize);
+    let n_layers = flag(&args, "--layers", 22usize);
+    let ffn = flag(&args, "--ffn", 6656usize);
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
+    let n_pairs: usize = flag(&args, "--n", 2000);
+    let seed: u64 = flag(&args, "--seed", 1234);
+    let max_add: i64 = flag(&args, "--max-add", 999);
+    let max_mul: i64 = flag(&args, "--max-mul", 99);
+    let max_new: usize = flag(&args, "--max-tokens", 32);
+    let out = flag_value(&args, "--out").expect("--out <file> is required");
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let tok = Tokenizer::from_file(std::path::Path::new(tok_path));
+    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn)
+        .with_kv_heads(kv_heads);
+    let mut seed_init = 1u64;
+    let model = TinyTransformer::new(cfg, device, |shape| {
+        seed_init = seed_init.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed_init, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed_init, 0.04)
+        }
+    });
+    xtrain_train::checkpoint::load_into(std::path::Path::new(ckpt.as_str()), &model.params())
+        .expect("load SFT checkpoint");
+
+    let gcfg = GenConfig {
+        max_add,
+        max_mul,
+        ops: vec![Op::Add, Op::Sub, Op::Mul],
+    };
+    let mut rng = seed.max(1);
+    let mut keys = HashSet::new();
+    let mut writer = std::io::BufWriter::new(std::fs::File::create(&out).expect("create out"));
+    let (mut written, mut skipped, mut attempts) = (0usize, 0usize, 0usize);
+
+    while written < n_pairs {
+        attempts += 1;
+        if attempts > n_pairs * 4 {
+            eprintln!("gen_dpo_pairs: stopping early at {written} pairs after {attempts} attempts");
+            break;
+        }
+        let p = gen_problem(&mut rng, &gcfg);
+        if !keys.insert(p.key()) {
+            continue;
+        }
+        let prompt_text = format!("User: {}\nAssistant:", p.question());
+        let ids: Vec<i32> = tok.encode(&prompt_text).into_iter().map(|t| t as i32).collect();
+        let out_ids = generate_greedy_cached(&model, device, &ids, max_new);
+        let cont = tok.decode(&out_ids[ids.len()..].iter().map(|&t| t as u32).collect::<Vec<_>>());
+        let seg = first_answer_segment(&cont).trim();
+        // A valid hard negative: a well-formed boxed answer that is WRONG.
+        if parse_boxed_answer(seg).is_some() && !check_answer(seg, p.answer()) {
+            writeln!(writer, "{}\t{}\t{}", p.question(), p.sft_answer(), seg).expect("write");
+            written += 1;
+        } else {
+            skipped += 1; // greedy was correct (~8%) or malformed → no clean negative
+        }
+    }
+    writer.flush().expect("flush");
+    println!(
+        "wrote {written} DPO pairs to {out} (skipped {skipped} no-negative; {attempts} attempts; \
+         chosen=gold, rejected=greedy-incorrect)"
+    );
+}

crates/xtrain-train/src/bin/train_dpo.rs

@@ -0,0 +1,233 @@
+//! DPO training on the verifiable arithmetic task (M3 / Stage P1).
+//!
+//! Loads the SFT checkpoint as the policy AND uses it as the frozen reference:
+//! reference logprobs `log πref(chosen)` / `log πref(rejected)` are **precomputed
+//! once** before any optimizer step (when policy == reference), then cached as
+//! constants — so only one model stays resident (the design's reference-logprob
+//! caching). Each step forwards the policy on the chosen and rejected completions,
+//! takes [`seq_logprob`] of each, and minimises [`dpo_loss`]; the two forwards
+//! share the policy params, so backward accumulates both branches' grads.
+//!
+//! Health metrics (per docs/18, the doc-13 "don't trust loss alone" lesson): the
+//! chosen−rejected **reward margin** and **preference accuracy** (margin > 0) — both
+//! should rise. The arithmetic-correctness payoff is measured separately by running
+//! `eval_arith` on the saved checkpoint.
+//!
+//!   train_dpo <tokenizer.json> <dpo.tsv> --init-ckpt <sft.ckpt> <arch flags> \
+//!       --beta 0.1 --steps 1000 --lr 5e-7 --ckpt <out.ckpt>
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("train_dpo: built without CUDA (no_cuda); run on a GPU host.");
+}
+
+#[cfg(not(no_cuda))]
+use xtrain_autodiff::ops;
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer, ids_tensor};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::Device;
+
+#[cfg(not(no_cuda))]
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+#[cfg(not(no_cuda))]
+fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}
+
+#[cfg(not(no_cuda))]
+fn flag_value(args: &[String], name: &str) -> Option<String> {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .cloned()
+}
+
+/// Frame a (question, completion) the same way the SFT loader does
+/// (`User: …\nAssistant:` prompt + ` {completion}\n<|endoftext|>`), then return the
+/// next-token (input, target) pair: input = tokens[..L-1], target = labels[1..L]
+/// with the prompt positions masked to -100 (only completion tokens supervised).
+#[cfg(not(no_cuda))]
+fn frame(
+    tok: &xserv_tokenizer::Tokenizer,
+    question: &str,
+    completion: &str,
+) -> (Vec<i32>, Vec<i32>) {
+    let prompt = format!("User: {question}\nAssistant:");
+    let answer = format!(" {completion}\n<|endoftext|>");
+    let p_ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
+    let a_ids: Vec<i32> = tok.encode(&answer).into_iter().map(|t| t as i32).collect();
+    let mut tokens = p_ids.clone();
+    tokens.extend_from_slice(&a_ids);
+    let mut labels = vec![-100i32; p_ids.len()];
+    labels.extend_from_slice(&a_ids);
+    let l = tokens.len();
+    (tokens[..l - 1].to_vec(), labels[1..l].to_vec())
+}
+
+/// Sequence logprob `Σ log πθ(completion)` of a framed (input, target) pair.
+#[cfg(not(no_cuda))]
+fn seq_lp(
+    model: &TinyTransformer,
+    device: Device,
+    input: &[i32],
+    target: &[i32],
+) -> xtrain_autodiff::tape::Var {
+    let logits = model.forward(&ids_tensor(input, device));
+    ops::seq_logprob(&logits, &ids_tensor(target, device))
+}
+
+#[cfg(not(no_cuda))]
+fn scalar(v: &xtrain_autodiff::tape::Var) -> f32 {
+    v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use xserv_tokenizer::Tokenizer;
+    use xtrain_optim::GpuAdamW;
+
+    let args: Vec<String> = std::env::args().collect();
+    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
+    let tok_path = positionals.first().expect("usage: train_dpo <tokenizer.json> <dpo.tsv> [flags]");
+    let tsv_path = positionals.get(1).expect("usage: train_dpo <tokenizer.json> <dpo.tsv> [flags]");
+
+    let n_heads = flag(&args, "--heads", 52usize);
+    let head_dim = flag(&args, "--head-dim", 32usize);
+    let n_layers = flag(&args, "--layers", 22usize);
+    let ffn = flag(&args, "--ffn", 6656usize);
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
+    let beta: f32 = flag(&args, "--beta", 0.1);
+    let steps: usize = flag(&args, "--steps", 1000);
+    let lr: f32 = flag(&args, "--lr", 5e-7);
+    let wd: f32 = flag(&args, "--wd", 0.0);
+    let clip: f32 = flag(&args, "--clip", 1.0);
+    let log_every: usize = flag(&args, "--log-every", 50);
+    let init_ckpt = flag_value(&args, "--init-ckpt").expect("--init-ckpt <sft.ckpt> is required");
+    let out_ckpt = flag_value(&args, "--ckpt").expect("--ckpt <out> is required");
+
+    // Load preference pairs: question<TAB>chosen<TAB>rejected.
+    let raw = std::fs::read_to_string(tsv_path).expect("read dpo tsv");
+    let pairs: Vec<(String, String, String)> = raw
+        .lines()
+        .filter(|l| !l.trim().is_empty())
+        .map(|l| {
+            let mut it = l.splitn(3, '\t');
+            let q = it.next().expect("question").to_string();
+            let c = it.next().expect("chosen").to_string();
+            let r = it.next().expect("rejected").to_string();
+            (q, c, r)
+        })
+        .collect();
+    assert!(!pairs.is_empty(), "no DPO pairs in {tsv_path}");
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let tok = Tokenizer::from_file(std::path::Path::new(tok_path.as_str()));
+    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn)
+        .with_kv_heads(kv_heads);
+    let mut seed_init = 1u64;
+    let model = TinyTransformer::new(cfg, device, |shape| {
+        seed_init = seed_init.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed_init, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed_init, 0.04)
+        }
+    });
+    xtrain_train::checkpoint::load_into(std::path::Path::new(&init_ckpt), &model.params())
+        .expect("load SFT checkpoint");
+    model.eval(); // DPO runs without dropout (deterministic logprobs)
+
+    // Pre-tokenize every pair once.
+    let framed: Vec<((Vec<i32>, Vec<i32>), (Vec<i32>, Vec<i32>))> = pairs
+        .iter()
+        .map(|(q, c, r)| (frame(&tok, q, c), frame(&tok, q, r)))
+        .collect();
+
+    // Reference logprobs: computed ONCE while policy == reference (SFT init), cached.
+    println!("precomputing reference logprobs for {} pairs…", framed.len());
+    let mut ref_c = Vec::with_capacity(framed.len());
+    let mut ref_r = Vec::with_capacity(framed.len());
+    for ((ci, ct), (ri, rt)) in &framed {
+        ref_c.push(scalar(&seq_lp(&model, device, ci, ct)));
+        ref_r.push(scalar(&seq_lp(&model, device, ri, rt)));
+    }
+
+    let params = model.params();
+    let mut opt = GpuAdamW::new(wd);
+    let n = framed.len();
+    // A fixed shuffle (LCG-strided) so steps sweep the dataset without bias.
+    let mut order: Vec<usize> = (0..n).collect();
+    let mut s = 0x9E3779B97F4A7C15u64;
+    for i in (1..n).rev() {
+        s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
+        let j = (s >> 33) as usize % (i + 1);
+        order.swap(i, j);
+    }
+
+    let start = std::time::Instant::now();
+    let (mut win_loss, mut win_margin, mut win_acc) = (0f32, 0f32, 0usize);
+    for step in 0..steps {
+        let i = order[step % n];
+        let ((ci, ct), (ri, rt)) = &framed[i];
+        let lpc = seq_lp(&model, device, ci, ct);
+        let lpr = seq_lp(&model, device, ri, rt);
+        let (lpc_v, lpr_v) = (scalar(&lpc), scalar(&lpr));
+        let margin = (lpc_v - ref_c[i]) - (lpr_v - ref_r[i]); // implicit reward margin
+        let loss = ops::dpo_loss(&lpc, &lpr, ref_c[i], ref_r[i], beta);
+        win_loss += scalar(&loss);
+        win_margin += margin;
+        win_acc += (margin > 0.0) as usize;
+
+        loss.backward();
+        let _ = xtrain_train::clip::clip_grad_norm_gpu(&params, clip, 1.0);
+        opt.step(lr, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+
+        if (step + 1) % log_every == 0 || step == steps - 1 {
+            let w = log_every.min(step + 1) as f32;
+            println!(
+                "step {:5}/{steps}: loss {:.4} | reward-margin {:+.4} | pref-acc {:.1}% | {:.1}s",
+                step + 1,
+                win_loss / w,
+                win_margin / w,
+                100.0 * win_acc as f32 / w,
+                start.elapsed().as_secs_f32(),
+            );
+            win_loss = 0.0;
+            win_margin = 0.0;
+            win_acc = 0;
+        }
+    }
+
+    xtrain_train::checkpoint::save(std::path::Path::new(&out_ckpt), &params).expect("save ckpt");
+    println!(
+        "DPO done: {} pairs, {steps} steps, beta {beta}, lr {lr:.1e} → {out_ckpt}",
+        framed.len()
+    );
+}

crates/xtrain-train/src/bin/train_grpo.rs

@@ -0,0 +1,294 @@
+//! GRPO training on the verifiable arithmetic task (M4 / Stage P3) — online,
+//! critic-free RL. The centerpiece: generation INSIDE the training loop.
+//!
+//! Each step: sample B prompts (fresh problems), roll out G completions per prompt
+//! (temperature sampling via the naive sampler — batched/cached rollout is the M2b/
+//! M4-perf follow-up), score each with the rule-based checker (reward ∈ {0,1}),
+//! compute the **group-relative advantage** `A_i = (r_i − mean) / (std + ε)` (no
+//! critic), then K inner clipped-PG epochs minimising [`clipped_pg_loss`] with a KL
+//! leash to the frozen reference (πref = the SFT checkpoint). Reward = pure 0/1
+//! correctness; the KL term (β) is what keeps format/coherence (the M3 collapse
+//! lesson — here it is an explicit leash, not just a hope).
+//!
+//! Health signal (the falsifiable "it learns"): **mean rollout reward must rise**
+//! (the RL analogue of T5's overfit-27/27). Held-out correctness is measured by
+//! eval_arith on the saved checkpoint.
+//!
+//!   train_grpo <tokenizer.json> --init-ckpt <sft.ckpt> <arch flags> \
+//!       --steps 200 --group 6 --prompts 8 --temp 1.0 --beta 0.04 --eps 0.2 \
+//!       --lr 1e-6 --max-add 20 --max-mul 9 --ckpt <out.ckpt>
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("train_grpo: built without CUDA (no_cuda); run on a GPU host.");
+}
+
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer, generate_cached_batch};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::{DType, Device};
+#[cfg(not(no_cuda))]
+use xtrain_train::grpo_batch::{PgSample, inner_pg_step_batched, per_token_logp_batched};
+#[cfg(not(no_cuda))]
+use xtrain_train::task::{check_answer, gen_problem, GenConfig, Op};
+
+#[cfg(not(no_cuda))]
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+#[cfg(not(no_cuda))]
+fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}
+
+#[cfg(not(no_cuda))]
+fn flag_value(args: &[String], name: &str) -> Option<String> {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .cloned()
+}
+
+#[cfg(not(no_cuda))]
+fn first_answer_segment(c: &str) -> &str {
+    let s = c.split("<|endoftext|>").next().unwrap_or(c);
+    s.split('\n').next().unwrap_or(s)
+}
+
+/// Build a model from the SFT checkpoint (bf16 compute to fit two 1B models). The
+/// policy enables activation recompute (T13) so its backward fits alongside the
+/// frozen reference + the Adam state; the reference only forwards (no backward).
+#[cfg(not(no_cuda))]
+fn load_model(cfg: Config, device: Device, ckpt: &str, recompute: bool) -> TinyTransformer {
+    let mut seed = 1u64;
+    let m = TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.04)
+        }
+    })
+    .with_compute_dtype(DType::BF16)
+    .with_recompute(recompute)
+    .with_flash(true);
+    xtrain_train::checkpoint::load_into(std::path::Path::new(ckpt), &m.params()).expect("load ckpt");
+    m.eval();
+    m
+}
+
+/// Frame (question, completion) like the SFT loader and return the next-token
+/// (input, target) pair (prompt masked to -100). Same as train_dpo.
+#[cfg(not(no_cuda))]
+fn frame(tok: &xserv_tokenizer::Tokenizer, question: &str, completion: &str) -> (Vec<i32>, Vec<i32>) {
+    let p_ids: Vec<i32> = tok
+        .encode(&format!("User: {question}\nAssistant:"))
+        .into_iter()
+        .map(|t| t as i32)
+        .collect();
+    let a_ids: Vec<i32> = tok
+        .encode(&format!(" {completion}\n<|endoftext|>"))
+        .into_iter()
+        .map(|t| t as i32)
+        .collect();
+    let mut tokens = p_ids.clone();
+    tokens.extend_from_slice(&a_ids);
+    let mut labels = vec![-100i32; p_ids.len()];
+    labels.extend_from_slice(&a_ids);
+    let l = tokens.len();
+    (tokens[..l - 1].to_vec(), labels[1..l].to_vec())
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use xserv_tokenizer::Tokenizer;
+    use xtrain_optim::GpuAdamW;
+
+    let args: Vec<String> = std::env::args().collect();
+    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
+    let tok_path = positionals.first().expect("usage: train_grpo <tokenizer.json> [flags]");
+
+    let n_heads = flag(&args, "--heads", 52usize);
+    let head_dim = flag(&args, "--head-dim", 32usize);
+    let n_layers = flag(&args, "--layers", 22usize);
+    let ffn = flag(&args, "--ffn", 6656usize);
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
+    let steps: usize = flag(&args, "--steps", 200);
+    let group: usize = flag(&args, "--group", 6);
+    let n_prompts: usize = flag(&args, "--prompts", 8);
+    let inner: usize = flag(&args, "--inner", 1);
+    // M2d: pack the step's N=B·G ragged samples into forward_batched chunks of this
+    // many samples (bounds the [chunk·Lmax, vocab] logits memory). Default = whole batch.
+    let micro: usize = flag(&args, "--micro", n_prompts * group.max(1));
+    let temp: f32 = flag(&args, "--temp", 1.0);
+    let beta: f32 = flag(&args, "--beta", 0.04);
+    let eps: f32 = flag(&args, "--eps", 0.2);
+    let lr: f32 = flag(&args, "--lr", 1e-6);
+    let clip: f32 = flag(&args, "--clip", 1.0);
+    let max_new: usize = flag(&args, "--max-tokens", 24);
+    let max_add: i64 = flag(&args, "--max-add", 20);
+    let max_mul: i64 = flag(&args, "--max-mul", 9);
+    let seed: u64 = flag(&args, "--seed", 20260630);
+    let log_every: usize = flag(&args, "--log-every", 20);
+    let init_ckpt = flag_value(&args, "--init-ckpt").expect("--init-ckpt <sft.ckpt> is required");
+    let out_ckpt = flag_value(&args, "--ckpt").expect("--ckpt <out> is required");
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let tok = Tokenizer::from_file(std::path::Path::new(tok_path.as_str()));
+    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
+    let policy = load_model(cfg, device, &init_ckpt, false); // flash keeps attn memory bounded
+    // Frozen πref for the KL leash — only resident when β>0 (a second 1B model is the
+    // memory long-pole; β=0 is pure PG and skips it, the gated degenerate).
+    let reference = if beta > 0.0 {
+        Some(load_model(cfg, device, &init_ckpt, false))
+    } else {
+        None
+    };
+
+    let gcfg = GenConfig {
+        max_add,
+        max_mul,
+        ops: vec![Op::Add, Op::Sub, Op::Mul],
+    };
+    let params = policy.params();
+    let mut opt = GpuAdamW::new(0.0);
+    let mut rng = seed.max(1);
+
+    let start = std::time::Instant::now();
+    let (mut win_reward, mut win_solved, mut win_n) = (0f32, 0usize, 0usize);
+    // Per-window phase timers (ms): rollout / capture / inner — to keep the step
+    // decomposition honest (M2d cut the training-side forwards 9×, so the question is
+    // what now dominates the step).
+    let (mut t_roll, mut t_cap, mut t_inner) = (0f32, 0f32, 0f32);
+    for step in 0..steps {
+        // ---- Rollout: B prompts × G completions, scored, group-advantage ----
+        // Collect ALL the step's framed samples first (input, target, adv), so the
+        // training-side forwards can be batched across the whole step (M2d) instead of
+        // run one ragged sequence at a time.
+        let t0 = std::time::Instant::now();
+        let mut raw: Vec<(Vec<i32>, Vec<i32>, f32)> = Vec::new();
+        for _ in 0..n_prompts {
+            let p = gen_problem(&mut rng, &gcfg);
+            let prompt_ids: Vec<i32> = tok
+                .encode(&format!("User: {}\nAssistant:", p.question()))
+                .into_iter()
+                .map(|t| t as i32)
+                .collect();
+            // M2b batched rollout: the G samples of this prompt decode in lockstep
+            // (one forward per step over the whole group → G× fewer kernel launches
+            // than G sequential single-seq rollouts; the M4 rollout long-pole fix).
+            let mut comps: Vec<(String, f32)> = Vec::with_capacity(group);
+            let outs = generate_cached_batch(&policy, device, &prompt_ids, group, max_new, temp, &mut rng);
+            for out in &outs {
+                let cont = tok.decode(&out[prompt_ids.len()..].iter().map(|&t| t as u32).collect::<Vec<_>>());
+                let seg = first_answer_segment(&cont).trim().to_string();
+                let r = if check_answer(&seg, p.answer()) { 1.0 } else { 0.0 };
+                comps.push((seg, r));
+            }
+            let mean = comps.iter().map(|c| c.1).sum::<f32>() / group as f32;
+            let var = comps.iter().map(|c| (c.1 - mean).powi(2)).sum::<f32>() / group as f32;
+            let std = var.sqrt();
+            win_reward += mean * group as f32;
+            win_solved += comps.iter().filter(|c| c.1 > 0.5).count();
+            win_n += group;
+            // A whole group with no reward variance gives zero advantage → skip
+            // (no learning signal, and avoids dividing by ~0).
+            if std < 1e-6 {
+                continue;
+            }
+            for (seg, r) in &comps {
+                let adv = (r - mean) / (std + 1e-4);
+                let (input, target) = frame(&tok, &p.question(), seg);
+                raw.push((input, target, adv));
+            }
+        }
+
+        t_roll += t0.elapsed().as_secs_f32() * 1e3;
+
+        // ---- Batched capture (M2d): logπ_old (policy) + logπ_ref (frozen) over ALL
+        // samples in forward_batched chunks, instead of one forward per sample. ----
+        if !raw.is_empty() {
+            let t1 = std::time::Instant::now();
+            let io: Vec<(Vec<i32>, Vec<i32>)> = raw.iter().map(|(i, t, _)| (i.clone(), t.clone())).collect();
+            let logp_old = per_token_logp_batched(&policy, device, &io, micro);
+            // β=0 ⇒ KL term drops ⇒ logp_ref unused; pass zeros (no reference model).
+            let logp_ref = match &reference {
+                Some(r) => per_token_logp_batched(r, device, &io, micro),
+                None => raw.iter().map(|(i, _, _)| vec![0.0; i.len()]).collect(),
+            };
+            let batch: Vec<PgSample> = raw
+                .iter()
+                .zip(logp_old)
+                .zip(logp_ref)
+                .map(|(((input, target, adv), lo), lr)| PgSample {
+                    input: input.clone(),
+                    target: target.clone(),
+                    adv: *adv,
+                    logp_old: lo,
+                    logp_ref: lr,
+                })
+                .collect();
+            t_cap += t1.elapsed().as_secs_f32() * 1e3;
+
+            // ---- K inner clipped-PG epochs, batched over the captured samples ----
+            let t2 = std::time::Instant::now();
+            for _ in 0..inner {
+                inner_pg_step_batched(&policy, device, &batch, eps, beta, micro);
+                let _ = xtrain_train::clip::clip_grad_norm_gpu(&params, clip, 1.0);
+                opt.step(lr, &params);
+                for p in &params {
+                    p.zero_grad();
+                }
+            }
+            t_inner += t2.elapsed().as_secs_f32() * 1e3;
+        }
+
+        if (step + 1) % log_every == 0 || step == steps - 1 {
+            let w = log_every.min(step + 1) as f32; // steps in this window
+            println!(
+                "step {:5}/{steps}: mean-reward {:.3} | solved {}/{} | {:.0}s | ms/step roll {:.0} cap {:.0} inner {:.0}",
+                step + 1,
+                win_reward / win_n.max(1) as f32,
+                win_solved,
+                win_n,
+                start.elapsed().as_secs_f32(),
+                t_roll / w,
+                t_cap / w,
+                t_inner / w,
+            );
+            win_reward = 0.0;
+            win_solved = 0;
+            win_n = 0;
+            t_roll = 0.0;
+            t_cap = 0.0;
+            t_inner = 0.0;
+            // Periodic save so a later OOM (naive rollout fragments the allocator —
+            // the long-pole the design doc flagged) still leaves an evaluatable ckpt.
+            xtrain_train::checkpoint::save(std::path::Path::new(&out_ckpt), &params).expect("save");
+        }
+    }
+
+    xtrain_train::checkpoint::save(std::path::Path::new(&out_ckpt), &params).expect("save ckpt");
+    println!("GRPO done: {steps} steps, G={group}, B={n_prompts}, beta {beta}, lr {lr:.1e} → {out_ckpt}");
+}

crates/xtrain-train/src/grpo_batch.rs

@@ -0,0 +1,162 @@
+//! Batched GRPO training-side forwards (post-training M2d). After M2b/M2c made the
+//! rollout cheap, the GRPO **step** is dominated by the per-sample full-sequence
+//! forwards: the `per_token_logp` captures (policy + reference) and the inner
+//! clipped-PG `forward`/`backward`s — each a single-sequence `forward` over a short
+//! ragged completion. This module packs the `N = B·G` ragged samples of a step into
+//! ONE `forward_batched`, amortising the per-launch overhead across N (the same win
+//! M2b gave the rollout).
+//!
+//! The enabling property: **right-padding is free under causal attention.** Pad each
+//! ragged completion on the RIGHT to the batch's `Lmax`; a real completion row is at
+//! an earlier position than the trailing pad, and causal masking forbids attending
+//! forward, so its logits are bit-identical to the unpadded single-sequence forward.
+//! The pad rows' own outputs are garbage but are masked out (`target = -100`).
+//!
+//! Both the looped (baseline) and batched paths live here so they share one source of
+//! truth — `bin/bench_grpo_batch` A/Bs them (timing + a closeness gate), and the
+//! per-row equivalence of the loss op is pinned by `clipped_pg_loss_batched_matches_looped`
+//! in `xtrain-autodiff/tests/autograd.rs`.
+
+#![cfg(not(no_cuda))]
+
+use xtrain_autodiff::ops;
+use xtrain_model::{TinyTransformer, ids_tensor};
+use xtrain_tensor::{Device, Tensor};
+
+/// One framed completion of a GRPO step: the next-token `(input, target)` pair
+/// (prompt positions masked to `-100` in `target`), its group-relative `adv`, and the
+/// per-position rollout-time / reference logprobs the clipped-PG loss needs.
+pub struct PgSample {
+    pub input: Vec<i32>,
+    pub target: Vec<i32>,
+    pub adv: f32,
+    pub logp_old: Vec<f32>,
+    pub logp_ref: Vec<f32>,
+}
+
+// ------------------------------- looped (baseline) -------------------------------
+
+/// Per-position `logπ(target_t)` of one framed `(input, target)` pair (= `−per_row`
+/// of cross_entropy; masked positions are 0). One single-sequence forward, no grad.
+pub fn per_token_logp(model: &TinyTransformer, device: Device, input: &[i32], target: &[i32]) -> Vec<f32> {
+    let logits = model.forward(&ids_tensor(input, device)).value();
+    let (_, per_row) = logits.cross_entropy(&ids_tensor(target, device));
+    per_row
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .iter()
+        .map(|p| -p)
+        .collect()
+}
+
+/// One inner clipped-PG epoch the looped way: per sample, a single-sequence forward +
+/// [`ops::clipped_pg_loss`] scaled by `1/N` + backward (grads accumulate on `model`'s
+/// params). Returns the summed scaled loss. Caller does clip + opt.step + zero_grad.
+pub fn inner_pg_step_looped(
+    model: &TinyTransformer,
+    device: Device,
+    batch: &[PgSample],
+    eps: f32,
+    beta: f32,
+) -> f32 {
+    let scale = 1.0 / batch.len() as f32;
+    let mut total = 0f32;
+    for s in batch {
+        let logits = model.forward(&ids_tensor(&s.input, device));
+        let loss = ops::clipped_pg_loss(&logits, &ids_tensor(&s.target, device), &s.logp_old, &s.logp_ref, s.adv, eps, beta);
+        let scaled = ops::scale(&loss, scale);
+        total += scaled.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+        scaled.backward();
+    }
+    total
+}
+
+// ------------------------------- batched (M2d) -----------------------------------
+
+/// Right-pad `m` ragged `i32` rows (each `< lmax` long) to `[m*lmax]` sequence-major,
+/// filling with `pad`. Used for both the id stream (pad = 0, arbitrary) and the target
+/// stream (pad = −100, ignored by cross_entropy).
+fn pack_i32(rows: &[&[i32]], lmax: usize, pad: i32) -> Vec<i32> {
+    let mut flat = vec![pad; rows.len() * lmax];
+    for (i, r) in rows.iter().enumerate() {
+        flat[i * lmax..i * lmax + r.len()].copy_from_slice(r);
+    }
+    flat
+}
+
+/// Batched [`per_token_logp`]: pack `samples` (each `(input, target)`) right-padded to
+/// `Lmax`, run ONE `forward_batched(batch = N)`, and slice each sample's `logπ` back to
+/// its real length. Equal to looping [`per_token_logp`] (right-pad is free under causal
+/// attention), to bf16 batch-reduction tolerance. `samples` are processed in chunks of
+/// `micro` (≥1) to bound the `[chunk*Lmax, vocab]` logits memory.
+pub fn per_token_logp_batched(
+    model: &TinyTransformer,
+    device: Device,
+    samples: &[(Vec<i32>, Vec<i32>)],
+    micro: usize,
+) -> Vec<Vec<f32>> {
+    let mut out = Vec::with_capacity(samples.len());
+    for chunk in samples.chunks(micro.max(1)) {
+        let m = chunk.len();
+        let lmax = chunk.iter().map(|(i, _)| i.len()).max().unwrap();
+        let ins: Vec<&[i32]> = chunk.iter().map(|(i, _)| i.as_slice()).collect();
+        let tgs: Vec<&[i32]> = chunk.iter().map(|(_, t)| t.as_slice()).collect();
+        let ids = Tensor::from_slice(&pack_i32(&ins, lmax, 0), &[m * lmax]).to_device(device);
+        let tgt = Tensor::from_slice(&pack_i32(&tgs, lmax, -100), &[m * lmax]).to_device(device);
+        let logits = model.forward_batched(&ids, m).value();
+        let (_, per_row) = logits.cross_entropy(&tgt);
+        let pr = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
+        for (i, (inp, _)) in chunk.iter().enumerate() {
+            let b = i * lmax;
+            out.push((0..inp.len()).map(|r| -pr[b + r]).collect());
+        }
+    }
+    out
+}
+
+/// One inner clipped-PG epoch, batched: pack the batch (in `micro`-sized chunks) and run
+/// ONE `forward_batched` + [`ops::clipped_pg_loss_batched`] + backward per chunk. The
+/// per-row `weight = 1/(N·n_s)` uses the GLOBAL `N = batch.len()` (not the chunk size),
+/// so chunked grad-accumulation reproduces the looped `Σ_s (1/N)(1/n_s)…` exactly.
+/// Returns the summed loss. Caller does clip + opt.step + zero_grad.
+pub fn inner_pg_step_batched(
+    model: &TinyTransformer,
+    device: Device,
+    batch: &[PgSample],
+    eps: f32,
+    beta: f32,
+    micro: usize,
+) -> f32 {
+    let inv_n = 1.0 / batch.len() as f32;
+    let mut total = 0f32;
+    for chunk in batch.chunks(micro.max(1)) {
+        let m = chunk.len();
+        let lmax = chunk.iter().map(|s| s.input.len()).max().unwrap();
+        let ins: Vec<&[i32]> = chunk.iter().map(|s| s.input.as_slice()).collect();
+        let tgs: Vec<&[i32]> = chunk.iter().map(|s| s.target.as_slice()).collect();
+        let ids = Tensor::from_slice(&pack_i32(&ins, lmax, 0), &[m * lmax]).to_device(device);
+        let tgt = Tensor::from_slice(&pack_i32(&tgs, lmax, -100), &[m * lmax]).to_device(device);
+
+        let mut logp_old = vec![0f32; m * lmax];
+        let mut logp_ref = vec![0f32; m * lmax];
+        let mut advantage = vec![0f32; m * lmax];
+        let mut weight = vec![0f32; m * lmax];
+        for (i, s) in chunk.iter().enumerate() {
+            let b = i * lmax;
+            let li = s.input.len();
+            logp_old[b..b + li].copy_from_slice(&s.logp_old);
+            logp_ref[b..b + li].copy_from_slice(&s.logp_ref);
+            let n_s = s.target.iter().filter(|&&t| t >= 0).count().max(1) as f32;
+            let w = inv_n / n_s; // = 1/(N · n_s)
+            for r in 0..lmax {
+                advantage[b + r] = s.adv;
+                weight[b + r] = w;
+            }
+        }
+        let logits = model.forward_batched(&ids, m);
+        let loss = ops::clipped_pg_loss_batched(&logits, &tgt, &logp_old, &logp_ref, &advantage, &weight, eps, beta);
+        total += loss.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+        loss.backward();
+    }
+    total
+}

crates/xtrain-train/src/lib.rs

@@ -15,6 +15,8 @@ pub mod task;
 #[cfg(not(no_cuda))]
 pub mod checkpoint;
 #[cfg(not(no_cuda))]
+pub mod grpo_batch;
+#[cfg(not(no_cuda))]
 pub mod sample;
 #[cfg(not(no_cuda))]
 mod train_loop;

crates/xtrain-train/tests/decode_batch.rs

@@ -0,0 +1,83 @@
+// M2b batched KV-cache decode — the token-identical gate.
+//
+// Batched decode rolls out G samples of one prompt in lockstep (one common decode
+// position each step, uniform RoPE via rope_pos, KV cache carrying a G dimension).
+// Under GREEDY decoding all G rows are deterministic and must each equal the
+// single-sequence greedy decode (generate_greedy_cached, itself gated token-
+// identical to the naive sampler). This pins that the G-way batching indexes each
+// sequence's K/V correctly (no cross-row contamination) and reproduces M2a exactly.
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{generate_cached_batch, generate_greedy_cached, Config, TinyTransformer};
+use xtrain_tensor::{DType, Device};
+
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+fn build(cfg: Config, device: Device) -> TinyTransformer {
+    let mut seed = 1u64;
+    TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.08)
+        }
+    })
+    .with_compute_dtype(DType::F32)
+}
+
+#[test]
+fn batched_greedy_decode_matches_single_seq() {
+    assert!(
+        device::device_count().expect("device count") > 0,
+        "no CUDA device"
+    );
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    // Real GQA (8 query / 2 kv heads → group 4) so repeat_kv(nh, batch=G) is exercised.
+    let cfg = Config::from_arch(48, 8, 16, 4, 256).with_kv_heads(2);
+    let model = build(cfg, device);
+    let prompt: Vec<i32> = vec![3, 9, 1, 14, 5];
+    let max_new = 24usize;
+    let g = 5usize;
+
+    let single = generate_greedy_cached(&model, device, &prompt, max_new);
+    let mut rng = 0u64;
+    let batched = generate_cached_batch(&model, device, &prompt, g, max_new, 0.0, &mut rng);
+
+    assert_eq!(batched.len(), g, "expected {g} sample rows");
+    for (row, seq) in batched.iter().enumerate() {
+        assert_eq!(
+            seq.len(),
+            single.len(),
+            "row {row} length {} vs single {}",
+            seq.len(),
+            single.len()
+        );
+        if seq != &single {
+            let first = seq.iter().zip(&single).position(|(a, b)| a != b).unwrap();
+            panic!(
+                "batched row {row} diverges from single-seq at index {first}: {:?} vs {:?}",
+                seq[first], single[first]
+            );
+        }
+    }
+    println!(
+        "batched decode OK: all {g} greedy rows token-identical to single-seq over {max_new} tokens"
+    );
+}

crates/xtrain-train/tests/decode_kv.rs

@@ -0,0 +1,94 @@
+// M2a KV-cache decode engine — the token-identical correctness gate.
+//
+// The centerpiece M2 invariant: greedy decode through the KV-cache incremental
+// engine (`xtrain_model::generate_greedy_cached`) must be TOKEN-IDENTICAL to the
+// naive full-recompute greedy (`xtrain_train::sample::generate` at temperature 0),
+// which re-runs the whole forward over the growing prefix each step. Same tokens ⇒
+// the cache + decode-time attention + RoPE-at-position reproduce the full forward.
+//
+// Numerics note: a randomly-initialised model has near-uniform logits, so argmax
+// can be fragile to ~1e-6 differences. This unit gate therefore runs in F32 (the
+// tightest path, and the dtype the eval harness actually uses) on a small model.
+// The headline gate on the trained v12 checkpoint (peaked logits → robust argmax)
+// is run on the GPU box and recorded in docs/18.
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, generate_greedy_cached};
+use xtrain_tensor::{DType, Device};
+
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+fn build(cfg: Config, device: Device, dtype: DType) -> TinyTransformer {
+    let mut seed = 1u64;
+    let m = TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.08)
+        }
+    });
+    m.with_compute_dtype(dtype)
+}
+
+// A real GQA config (8 query / 2 kv heads → group 4) to exercise repeat_kv in the
+// decode path; head_dim 16, dim 128, 4 layers.
+fn gqa_cfg() -> Config {
+    Config::from_arch(48, 8, 16, 4, 256).with_kv_heads(2)
+}
+
+#[test]
+fn kv_cache_decode_is_token_identical_to_naive_f32() {
+    assert!(
+        device::device_count().expect("device count") > 0,
+        "no CUDA device"
+    );
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let model = build(gqa_cfg(), device, DType::F32);
+    let prompt: Vec<i32> = vec![1, 5, 9, 13, 2, 7];
+    let max_new = 24usize;
+
+    let mut rng = 7u64;
+    let naive = xtrain_train::sample::generate(&model, device, &prompt, max_new, 0.0, &mut rng);
+    let cached = generate_greedy_cached(&model, device, &prompt, max_new);
+
+    assert_eq!(
+        naive.len(),
+        cached.len(),
+        "length mismatch: naive {} vs cached {}",
+        naive.len(),
+        cached.len()
+    );
+    if naive != cached {
+        // Report the first divergence for debugging.
+        let first = naive
+            .iter()
+            .zip(&cached)
+            .position(|(a, b)| a != b)
+            .unwrap();
+        panic!(
+            "token divergence at index {first}: naive={:?} cached={:?}\nnaive  ={naive:?}\ncached ={cached:?}",
+            naive[first], cached[first]
+        );
+    }
+    println!(
+        "KV-cache decode token-identical to naive over {} generated tokens (F32, GQA 8/2)",
+        max_new
+    );
+}

csrc/ops/nn.cu

@@ -242,6 +242,96 @@ void launch_rope_f32(const float* x, float* y, int tokens, int heads,
     rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, period);
 }
 
+// RoPE at an absolute position offset (KV-cache decode-time, forward only). Same
+// rotate_half as rope_k, but row `tok`'s position is `pos0 + tok` (no modulo) —
+// a single new decode token sits at absolute position pos0. The training rope_k
+// (position = tok % period) is left untouched, so this adds no training-path risk.
+__global__ void rope_at_k(const float* x, float* y, int heads, int head_dim,
+                          float theta, int pos0) {
+    int tok = blockIdx.x;
+    int head = blockIdx.y;
+    int half = head_dim / 2;
+    int i = threadIdx.x;
+    if (i >= half) return;
+    int pos = pos0 + tok;
+    float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
+    float angle = (float)pos * freq;
+    float c = cosf(angle), sn = sinf(angle);
+    int base = (tok * heads + head) * head_dim;
+    float x0 = x[base + i], x1 = x[base + i + half];
+    y[base + i]        = x0 * c - x1 * sn;
+    y[base + i + half] = x1 * c + x0 * sn;
+}
+void launch_rope_at_f32(const float* x, float* y, int tokens, int heads,
+                        int head_dim, float theta, int pos0, void* s) {
+    dim3 grid(tokens, heads);
+    int blk = head_dim / 2;
+    rope_at_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, pos0);
+}
+
+// RoPE with a PER-ROW absolute position (batched KV-cache decode, M2b): row `tok`'s
+// position is `positions[tok]` (an i32 per token). For G-way batched decode all G
+// rows share one decode position; for ragged batches each row carries its own.
+// Forward only; the training rope_k is untouched.
+__global__ void rope_pos_k(const float* x, const int* positions, float* y,
+                           int heads, int head_dim, float theta) {
+    int tok = blockIdx.x;
+    int head = blockIdx.y;
+    int half = head_dim / 2;
+    int i = threadIdx.x;
+    if (i >= half) return;
+    int pos = positions[tok];
+    float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
+    float angle = (float)pos * freq;
+    float c = cosf(angle), sn = sinf(angle);
+    int base = (tok * heads + head) * head_dim;
+    float x0 = x[base + i], x1 = x[base + i + half];
+    y[base + i]        = x0 * c - x1 * sn;
+    y[base + i + half] = x1 * c + x0 * sn;
+}
+void launch_rope_pos_f32(const float* x, const int* positions, float* y,
+                         int tokens, int heads, int head_dim, float theta, void* s) {
+    dim3 grid(tokens, heads);
+    int blk = head_dim / 2;
+    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.
+__global__ void scale_rows_k(const float* x, const float* s, float* y,
+                             int rows, int cols) {
+    int r = blockIdx.x;
+    float sr = s[r];
+    for (int c = threadIdx.x; c < cols; c += blockDim.x)
+        y[r * cols + c] = x[r * cols + c] * sr;
+}
+void launch_scale_rows_f32(const float* x, const float* s, float* y,
+                           int rows, int cols, void* st) {
+    int blk = cols < 1024 ? cols : 1024;
+    if (blk < 32) blk = 32;
+    scale_rows_k<<<rows, blk, 0, (cudaStream_t)st>>>(x, s, y, rows, cols);
+}
+
 __global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim,
                          float theta, int period) {
     int tok = blockIdx.x;

docs/18-post-training-rl-sft.md

@@ -360,3 +360,270 @@ gap is exactly what the verifiable reward in M3 (DPO) / M4 (GRPO) is built to cl
 held-out correct > 0 confirms the checker + eval harness score real matches (not just format).
 M1 delivers the format floor + the reusable task spec / checker / eval harness — not arithmetic
 skill, which is downstream by design.
+
+### M2a — KV-cache incremental-decode engine (single sequence, landed)
+
+The decode engine (D3, built up front) that replaces the naive sampler — which re-runs the
+full forward over the growing prefix every step (O(t²), a fresh autograd graph per token). Two
+forward-only primitives + a raw-Tensor per-token block forward, each gated in isolation.
+
+**Primitives (`xtrain-tensor`, both forward-only):**
+- `Tensor::rope_at(theta, pos0)` — RoPE at a token's *absolute* position (`pos = pos0 + row`,
+  no modulo), vs the training `rope` (`pos = row % period`) which is left untouched (new CUDA
+  kernel `rope_at_k` → no training-path risk). Cached K is stored post-RoPE, so it must match
+  what the full forward produced at that position. **Gate:** bit-identical to the full-sequence
+  rope's row `t` (`integration::rope_at_matches_full_rope_row`).
+- `Tensor::decode_attention(k, v, scale)` — single-query × cached-K/V SDPA (`[bh,1,hd]` vs
+  `[bh,t,hd]`, no causal mask: the one query sees all cached keys). Composed from the existing
+  strided batched GEMM + plain softmax — **no new kernel**. **Gate:** equals the full causal
+  attention's last query row, max |Δ| 6e-8 (`integration::decode_attention_matches_…`).
+
+**Engine (`xtrain-model/src/decode.rs`, `generate_greedy_cached`):** per-layer K/V cache +
+single-token incremental forward. Prefill = the first `prompt.len()` decode steps (one code
+path). Mirrors `model::block_forward` at the raw-Tensor level (no autograd tape — inference
+needs no grads), pulling weights via the public `params()` stable order (no model-internal
+visibility changes). The cache is host-accumulated token-major f32, rebuilt per step — the
+honest M2a baseline; M2b moves it device-side + adds batched ragged decode.
+
+**Gate (the M2 centerpiece — token-identical):** KV-cache greedy decode is byte-for-byte the
+same token sequence as the naive full-recompute greedy. Verified two ways:
+- `xtrain-train/tests/decode_kv.rs` — small GQA model (8 query / 2 kv heads), F32, 24 generated
+  tokens, exact token-equality. (Unit gate runs F32: a random model's near-uniform logits make
+  argmax fragile to ~1e-6, so the tightest path is used; the trained model below has peaked
+  logits → robust.)
+- v12 1.05B SFT checkpoint: `eval_arith --cached` produces the **identical** eval outcome to the
+  naive run (format 100/100, correct 8/100) and byte-identical completions.
+
+**Throughput baseline (v12 1.05B, batch 1, F32, profile-first — measured, not assumed):** the
+cache win is **sequence-length-dependent**, which is the honest systems finding here:
+
+| max_new | naive | kv-cache | note |
+|---------|-------|----------|------|
+| 32      | 108 tok/s | 111 tok/s | ~1.0× — both **launch/overhead-bound** at short seq |
+| 128     | 69 tok/s  | **133 tok/s** | **~1.9×** — naive's O(t²) recompute starts to bite |
+| 256     | **OOM**   | 129 tok/s | naive rebuilds the O(seq²) graph every step → OOM |
+
+Cached throughput stays ~constant (O(1)/token compute + constant memory); naive **decays**
+(108→69 tok/s, O(t)/token) and eventually **OOMs** (the full autograd graph per step). So at the
+short arithmetic-eval lengths the cache is overhead-bound and gives ~nothing — it matters for
+**long rollouts** (DPO pair-generation, GRPO completions), exactly where M3/M4 use it. (M2a's
+per-layer host round-trip is part of why short-seq is overhead-bound; M2b's device-side cache
+targets it.) This is the same measure-first lesson as T17 (process-per-GPU throughput-neutral):
+the win is real but only in the regime that actually stresses the bottleneck.
+
+### M3 — DPO (offline preference optimization, landed; honest negative result)
+
+The first real alignment method. Infra landed and gated; the empirical finding is that DPO
+**does not improve held-out arithmetic correctness on this task** — a genuine, on-theme negative
+result (the design doc's "RL is finicky" risk, made concrete).
+
+**Two new autograd ops (`xtrain-autodiff`, both reuse the CE kernel — no new CUDA):**
+- `seq_logprob(logits, target)` = `Σ log πθ(target)` over non-ignored positions (the per-
+  sequence logprob DPO compares). `= −Σ per_row` of cross_entropy (ignored rows already 0, like
+  SFT masking); backward = `cross_entropy_backward(probs, target, −upstream)` (SUM, no mean).
+  **Gate:** finite-diff grad-check with a `-100` completion mask.
+- `dpo_loss(lpθ_chosen, lpθ_rejected, lpref_chosen, lpref_rejected, β)` = `−log σ(Δ)` with the
+  two policy logprobs as parents (ref logprobs constant). **Gate:** grad-check both parents +
+  degenerate points (policy==ref ⇒ Δ=0, L=log2, grads ∓β/2; β=0 ⇒ grads 0).
+
+**Pair construction (`gen_dpo_pairs`, aligned decision):** chosen = gold answer; rejected = the
+SFT model's own **greedy** (KV-cache engine, M2a) completion when it's a format-valid WRONG
+boxed answer — a hard negative in the model's distribution. Since SFT is ~8% correct (M1),
+greedy is wrong ~92% of the time, so this is fast and deterministic; ~8% of prompts are skipped
+(greedy correct). 1500 pairs generated (158 skipped) in ~8 min.
+
+**Training (`train_dpo`):** loads the SFT ckpt as policy AND frozen reference; **precomputes the
+reference logprobs once** (while policy == reference) and caches them — one resident model. Each
+step forwards the policy on chosen + rejected, `seq_logprob` each, minimises `dpo_loss`; the two
+forwards share params so backward accumulates both branches. Loss **starts at exactly log2**
+(Δ=0 at init) — a built-in correctness check that fired correctly. Tracks reward margin +
+preference accuracy.
+
+**Result (v12 1.05B, 1500 pairs, β=0.1; 100 held-out prompts, vs the SFT baseline format
+100/100, correct 8/100):**
+
+| run                       | reward margin | pref-acc | format | correct |
+|---------------------------|---------------|----------|--------|---------|
+| SFT (baseline)            | —             | —        | 100/100 | 8/100  |
+| DPO lr 5e-7 × 300         | +0.78         | ~82%     | 100/100 | 7/100  |
+| DPO lr 5e-7 × 800         | +1.25         | ~82%     | 100/100 | 5/100  |
+| DPO lr 1e-6 × 2000        | **+34.2**     | ~76%     | **0/100** | 0/100 |
+
+The reward margin and preference accuracy rise cleanly (the loss IS being optimized — the infra
+is correct), but the implicit reward **does not transfer to held-out correctness**: it stays
+~5–8% (all within the ~2.7% std-error of 100 prompts — statistically flat), and pushing harder
+**over-optimizes to collapse** (margin +34 = huge KL from the reference → the model emits
+garbage, `46 * 80 = CRAFTIE SERIES SERIES…`, format 0%).
+
+**The lesson (why):** chosen and rejected differ only in the final number tokens, so DPO raises
+`log p(correct) − log p(wrong)` for the *specific* training pairs — it **reweights the existing
+distribution, it does not install the capability**. The base model has no arithmetic algorithm,
+so preferring correct-vs-wrong final answers on seen pairs cannot generalize to unseen problems;
+and the only way to drive the margin far is to globally distort the distribution → incoherence.
+**DPO works when the chosen is already plausible under the policy; it cannot manufacture
+knowledge the model lacks.** This is the precise motivation for **M4 GRPO**: optimize the *actual
+verifiable reward* online (sample → check → reinforce what is genuinely correct), rather than a
+fixed-pair proxy — though GRPO faces the same 8%-correct sparsity, so whether it moves the metric
+is M4's open question. Gate met for M3 = the infra is correct (op grad-checks, log2-at-init,
+margin/acc rise); the correctness flatness is the reported finding, not a bug.
+
+### M4 — GRPO (online RL, critic-free, landed; infra + two honest systems walls)
+
+The centerpiece: generation INSIDE the training loop. Infra built and gated; the run surfaces
+two concrete systems findings (the memory long-pole + the rollout long-pole, both flagged in the
+design doc's Risks) and the same capability wall as M3.
+
+**Task made learnable first (per the aligned decision "easier task → then M4"):** the v12 SFT
+model scores ~8% on the hard task *and* on easy problems — it learned format, not arithmetic. So
+the easy task (operands ≤20, ops `+ − ×`) was re-SFT'd from the v12 base → **held-out 18.7%**
+(100% format), a baseline with reward variance for GRPO. Note: even easy arithmetic plateaus at
+~19% held-out (250 vs 600 SFT steps identical) — a 1B web-text model does not generalize the
+add/sub algorithm from ~550 examples; it memorizes train (982 total problems, 550 seen).
+
+**New op (`xtrain-autodiff`, reuses the CE kernel + one new primitive):**
+- `clipped_pg_loss(logits, target, logp_old, logp_ref, A, ε, β)` — per completion token
+  `ρ_t = exp(logπθ_t − logp_old_t)`, `L = −mean min(ρA, clip(ρ,1±ε)A) + β·mean KL` (k3), masked
+  to completion tokens. Backward reuses `(probs − onehot)` + `scale_rows` (a new ~5-line per-row
+  scale kernel — the per-token coefficient varies, which CE-backward's single scalar can't
+  express). **Gate:** grad-check the active PG path + the A=0 (KL-only) path; degenerate value
+  checks ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL.
+
+**Loop (`train_grpo`):** per step — sample B prompts, roll out G completions each, score (reward
+0/1), group-relative advantage `A=(r−mean)/(std+ε)` (no critic; all-correct/all-wrong groups
+skipped — zero advantage), capture `logπθ_old`/`logπref` per token, K inner clipped-PG epochs.
+Rollout uses the M2 KV-cache engine with **temperature sampling** (added in M4): single-row
+`[1,vocab]` logits per step vs the naive sampler's `[seq,vocab]`.
+
+**Systems wall #1 — memory (the design doc's "two/three resident models"):** KL-leash GRPO needs
+policy + frozen reference, two 1.05B fp32-master models + AdamW m/v ≈ 21 GB fixed + training
+activations → unreliably OOMs on a 32 GB 5090 (fragmentation tips it over). To get a completing
+run, `β=0` (pure PG) drops the reference model (−4.2 GB). So the *principled* KL-leash version is
+memory-bound at this model size on this hardware — a real, reported constraint, not a bug.
+
+**Systems wall #2 — rollout (the design doc's "rollout is the long pole"):** the naive sampler's
+growing `[seq,vocab]` allocations fragment the caching allocator over a long rollout → OOM. The
+cached temperature rollout (single-row logits) is lighter; but single-sequence cached decode is
+slow (the M2a host-round-trip), so rollout still dominates wall-clock (~16 s/step at G=6·B=6).
+Batched ragged decode (M2b) is the real fix and is deferred to where it is load-bearing.
+
+**Result (easy task, β=0, G=6·B=6, 40 steps, lr 5e-7; 150 held-out, vs SFT 28/150 = 18.7%):**
+mean rollout reward fluctuates ~0.58–0.81 (noisy, inflated by train-set overlap in the sampled
+problems); **format stays 100/100** (no collapse even without the KL leash, at this gentle lr);
+**held-out 30/150 = 20.0%** — `+1.3 pp`, within the ~3% std-error of 150 prompts, i.e.
+**statistically flat**, the same wall as M3 DPO.
+
+**The consistent M3+M4 lesson:** on a task where the base model lacks the underlying capability,
+**neither offline preference optimization (DPO) nor online RL (GRPO) moves held-out correctness**
+— each optimizes its objective (margin / reward) on the *training distribution* it can reach
+(here inflated by memorization), but cannot install a *generalizable* algorithm the model never
+had. RL reinforces what the model already does; it does not teach arithmetic. Gate met for M4 =
+the infra is correct (PG/KL grad-checks + degenerate checks, the loop runs, reward signal + KL
+leash wired, format held); the held-out flatness + the two memory/throughput walls are the
+reported findings. The honest end-state of the post-training arc: **a complete, correctness-gated
+SFT → KV-cache → DPO → GRPO stack** — the infrastructure learned in full, with measured, honest
+limits on what alignment can do for a capability the base model lacks.
+
+### M2b — batched KV-cache decode (landed; completes the M2 engine, fixes the rollout long-pole)
+
+Built after M4 (where the rollout long-pole bit hardest): decode the **G samples of one prompt in
+lockstep** — one forward per step over the whole group → G× fewer kernel launches, the deferred
+fix from M2a.
+
+**One new primitive:** `rope_pos(x, positions[])` — RoPE with a *per-row* absolute position (new
+forward-only kernel), since the G batched rows share one decode position (M2a's `rope_at` does
+`pos0 + row`, wrong for a batch at a single position). **Gate:** bit-identical to the full rope
+for positions `[0..n]`, and to `rope_at(P)` per row for a uniform `P`.
+
+**Engine (`generate_cached_batch`):** `BatchKVCache` carries a G dimension (`[T, G·num_kv, hd]`
+host-accumulated → `[G·num_kv, T, hd]`); the batched `decode_step` threads G through embed /
+projections / QK-norm / `rope_pos` / cache. Two M2a pieces drop in unchanged: `decode_attention`
+is already batch-agnostic (`bh = G·nh`), and `repeat_kv(nh, batch=G)` broadcasts per group. No
+finished-mask (all G generate `max_new`; the caller cuts at EOS) and no ragged-length prompts yet
+— both perf-only follow-ups.
+
+**Gate (token-identical):** all G **greedy** rows are byte-identical to the single-sequence decode
+(`tests/decode_batch.rs`, 8 query / 2 kv heads → exercises the `repeat_kv` batching) — pins that
+G-way batching indexes each sequence's K/V with no cross-row contamination.
+
+**Throughput (v12 1.05B, G=6·B=6, easy task, rollout wired into `train_grpo`):** ~8.5 s/step vs
+~14–16 s/step for the single-seq cached rollout — **~1.7×**, rollout-inclusive. Short of the full
+G× because (a) the per-token-logp forwards + the PG update also cost, and (b) the M2a per-layer
+**host round-trip** is still there (now G× the data in one transfer, not removed). The full
+device-side cache (no host round-trip) is the remaining decode-engine optimization. Batching also
+**stabilises memory**: one batched forward per step vs G separate allocations that fragmented the
+caching allocator (the M4 OOM). So M2b closes the decode-engine milestone (M2a single-seq + M2b
+batched) and turns the rollout long-pole from "OOM/unbounded" into a bounded ~1.7× win — measured,
+with the device-cache as the named next lever.
+
+### M2c — device-side KV cache (landed; the bottleneck moved, a profile-first finding)
+
+The named M2b follow-up: keep K/V on the GPU (`[bh,T,hd]`, an `Option<Tensor>` per layer) and
+grow it by one token per step via a new `cat_seq` kernel (concat along the seq dim) — removing 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 (cleaner than the host `Vec` + rebuild).
+
+**Gates hold:** `cat_seq == host concat`; `decode_kv` single-seq + `decode_batch` G-way both still
+**token-identical**; GQA training path unaffected.
+
+**The finding (why this is a measure-first lesson, not a speedup story):** 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** (2 forwards/sample) and the
+**PG-update** forwards+backwards (`model.forward`, full-sequence, per sample) now dominate. So the
+long pole **shifted** from the rollout to the training-side forwards (cf. T11/T17/M2a: profile
+before optimizing — the bottleneck you fixed is not the one that remains). The device cache is
+still a real, correctness-gated improvement (cleaner code, less PCIe, ~10% decode); the honest
+headline is that the *next* decode lever is **ragged batched prefill of the per-sample forwards**,
+not the cache. The M2 decode engine is now M2a (single-seq) + M2b (batched) + M2c (device cache),
+all token-identical-gated; the post-training stack remains complete with its bottleneck mapped.
+
+### M2d — batch the GRPO training-side forwards (landed; the lever M2c named, + a decomposition correction)
+
+M2c named the next lever: **ragged batched prefill of the per-sample training-side forwards**. Those
+forwards are the two phases that, per step, run one single-sequence `forward` per sample: the
+`per_token_logp` **captures** (logπ_old policy + logπ_ref reference) and the inner **clipped-PG**
+forward/backwards. M2d packs all `N = B·G` ragged samples of a step into ONE `forward_batched`.
+
+**The enabling property — right-padding is free under causal attention.** Pad each ragged completion
+on the RIGHT to the batch's `Lmax`. A real completion row sits at an earlier position than the
+trailing pad, and causal masking forbids attending forward, so its logits are **bit-identical** to
+the unpadded single-sequence forward; the pad rows are garbage but masked out (`target = -100`). This
+is exactly why training engines pad-and-mask rather than run ragged. Two new pieces:
+- `per_token_logp_batched` (`crates/xtrain-train/src/grpo_batch.rs`): right-pad → one
+  `forward_batched(batch = N)` → slice each sample's logπ back to its real length.
+- `ops::clipped_pg_loss_batched` (`crates/xtrain-autodiff/src/ops.rs`): like the per-sample
+  `clipped_pg_loss`, but takes **per-row** `advantage[t]` (the owning sample's `A`) and **per-row**
+  `weight[t]` (the full normaliser; the caller passes `1/(N·n_s)`). It does NOT compute its own
+  `1/n_tokens`, so folding `weight = 1/(N·n_s)` reproduces the looped `Σ_s (1/N)(1/n_s)…`
+  **bit-for-bit** (the per-row CE backward is row-local). A `--micro` knob packs in chunks to bound
+  the `[chunk·Lmax, vocab]` logits memory; the weight uses the GLOBAL `N`, so chunked
+  grad-accumulation is exact. Both `train_grpo` and the bench call these shared helpers.
+
+**Correctness gates (exact, not bf16-noisy):**
+- `xtrain-model::forward_batched_ragged_matches_looped` — forward_batched on right-padded ragged
+  sequences == per-sequence single-seq forward on the real rows, **max|Δlogit| = 3.7e-7 (fp32) and
+  0.0 (bf16)**, both composed + flash. Pins "right-pad is free".
+- `xtrain-autodiff::clipped_pg_loss_batched_matches_looped` — batched op == looped
+  `Σ_s (1/N)·clipped_pg_loss_s`, **loss Δ=1.5e-8, grad max|Δ|=7.5e-9 (f32)**.
+Composed, these prove the batched GRPO step == the looped step. End-to-end: a short SFT (v12 base,
+150 steps, arith) → `train_grpo` 12 steps runs clean — **no OOM** (1B master + AdamW + batched
+activations fit with `micro=16`), mean-reward rises, the batched inner executes.
+
+**Throughput (bench `bin/bench_grpo_batch`, v12 1.05B, N=48 ragged, micro=16, β=0, weight-independent):**
+
+| phase (per step)        | looped (single-seq) | batched (M2d) | speedup |
+|-------------------------|---------------------|---------------|---------|
+| capture `per_token_logp`| 622 ms              | 71 ms         | 8.7×    |
+| inner clipped-PG fwd+bwd| 1907 ms             | 208 ms        | 9.2×    |
+| **training forwards**   | **2526 ms**         | **280 ms**    | **9.0×**|
+
+**The decomposition correction (the honest finding).** M2c claimed "the per-sample training forwards
+now dominate the step." The clean per-component bench falsifies the strong form: the training
+forwards were **~2.5 s of the ~8.5 s step (~30%)** — substantial and worth the 9× win, but the
+**rollout (`generate_cached_batch`, ~6 s) was always the larger share.** After M2d cuts the training
+forwards to ~0.28 s, the step is **~95% rollout** — the long pole has swung back to the rollout. So
+M2d removes the training-forward overhang (a real, exactly-gated 9× on its component), and re-confirms
+the same measure-first lesson one more time: the next **step-level** lever is **full B×G rollout
+batching** — today only the `G` samples of each prompt decode in lockstep (M2b); the `B` prompts are
+still sequential. M2d closes the "ragged batched per-sample forwards" lever M2c named; the post-
+training stack stays complete, now with the step decomposition measured, not asserted.

docs/evolution.md

@@ -97,6 +97,18 @@ Phase 1/2 把**预训练全栈**学完后，Phase 3 转向**后训练 infra**（
 
 **M1（SFT task baseline，已落地）**：可验证算术任务 + 数据生成器 + 评分器一套，host-side 9/9 单测过（masking、SFT-target 自洽 2000 样、parser 边界、种子确定性）。dash5 单卡从 v12 基座 SFT（loss 4.68→~0.34，best val 0.386）。**100 留出题 eval：格式 `\boxed{}` 习得率 base 0% → SFT 100%；算术正确率 8%。**——SFT 只买**格式**（0%→100% 干净落地），算术正确性是 base 模型本身弱项（如 `46*80` 框成 3380），正是 M3/M4 的可验证 reward 要去补的残差。一条诚实账：M1 用的是**朴素无 KV-cache 采样器**（每 token 全量 forward），100 题已经很慢——这正是 M2 解码引擎前置的动机。
 
+**M2a（KV-cache 增量解码引擎，单序列，已落地）**：两个 forward-only 原语 + 裸 Tensor 逐 token block forward，各自隔离闸门。`rope_at`（绝对位置 RoPE，新 kernel，不动训练 `rope` → 训练路径零风险）逐位等于全序列 rope 的对应行；`decode_attention`（单 query × cached-K/V，由现成 strided-gemm + 普通 softmax 组合，**零新 kernel**）等于全 causal attention 末行（max|Δ| 6e-8）。引擎 `generate_greedy_cached` 镜像 `block_forward` 在 Tensor 层（无 autograd tape，推理不需梯度），靠**公开 `params()` 稳定顺序**拿权重（零 model 可见性改动）。**核心闸门 = token-identical**：与朴素全重算贪心逐 token 一致（小 GQA 单测 + v12 1.05B 上 cached eval 与 naive **逐字节相同**：format 100/100, correct 8/100）。**吞吐 baseline（v12, batch1, F32，profile-first 实测）= cache 收益随序列长度而定**：max_new 32 ≈ 持平（108 vs 111，短序列 launch 开销 bound）、128 **~1.9×**（69 vs 133）、256 naive **OOM** vs cached 129 tok/s。cached 吞吐**近恒定**（O(1)/token + 恒定显存），naive **衰减**（O(t)/token，O(seq²) 图 → OOM）。⇒ 短 eval prompt overhead-bound、cache 几乎无收益，真正受益的是**长 rollout**（DPO 造对 / GRPO completion）——与 T17（process-per-GPU 吞吐中性）同一条 measure-first 教训：收益真实，但只在真正压到瓶颈的 regime 里。M2a 的 per-layer 主机往返是短序列 overhead-bound 的一部分原因，M2b（device 端 cache + 批量 ragged）针对它。
+
+**M3（DPO，离线偏好优化，已落地 + 诚实负结果）**：两个复用 CE kernel 的新算子（零新 CUDA）——`seq_logprob`（Σ log πθ over 非 mask 位，反向 = CE_backward 取负求和；grad-check + mask）、`dpo_loss`（−log σ(Δ)，双 policy logprob 父节点；grad-check + 退化 Δ=0→log2/∓β·½、β=0→0）。造对（`gen_dpo_pairs`）= chosen=gold、rejected=SFT 自己 greedy（用 M2a 引擎）的格式合法**错误**答案（8% greedy 答对的跳过）。训练（`train_dpo`）把 SFT ckpt 同时作 policy 和冻结 reference，**一次性预算 reference logprob 并缓存**（单模型驻留），每步 policy forward chosen+rejected → seq_logprob → dpo_loss，两 forward 共享 param 累积梯度；**loss 起步恰好 log2**（Δ=0 内置校验）。**结果（v12, 1500 对, β0.1；100 留出题 vs SFT 8/100）**：reward-margin 与 pref-acc 干净上升（loss 被正确优化、infra 对），但**不转化为 held-out 正确率**——lr5e-7×300→7%、×800→5%、lr1e-6×2000→margin+34 **崩溃**（0% 格式、输出垃圾），三档都在 100 题 ~2.7% 标准误内 = 统计持平。**教训**：chosen/rejected 只差最终数字 token，DPO 提升的是**特定训练对的 token 偏好、reweight 现有分布,不 install 能力**；base 模型没有算术算法,偏好优化不泛化,推狠了只是全局扭曲分布→不连贯。**DPO 在 chosen 本就 plausible 时有效,不能凭空造模型没有的知识**——这正是 M4 GRPO 的动机:在线优化**真实可验证 reward**(采样→check→强化真正对的)而非固定对的 proxy(但 GRPO 同样面对 8% 稀疏,能否抬动指标是 M4 的 open question)。与 v8/T17 同源的诚实账：跑通+闸门齐全,负结果如实记。
+
+**M4（GRPO,在线 critic-free RL,已落地 + 两道诚实系统墙 + 一致负结果）**：新算子 `clipped_pg_loss`（per-token ρ + clip + k3 KL,反向用新增 `scale_rows` per-row 缩放 kernel;grad-check active+A=0 路径 + 退化 ε→∞ vanilla/β=0 无KL)。环 `train_grpo`:采 B prompt × rollout G → checker reward 0/1 → group-relative advantage `(r−mean)/(std+ε)`(无 critic,全对/全错组跳过)→ 存 πθ_old/πref per-token → K 内层 clipped-PG。rollout 用 **M2 引擎 + 新加的 temperature 采样**（单行 logits 比 naive `[seq,vocab]` 轻)。**先把任务改简单**:v12 SFT 在硬/易题都 ~8-9%(只会格式不会算术)→ 在 easy(操作数≤20)上从 v12 base 重训 SFT → held-out **18.7%**;但 250/600 步同样 18.7% = 1B web-text 模型从 ~550 例**不泛化加减法、只记 train**。**两道系统墙(设计文档 Risks 预言)**:① 显存——KL-leash 要 policy+reference 两个 1B fp32-master+Adam≈21GB,加激活在 32GB 5090 上不稳定 OOM → 只能 `β=0`(去掉 reference)跑完;② rollout 长杆——naive 采样增长序列撑碎 allocator,cached 采样更轻但单序列慢仍主导墙钟(~16s/step)。**结果**(easy, β=0, G6·B6, 40步, lr5e-7;150 留出 vs SFT 18.7%):reward 噪声 ~0.58-0.81(被 train 重叠抬),**format 100/100 不崩**(温和 lr 下 β=0 也没崩),**held-out 20.0%**(+1.3pp,~3% 标准误内 = 统计持平)。**M3+M4 一致教训**:模型缺底层能力时,离线偏好(DPO)和在线 RL(GRPO)**都不抬 held-out**——各自在能触及的训练分布上优化目标(被记忆抬高),装不进可泛化算法;**RL 强化模型已会的,不教算术**。**后训练弧诚实终态 = 一套完整、闸门齐全的 SFT → KV-cache → DPO → GRPO 栈**,infra 学全,并测得对齐对"base 缺失能力"能做什么的诚实边界。
+
+**M2b（批量 KV-cache 解码,已落地,补全 M2 引擎 + 修 rollout 长杆）**：M4 后补的 rollout 长杆修复——一个 prompt 的 **G 个样本同步解码**(每步一次 forward 跑整组 → G× 更少 kernel 启动)。一个新原语 `rope_pos`(逐 row 绝对位置 kernel,G 行共享一个解码位置;闸门 = `[0..n]` 逐位等于全 rope、统一 P 逐行等于 `rope_at(P)`,bit-identical)。引擎 `generate_cached_batch`:`BatchKVCache` 带 G 维,批量 `decode_step` 把 G 贯穿 embed/proj/QK-norm/`rope_pos`/cache;**M2a 两件零改动复用**——`decode_attention` 本就 batch-agnostic(bh=G·nh)、`repeat_kv(nh,batch=G)` 按组广播。闸门 = G 个贪心行逐字节等于单序列(`tests/decode_batch.rs`,8q/2kv 头练 repeat_kv 批量)。**吞吐**(v12, G6·B6, 接进 train_grpo):**~8.5s/step vs 单序列 ~14-16s/step ≈ 1.7×**(rollout-inclusive;未到满 G× 因 per_token_logp + PG 更新也占时间、M2a 主机往返还在);且**显存更稳**(一次批量 forward vs G 次分配撑碎 allocator 的 M4 OOM)。⇒ M2 引擎闭环(M2a 单序列 + M2b 批量),rollout 长杆从"OOM/无界"变成有界 ~1.7× 收益,device 端 cache 是点名的下一杠杆。
+
+**M2c（device 端 KV cache,已落地,瓶颈转移的 profile-first 发现）**：K/V 留 device 为 `[bh,T,hd]`(每层 `Option<Tensor>`),每步用新 `cat_seq` kernel(沿 seq 拼接)append 一个 token——去掉 M2a/M2b 每层**主机往返** + `transpose_3d01`,单序列和批量都重构到它(比 host Vec+rebuild 干净)。闸门全保:`cat_seq`==host concat、decode_kv 单序列 + decode_batch 批量仍 **token-identical**、GQA 训练路径不受影响。**发现(measure-first 的点,不是加速故事)**:去掉主机往返让**纯单序列解码 +10%**(133→147 tok/s@128),但 **GRPO step 不动**(~8.5s/step)——因为 M2b 批量化后 rollout 已不是 step 瓶颈,**per-sample `per_token_logp` 捕获(2×/样本)+ PG 更新 forward/backward(全序列 `model.forward`)成了主导**。长杆从 rollout **转移**到训练侧 forward(同 T11/T17/M2a:profile 后再动手——你修的不是剩下的瓶颈)。device cache 仍是真实、闸门齐全的改进(更干净、少 PCIe、解码 +10%),但下一杠杆是 **per-sample forward 的 ragged 批量**而非 cache。M2 引擎现 = M2a(单序列)+ M2b(批量)+ M2c(device cache),全 token-identical-gated;后训练栈完整、瓶颈已测绘。
+
+**M2d（批量 GRPO 训练侧 forward,已落地,M2c 点名的杠杆 + 一处 decomposition 纠正）**：M2c 点名的下一杠杆——把每步 `N=B·G` 个 ragged 样本的训练侧 forward(`per_token_logp` 捕获 + inner clipped-PG fwd/bwd)打包进**一次 `forward_batched`**。**使能性质 = causal 下右 padding 免费**:真 completion 行位置早于尾部 pad,causal 禁止前向 attend,故真行 logits 与单序列 forward **逐位相同**,pad 行垃圾被 `target=-100` 屏蔽——这正是训练引擎 pad-and-mask 而非跑 ragged 的原因。两件新东西:`per_token_logp_batched`(右 pad → 一次 `forward_batched(N)` → 按真长切片)、`ops::clipped_pg_loss_batched`(per-row `advantage[t]` + per-row `weight[t]`,caller 传 `1/(N·n_s)`,op 不再自算 `1/n_tokens` → 折进 weight 即与 looped `Σ_s (1/N)(1/n_s)…` **逐位等价**;`--micro` 分块界定 `[chunk·Lmax,vocab]` logits 显存,weight 用全局 N 故分块梯度累积精确)。**两道精确闸门**:`forward_batched_ragged_matches_looped`(右 pad 批量 forward == 单序列,fp32 max|Δ|=3.7e-7、bf16 **0.0**,composed+flash)+ `clipped_pg_loss_batched_matches_looped`(批量 op == looped,loss Δ=1.5e-8/grad 7.5e-9,f32),复合即证端到端等价;端到端短 SFT→`train_grpo` 12 步**不 OOM**(1B master+AdamW+批量激活 micro=16 容得下)、批量 inner 执行。**吞吐(bench,v12 1.05B,N=48,micro16,权重无关)**:capture 622→71ms(8.7×)、inner 1907→208ms(9.2×)、**训练侧 forward 合计 2526→280ms(9.0×)**。**Decomposition 纠正(诚实发现)**:M2c 说"训练侧 forward 主导 step",干净分量 bench 证伪强形式——训练侧 forward 是 **~8.5s step 里的 ~2.5s(~30%)**,可观、值这 9×,但 **rollout(`generate_cached_batch` ~6s)一直是更大头**;M2d 把训练侧砍到 ~0.28s 后,step **~95% 是 rollout**,长杆又摆回 rollout。⇒ M2d 拔掉训练侧 forward 这块 overhang(分量级精确 9×),再次印证 measure-first:**step 级下一杠杆 = 全 B×G rollout 批量**(今天只有每 prompt 的 G 同步、B 个 prompt 仍串行)。后训练栈保持完整,step decomposition 现为**实测**而非断言。
+
 ## 四、perf 杠杆台账（详见 [known-issues.md](known-issues.md)）
 
 - **已修**：KI-1 单序列 launch-bound（T10）· KI-5 per-op cudaMalloc 串行（T11）· KI-2 bf16/OOM（T12）· KI-3 激活重计算（T13，解锁 dim1024，v8 用上）。