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).

README.md

@@ -6,10 +6,16 @@ inference side). A learning project: hand-write the entire training-systems stac
 gradient checkpointing), then use it to run a multi-version **scaling study** that maps
 the data-vs-capacity frontier for a tiny model.
 
-> **Status: complete.** From-scratch full stack (phases T1–T13) + an 8-version scaling
-> ladder (v0–v8). Trains a Qwen3-compatible LM whose weights load into **xserv** and
-> generate **token-identical** output. This README is the capstone; per-topic detail
-> lives in [`docs/`](docs/).
+> **Status: complete — three phases.**
+> **Phase 1** = the from-scratch full stack (T1–T13) + an 8-version scaling study (v0–v8):
+> hand-write the whole training-systems stack, then map the data-vs-capacity frontier.
+> **Phase 2** = systems-stack depth (T14–T18): hand-write the five deferred training-stack
+> features — fused flash-attention, real GQA, gradient accumulation, process-per-GPU DDP,
+> dropout. **Phase 3** = one Chinchilla-style double-axis run (v9): dim1280 true-GQA +
+> 6.01B FineWeb tokens, validating the v8 conclusion that data and capacity must scale
+> together. Trains Qwen3-compatible LMs whose weights load into **xserv**; deterministic
+> gates stay byte-identical, while large BF16 checkpoints are served and checked for
+> prompt-level drift. This README is the capstone; per-topic detail lives in [`docs/`](docs/).
 
 ---
 
@@ -22,20 +28,23 @@ borrows, the rest hand-written CUDA + Rust:
 |---|---|
 | `xtrain-cuda` | CUDA Runtime FFI, RAII `GpuBuffer`, **caching/pool allocator**, cuBLAS (sgemm + bf16 GemmEx) bindings |
 | `xtrain-tensor` | tensor (dtype/shape/strides/storage), elementwise + transpose + embedding kernels |
-| `xtrain-autodiff` | **tape autograd engine** (grad accumulation), per-op backward, finite-diff grad-check, **checkpoint** (recompute) primitive |
-| `xtrain-model` | tiny **Qwen3-style** transformer (RoPE + RMSNorm + QK-norm + SwiGLU), batched forward |
+| `xtrain-autodiff` | **tape autograd engine** (grad accumulation), per-op backward, finite-diff grad-check, **checkpoint** (recompute) primitive, **fused flash-attention** (online-softmax) fwd/bwd, **`repeat_kv`** broadcast (GQA), **`dropout`** (counter-based device RNG + mask) |
+| `xtrain-model` | tiny **Qwen3-style** transformer (RoPE + RMSNorm + QK-norm + SwiGLU), batched forward, **GQA** (`num_kv_heads<num_heads`), residual/MLP **dropout** |
 | `xtrain-optim` | hand-written **AdamW** (host + GPU kernels) |
-| `xtrain-train` | training loop, LR schedule, grad clip, checkpoint, BPE corpus + cache, samplers, safetensors export |
-| `xtrain-distributed` | **NCCL DDP** (thread-per-GPU, all-reduce) |
+| `xtrain-train` | training loop, LR schedule, grad clip, **gradient accumulation**, checkpoint, BPE corpus + cache, samplers, safetensors export |
+| `xtrain-distributed` | **NCCL DDP** (thread-per-GPU + torchrun-style process-per-GPU launcher / cross-process `ncclUniqueId`, all-reduce) |
 
 Every op's backward is verified against **finite differences** and against **PyTorch**
 (forward + per-parameter grads, batch > 1). Trained weights export to HF-safetensors and
-load into xserv (Qwen3, BF16) producing token-identical greedy output — the closed loop.
+load into xserv (Qwen3, BF16); deterministic fixtures produce token-identical greedy output,
+and large checkpoints are validated end-to-end in the serving path.
 
-## The build journey — phases T1–T13
+## The build journey — Phase 1 (T1–T13) + Phase 2 (T14–T18)
 
 Each phase: design doc + implementation + tests + a scoped commit (see [`docs/`](docs/) and
-[`docs/evolution.md`](docs/evolution.md) for the per-axis changelog).
+[`docs/evolution.md`](docs/evolution.md) for the per-axis changelog). **Phase 1 (T1–T13)**
+hand-built the stack and fixed the four real bottlenecks; **Phase 2 (T14–T18)** went back to
+hand-write five deferred training-stack features — see the Phase-2 summary below the table.
 
 | phase | what | result |
 |---|---|---|
@@ -50,12 +59,57 @@ Each phase: design doc + implementation + tests + a scoped commit (see [`docs/`]
 | **T11** | **device caching allocator** (fixes KI-5) | single-GPU 2.3×; **8-GPU 461K tok/s** |
 | **T12** | **bf16 mixed precision** (fp32 master, fixes KI-2) | dim768 OOM solved; −29% mem |
 | **T13** | **activation recompute** / checkpointing (fixes KI-3) | dim1024 fits; grads bit-identical |
+| **T14** | **fused flash-attention** kernel (online softmax, no materialized N×N; opt-in `--flash`) | peak mem −16%@1k / −23%@2k seq; flash==composed (grads/PyTorch) |
+| **T15** | **grouped-query attention** (`num_kv_heads<num_heads`; `repeat_kv` broadcast feeds both SDPA paths; backward sums each kv head's group; `--kv-heads`) | repeat_kv grad-check + **group=1 bit-identical to MHA**; GQA flash==composed; PyTorch GQA B>1; **xserv closed loop with real `num_key_value_heads`** token-identical |
+| **T16** | **gradient accumulation** (`--accum-steps`; DDP all-reduces only at the boundary) | equiv to N× big batch (grad 3.8e-5); same effective-64 batch 27.7GB→7.2GB (−74%) |
+| **T17** | **process-per-GPU** DDP (torchrun-style: 1 worker process / CUDA context per GPU; launcher mints `ncclUniqueId` → hex env injection; `train_rank` reused unchanged; thread-per-GPU path kept) | proc==thread loss 1.5e-7, cross-rank 1.2e-7, xserv md5 identical · **measured no-op on throughput**: thread 5.27× vs proc 5.31×@8 (8 GPUs 95–99% util) → residual non-linearity is NCCL/PCIe, *not* CUDA-context serialization (falsifies the old KI-5 hypothesis) |
 | **T18** | **dropout** (hand counter-based device RNG + mask, inverted scaling, train/eval switch) | fixed-seed grad-check; **p=0 bit-identical**; recompute-safe |
 
 The four performance fixes (T10–T13) each removed a real bottleneck — see
-[`docs/known-issues.md`](docs/known-issues.md).
+[`docs/known-issues.md`](docs/known-issues.md) — which is where **Phase 1** closed.
 
-## The scaling study — v0 → v8
+## Phase 2 — systems-stack depth (T14–T18)
+
+Phase 1 fixed bottlenecks; Phase 2 went back to hand-write the five training-stack features
+that had been **explicitly deferred** earlier (project's actual goal = learn the whole stack).
+Each is opt-in, kept the default path **bit-identical**, and held a **hard correctness gate**:
+
+- **T14 · fused flash-attention** ([`docs/13-flash-attention.md`](docs/13-flash-attention.md)) —
+  a single hand-written kernel: **online (streaming) softmax, tiled over KV, never materializes
+  the `N×N` scores**; flash-style backward recomputes scores + the `D=ΣdO·O` Jacobian
+  simplification for dQ/dK/dV. Opt-in `--flash`, default off. **The win is memory, not
+  wall-clock**: peak activation **−16%@seq1024 / −23%@seq2048** (grows with seq, since the
+  `N×N` never lands), but **~2.3× slower** at head-dim 64 (a hand kernel can't beat cuBLAS
+  tensor-cores on a small head). Gate: flash == composed (loss rel `0.0`, grad `4.4e-5`),
+  PyTorch B>1 `7.9e-6`.
+- **T15 · real GQA** ([`docs/14-gqa.md`](docs/14-gqa.md)) — `num_kv_heads < num_heads` via a new
+  `repeat_kv` **broadcast op** that copies K/V `group = nh/num_kv` times to feed **both**
+  (composed + flash) SDPA paths **unchanged**; its **backward is a deterministic group-sum**
+  (no atomics) collapsing each kv head's query-head group. Gate: `repeat_kv` grad-check +
+  **group=1 bit-identical to MHA** (regression guard); **xserv closed loop with real
+  `num_key_value_heads`** token-identical.
+- **T16 · gradient accumulation** ([`docs/15-grad-accum.md`](docs/15-grad-accum.md)) — N
+  micro-steps scaled by `1/N` accumulate on the tape, then one AdamW step; DDP **all-reduces
+  only at the accumulation boundary**. Decouples effective batch from activation memory: same
+  effective batch 64, big-batch **27.7GB (OOM)** → accum 4×16 **7.2GB (−74%)**. Gate: `accum=N`
+  ≡ one N× batch (grad `3.8e-5`); `accum=1` bit-identical.
+- **T18 · dropout** ([`docs/17-dropout.md`](docs/17-dropout.md)) — a **stateless counter-based
+  device RNG** (Philox-style bit-mix) → Bernoulli mask, inverted `1/(1−p)` scaling in train,
+  identity in eval; wired at the two residual sites (attn-out, mlp-out). Stateless RNG is what
+  makes it **compose bit-exactly with T13 activation recompute** — the backward re-run
+  regenerates the *same* mask from `(seed, index)`. Gate: fixed-seed grad-check; **p=0
+  bit-identical**.
+- **T17 · process-per-GPU** ([`docs/16-process-per-gpu.md`](docs/16-process-per-gpu.md)) — a
+  torchrun-style launcher: one worker process + CUDA context per GPU, the launcher mints one
+  `ncclUniqueId` and **hex-injects it into each child's env** (no shared FS/TCP, no race); the
+  worker reuses the T8 `train_rank` **unchanged**. Built and **correct** (proc vs thread loss
+  `1.5e-7`, cross-rank `1.2e-7`, xserv md5 identical) — but **measured throughput-neutral**:
+  8-GPU thread **491K (5.27×)** vs proc **493K (5.31×)**, `<1%`. This **falsifies** the
+  long-standing KI-5/T11 hypothesis that thread-per-GPU's shared CUDA context caused the
+  residual ~5×@8; with all 8 GPUs at 95–99% util, the residual is the **NCCL all-reduce + PCIe
+  topology wall**, not context serialization. The third profile-first falsification (see below).
+
+## The scaling study — v0 → v10
 
 Same Qwen3-style architecture throughout; we scaled **dim** and **data** and read out val
 loss (full per-run detail in [`docs/runs/`](docs/runs/)).
@@ -68,11 +122,13 @@ loss (full per-run detail in [`docs/runs/`](docs/runs/)).
 | v6 | FineWeb-edu 1.02ep | 768 / 127M | 3.07\* | **corpus swap → graduates to real text** |
 | v7 | FineWeb-edu 1.45ep | 768 / 127M | 3.01\* | same subset, more epochs → near-ceiling |
 | **v8** | FineWeb-edu 1.05ep | **1024 / 226M** | **2.98\*** | **capacity → helps** |
+| **v9** | FineWeb-edu 6.01B / ~1ep | **1280 / 357M + GQA** | **2.89\*** | **data + capacity → helps** |
+| **v10** | FineWeb-edu 6.76B / ~1ep | **1280 / 357M + GQA** | **2.88\*** | **data-only top-up → small gain** |
 
 \* FineWeb-edu val is a different (harder) distribution — **not comparable** to the
 TinyStories val of v0–v5. Judge v6+ by sample quality + transfer, not the number.
 
-### Three findings
+### Four findings
 
 1. **Data volume saturates.** TinyStories at dim768: 3.5× more tokens (v4→v5) bought only
    −5% val, curve flat. The narrow synthetic corpus is exhausted at this model size.
@@ -81,10 +137,18 @@ TinyStories val of v0–v5. Judge v6+ by sample quality + transfer, not the numb
    historical/scientific expository prose. (Cost: TinyStories transfer val 1.11 → 2.75.)
 3. **Capacity helps.** v8 (dim1024, ~1 epoch) beats both v6 (dim768, same epoch, by 0.085)
    and v7 (dim768, *more* data, by 0.035) → the dim768 runs were partly capacity-limited.
+4. **Double-axis scale helps.** v9 scales both axes (dim1280/core357M + 6.01B FineWeb tokens)
+   and beats v8 by another 0.095 val loss (~3.2%). The direction is validated, but the gain is
+   still incremental and greedy decoding still repeats.
+5. **Moving validation tails must stop.** v10 added one more FineWeb shard and got moving-tail
+   val 2.8816, but appending data moves the held-out tail. A fixed eval v1 was created from the
+   shard010 tail: v6/v7/v8/v9/v10 = 3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814. Future runs
+   should report this fixed eval first.
 
-**Meta-finding:** every *single*-axis lever (data volume, corpus breadth, capacity) is now
-worth only **~3%**. Per the Chinchilla lesson, further gains require scaling **data and
-capacity together** — single-axis moves are exhausted.
+**Meta-finding:** every lever is now in the **~3% or smaller** regime. Single-axis moves were
+exhausted by v8; v9 confirms Chinchilla-style double-axis scale works; v10 shows a data-only
+top-up mostly adapts to the new shard. The next useful run should change model/context, not just
+append another shard.
 
 ## Efficiency — throughput & MFU
 
@@ -103,25 +167,30 @@ versions — a fixed-MFU estimate is off by up to ~100× for the early launch-bo
 
 ## Engineering lessons
 
-- **Profile before optimizing.** Two "known" perf fixes were *falsified by measurement*
-  before being shipped: "bigger batch fixes DDP scaling" (real cause: single-seq
-  launch-bound → T10) and "bucket the all-reduce" (real cause: per-op `cudaMalloc`
-  serialization → T11 caching allocator). Both would have been no-ops; both got reverted +
-  re-diagnosed instead of shipped.
-- **Honest correctness.** QK-norm was *added* to match xserv's Qwen3 (not faked); every perf
-  change kept a hard correctness gate (recompute grads bit-identical; bf16 keeps the fp32
-  path untouched; the full grad-check / PyTorch / DDP / xserv suite must stay green).
-- **The closed loop matters.** Exporting to xserv and checking token-identical greedy output
-  caught real bugs and proved the whole stack end-to-end.
+- **Profile before optimizing.** *Three* "known" fixes were *falsified by measurement*: (1)
+  "bigger batch fixes DDP scaling" (real cause: single-seq launch-bound → T10); (2) "bucket the
+  all-reduce" (real cause: per-op `cudaMalloc` serialization → T11 caching allocator); and (3)
+  "process-per-GPU would fix the residual ~5×@8" (T17 — built the torchrun-style launcher and
+  measured it **throughput-neutral**: the residual is the NCCL/PCIe communication wall, not
+  shared-context serialization). All three would have been no-ops; each got measured and either
+  reverted or recorded as a deliberate negative result instead of shipped on faith.
+- **Honest correctness.** QK-norm was *added* to match xserv's Qwen3 (not faked); every change
+  kept a hard correctness gate, and **no tolerance was ever loosened to go green**. Phase 2 held
+  the line: flash == composed SDPA (grads/PyTorch), GQA group=1 bit-identical to MHA, gradient
+  accumulation `accum=1` bit-identical, dropout p=0 bit-identical *and* dropout × recompute
+  bit-exact, the default path unchanged on every feature, and the **xserv closed-loop md5
+  byte-identical (`b04fc9f9`) throughout the deterministic gates**.
+- **The closed loop matters.** Exporting to xserv and checking generated continuations caught
+  real bugs and proved the whole stack end-to-end.
 
 ## Running it
 
 Everything trains on a remote 8× RTX 5090 box; model artifacts live in a registry
-(`tiny-models/v0…v8`). Serve any trained version in xserv:
+(`tiny-models/v0…v10`). Serve any trained version in xserv:
 
 ```bash
 # on the GPU box
-cargo run -p xserv-model --release --bin xserv-cli -- <registry>/v8-fineweb-edu-dim1024 --max-tokens 100
+cargo run -p xserv-model --release --bin xserv-cli -- <registry>/v10-fineweb-edu-dim1280-gqa-data6765 --max-tokens 100
 # then type a prompt, e.g.  In science,
 ```
 
@@ -136,6 +205,6 @@ cargo test --workspace            # autograd grad-checks, PyTorch parity, DDP, e
 ## Doc index
 
 - [`docs/evolution.md`](docs/evolution.md) — per-milestone changes across algorithm / architecture / infra / dataset.
-- [`docs/runs/README.md`](docs/runs/README.md) — the v0–v8 comparison; [`docs/runs/0N-*.md`](docs/runs/) — per-run detail.
-- [`docs/00-*` … `12-*`](docs/) — per-phase design docs (build chain → tensor → autograd → transformer → training → perf → distributed → export → batched → allocator → bf16 → recompute).
-- [`docs/known-issues.md`](docs/known-issues.md) — perf backlog (KI-1/2/3/5 fixed; KI-4 + process-per-GPU open).
+- [`docs/runs/README.md`](docs/runs/README.md) — the v0–v10 comparison; [`docs/runs/0N-*.md`](docs/runs/) — per-run detail.
+- [`docs/00-*` … `17-*`](docs/) — per-phase design docs (build chain → tensor → autograd → transformer → training → perf → distributed → export → batched → allocator → bf16 → recompute → flash-attention → GQA → grad-accum → process-per-GPU → dropout).
+- [`docs/known-issues.md`](docs/known-issues.md) — perf backlog (KI-1/2/3/5 fixed; process-per-GPU CLOSED = measured no-op; KI-4 = accepted modeling tradeoff).

crates/xtrain-autodiff/src/ops.rs

@@ -350,8 +350,56 @@ pub fn attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
     )
 }
 
+/// Fused FLASH causal scaled-dot-product attention (Phase T14). Same interface as
+/// [`attention`] (`q`,`k`,`v` each `[bh, seq, head_dim]`), but the forward is a
+/// SINGLE fused kernel with an online softmax over KV tiles — the `[bh,seq,seq]`
+/// score matrix is NEVER materialized, and backward caches only the per-row
+/// logsumexp (O(N)) instead of the whole probs (O(N²)). Mathematically the same
+/// SDPA, so it matches the composed [`attention`] within fp/bf16 tolerance.
+/// Opt-in via the model's `--flash` flag; the composed path stays the default.
+pub fn flash_attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
+    let (out, lse) = q.value().flash_attention(&k.value(), &v.value(), scale);
+    let out_cache = out.clone();
+    Var::from_op(
+        out,
+        vec![q.clone(), k.clone(), v.clone()],
+        Box::new(move |dout, parents| {
+            let q = parents[0].value();
+            let k = parents[1].value();
+            let v = parents[2].value();
+            let (dq, dk, dv) =
+                Tensor::flash_attention_backward(&q, &k, &v, &out_cache, &lse, dout, scale);
+            Var::push_grad(&parents[0], dq);
+            Var::push_grad(&parents[1], dk);
+            Var::push_grad(&parents[2], dv);
+        }),
+    )
+}
+
+/// GQA repeat_kv head broadcast (Phase T15). `kv`:[batch·num_kv, seq, head_dim]
+/// (a K or V tensor) → `[batch·nh, seq, head_dim]`, each KV head broadcast to its
+/// `group = nh/num_kv` query heads (qh ← kv head qh/group, contiguous groups —
+/// matches xserv's repeat_kv). Feeds the unchanged composed/flash SDPA so GQA is
+/// "free" for both. Backward SUMS the `group` query heads sharing each KV head back
+/// onto it (the multi-group grad accumulation). `nh == num_kv` (group 1) is identity
+/// → bit-identical to the MHA path. `batch` lets the op recover num_kv from kv's bh.
+pub fn repeat_kv(kv: &Var, nh: usize, batch: usize) -> Var {
+    let bh_kv = kv.value().shape()[0];
+    let num_kv = bh_kv / batch;
+    let out = kv.value().repeat_kv(nh, batch);
+    Var::from_op(
+        out,
+        vec![kv.clone()],
+        Box::new(move |dout, parents| {
+            let din = Tensor::repeat_kv_backward(dout, num_kv, batch);
+            Var::push_grad(&parents[0], din);
+        }),
+    )
+}
+
 /// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
-/// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
+/// row. Negative targets are ignored, which is useful for assistant-only SFT
+/// masks. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/valid_rows`,
 /// scaled by the upstream scalar grad.
 pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
     // CE math is fp32 (cross_entropy upcasts bf16 logits internally + caches fp32
@@ -360,10 +408,22 @@ pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
     // fp32 logits buffer) is a real activation-memory saving at large vocab.
     let logit_dtype = x.value().dtype();
     let (probs, per_row) = x.value().cross_entropy(target);
-    let rows = x.value().shape()[0];
+    let cols = x.value().shape()[1] as i32;
+    let target_host = target.to_device(xtrain_tensor::Device::Cpu);
+    let valid_rows = target_host
+        .as_slice::<i32>()
+        .iter()
+        .filter(|&&t| {
+            if t >= cols {
+                panic!("cross_entropy target {t} out of vocab range {cols}");
+            }
+            t >= 0
+        })
+        .count()
+        .max(1);
     // Mean loss as a host scalar wrapped back into a [1] tensor.
     let mean = per_row.to_device(xtrain_tensor::Device::Cpu);
-    let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / rows as f32;
+    let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / valid_rows as f32;
     let loss = Tensor::from_slice(&[mean_val], &[1]).to_device(x.value().device());
 
     let target = target.clone();
@@ -373,9 +433,251 @@ pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
         Box::new(move |d, parents| {
             // `d` is the scalar upstream grad (1.0 when this is the loss root).
             let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
-            let scale = upstream / rows as f32;
+            let scale = upstream / valid_rows as f32;
             let dx = Tensor::cross_entropy_backward(&probs, &target, scale);
             Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
         }),
     )
 }
+
+/// 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

@@ -625,6 +625,245 @@ fn attention_batched_bwd() {
     );
 }
 
+// ---- fused FLASH causal attention (the T14 op) ----
+// Same structure + dimensions as attention_batched_bwd (bh=2,seq=5,hd=6), but
+// exercises ops::flash_attention. Grad-check dq/dk/dv against finite-diff of
+// L=sum(W∘out). This is the SINGLE-tile regime (seq<FA_TILE=32), matching the
+// trusted composed grad-check's clean near-zero behavior; the MULTI-tile online-
+// softmax path (seq>FA_TILE) is validated against the already-grad-checked
+// composed backward by `flash_bwd_matches_composed_bwd` (seq=40) — sharper than
+// finite-diff, which is unreliable on the near-zero grad elements a long softmax
+// produces.
+#[test]
+fn flash_attention_batched_bwd() {
+    require_gpu();
+    let (bh, seq, hd) = (2, 5, 6);
+    let n = bh * seq * hd;
+    let scale = 1.0 / (hd as f32).sqrt();
+    // Scale Q/K up so the softmax is non-uniform (sharper attention) → the dQ/dK
+    // gradients are well-conditioned, not the near-zero saddle values a uniform
+    // softmax produces (those make central finite-diff give spurious 0.0 / sign
+    // flips that aren't backward bugs — cf. flash_bwd_matches_composed_bwd).
+    let q_h: Vec<f32> = fill(n, 241).iter().map(|v| v * 2.5).collect();
+    let k_h: Vec<f32> = fill(n, 242).iter().map(|v| v * 2.5).collect();
+    let v_h = fill(n, 243);
+    let w = fill(n, 244);
+
+    let q = Var::leaf(cuda(&q_h, &[bh, seq, hd]));
+    let k = Var::leaf(cuda(&k_h, &[bh, seq, hd]));
+    let v = Var::leaf(cuda(&v_h, &[bh, seq, hd]));
+    let out = ops::flash_attention(&q, &k, &v, scale);
+    scalar_loss(&out, &w).backward();
+
+    let dq = q.grad().unwrap().to_device(Device::Cpu);
+    let dk = k.grad().unwrap().to_device(Device::Cpu);
+    let dv = v.grad().unwrap().to_device(Device::Cpu);
+
+    let fwd = move |qh: &[f32], kh: &[f32], vh: &[f32]| -> f32 {
+        let qv = cuda(qh, &[bh, seq, hd]);
+        let kv = cuda(kh, &[bh, seq, hd]);
+        let vv = cuda(vh, &[bh, seq, hd]);
+        let (o, _) = qv.flash_attention(&kv, &vv, scale);
+        weighted_sum(&o, &w)
+    };
+    // Attention dQ/dK carry softmax curvature; for the small grad magnitudes here
+    // a larger eps (2e-3) cuts the f32 rounding term (∝|L|/eps) that dominates the
+    // O(eps²) truncation on a ~4e-4 grad. (dV is exactly linear → cfg_linear.)
+    let cfg_attn = GradCheckConfig {
+        eps: 2e-3,
+        rel_tol: 3e-2,
+        atol: 1e-3,
+    };
+    let (kf, vf, ff) = (k_h.clone(), v_h.clone(), fwd.clone());
+    let lq = move |x: &[f32], _s: &[usize]| ff(x, &kf, &vf);
+    report(
+        "flash dQ",
+        &grad_check(&q_h, &[bh, seq, hd], &lq, dq.as_slice::<f32>(), cfg_attn),
+    );
+    let (qf, vf, ff) = (q_h.clone(), v_h.clone(), fwd.clone());
+    let lk = move |x: &[f32], _s: &[usize]| ff(&qf, x, &vf);
+    report(
+        "flash dK",
+        &grad_check(&k_h, &[bh, seq, hd], &lk, dk.as_slice::<f32>(), cfg_attn),
+    );
+    let (qf, kf, ff) = (q_h.clone(), k_h.clone(), fwd.clone());
+    let lv = move |x: &[f32], _s: &[usize]| ff(&qf, &kf, x);
+    report(
+        "flash dV",
+        &grad_check(
+            &v_h,
+            &[bh, seq, hd],
+            &lv,
+            dv.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}
+
+// flash forward must equal the composed attention forward (same SDPA math).
+#[test]
+fn flash_matches_composed_fwd() {
+    require_gpu();
+    let (bh, seq, hd) = (2, 40, 16);
+    let n = bh * seq * hd;
+    let scale = 1.0 / (hd as f32).sqrt();
+    let q = cuda(&fill(n, 341), &[bh, seq, hd]);
+    let k = cuda(&fill(n, 342), &[bh, seq, hd]);
+    let v = cuda(&fill(n, 343), &[bh, seq, hd]);
+    let (oc, _) = q.attention(&k, &v, scale);
+    let (of, _) = q.flash_attention(&k, &v, scale);
+    let oc = oc.to_device(Device::Cpu);
+    let of = of.to_device(Device::Cpu);
+    let max_rel = oc
+        .as_slice::<f32>()
+        .iter()
+        .zip(of.as_slice::<f32>())
+        .map(|(c, f)| (c - f).abs() / (c.abs() + 1e-6))
+        .fold(0.0f32, f32::max);
+    println!("flash-vs-composed fwd max rel: {max_rel:.3e}");
+    assert!(
+        max_rel < 1e-4,
+        "flash fwd diverges from composed: {max_rel:.3e}"
+    );
+}
+
+// flash backward must equal the (already grad-checked) composed backward. This is
+// a sharper test than finite-diff: both share the trusted composed forward as the
+// reference, so it isolates the flash bwd dQ/dK/dV math from finite-diff noise on
+// near-zero gradient elements.
+#[test]
+fn flash_bwd_matches_composed_bwd() {
+    require_gpu();
+    let (bh, seq, hd) = (2, 40, 16);
+    let n = bh * seq * hd;
+    let scale = 1.0 / (hd as f32).sqrt();
+    let q_h = fill(n, 441);
+    let k_h = fill(n, 442);
+    let v_h = fill(n, 443);
+    let w = fill(n, 444);
+
+    let run = |flash: bool| -> (Vec<f32>, Vec<f32>, Vec<f32>) {
+        let q = Var::leaf(cuda(&q_h, &[bh, seq, hd]));
+        let k = Var::leaf(cuda(&k_h, &[bh, seq, hd]));
+        let v = Var::leaf(cuda(&v_h, &[bh, seq, hd]));
+        let out = if flash {
+            ops::flash_attention(&q, &k, &v, scale)
+        } else {
+            ops::attention(&q, &k, &v, scale)
+        };
+        scalar_loss(&out, &w).backward();
+        let g = |x: &Var| {
+            x.grad()
+                .unwrap()
+                .to_device(Device::Cpu)
+                .as_slice::<f32>()
+                .to_vec()
+        };
+        (g(&q), g(&k), g(&v))
+    };
+    let (cq, ck, cv) = run(false);
+    let (fq, fk, fv) = run(true);
+    let maxrel = |a: &[f32], b: &[f32]| -> f32 {
+        a.iter()
+            .zip(b)
+            .map(|(x, y)| (x - y).abs() / (x.abs() + y.abs() + 1e-4))
+            .fold(0.0f32, f32::max)
+    };
+    let (rq, rk, rv) = (maxrel(&cq, &fq), maxrel(&ck, &fk), maxrel(&cv, &fv));
+    println!("flash-vs-composed bwd max rel: dQ {rq:.3e} dK {rk:.3e} dV {rv:.3e}");
+    assert!(rq < 2e-2, "dQ diverges: {rq:.3e}");
+    assert!(rk < 2e-2, "dK diverges: {rk:.3e}");
+    assert!(rv < 2e-2, "dV diverges: {rv:.3e}");
+}
+
+// ---- GQA repeat_kv head broadcast (Phase T15) ----
+//
+// repeat_kv expands K/V from [batch·num_kv, seq, hd] to [batch·nh, seq, hd]; each
+// kv head is broadcast to its `group = nh/num_kv` query heads. The forward is a
+// gather (a linear map), so finite-diff is clean. The CRITICAL gate is the
+// BACKWARD: a kv head receives the SUM of the `group` query heads sharing it —
+// the multi-group-to-one grad accumulation GQA correctness hinges on. We grad-check
+// din against finite-diff of L = sum(W∘out) with group>1, plus assert the forward
+// actually broadcasts and that group==1 is exact identity.
+#[test]
+fn repeat_kv_grad() {
+    require_gpu();
+    // batch 2, num_kv 2 → bh_kv 4 input rows; nh 6 → group 3, bh_q 12 output rows.
+    let (batch, num_kv, nh, seq, hd) = (2usize, 2usize, 6usize, 4usize, 5usize);
+    let n_in = batch * num_kv * seq * hd;
+    let n_out = batch * nh * seq * hd;
+    let x_h = fill(n_in, 711);
+    let w = fill(n_out, 712);
+
+    let kv = Var::leaf(cuda(&x_h, &[batch * num_kv, seq, hd]));
+    let out = ops::repeat_kv(&kv, nh, batch);
+    assert_eq!(out.value().shape(), &[batch * nh, seq, hd]);
+
+    // Forward sanity: out head (b·nh + qh) must equal in head (b·num_kv + qh/group).
+    let group = nh / num_kv;
+    let out_h = out
+        .value()
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec();
+    let row = seq * hd;
+    for b in 0..batch {
+        for qh in 0..nh {
+            let kvh = qh / group;
+            let o0 = (b * nh + qh) * row;
+            let i0 = (b * num_kv + kvh) * row;
+            for e in 0..row {
+                assert_eq!(out_h[o0 + e], x_h[i0 + e], "repeat_kv fwd mismatch");
+            }
+        }
+    }
+
+    scalar_loss(&out, &w).backward();
+    let din = kv.grad().unwrap().to_device(Device::Cpu);
+
+    let fwd = move |xh: &[f32], _s: &[usize]| -> f32 {
+        let kv = cuda(xh, &[batch * num_kv, seq, hd]);
+        let o = kv.repeat_kv(nh, batch);
+        weighted_sum(&o, &w)
+    };
+    // repeat_kv is exactly linear (gather/sum), so the linear-op tolerances apply.
+    report(
+        "repeat_kv din",
+        &grad_check(
+            &x_h,
+            &[batch * num_kv, seq, hd],
+            &fwd,
+            din.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}
+
+// group==1 (num_kv == nh) must be a bit-exact identity in BOTH directions — this is
+// the regression guard that makes the MHA path (kv_heads == n_heads) unchanged.
+#[test]
+fn repeat_kv_identity_group1() {
+    require_gpu();
+    let (batch, nh, seq, hd) = (2usize, 3usize, 4usize, 5usize);
+    let n = batch * nh * seq * hd;
+    let x_h = fill(n, 721);
+    let w = fill(n, 722);
+    let kv = Var::leaf(cuda(&x_h, &[batch * nh, seq, hd]));
+    let out = ops::repeat_kv(&kv, nh, batch); // group 1
+    let out_h = out
+        .value()
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec();
+    assert_eq!(out_h, x_h, "group-1 repeat_kv fwd must be identity");
+    scalar_loss(&out, &w).backward();
+    let din = kv.grad().unwrap().to_device(Device::Cpu);
+    // dL/din = w exactly (identity forward → grad passes through unchanged).
+    for (g, expect) in din.as_slice::<f32>().iter().zip(&w) {
+        assert_eq!(*g, *expect, "group-1 repeat_kv bwd must be identity");
+    }
+}
+
 // ---- dropout (Phase T18) ----
 //
 // Fixed-seed finite-diff grad-check. Under a fixed `seed` the mask is constant
@@ -676,9 +915,17 @@ fn dropout_expectation_and_keep_rate() {
         let (out, mask) = x.dropout(p, 0x5EED_0000 + t as u64);
         let out_h = out.to_device(Device::Cpu);
         let mask_h = mask.to_device(Device::Cpu);
-        let mean_out: f64 =
-            out_h.as_slice::<f32>().iter().map(|&v| v as f64).sum::<f64>() / n as f64;
-        let kept = mask_h.as_slice::<f32>().iter().filter(|&&m| m != 0.0).count();
+        let mean_out: f64 = out_h
+            .as_slice::<f32>()
+            .iter()
+            .map(|&v| v as f64)
+            .sum::<f64>()
+            / n as f64;
+        let kept = mask_h
+            .as_slice::<f32>()
+            .iter()
+            .filter(|&&m| m != 0.0)
+            .count();
         mean_out_acc += mean_out;
         keep_acc += kept as f64 / n as f64;
     }
@@ -758,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/build.rs

@@ -36,6 +36,8 @@ fn main() {
         .file("../../csrc/ops/model.cu")
         .file("../../csrc/ops/optim.cu")
         .file("../../csrc/ops/attention.cu")
+        .file("../../csrc/ops/flash_attention.cu")
+        .file("../../csrc/ops/repeat_kv.cu")
         .file("../../csrc/ops/cast.cu")
         .file("../../csrc/ops/dropout.cu")
         .compile("xtrain_cuda_kernels");

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,
@@ -243,6 +288,90 @@ unsafe extern "C" {
     );
 }
 
+// Fused flash-attention (csrc/ops/flash_attention.cu, Phase T14). A SINGLE kernel
+// each for forward/backward that streams over KV tiles with an online softmax and
+// NEVER materializes the [bh,S,S] score matrix. Q/K/V/out are [bh,S,hd] row-major
+// F32; the forward saves only the per-row logsumexp `l` ([bh*S], O(N)) for backward.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    // Forward: o[bh,S,hd] = softmax(causal(Q·Kᵀ·scale))·V, online over KV tiles.
+    // Also writes l[bh*S] = per-row logsumexp (saved for backward, not the scores).
+    #[allow(clippy::too_many_arguments)]
+    pub fn launch_flash_attention_fwd_f32(
+        q: *const f32,
+        k: *const f32,
+        v: *const f32,
+        o: *mut f32,
+        l: *mut f32,
+        bh: i32,
+        seq: i32,
+        hd: i32,
+        scale: f32,
+        s: CudaStream,
+    );
+    // Per-row D[i]=Σ_d dO[i,d]·O[i,d] over `rows`=bh*S rows of width `hd`. Must run
+    // before the backward kernel (which takes the precomputed D, not O).
+    pub fn launch_flash_attention_rowdot_f32(
+        d_o: *const f32,
+        o: *const f32,
+        d_d: *mut f32,
+        rows: i32,
+        hd: i32,
+        s: CudaStream,
+    );
+    // Backward: recomputes scores from Q/K/V + saved logsumexp `l` (NO cached probs)
+    // and the precomputed `d_d` (= D), produces dq/dk/dv. dq/dk/dv must be PRE-ZEROED
+    // (dk/dv are accumulated across query rows via atomicAdd).
+    #[allow(clippy::too_many_arguments)]
+    pub fn launch_flash_attention_bwd_f32(
+        q: *const f32,
+        k: *const f32,
+        v: *const f32,
+        d_o: *const f32,
+        l: *const f32,
+        d_d: *mut f32,
+        dq: *mut f32,
+        dk: *mut f32,
+        dv: *mut f32,
+        bh: i32,
+        seq: i32,
+        hd: i32,
+        scale: f32,
+        s: CudaStream,
+    );
+}
+
+// GQA repeat_kv head broadcast (csrc/ops/repeat_kv.cu, Phase T15). Expands a K/V
+// tensor from [batch·num_kv, S, hd] to the full [batch·nh, S, hd] so the SDPA
+// (composed or flash, both untouched) sees a full set of heads. Forward gathers
+// (out head qh ← kv head qh/group, group = nh/num_kv); backward sums the `group`
+// query heads sharing each kv head (deterministic, no atomics). All F32.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    // Forward: out[b·nh+qh] = in[b·num_kv + qh/group], per [S,hd] head block.
+    pub fn launch_repeat_kv_fwd_f32(
+        input: *const f32,
+        out: *mut f32,
+        batch: i32,
+        nh: i32,
+        num_kv: i32,
+        seq: i32,
+        hd: i32,
+        s: CudaStream,
+    );
+    // Backward: din[b·num_kv+kvh] = Σ_{r<group} dout[b·nh + kvh·group + r].
+    pub fn launch_repeat_kv_bwd_f32(
+        dout: *const f32,
+        din: *mut f32,
+        batch: i32,
+        nh: i32,
+        num_kv: i32,
+        seq: i32,
+        hd: i32,
+        s: CudaStream,
+    );
+}
+
 // GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
 // the global grad-norm reduction + in-place rescale (Phase T7).
 #[cfg(not(no_cuda))]

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

@@ -61,6 +61,8 @@ fn main() {
     let head_dim = flag(&args, "--head-dim", 16usize);
     let n_layers = flag(&args, "--layers", 4usize);
     let ffn = flag(&args, "--ffn", 64usize);
+    // GQA (Phase T15): num K/V heads (must divide --heads). Default = --heads (MHA).
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
     // `--dim` is informational; dim is always n_heads*head_dim. Warn on mismatch.
     let dim_flag = flag(&args, "--dim", 0usize);
     if dim_flag != 0 && dim_flag != n_heads * head_dim {
@@ -74,6 +76,10 @@ fn main() {
     // Optimization knobs (mirror bin/train).
     let steps: usize = flag(&args, "--steps", 100);
     let batch: usize = flag(&args, "--batch", 16);
+    // Micro-batch gradient accumulation (Phase T16): effective global batch =
+    // accum_steps × batch, all-reducing only at the accumulation boundary. Default
+    // 1 = no accumulation (bit-identical to the pre-T16 DDP path).
+    let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
     let seq_len: usize = flag(&args, "--seq", 64);
     let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
     let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
@@ -82,6 +88,13 @@ 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);
+    let sft_tsv = args.iter().any(|a| a == "--sft-tsv");
+    // Dropout (Phase T18/T21): residual-path dropout prob, active at training time
+    // only (inverted scaling), identity at eval/sampling/export. Default 0 = off
+    // (forward graph bit-identical to the no-dropout path). Mirrors bin/train; the
+    // train_rank loop calls model.train() each step so dropout is actually live
+    // under DDP (T21 wired this — the launcher previously never set training mode).
+    let dropout: f32 = flag(&args, "--dropout", 0.0f32);
     // bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
     // activations. Opt-in; default fp32 reproduces v0–v4 numerics.
     let bf16 = args.iter().any(|a| a == "--bf16");
@@ -89,11 +102,19 @@ fn main() {
     // rank checkpoints its own forward/backward; exact grads, lower peak activation
     // memory (lets dim1024 batch32 fit). Opt-in; default off.
     let recompute = args.iter().any(|a| a == "--recompute");
+    // Fused flash-attention (Phase T14): single fused SDPA kernel, online softmax,
+    // no materialized [bh,S,S] scores. Opt-in; default off keeps the composed path.
+    let flash = args.iter().any(|a| a == "--flash");
     let ckpt: Option<PathBuf> = args
         .iter()
         .position(|a| a == "--ckpt")
         .and_then(|i| args.get(i + 1))
         .map(PathBuf::from);
+    let init_ckpt: Option<PathBuf> = args
+        .iter()
+        .position(|a| a == "--init-ckpt")
+        .and_then(|i| args.get(i + 1))
+        .map(PathBuf::from);
 
     // Use every visible GPU as a rank (CUDA_VISIBLE_DEVICES selects the set;
     // device ordinals are 0..count within it).
@@ -114,12 +135,19 @@ fn main() {
     );
 
     // Reuse the cached token-id stream (v1's u16 cache); never re-tokenize 2GB.
-    let corpus = Corpus::load_cached(&tok_path, &corpus_path);
+    let corpus = if sft_tsv {
+        Corpus::load_sft_tsv_cached(&tok_path, &corpus_path)
+    } else {
+        Corpus::load_cached(&tok_path, &corpus_path)
+    };
     println!(
         "corpus: {} tokens, vocab {}",
         corpus.len(),
         corpus.vocab_size
     );
+    if sft_tsv {
+        println!("SFT TSV: ON (assistant-only loss via ignore-index labels)");
+    }
     let vocab = corpus.vocab_size;
     // Hold out a tail slice for validation (rank 0 evaluates on it).
     let (train_corpus, valid) = if val_tokens > 0 {
@@ -130,13 +158,16 @@ fn main() {
         (corpus, None)
     };
 
-    let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
+    let mut cfg =
+        Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
+    cfg.dropout = dropout;
     println!(
-        "model: dim {} layers {} heads {} head_dim {} ffn {} → core {:.3}M params \
+        "model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
          (+ embed/lm {:.2}M = {:.2}M total)",
         cfg.dim,
         cfg.n_layers,
         cfg.n_heads,
+        cfg.num_kv_heads,
         cfg.head_dim,
         cfg.ffn_hidden,
         cfg.core_params() as f32 / 1e6,
@@ -147,6 +178,7 @@ fn main() {
     let dcfg = DdpConfig {
         seq_len,
         batch_size: batch,
+        accum_steps,
         steps,
         schedule: LrSchedule {
             max_lr,
@@ -164,8 +196,9 @@ fn main() {
     };
 
     println!(
-        "training: {steps} steps, seq {seq_len}, global batch {batch}, lr {max_lr:.1e}→{min_lr:.1e}, \
-         eval every {eval_every}"
+        "training: {steps} steps, seq {seq_len}, global batch {batch} × accum {accum_steps} = \
+         effective global batch {}, lr {max_lr:.1e}→{min_lr:.1e}, eval every {eval_every}",
+        batch * accum_steps
     );
 
     if bf16 {
@@ -174,6 +207,16 @@ fn main() {
     if recompute {
         println!("activation recompute: ON (per-block gradient checkpointing)");
     }
+    if 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)");
+    }
+    if let Some(path) = &init_ckpt {
+        println!("init checkpoint: {}", path.display());
+    }
+    let init_ckpt_for_ranks = init_ckpt.clone();
     let results = launch(
         &devices,
         &train_corpus,
@@ -187,6 +230,13 @@ fn main() {
             if recompute {
                 m = m.with_recompute(true);
             }
+            if flash {
+                m = m.with_flash(true);
+            }
+            if let Some(path) = &init_ckpt_for_ranks {
+                xtrain_train::checkpoint::load_into(path, &m.params())
+                    .expect("load init checkpoint");
+            }
             m
         },
     );

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

@@ -0,0 +1,215 @@
+//! Process-per-GPU DDP launcher / worker (Phase T17, torchrun-style).
+//!
+//! ONE binary, two modes (it self-detects via `XTRAIN_RANK`):
+//!  - **launcher** (env unset): mints the NCCL `ncclUniqueId`, then spawns one
+//!    WORKER process per visible GPU, re-execing this same binary with the same
+//!    argv plus `XTRAIN_{RANK,WORLD,LOCAL_RANK,NCCL_ID}` env, and waits for them.
+//!  - **worker** (`XTRAIN_RANK` set): binds its GPU (→ its own CUDA context),
+//!    inits NCCL with the launcher-supplied id, builds its model, runs
+//!    `train_rank` — the T8 training step reused UNCHANGED.
+//!
+//! Versus `train_ddp` (thread-per-GPU, kept as the regression baseline) the ONLY
+//! difference is the launch model + cross-process UniqueId bootstrap. CLI flags
+//! 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
+//!   CUDA_VISIBLE_DEVICES=0,1,2,3 cargo run -p xtrain-distributed --release \
+//!       --bin train_ddp_mp -- /opt/wjh/models/gpt2/tokenizer.json \
+//!       data/tinystories-valid-3mb.txt \
+//!       --dim 384 --heads 12 --head-dim 32 --layers 12 --ffn 1536 \
+//!       --steps 200 --batch 128 --seq 256
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("train_ddp_mp: built without CUDA (no_cuda); run on a GPU host (dash5).");
+}
+
+#[cfg(not(no_cuda))]
+use std::path::PathBuf;
+
+// A flag like `--dim 384`: scan argv for `name`, parse the following token.
+#[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 main() {
+    use xtrain_cuda::device;
+    use xtrain_distributed::DdpConfig;
+    use xtrain_distributed::proc::{ModelOpts, launch_processes, run_worker, worker_env};
+    use xtrain_model::Config;
+    use xtrain_train::data::Corpus;
+    use xtrain_train::schedule::LrSchedule;
+
+    let args: Vec<String> = std::env::args().collect();
+
+    // ── Launcher mode: no XTRAIN_RANK in env → spawn one worker per visible GPU.
+    let env = worker_env();
+    if env.is_none() {
+        let count = device::device_count().expect("device_count");
+        assert!(count > 0, "no CUDA device visible");
+        let world = count as usize;
+        // Forward the full argv (minus argv[0]) to each worker verbatim.
+        let extra: Vec<String> = args[1..].to_vec();
+        println!("DDP (process-per-GPU): launching {world} worker processes (one per visible GPU)");
+        match launch_processes(world, &extra) {
+            Ok(()) => {}
+            Err(e) => {
+                eprintln!("launcher: {e}");
+                std::process::exit(1);
+            }
+        }
+        return;
+    }
+    let env = env.unwrap();
+
+    // ── Worker mode: build config from the forwarded argv, then train this rank.
+    // First two non-flag positionals: tokenizer.json, corpus.txt.
+    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
+    let tok_path = positionals
+        .first()
+        .map(|s| PathBuf::from(s.as_str()))
+        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
+    let corpus_path = positionals
+        .get(1)
+        .map(|s| PathBuf::from(s.as_str()))
+        .unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));
+
+    // Architecture (scaling-ladder rung). Defaults = v0-baseline tiny config.
+    let n_heads = flag(&args, "--heads", 2usize);
+    let head_dim = flag(&args, "--head-dim", 16usize);
+    let n_layers = flag(&args, "--layers", 4usize);
+    let ffn = flag(&args, "--ffn", 64usize);
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
+    let dim_flag = flag(&args, "--dim", 0usize);
+    if dim_flag != 0 && dim_flag != n_heads * head_dim {
+        eprintln!(
+            "warning: --dim {dim_flag} != heads*head_dim {}; using {}",
+            n_heads * head_dim,
+            n_heads * head_dim
+        );
+    }
+
+    // Optimization knobs (mirror train_ddp).
+    let steps: usize = flag(&args, "--steps", 100);
+    let batch: usize = flag(&args, "--batch", 16);
+    let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
+    let seq_len: usize = flag(&args, "--seq", 64);
+    let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
+    let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
+    let weight_decay: f32 = flag(&args, "--wd", 0.1);
+    let max_grad_norm: f32 = flag(&args, "--clip", 1.0);
+    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"),
+        flash: args.iter().any(|a| a == "--flash"),
+    };
+    let ckpt: Option<PathBuf> = args
+        .iter()
+        .position(|a| a == "--ckpt")
+        .and_then(|i| args.get(i + 1))
+        .map(PathBuf::from);
+
+    assert_eq!(
+        batch % env.world,
+        0,
+        "global batch {batch} not divisible by world {}",
+        env.world
+    );
+
+    // Each worker loads the corpus independently (read-only u16 cache hit → cheap).
+    let corpus = Corpus::load_cached(&tok_path, &corpus_path);
+    let vocab = corpus.vocab_size;
+    let (train_corpus, valid): (Corpus, Option<Corpus>) = if val_tokens > 0 {
+        let (t, v) = corpus.split_tail(val_tokens);
+        (t, Some(v))
+    } else {
+        (corpus, None)
+    };
+
+    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!(
+            "model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
+             (+ embed/lm {:.2}M = {:.2}M total) | world={} mode=process-per-GPU",
+            cfg.dim,
+            cfg.n_layers,
+            cfg.n_heads,
+            cfg.num_kv_heads,
+            cfg.head_dim,
+            cfg.ffn_hidden,
+            cfg.core_params() as f32 / 1e6,
+            (cfg.num_params() - cfg.core_params()) as f32 / 1e6,
+            cfg.num_params() as f32 / 1e6,
+            env.world,
+        );
+        if opts.bf16 {
+            println!("bf16 mixed precision: ON (fp32 master weights)");
+        }
+        if opts.recompute {
+            println!("activation recompute: ON (per-block gradient checkpointing)");
+        }
+        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 {
+        seq_len,
+        batch_size: batch,
+        accum_steps,
+        steps,
+        schedule: LrSchedule {
+            max_lr,
+            min_lr,
+            warmup: (steps / 20).max(5),
+            total: steps,
+        },
+        weight_decay,
+        max_grad_norm,
+        log_every: 50,
+        seed: 42,
+        eval_every,
+        eval_batches,
+        ckpt_path: ckpt.clone(),
+    };
+
+    let res = run_worker(&env, cfg, opts, &train_corpus, valid.as_ref(), &dcfg);
+
+    if env.rank == 0 {
+        let start = res.losses.first().copied().unwrap_or(0.0);
+        let end = res.losses.last().copied().unwrap_or(0.0);
+        println!("train loss: start {start:.4} → end {end:.4}");
+        if let Some(best) = res.best_val {
+            println!("best val loss: {best:.4}");
+        }
+        if let Some((s, v)) = res.evals.last() {
+            println!("final val loss (step {s}): {v:.4}");
+        }
+        if let Some(path) = &ckpt {
+            println!("best-val checkpoint → {}", path.display());
+        }
+    }
+}

crates/xtrain-distributed/src/ddp.rs

@@ -35,6 +35,13 @@ pub struct DdpConfig {
     pub seq_len: usize,
     /// Global batch size; must be divisible by the world size.
     pub batch_size: usize,
+    /// Micro-batch gradient accumulation (Phase T16): each optimizer step
+    /// accumulates grads over `accum_steps` micro-batches, giving an EFFECTIVE
+    /// global batch of `accum_steps × batch_size`. The cross-rank all-reduce
+    /// fires ONLY at the accumulation boundary (after the last micro-step) —
+    /// intermediate micro-steps skip the NCCL collective entirely. `1` = no
+    /// accumulation (bit-identical to the pre-T16 DDP path).
+    pub accum_steps: usize,
     pub steps: usize,
     pub schedule: LrSchedule,
     pub weight_decay: f32,
@@ -96,6 +103,7 @@ pub fn train_rank(
     // (sum across ranks, /world) then gives Σ_global/(world·b_local) = Σ_global/
     // B_global — already the global-batch mean — so the clip pre-scale is 1.0.
     let batch_local = cfg.batch_size / ctx.world;
+    let accum = cfg.accum_steps.max(1);
     let start = Instant::now();
     let mut tokens_seen: u64 = 0;
     // Rank 0 owns the held-out eval + best-val checkpoint (params are identical
@@ -105,36 +113,59 @@ pub fn train_rank(
     for step in 0..cfg.steps {
         let lr = cfg.schedule.lr(step);
 
-        // Draw the whole global batch from the shared RNG (same on every rank);
-        // collect only this rank's shard (global index % world == rank) and run it
-        // as ONE batched forward/backward. The union of shards == the single-GPU
-        // batch; each rank's backward yields its local mean (Σ_local / b_local).
-        let mut inputs = Vec::with_capacity(batch_local);
-        let mut targets_v = Vec::with_capacity(batch_local);
-        for i in 0..cfg.batch_size {
-            let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
-            if i % ctx.world == ctx.rank {
-                inputs.push(input);
-                targets_v.push(target);
+        // Accumulate grads over `accum` micro-batches, then ONE optimizer step
+        // (Phase T16). Per micro-batch: draw the whole micro global batch from the
+        // shared RNG (same on every rank), keep only this rank's shard (global index
+        // % world == rank), run it as ONE batched forward/backward. Each micro-loss
+        // is scaled by 1/accum before backward (the tape SUM-accumulates the scaled
+        // grads across the `accum` micro-backwards) so the boundary grad equals a
+        // single step over an `accum × batch_size` global batch. `accum == 1` skips
+        // the scale → bit-identical to the pre-T16 DDP path. The cross-rank
+        // all-reduce fires ONLY after the last micro-step (intermediate micro-steps
+        // are local-only, no NCCL).
+        let mut local_sum = 0.0f32; // Σ over micro of (local_mean · b_local)
+        // Training mode → dropout active (T18; no-op when cfg.dropout == 0). Set
+        // each step so it is restored after a periodic eval flips the model to eval
+        // mode (eval_loss calls model.eval() and does not restore). Mirrors the
+        // single-GPU loop's train/eval discipline — without this, DDP forwards run
+        // in the default eval (identity) mode and --dropout is silently ignored
+        // (the T21 launcher-wiring gap the V9-PILOT caught). Each micro-step's
+        // forward bumps the per-step seed → fresh masks.
+        model.train();
+        for _ in 0..accum {
+            let mut inputs = Vec::with_capacity(batch_local);
+            let mut targets_v = Vec::with_capacity(batch_local);
+            for i in 0..cfg.batch_size {
+                let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
+                if i % ctx.world == ctx.rank {
+                    inputs.push(input);
+                    targets_v.push(target);
+                }
             }
+            let ids = batched_ids_tensor(&inputs, device);
+            let targets = batched_ids_tensor(&targets_v, device);
+            let loss = model.loss_batched(&ids, &targets, batch_local);
+            local_sum += read_scalar(&loss) * batch_local as f32; // local mean·b_local
+            if accum == 1 {
+                loss.backward();
+            } else {
+                xtrain_autodiff::ops::scale(&loss, 1.0 / accum as f32).backward();
+            }
+            tokens_seen += (batch_local * cfg.seq_len) as u64;
         }
-        let ids = batched_ids_tensor(&inputs, device);
-        let targets = batched_ids_tensor(&targets_v, device);
-        let loss = model.loss_batched(&ids, &targets, batch_local);
-        let local_mean = read_scalar(&loss); // Σ_local / b_local
-        loss.backward();
-        tokens_seen += (batch_local * cfg.seq_len) as u64;
 
-        // AllReduce(sum) + /world the grads → every rank holds Σ_global/B_global
-        // (local means summed over ranks, /world = global mean). See note above.
+        // Accumulation boundary: ONE AllReduce(sum) + /world over the accumulated
+        // grads → every rank holds the effective-batch (accum·B_global) mean grad
+        // (the per-micro 1/accum scaling is already baked into each backward; the
+        // /world here is orthogonal to accum). Intermediate micro-steps issued NO
+        // NCCL — only this single boundary collective per optimizer step.
         ctx.all_reduce_average_grads(&params);
-        // Reported loss = global mean: sum the per-rank local sums (= mean·b_local)
-        // across ranks, /B_global. With equal b_local this is mean over ranks.
-        let step_loss =
-            all_reduce_loss(ctx, local_mean * batch_local as f32) / cfg.batch_size as f32;
+        // Reported loss = effective-batch mean: AllReduce(sum) the per-rank local
+        // sums across ranks, /(accum·B_global).
+        let step_loss = all_reduce_loss(ctx, local_sum) / (accum * cfg.batch_size) as f32;
         losses.push(step_loss);
 
-        // Grads are already the global-batch mean — just clip (pre-scale 1.0).
+        // Grads are already the effective-batch mean — just clip (pre-scale 1.0).
         let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
         opt.step(lr, &params);
         for p in &params {

crates/xtrain-distributed/src/lib.rs

@@ -18,8 +18,13 @@
 
 pub mod ddp;
 pub mod ffi;
+pub mod proc;
 
 pub use ddp::{DdpConfig, DdpResult, build_model, launch, train_rank};
+pub use proc::{
+    ModelOpts, WorkerEnv, build_worker_model, hex_decode_unique_id, hex_encode_unique_id,
+    launch_processes, run_worker, worker_env,
+};
 
 use std::ffi::c_void;
 

crates/xtrain-distributed/src/proc.rs

@@ -0,0 +1,200 @@
+//! Process-per-GPU DDP launcher + worker (Phase T17, torchrun-style).
+//!
+//! T8's DDP is single-process, thread-per-GPU: N rank threads share ONE CUDA
+//! primary context, so much of the driver work (kernel launch, cuBLAS handle,
+//! stream queueing) serializes at the context level — the residual ~5×@8
+//! non-linearity left after T11's allocator fix (see docs/10 / KI-5).
+//!
+//! Process-per-GPU gives each rank its OWN OS process and OWN CUDA context, so
+//! those driver calls no longer queue in a shared context. Only the LAUNCH model
+//! and the cross-process NCCL bootstrap change; the training step
+//! (`train_rank` → grad all-reduce → local AdamW) and the consistency argument
+//! are reused from T8 UNCHANGED.
+//!
+//! UniqueId rendezvous: the LAUNCHER (the common parent of every worker) mints
+//! the `ncclUniqueId` once, hex-encodes it, and injects it into each worker's env
+//! at spawn time. No shared file / TCP server / polling — the id is atomically
+//! present before the child exists, so there is no "id not ready yet" race. This
+//! is the simplest single-node mechanism (see docs/16).
+
+use std::path::PathBuf;
+use std::process::{Command, Stdio};
+
+use xtrain_model::{Config, TinyTransformer};
+use xtrain_tensor::{DType, Device};
+use xtrain_train::data::Corpus;
+
+use crate::ddp::{DdpConfig, DdpResult, build_model, train_rank};
+use crate::ffi::NcclUniqueId;
+use crate::{DdpContext, get_unique_id};
+
+// Env keys the launcher sets on every spawned worker (torchrun-style: a worker
+// detects its role by the presence of `XTRAIN_RANK`).
+pub const ENV_RANK: &str = "XTRAIN_RANK";
+pub const ENV_WORLD: &str = "XTRAIN_WORLD";
+pub const ENV_LOCAL_RANK: &str = "XTRAIN_LOCAL_RANK";
+pub const ENV_NCCL_ID: &str = "XTRAIN_NCCL_ID";
+
+/// Hex-encode the 128-byte `ncclUniqueId` for env transport (128 B → 256 chars,
+/// well under any env-var length limit). `c_char` is signed on this target, so
+/// reinterpret the bytes as `u8` first.
+pub fn hex_encode_unique_id(id: &NcclUniqueId) -> String {
+    let mut s = String::with_capacity(256);
+    for &b in &id.internal {
+        s.push_str(&format!("{:02x}", b as u8));
+    }
+    s
+}
+
+/// Inverse of [`hex_encode_unique_id`]: parse 256 hex chars back into the
+/// 128-byte opaque blob. Panics on malformed input (the launcher always writes a
+/// well-formed value, so a bad value means a corrupted env).
+pub fn hex_decode_unique_id(hex: &str) -> NcclUniqueId {
+    assert_eq!(
+        hex.len(),
+        256,
+        "NCCL id hex must be 256 chars, got {}",
+        hex.len()
+    );
+    let mut id = NcclUniqueId::default();
+    for (i, slot) in id.internal.iter_mut().enumerate() {
+        let byte = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16).expect("NCCL id hex byte parse");
+        *slot = byte as std::os::raw::c_char;
+    }
+    id
+}
+
+/// Spawn `world` worker processes (re-exec of the current binary with the same
+/// argv), each pinned to one GPU via `XTRAIN_LOCAL_RANK`, and wait for all of
+/// them. The launcher mints the `ncclUniqueId` and injects it (hex) into every
+/// worker's env, so the cross-process NCCL bootstrap needs no shared file/TCP.
+///
+/// Returns `Ok(())` iff every worker exits 0; otherwise an error naming the first
+/// failing rank (so the caller — `main` / a test — can propagate a non-zero exit).
+/// `extra_args` is forwarded to each worker verbatim (so all training hyper-params
+/// pass straight through); the workers inherit the launcher's env (incl.
+/// `CUDA_VISIBLE_DEVICES`) plus the four `XTRAIN_*` keys.
+pub fn launch_processes(world: usize, extra_args: &[String]) -> Result<(), String> {
+    let exe = std::env::current_exe().map_err(|e| format!("current_exe: {e}"))?;
+    let id = get_unique_id();
+    let id_hex = hex_encode_unique_id(&id);
+
+    let mut children = Vec::with_capacity(world);
+    for rank in 0..world {
+        let child = Command::new(&exe)
+            .args(extra_args)
+            .env(ENV_RANK, rank.to_string())
+            .env(ENV_WORLD, world.to_string())
+            // Single node: local rank == global rank == device ordinal within the
+            // visible set. (Multi-node would split these; see docs/16 follow-up.)
+            .env(ENV_LOCAL_RANK, rank.to_string())
+            .env(ENV_NCCL_ID, &id_hex)
+            // Workers inherit stdout/stderr so rank 0's training log surfaces.
+            .stdout(Stdio::inherit())
+            .stderr(Stdio::inherit())
+            .spawn()
+            .map_err(|e| format!("spawn worker rank {rank}: {e}"))?;
+        children.push((rank, child));
+    }
+
+    let mut first_err: Option<String> = None;
+    for (rank, mut child) in children {
+        let status = child
+            .wait()
+            .map_err(|e| format!("wait worker rank {rank}: {e}"))?;
+        if !status.success() && first_err.is_none() {
+            first_err = Some(format!("worker rank {rank} exited with {status}"));
+        }
+    }
+    match first_err {
+        Some(e) => Err(e),
+        None => Ok(()),
+    }
+}
+
+/// The four `XTRAIN_*` values a worker reads from its env. Present iff this
+/// process was spawned by [`launch_processes`].
+pub struct WorkerEnv {
+    pub rank: usize,
+    pub world: usize,
+    pub local_rank: u32,
+    pub id: NcclUniqueId,
+}
+
+/// Read the worker env if this process is a spawned worker (i.e. `XTRAIN_RANK`
+/// is set), else `None` (this process is the launcher).
+pub fn worker_env() -> Option<WorkerEnv> {
+    let rank: usize = std::env::var(ENV_RANK).ok()?.parse().ok()?;
+    let world: usize = std::env::var(ENV_WORLD)
+        .expect("XTRAIN_WORLD set with XTRAIN_RANK")
+        .parse()
+        .expect("XTRAIN_WORLD parse");
+    let local_rank: u32 = std::env::var(ENV_LOCAL_RANK)
+        .expect("XTRAIN_LOCAL_RANK set with XTRAIN_RANK")
+        .parse()
+        .expect("XTRAIN_LOCAL_RANK parse");
+    let id_hex = std::env::var(ENV_NCCL_ID).expect("XTRAIN_NCCL_ID set with XTRAIN_RANK");
+    let id = hex_decode_unique_id(&id_hex);
+    Some(WorkerEnv {
+        rank,
+        world,
+        local_rank,
+        id,
+    })
+}
+
+/// Per-worker model construction knobs (the opt-in feature flags the launcher
+/// forwards). Mirrors the closure `train_ddp` passes to the thread-per-GPU
+/// `launch`, but here it runs once in this worker's own process/context.
+#[derive(Clone, Copy, Default)]
+pub struct ModelOpts {
+    pub bf16: bool,
+    pub recompute: bool,
+    pub flash: bool,
+}
+
+/// Run this worker: bind its GPU (→ its own CUDA context), init NCCL with the
+/// launcher-supplied id, build its model with the deterministic init (same as
+/// every rank + the single-GPU baseline), and run `train_rank`. Reuses the T8
+/// training step verbatim — the only difference from thread-per-GPU is how this
+/// rank was started and how it got the `UniqueId`.
+///
+/// `valid` is the held-out corpus for rank 0's periodic eval (pass `None` on
+/// other ranks or when `cfg.eval_every == 0`).
+pub fn run_worker(
+    env: &WorkerEnv,
+    cfg: Config,
+    opts: ModelOpts,
+    corpus: &Corpus,
+    valid: Option<&Corpus>,
+    dcfg: &DdpConfig,
+) -> DdpResult {
+    // Binding the device here establishes this process's own CUDA primary context.
+    let ctx = DdpContext::init(env.rank, env.world, env.id, env.local_rank);
+    let device = Device::Cuda(env.local_rank);
+    let model = build_worker_model(cfg, opts, device);
+    let v = if env.rank == 0 { valid } else { None };
+    train_rank(&ctx, &model, device, corpus, v, dcfg)
+}
+
+/// Build the worker's model with the deterministic `build_model` init + the
+/// opt-in feature flags. Shared by `run_worker` and the test worker.
+pub fn build_worker_model(cfg: Config, opts: ModelOpts, device: Device) -> TinyTransformer {
+    let mut m = build_model(cfg, device);
+    if opts.bf16 {
+        m = m.with_compute_dtype(DType::BF16);
+    }
+    if opts.recompute {
+        m = m.with_recompute(true);
+    }
+    if opts.flash {
+        m = m.with_flash(true);
+    }
+    m
+}
+
+/// Convenience: the directory tests/bins can stash per-rank result dumps in
+/// (a worker writes its loss/params there; the launching test reads them back).
+pub fn rank_dump_path(dir: &std::path::Path, rank: usize) -> PathBuf {
+    dir.join(format!("rank{rank}.dump"))
+}

crates/xtrain-distributed/tests/ddp_correctness.rs

@@ -27,6 +27,7 @@ fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
         .collect();
     Corpus {
         tokens,
+        labels: None,
         vocab_size: vocab,
     }
 }
@@ -38,6 +39,53 @@ fn test_config(vocab: usize) -> Config {
     cfg
 }
 
+/// Run `cfg`/`dcfg` as a DDP job over `devices` (the same launcher path as
+/// production — `DdpContext::init` + `train_rank` per rank) and return rank 0's
+/// (loss trace, final params on host, final `is_training()` flag). `cfg` carries
+/// the dropout prob; `dcfg` carries the loop knobs. Caller asserts.
+///
+/// `world == 1` is the deterministic path: `all_reduce_average_grads` short-circuits
+/// (no NCCL collective), so the run is bit-reproducible — used for the bit-identity
+/// gate. `world >= 2` exercises the real cross-rank NCCL all-reduce, which is not
+/// bit-reproducible run-to-run on this PCIe box (KI-5), so those gates use the same
+/// ULP/relative tolerances as the rest of this file.
+fn run_ddp(
+    devices: &[u32],
+    cfg: Config,
+    corpus: &Corpus,
+    valid: Option<&Corpus>,
+    dcfg: &DdpConfig,
+) -> (Vec<f32>, Vec<Vec<f32>>, bool) {
+    let world = devices.len();
+    let id = get_unique_id();
+    let results: Vec<(Vec<f32>, Vec<Vec<f32>>, bool)> = std::thread::scope(|s| {
+        let handles: Vec<_> = devices
+            .iter()
+            .enumerate()
+            .map(|(rank, &dev)| {
+                let dcfg = dcfg.clone();
+                let corpus = &corpus;
+                s.spawn(move || {
+                    let ctx = DdpContext::init(rank, world, id, dev);
+                    let device = Device::Cuda(dev);
+                    let model = build_model(cfg, device);
+                    // Only rank 0 holds the val corpus (mirrors launch()).
+                    let v = if rank == 0 { valid } else { None };
+                    let res = train_rank(&ctx, &model, device, corpus, v, &dcfg);
+                    let host = model
+                        .params()
+                        .iter()
+                        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+                        .collect::<Vec<_>>();
+                    (res.losses, host, model.is_training())
+                })
+            })
+            .collect();
+        handles.into_iter().map(|h| h.join().unwrap()).collect()
+    });
+    results.into_iter().next().unwrap()
+}
+
 // Single-GPU baseline: the SAME loop as the DDP rank but world=1, so the global
 // batch is processed on one device. Returns (loss trace, final params on host).
 fn run_single_gpu(cfg: Config, corpus: &Corpus, dcfg: &DdpConfig) -> (Vec<f32>, Vec<Vec<f32>>) {
@@ -94,6 +142,7 @@ fn ddp_matches_single_gpu_and_params_consistent() {
     let dcfg = DdpConfig {
         seq_len: 32,
         batch_size: 8, // global; 4 per rank with world=2
+        accum_steps: 1,
         steps,
         schedule: LrSchedule {
             max_lr: 3e-3,
@@ -195,6 +244,127 @@ fn ddp_matches_single_gpu_and_params_consistent() {
     assert!(max_sdiff < 1e-2, "DDP params diverged from single-GPU");
 }
 
+#[test]
+fn ddp_with_accum_matches_single_gpu_big_batch() {
+    // T16: DDP + gradient accumulation must match a single-GPU big-batch baseline
+    // of the SAME effective batch. world=2, accum=2, per-rank micro-batch 2 →
+    // effective global batch = world·accum·b_local = 2·2·2 = 8. Compared against a
+    // single-GPU run with batch 8, accum 1 (the big-batch baseline). The all-reduce
+    // fires only at the accumulation boundary (once per optimizer step, not per
+    // micro-step) — enforced by the train_rank implementation; the load-bearing
+    // gate here is that loss + final params still match the big-batch baseline.
+    let world = 2usize;
+    if device::device_count().unwrap_or(0) < world as i32 {
+        eprintln!("skip: need >= {world} GPUs");
+        return;
+    }
+
+    let vocab = 64usize;
+    let cfg = test_config(vocab);
+    let corpus = synth_corpus(vocab, 4096);
+    let steps = 20usize;
+    let effective_batch = 8usize; // world(2) · accum(2) · b_local(2)
+    let sched = LrSchedule {
+        max_lr: 3e-3,
+        min_lr: 3e-4,
+        warmup: 3,
+        total: steps,
+    };
+
+    // Single-GPU big-batch baseline: world=1, accum=1, batch = effective_batch.
+    let baseline_cfg = DdpConfig {
+        seq_len: 32,
+        batch_size: effective_batch,
+        accum_steps: 1,
+        steps,
+        schedule: sched,
+        weight_decay: 0.1,
+        max_grad_norm: 1.0,
+        log_every: 1_000_000,
+        seed: 7,
+        eval_every: 0,
+        eval_batches: 0,
+        ckpt_path: None,
+    };
+    let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &baseline_cfg);
+
+    // DDP + accumulation: world=2, accum=2 → per-rank micro-batch = batch/world = 2.
+    let ddp_cfg = DdpConfig {
+        batch_size: effective_batch / 2, // per-step global batch; ×accum = effective
+        accum_steps: 2,
+        ..baseline_cfg
+    };
+    let devices = [0u32, 1u32];
+    let id = get_unique_id();
+    let results: Vec<(Vec<f32>, Vec<Vec<f32>>)> = std::thread::scope(|s| {
+        let handles: Vec<_> = devices
+            .iter()
+            .enumerate()
+            .map(|(rank, &dev)| {
+                let ddp_cfg = ddp_cfg.clone();
+                let corpus = &corpus;
+                s.spawn(move || {
+                    let ctx = DdpContext::init(rank, world, id, dev);
+                    let device = Device::Cuda(dev);
+                    let model = build_model(cfg, device);
+                    let res = train_rank(&ctx, &model, device, corpus, None, &ddp_cfg);
+                    let host = model
+                        .params()
+                        .iter()
+                        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+                        .collect::<Vec<_>>();
+                    (res.losses, host)
+                })
+            })
+            .collect();
+        handles.into_iter().map(|h| h.join().unwrap()).collect()
+    });
+
+    let (ddp_losses, ddp_p0) = &results[0];
+    let (_, ddp_p1) = &results[1];
+
+    // (a) Loss trajectory matches the single-GPU big-batch baseline.
+    let mut max_rel = 0.0f32;
+    for (s, d) in single_losses.iter().zip(ddp_losses) {
+        max_rel = max_rel.max((s - d).abs() / s.abs().max(1e-6));
+    }
+    println!(
+        "DDP+accum(w2·a2·b2) vs single-GPU big-batch(8): single[last]={:.6} ddp[last]={:.6} max_rel={max_rel:.2e}",
+        single_losses.last().unwrap(),
+        ddp_losses.last().unwrap()
+    );
+    assert!(
+        max_rel < 1e-3,
+        "DDP+accum loss diverged from big-batch baseline: {max_rel:.3e}"
+    );
+
+    // (b) Cross-rank parameter agreement (same KI-5 ULP tolerance as the base test).
+    let mut max_pdiff = 0.0f32;
+    for (a, b) in ddp_p0.iter().zip(ddp_p1) {
+        for (x, y) in a.iter().zip(b) {
+            max_pdiff = max_pdiff.max((x - y).abs());
+        }
+    }
+    println!("DDP+accum cross-rank max |param diff| = {max_pdiff:.3e}");
+    assert!(
+        max_pdiff < 1e-6,
+        "ranks' params drifted apart: {max_pdiff:.3e}"
+    );
+
+    // (c) Final params match single-GPU big-batch within fp tolerance.
+    let mut max_sdiff = 0.0f32;
+    for (a, b) in ddp_p0.iter().zip(&single_params) {
+        for (x, y) in a.iter().zip(b) {
+            max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
+        }
+    }
+    println!("DDP+accum vs single-GPU big-batch max rel |param diff| = {max_sdiff:.3e}");
+    assert!(
+        max_sdiff < 1e-2,
+        "DDP+accum params diverged from big-batch baseline"
+    );
+}
+
 #[test]
 fn ddp_throughput_scaling() {
     let max_gpus = device::device_count().unwrap_or(0) as usize;
@@ -230,6 +400,7 @@ fn ddp_throughput_scaling() {
         let dcfg = DdpConfig {
             seq_len,
             batch_size: per_gpu_batch * world,
+            accum_steps: 1,
             steps,
             schedule: LrSchedule {
                 max_lr: 1e-3,
@@ -263,3 +434,181 @@ fn ddp_throughput_scaling() {
         );
     }
 }
+
+/// T21 regression: prove dropout is actually LIVE under DDP (with `p>0`), and that
+/// `p=0` is bit-identical to the no-dropout path. Guards the V9-PILOT launcher-
+/// wiring gap — `train_ddp` had no `--dropout` flag and `train_rank` never called
+/// `model.train()`, so under DDP every forward ran in the default eval mode and
+/// dropout was a silent identity regardless of config. Op/single-GPU tests never
+/// exercised dropout-under-DDP, so it slipped through; this test runs the REAL
+/// launcher path (`DdpContext::init` + `train_rank`).
+///
+/// On the pre-T21 code, both load-bearing gates FAIL: GATE B (p>0 trace would be
+/// bit-identical to p=0 — model stuck in eval mode → dropout is identity) and GATE C
+/// (`is_training()` would be false after the run).
+///
+/// p=0 regression (GATE A) is checked at `world=1`, ONE step, where the NCCL
+/// all-reduce short-circuits: the p=0 FORWARD is byte-identical to no-dropout so the
+/// loss is BIT-IDENTICAL (== 0.0), and the post-step params match within the engine's
+/// atomicAdd backward-reduction ULP floor (< 1e-7, dropout-independent — the
+/// fresh-train md5 caveat). The cross-rank NCCL all-reduce (`world>=2`) is not
+/// bit-reproducible run-to-run on this PCIe box (KI-5, observed ≤~2.4e-7), so the
+/// `world=2` p=0-vs-no-dropout check (GATE A2) uses the same KI-5 ULP tolerance as the
+/// rest of this file. GATE B's live-dropout signal (>1e-3) sits ~4 orders of magnitude
+/// above every noise floor here, so it carries the load.
+#[test]
+fn ddp_dropout_is_live_and_p0_bit_identical() {
+    if device::device_count().unwrap_or(0) < 2 {
+        eprintln!("skip: need >= 2 GPUs");
+        return;
+    }
+
+    let vocab = 64usize;
+    let corpus = synth_corpus(vocab, 4096);
+    let steps = 20usize;
+    // eval_every < steps so a periodic eval fires MID-run (flipping the model to
+    // eval mode via eval_loss → model.eval()). The per-step model.train() must
+    // restore training mode so dropout stays live across the eval boundary — this is
+    // exactly the train/eval discipline the pilot called out. A held-out slice gives
+    // rank 0 something to eval on.
+    let valid = synth_corpus(vocab, 512);
+    let base_dcfg = DdpConfig {
+        seq_len: 32,
+        batch_size: 8, // global; 4 per rank with world=2
+        accum_steps: 1,
+        steps,
+        schedule: LrSchedule {
+            max_lr: 3e-3,
+            min_lr: 3e-4,
+            warmup: 3,
+            total: steps,
+        },
+        weight_decay: 0.1,
+        max_grad_norm: 1.0,
+        log_every: 1_000_000, // silence per-step logging
+        seed: 7,
+        eval_every: 7, // fires at steps 6, 13, 19 — flips to eval mode mid-run
+        eval_batches: 4,
+        ckpt_path: None,
+    };
+
+    // --- GATE A: p=0 == no-dropout at world=1, ONE step (the deterministic scope). ---
+    // The regression guard for `--dropout 0`. ops::dropout(p=0) returns x.clone() (a
+    // graph no-op) regardless of training mode, so the p=0 FORWARD graph is byte-for-
+    // byte the no-dropout forward → loss[0] must be BIT-IDENTICAL (the load-bearing
+    // claim, asserted == 0.0). At world=1 the NCCL all-reduce short-circuits, and one
+    // step has no optimizer-state compounding; the only residual non-determinism is
+    // the engine's atomicAdd backward-reduction ORDER (the documented fresh-train md5
+    // caveat — dropout-INDEPENDENT, present with or without the dropout op), which
+    // moves the post-step params by a single grad ULP. So params are checked against
+    // that tight reduction floor (< 1e-7), the same nature as the cross-rank KI-5
+    // tolerance used elsewhere in this file — not a dropout signal. GATE B (live) has
+    // a >1e-3 signal, ~4 orders of magnitude above this floor, so it carries the load.
+    let d1 = [0u32];
+    let dcfg_1step = DdpConfig {
+        steps: 1,
+        eval_every: 0,
+        ..base_dcfg.clone()
+    };
+    let cfg_nodrop = test_config(vocab); // cfg.dropout defaults to 0.0
+    assert_eq!(cfg_nodrop.dropout, 0.0, "baseline cfg must have dropout 0");
+    let mut cfg_p0 = test_config(vocab);
+    cfg_p0.dropout = 0.0; // explicitly set p=0 — must not perturb anything
+    let (loss_nd1, params_nd1, _) = run_ddp(&d1, cfg_nodrop, &corpus, None, &dcfg_1step);
+    let (loss_p01, params_p01, _) = run_ddp(&d1, cfg_p0, &corpus, None, &dcfg_1step);
+    let max_loss_diff_1 = (loss_nd1[0] - loss_p01[0]).abs();
+    let max_param_diff_1 = params_nd1
+        .iter()
+        .zip(&params_p01)
+        .flat_map(|(a, b)| a.iter().zip(b).map(|(x, y)| (x - y).abs()))
+        .fold(0.0f32, f32::max);
+    println!(
+        "T21 GATE A (world=1, 1 step, p=0 vs no-dropout): |loss diff| = {max_loss_diff_1:.3e} \
+         (bit-identical forward), max |param diff| = {max_param_diff_1:.3e} (atomicAdd floor)"
+    );
+    assert_eq!(
+        max_loss_diff_1, 0.0,
+        "world=1 p=0 forward loss not bit-identical to no-dropout path"
+    );
+    assert!(
+        max_param_diff_1 < 1e-7,
+        "world=1 p=0 post-step params diverged from no-dropout beyond the atomicAdd \
+         reduction floor: {max_param_diff_1:.3e}"
+    );
+
+    // --- world=2 runs: real cross-rank NCCL all-reduce (the production path). ---
+    let d2 = [0u32, 1u32];
+    let mut cfg_p0_w2 = test_config(vocab);
+    cfg_p0_w2.dropout = 0.0;
+    let mut cfg_p_w2 = test_config(vocab);
+    cfg_p_w2.dropout = 0.2;
+    let (loss_p0_2, _params_p0_2, _) = run_ddp(&d2, cfg_p0_w2, &corpus, Some(&valid), &base_dcfg);
+    let (loss_p_2, _params_p_2, _) = run_ddp(&d2, cfg_p_w2, &corpus, Some(&valid), &base_dcfg);
+
+    // GATE A2 — under DDP (world=2), p=0 matches a separate no-dropout baseline within
+    // NCCL's run-to-run ULP noise (KI-5; the all-reduce is not bit-reproducible). This
+    // confirms enabling dropout=0 doesn't perturb the DDP path beyond that noise floor.
+    let (loss_nd_2, _, _) = run_ddp(&d2, test_config(vocab), &corpus, Some(&valid), &base_dcfg);
+    let max_loss_diff_2 = loss_nd_2
+        .iter()
+        .zip(&loss_p0_2)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0.0f32, f32::max);
+    println!(
+        "T21 GATE A2 (world=2 p=0 vs no-dropout, KI-5 noise): max |loss diff| = {max_loss_diff_2:.3e}"
+    );
+    assert!(
+        max_loss_diff_2 < 1e-6,
+        "world=2 p=0 diverged from no-dropout beyond NCCL noise: {max_loss_diff_2:.3e}"
+    );
+
+    // GATE B — dropout is LIVE with p>0 under DDP. If model.train() were not wired
+    // (the pre-T21 bug), the model would stay in eval mode and the p=0.2 forward would
+    // be IDENTITY → loss trace bit-identical to p=0 (diff at the ~1e-7 NCCL noise
+    // floor). A difference orders of magnitude above that proves dropout masks are
+    // actually applied during the training forward — and that they survive the mid-run
+    // eval flips (model.train() is re-asserted each step). Inverted scaling + masking
+    // perturbs every step, so the gap is large (>1e-3 ≫ KI-5 noise ~2.4e-7).
+    let max_live_diff = loss_p0_2
+        .iter()
+        .zip(&loss_p_2)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0.0f32, f32::max);
+    println!(
+        "T21 GATE B (dropout live, world=2): p0[last]={:.6} p0.2[last]={:.6} max |loss diff| = {max_live_diff:.3e}",
+        loss_p0_2.last().unwrap(),
+        loss_p_2.last().unwrap()
+    );
+    assert!(
+        max_live_diff > 1e-3,
+        "p=0.2 DDP loss trace matches p=0 — dropout is NOT live under DDP \
+         (model.train() not wired): max |loss diff| {max_live_diff:.3e}"
+    );
+
+    // No NaN/Inf in the p>0 run (dropout converges normally under DDP).
+    assert!(
+        loss_p_2.iter().all(|l| l.is_finite()),
+        "p=0.2 DDP loss has non-finite values"
+    );
+
+    // GATE C — train_rank actually sets TRAINING mode (direct, complementary proof of
+    // model.train() being wired). Use a dedicated short run with eval_every=0 so no
+    // eval fires: a model that finishes a training step in training mode proves
+    // train_rank called model.train(). (With eval enabled, eval_loss → model.eval()
+    // runs LAST on the final step and legitimately leaves the model in eval mode —
+    // same as the single-GPU loop — so is_training() after an eval-enabled run reflects
+    // the final eval, not the training-mode wiring. GATE B already proves dropout
+    // survives the mid-run eval flips via the per-step model.train() restore.) On the
+    // pre-T21 code is_training() stays false (model never left the default eval mode).
+    let dcfg_noeval = DdpConfig {
+        steps: 2,
+        eval_every: 0,
+        ..base_dcfg.clone()
+    };
+    let (_, _, train_flag) = run_ddp(&d1, cfg_p_w2, &corpus, None, &dcfg_noeval);
+    assert!(
+        train_flag,
+        "model not in training mode after a no-eval DDP run — model.train() not wired in train_rank"
+    );
+    println!("T21 GATE C (train_rank sets training mode): is_training() == true ✅");
+}

crates/xtrain-distributed/tests/ddp_proc.rs

@@ -0,0 +1,409 @@
+//! Process-per-GPU DDP acceptance (Phase T17). Gated to a GPU host; skips when
+//! fewer than 2 GPUs. Run with `--test-threads=1` (distributed tests deadlock if
+//! they contend for the same GPUs in parallel — known harness property).
+//!
+//! Self-launching: the test binary detects WORKER mode via `XTRAIN_RANK` (set by
+//! `launch_processes`). In worker mode it runs `run_worker` on a synthetic corpus
+//! and dumps its per-step loss trace + final params to a per-rank file; in normal
+//! mode it is the launcher — it runs the single-GPU baseline, spawns N worker
+//! processes (re-execing itself), reads their dumps back, and asserts:
+//!   (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))]
+
+use std::io::Write;
+use std::path::Path;
+
+use xtrain_cuda::device;
+use xtrain_distributed::proc::{launch_processes, rank_dump_path, worker_env};
+use xtrain_distributed::{DdpConfig, DdpContext, build_model, train_rank};
+use xtrain_model::{Config, batched_ids_tensor};
+use xtrain_optim::GpuAdamW;
+use xtrain_tensor::Device;
+use xtrain_train::clip::clip_grad_norm_gpu;
+use xtrain_train::data::Corpus;
+use xtrain_train::schedule::LrSchedule;
+
+// ── Shared fixture (identical on launcher + every worker, so they agree) ──────
+
+const VOCAB: usize = 64;
+const STEPS: usize = 20;
+
+fn synth_corpus() -> Corpus {
+    let tokens: Vec<i32> = (0..4096)
+        .map(|i| (i * 7 + 3) as i32 % VOCAB as i32)
+        .collect();
+    Corpus {
+        tokens,
+        labels: None,
+        vocab_size: VOCAB,
+    }
+}
+
+fn test_config() -> Config {
+    let mut cfg = Config::tiny();
+    cfg.vocab = VOCAB;
+    cfg.n_layers = 2;
+    cfg
+}
+
+fn dcfg(batch_size: usize) -> DdpConfig {
+    DdpConfig {
+        seq_len: 32,
+        batch_size,
+        accum_steps: 1,
+        steps: STEPS,
+        schedule: LrSchedule {
+            max_lr: 3e-3,
+            min_lr: 3e-4,
+            warmup: 3,
+            total: STEPS,
+        },
+        weight_decay: 0.1,
+        max_grad_norm: 1.0,
+        log_every: 1_000_000,
+        seed: 7,
+        eval_every: 0,
+        eval_batches: 0,
+        ckpt_path: None,
+    }
+}
+
+// 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() {
+    let Some(env) = worker_env() else { return };
+    let dump_dir = std::env::var(ENV_DUMP_DIR).expect("dump dir env");
+    // This is the worker body `run_worker` performs in production (init ctx →
+    // build deterministic model → train_rank). We train ONCE inline so we can dump
+    // both the loss trace AND the final params for the launcher to check; the
+    // 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);
+    // 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,
+        device,
+        &synth_corpus(),
+        None,
+        &dcfg(GLOBAL_BATCH),
+    );
+    let params: Vec<Vec<f32>> = model
+        .params()
+        .iter()
+        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+        .collect();
+    write_dump(&dump_dir, env.rank, &res.losses, &params);
+    std::process::exit(0);
+}
+
+fn write_dump(dir: &str, rank: usize, losses: &[f32], params: &[Vec<f32>]) {
+    let path = rank_dump_path(Path::new(dir), rank);
+    let mut f = std::fs::File::create(&path).expect("create dump");
+    // Line 1: losses (space-separated). Following lines: one param tensor each.
+    let loss_line: Vec<String> = losses.iter().map(|x| format!("{x:.8e}")).collect();
+    writeln!(f, "{}", loss_line.join(" ")).unwrap();
+    for p in params {
+        let line: Vec<String> = p.iter().map(|x| format!("{x:.8e}")).collect();
+        writeln!(f, "{}", line.join(" ")).unwrap();
+    }
+}
+
+fn read_dump(dir: &str, rank: usize) -> (Vec<f32>, Vec<Vec<f32>>) {
+    let path = rank_dump_path(Path::new(dir), rank);
+    let text = std::fs::read_to_string(&path).expect("read dump");
+    let mut lines = text.lines();
+    let losses: Vec<f32> = lines
+        .next()
+        .unwrap()
+        .split_whitespace()
+        .map(|s| s.parse().unwrap())
+        .collect();
+    let params: Vec<Vec<f32>> = lines
+        .map(|l| l.split_whitespace().map(|s| s.parse().unwrap()).collect())
+        .collect();
+    (losses, params)
+}
+
+// ── Single-GPU baseline (same loop as the DDP rank, world=1) ──────────────────
+
+fn run_single_gpu(cfg: Config, corpus: &Corpus, d: &DdpConfig) -> (Vec<f32>, Vec<Vec<f32>>) {
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+    let model = build_model(cfg, device);
+    let params = model.params();
+    let mut opt = GpuAdamW::new(d.weight_decay);
+    let mut rng = d.seed;
+    let mut losses = Vec::new();
+    for step in 0..d.steps {
+        let lr = d.schedule.lr(step);
+        let mut inputs = Vec::with_capacity(d.batch_size);
+        let mut targets_v = Vec::with_capacity(d.batch_size);
+        for _ in 0..d.batch_size {
+            let (input, target) = corpus.sample(d.seq_len, &mut rng);
+            inputs.push(input);
+            targets_v.push(target);
+        }
+        let ids = batched_ids_tensor(&inputs, device);
+        let targets = batched_ids_tensor(&targets_v, device);
+        let loss = model.loss_batched(&ids, &targets, d.batch_size);
+        losses.push(loss.value().to_device(Device::Cpu).as_slice::<f32>()[0]);
+        loss.backward();
+        clip_grad_norm_gpu(&params, d.max_grad_norm, 1.0);
+        opt.step(lr, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+    }
+    let host = params
+        .iter()
+        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+        .collect();
+    (losses, host)
+}
+
+// ── The test (launcher mode) ──────────────────────────────────────────────────
+
+#[test]
+fn proc_per_gpu_matches_single_gpu_and_thread_path() {
+    // If this process was spawned as a worker, do the worker job and exit before
+    // the test framework runs anything else.
+    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 cfg = test_config();
+    let corpus = synth_corpus();
+    let d = dcfg(GLOBAL_BATCH);
+
+    // (1) Single-GPU baseline over the global batch.
+    let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &d);
+
+    // (2) Thread-per-GPU path (T8 `launch`) — the regression baseline to match.
+    let thread_results =
+        xtrain_distributed::launch(&[0u32, 1u32], &corpus, None, &d, move |device| {
+            build_model(cfg, device)
+        });
+    let thread_losses = &thread_results[0].losses;
+
+    // (3) Process-per-GPU: spawn 2 worker processes (re-exec of this test binary),
+    //     each dumps its loss trace + final params to a temp dir.
+    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. 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
+    // worker process does the worker job and exits without touching other tests.
+    let worker_args = [
+        "--exact".to_string(),
+        "proc_per_gpu_matches_single_gpu_and_thread_path".to_string(),
+        "--test-threads=1".to_string(),
+        "--nocapture".to_string(),
+    ];
+    launch_processes(world, &worker_args).expect("worker processes failed");
+
+    let (proc_losses0, proc_p0) = read_dump(dump_dir.to_str().unwrap(), 0);
+    let (_proc_losses1, proc_p1) = read_dump(dump_dir.to_str().unwrap(), 1);
+
+    // (a) process-per-GPU loss matches single-GPU.
+    let max_rel_single = max_rel(&single_losses, &proc_losses0);
+    println!(
+        "proc-per-GPU vs single-GPU loss: single[last]={:.6} proc[last]={:.6} max_rel={max_rel_single:.2e}",
+        single_losses.last().unwrap(),
+        proc_losses0.last().unwrap()
+    );
+    assert!(
+        max_rel_single < 1e-3,
+        "proc-per-GPU loss diverged from single-GPU: {max_rel_single:.3e}"
+    );
+
+    // (c) process-per-GPU loss matches the thread-per-GPU path.
+    let max_rel_thread = max_rel(thread_losses, &proc_losses0);
+    println!(
+        "proc-per-GPU vs thread-per-GPU loss: thread[last]={:.6} proc[last]={:.6} max_rel={max_rel_thread:.2e}",
+        thread_losses.last().unwrap(),
+        proc_losses0.last().unwrap()
+    );
+    assert!(
+        max_rel_thread < 1e-3,
+        "proc-per-GPU loss diverged from thread-per-GPU: {max_rel_thread:.3e}"
+    );
+
+    // (b) cross-rank parameter agreement (KI-5 ULP tolerance).
+    let mut max_pdiff = 0.0f32;
+    for (a, b) in proc_p0.iter().zip(&proc_p1) {
+        for (x, y) in a.iter().zip(b) {
+            max_pdiff = max_pdiff.max((x - y).abs());
+        }
+    }
+    println!("proc-per-GPU cross-rank max |param diff| = {max_pdiff:.3e}");
+    assert!(
+        max_pdiff < 1e-6,
+        "ranks' params drifted apart: {max_pdiff:.3e}"
+    );
+
+    // Bonus sanity: proc-per-GPU final params vs single-GPU within fp tolerance.
+    let mut max_sdiff = 0.0f32;
+    for (a, b) in proc_p0.iter().zip(&single_params) {
+        for (x, y) in a.iter().zip(b) {
+            max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
+        }
+    }
+    println!("proc-per-GPU vs single-GPU max rel |param diff| = {max_sdiff:.3e}");
+    assert!(
+        max_sdiff < 1e-2,
+        "proc-per-GPU params diverged from single-GPU"
+    );
+
+    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)
+        .map(|(s, d)| (s - d).abs() / s.abs().max(1e-6))
+        .fold(0.0f32, f32::max)
+}

crates/xtrain-model/src/config.rs

@@ -10,8 +10,15 @@ pub struct Config {
     pub dim: usize,
     /// Number of decoder blocks.
     pub n_layers: usize,
-    /// Number of attention heads.
+    /// Number of attention (query) heads.
     pub n_heads: usize,
+    /// Number of key/value heads (Phase T15, GQA). Each KV head is shared by a
+    /// group of `n_heads / num_kv_heads` query heads (repeat_kv). Must divide
+    /// `n_heads`. `num_kv_heads == n_heads` (the default) = MHA, bit-identical to
+    /// the pre-T15 path; `num_kv_heads < n_heads` = real grouped-query attention,
+    /// shrinking the K/V projections to `num_kv_heads * head_dim` and exported as a
+    /// real `num_key_value_heads`.
+    pub num_kv_heads: usize,
     /// Per-head dimension (`dim / n_heads`).
     pub head_dim: usize,
     /// SwiGLU hidden width (gate/up project to this, down projects back).
@@ -37,6 +44,7 @@ impl Config {
             dim: n_heads * head_dim,
             n_layers: 2,
             n_heads,
+            num_kv_heads: n_heads, // default = MHA
             head_dim,
             ffn_hidden: 64,
             eps: 1e-5,
@@ -62,6 +70,7 @@ impl Config {
             dim: n_heads * head_dim,
             n_layers,
             n_heads,
+            num_kv_heads: n_heads, // default = MHA; set via with_kv_heads for GQA
             head_dim,
             ffn_hidden,
             eps: 1e-5,
@@ -70,6 +79,27 @@ impl Config {
         }
     }
 
+    /// Set the number of K/V heads (Phase T15, GQA). Builder-style so existing
+    /// `from_arch` call sites stay MHA unless they opt in. Asserts `num_kv_heads`
+    /// divides `n_heads`.
+    pub fn with_kv_heads(mut self, num_kv_heads: usize) -> Self {
+        assert!(num_kv_heads > 0, "num_kv_heads must be > 0");
+        assert_eq!(
+            self.n_heads % num_kv_heads,
+            0,
+            "n_heads {} not divisible by num_kv_heads {num_kv_heads}",
+            self.n_heads
+        );
+        self.num_kv_heads = num_kv_heads;
+        self
+    }
+
+    /// KV projection width (`num_kv_heads * head_dim`). For GQA this is smaller than
+    /// `dim`; for MHA it equals `dim`.
+    pub fn kv_dim(&self) -> usize {
+        self.num_kv_heads * self.head_dim
+    }
+
     /// Transformer-core parameter count: everything except the token embedding and
     /// the LM head (the two `vocab × dim` tables). This is the figure the scaling
     /// ladder is sized against — the 50257-vocab embed+lm_head adds a fixed ~25M on
@@ -82,7 +112,8 @@ impl Config {
     pub fn num_params(&self) -> usize {
         let per_layer = 2 * self.dim                       // 2 rmsnorm gammas
             + 2 * self.head_dim                             // q/k per-head norm gammas
-            + 3 * self.dim * self.dim                       // q/k/v proj
+            + self.dim * self.dim                           // q proj [dim,dim]
+            + 2 * self.dim * self.kv_dim()                  // k/v proj [dim,kv_dim] (GQA: smaller)
             + self.dim * self.dim                           // out proj
             + 2 * self.dim * self.ffn_hidden                // gate/up proj
             + self.ffn_hidden * self.dim; // down proj

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/src/model.rs

@@ -13,8 +13,8 @@ use xtrain_tensor::{DType, Device, Tensor};
 struct Block {
     attn_norm: Var, // [dim]
     wq: Var,        // [dim, dim]
-    wk: Var,        // [dim, dim]
-    wv: Var,        // [dim, dim]
+    wk: Var,        // [dim, kv_dim] — kv_dim = num_kv_heads·head_dim (GQA; = dim for MHA)
+    wv: Var,        // [dim, kv_dim]
     q_norm: Var,    // [head_dim] — per-head QK-norm (Qwen3-style)
     k_norm: Var,    // [head_dim]
     wo: Var,        // [dim, dim]
@@ -49,6 +49,14 @@ pub struct TinyTransformer {
     /// existing numerics are bit-identical; recompute is mathematically exact, so
     /// grads match the non-checkpointed path within fp tolerance.
     recompute: bool,
+    /// Fused flash-attention (Phase T14). When `true`, the SDPA core runs through
+    /// the hand-written single fused kernel ([`ops::flash_attention`]): online
+    /// softmax over KV tiles, the `[bh,seq,seq]` score matrix NEVER materialized,
+    /// backward caches only the O(N) logsumexp. Default `false` = the composed T10
+    /// path (`cublasSgemmStridedBatched`×2 + causal-softmax kernel, O(N²) probs),
+    /// so the default graph is unchanged. Mathematically the same SDPA → grads/loss
+    /// match the composed path within fp/bf16 tolerance. Opt-in via `--flash`.
+    use_flash: bool,
     /// Training mode for dropout (Phase T18). `true` → the attn/MLP sub-block
     /// outputs pass through `ops::dropout` (with `cfg.dropout` and a per-step,
     /// per-site seed); `false` (default) → dropout is identity (eval/sampling/
@@ -83,8 +91,9 @@ impl TinyTransformer {
             .map(|_| Block {
                 attn_norm: mk(&[cfg.dim]),
                 wq: mk(&[cfg.dim, cfg.dim]),
-                wk: mk(&[cfg.dim, cfg.dim]),
-                wv: mk(&[cfg.dim, cfg.dim]),
+                // GQA (T15): K/V project to num_kv_heads·head_dim (= dim when MHA).
+                wk: mk(&[cfg.dim, cfg.kv_dim()]),
+                wv: mk(&[cfg.dim, cfg.kv_dim()]),
                 q_norm: mk(&[cfg.head_dim]),
                 k_norm: mk(&[cfg.head_dim]),
                 wo: mk(&[cfg.dim, cfg.dim]),
@@ -105,6 +114,7 @@ impl TinyTransformer {
             lm_head,
             compute_dtype: DType::F32,
             recompute: false,
+            use_flash: false,
             training: Cell::new(false),
             step_seed: Cell::new(0),
         }
@@ -144,6 +154,19 @@ impl TinyTransformer {
         self.recompute
     }
 
+    /// Enable the fused flash-attention SDPA core (Phase T14). Builder-style and
+    /// opt-in; default off keeps the composed T10 path (so the default graph is
+    /// unchanged). On, the SDPA runs through [`ops::flash_attention`] — same SDPA
+    /// math, online softmax, no materialized `[bh,seq,seq]` scores.
+    pub fn with_flash(mut self, use_flash: bool) -> Self {
+        self.use_flash = use_flash;
+        self
+    }
+
+    pub fn use_flash(&self) -> bool {
+        self.use_flash
+    }
+
     /// Switch to training mode (Phase T18): dropout (if `cfg.dropout > 0`) is
     /// active in subsequent forwards. The training loop calls this before stepping.
     pub fn train(&self) {
@@ -251,15 +274,19 @@ impl TinyTransformer {
                 // the block forward is re-run in backward to recover the grads. The
                 // segment fn captures only `Copy` config (no borrow of `self`) and
                 // receives the block's params via the slice, in `block_params` order.
-                let (cfg, cdt) = (self.cfg, self.compute_dtype);
+                // `flash` is captured too → the recompute segment also runs flash;
+                // `dropout_p`/`block_seed` are captured so the recompute re-derives
+                // the same per-site dropout masks (counter-based RNG, exact).
+                let (cfg, cdt, flash) = (self.cfg, self.compute_dtype, self.use_flash);
                 let seg = move |x: &Var, p: &[Var]| {
-                    block_forward(cfg, cdt, batch, seq, dropout_p, block_seed, x, p)
+                    block_forward(cfg, cdt, flash, batch, seq, dropout_p, block_seed, x, p)
                 };
                 xtrain_autodiff::checkpoint::checkpoint(seg, &h, &b.block_params())
             } else {
                 block_forward(
                     self.cfg,
                     self.compute_dtype,
+                    self.use_flash,
                     batch,
                     seq,
                     dropout_p,
@@ -342,7 +369,8 @@ fn norm_gamma(cdt: DType, gamma: &Var) -> Var {
 
 /// One transformer block's forward: pre-norm + multi-head causal attention +
 /// (T18) dropout + residual, then pre-norm + SwiGLU MLP + dropout + residual.
-/// Pure in `(cfg, cdt, batch, seq, dropout_p, block_seed, input, params)` (no
+/// Attention runs the composed or fused-flash (T14) SDPA per `flash`. Pure in
+/// `(cfg, cdt, flash, batch, seq, dropout_p, block_seed, input, params)` (no
 /// `&self`, all `Copy`) so it can be the segment fn of
 /// [`xtrain_autodiff::checkpoint`] for activation recomputation (T13) — the
 /// recompute re-derives the same per-site seeds, so the dropout masks are
@@ -353,6 +381,7 @@ fn norm_gamma(cdt: DType, gamma: &Var) -> Var {
 fn block_forward(
     cfg: Config,
     cdt: DType,
+    flash: bool,
     batch: usize,
     seq: usize,
     dropout_p: f32,
@@ -372,7 +401,7 @@ fn block_forward(
     // --- Attention sub-block (pre-norm + dropout + residual) ---
     let normed = ops::rms_norm(h, &norm_gamma(cdt, attn_norm), cfg.eps);
     let attn = attention(
-        cfg, cdt, batch, seq, &normed, wq, wk, wv, q_norm, k_norm, wo,
+        cfg, cdt, flash, batch, seq, &normed, wq, wk, wv, q_norm, k_norm, wo,
     );
     let attn = ops::dropout(&attn, dropout_p, attn_seed);
     let h = ops::add(h, &attn);
@@ -395,6 +424,7 @@ fn block_forward(
 fn attention(
     cfg: Config,
     cdt: DType,
+    flash: bool,
     batch: usize,
     seq: usize,
     x: &Var,
@@ -406,41 +436,56 @@ fn attention(
     wo: &Var,
 ) -> Var {
     let (nh, hd) = (cfg.n_heads, cfg.head_dim);
+    let num_kv = cfg.num_kv_heads; // GQA (T15): K/V have fewer heads than Q
     let total = batch * seq;
-    let bh = batch * nh;
     let scale = 1.0 / (hd as f32).sqrt();
 
-    // Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
-    //   [B*S,dim] @ [dim,dim] = [B*S,dim]
-    //   reshape   [B*S, nh, hd]
+    // Project, qk-norm + RoPE, then lay out as a batched [B*heads, seq, hd] tensor.
+    // `heads` = nh for Q, num_kv for K/V (GQA; equal for MHA).
+    //   [B*S,dim] @ [dim,heads*hd] = [B*S, heads*hd]
+    //   reshape   [B*S, heads, hd]
     //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
-    //   rope      [B*S, nh, hd] with per-sequence position (period = seq)
-    //   reshape   [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
-    let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
-        let r = ops::reshape(&proj, &[total, nh, hd]);
+    //   rope      [B*S, heads, hd] with per-sequence position (period = seq)
+    //   reshape   [B, S, heads, hd] → transpose(1,2) → [B, heads, S, hd] → [B*heads, S, hd]
+    let to_bh = |proj: Var, heads: usize, norm: Option<&Var>| -> Var {
+        let r = ops::reshape(&proj, &[total, heads, hd]);
         let r = match norm {
-            // Per-head RMSNorm: flatten the (B*S,nh) head rows, norm over hd,
+            // Per-head RMSNorm: flatten the (B*S,heads) head rows, norm over hd,
             // restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
             Some(gamma) => {
-                let flat = ops::reshape(&r, &[total * nh, hd]);
+                let flat = ops::reshape(&r, &[total * heads, hd]);
                 let normed = ops::rms_norm(&flat, &norm_gamma(cdt, gamma), cfg.eps);
-                let r = ops::reshape(&normed, &[total, nh, hd]);
+                let r = ops::reshape(&normed, &[total, heads, hd]);
                 ops::rope(&r, cfg.rope_theta, seq)
             }
             None => r,
         };
-        let r = ops::reshape(&r, &[batch, seq, nh, hd]);
-        let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
-        ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
+        let r = ops::reshape(&r, &[batch, seq, heads, hd]);
+        let t = ops::transpose_4d12(&r); // [B, heads, S, hd]
+        ops::reshape(&t, &[batch * heads, seq, hd]) // [B*heads, S, hd]
     };
 
-    let q = to_bh(linear(cdt, x, wq), Some(q_norm));
-    let k = to_bh(linear(cdt, x, wk), Some(k_norm));
-    let v = to_bh(linear(cdt, x, wv), None);
+    let q = to_bh(linear(cdt, x, wq), nh, Some(q_norm));
+    // K/V are laid out with num_kv heads, then repeat_kv-broadcast to nh heads so
+    // the SDPA below (composed or flash, both unchanged) sees a full head set. The
+    // broadcast's backward sums each KV head's group of query-head grads (GQA). For
+    // MHA (num_kv == nh) repeat_kv is identity → bit-identical to the pre-T15 path.
+    let k = to_bh(linear(cdt, x, wk), num_kv, Some(k_norm));
+    let v = to_bh(linear(cdt, x, wv), num_kv, None);
+    let (k, v) = if num_kv == nh {
+        (k, v)
+    } else {
+        (ops::repeat_kv(&k, nh, batch), ops::repeat_kv(&v, nh, batch))
+    };
 
-    // Fused batched causal SDPA over all B*nh (sequence,head) blocks at once
-    // (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
-    let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]
+    // Causal SDPA over all B*nh (sequence,head) blocks. `flash` (T14) picks the
+    // single fused flash kernel (online softmax, no materialized [bh,S,S] scores);
+    // otherwise the composed T10 path (2 batched GEMMs + 1 causal-softmax kernel).
+    let out = if flash {
+        ops::flash_attention(&q, &k, &v, scale) // [B*nh, S, hd]
+    } else {
+        ops::attention(&q, &k, &v, scale) // [B*nh, S, hd]
+    };
 
     // Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
     // [B,S,nh,hd] → [B*S, dim].

crates/xtrain-model/tests/dropout.rs

@@ -66,10 +66,18 @@ fn tiny_cfg(dropout: f32) -> Config {
 fn batch_data(cfg: &Config, device: Device) -> (xtrain_tensor::Tensor, xtrain_tensor::Tensor) {
     let (batch, seq) = (3usize, 6usize);
     let seqs: Vec<Vec<i32>> = (0..batch)
-        .map(|b| (0..seq).map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32).collect())
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
+                .collect()
+        })
         .collect();
     let tgts: Vec<Vec<i32>> = (0..batch)
-        .map(|b| (0..seq).map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32).collect())
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
+                .collect()
+        })
         .collect();
     (
         batched_ids_tensor(&seqs, device),
@@ -94,7 +102,11 @@ fn fwd_bwd(
     let loss = m.loss_batched(ids, tgt, batch);
     let loss_val = host(&loss.value())[0];
     loss.backward();
-    let grads: Vec<Vec<f32>> = m.params().iter().map(|p| host(&p.grad().unwrap())).collect();
+    let grads: Vec<Vec<f32>> = m
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().unwrap()))
+        .collect();
     (logits, loss_val, grads)
 }
 
@@ -186,7 +198,9 @@ fn recompute_with_dropout(dtype: DType, grad_tol: f32) {
 
     // Both models: same init, train mode, p=0.2. step_seed starts at 0 and bumps
     // to 1 on the first training forward in BOTH, so they draw the same masks.
-    let off = build(cfg, device).with_compute_dtype(dtype).with_training(true);
+    let off = build(cfg, device)
+        .with_compute_dtype(dtype)
+        .with_training(true);
     let on = build(cfg, device)
         .with_compute_dtype(dtype)
         .with_recompute(true)
@@ -194,11 +208,19 @@ fn recompute_with_dropout(dtype: DType, grad_tol: f32) {
 
     let off_loss = off.loss_batched(&ids, &tgt, batch);
     off_loss.backward();
-    let off_grads: Vec<Vec<f32>> = off.params().iter().map(|p| host(&p.grad().unwrap())).collect();
+    let off_grads: Vec<Vec<f32>> = off
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().unwrap()))
+        .collect();
 
     let on_loss = on.loss_batched(&ids, &tgt, batch);
     on_loss.backward();
-    let on_grads: Vec<Vec<f32>> = on.params().iter().map(|p| host(&p.grad().unwrap())).collect();
+    let on_grads: Vec<Vec<f32>> = on
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().unwrap()))
+        .collect();
 
     let mut max_rel = 0.0f32;
     for (a, b) in off_grads.iter().flatten().zip(on_grads.iter().flatten()) {
@@ -220,3 +242,81 @@ fn dropout_recompute_matches_fp32() {
 fn dropout_recompute_matches_bf16() {
     recompute_with_dropout(DType::BF16, 5e-3);
 }
+
+// --- Cross-feature gate (Phase-2 integration): flash (T14) + dropout (T18)
+// together in the SAME model still grad-checks. Build two identical models, both
+// in train mode with p=0.2 (so dropout fires), one with `--flash` on, one off.
+// The dropout site seeds are a pure function of (step_seed, layer, site) and are
+// INDEPENDENT of flash, so both models draw the SAME masks on their first training
+// forward → the only difference is the SDPA reduction order. Assert logits/loss/
+// grads match within the flash-vs-composed tolerance and are finite. This is the
+// orthogonality check for the two Phase-2 features landing together.
+#[test]
+fn flash_plus_dropout_grad_check_fp32() {
+    let device = require_gpu();
+    let batch = 3;
+    // seq=40 > FA_TILE=32 exercises flash's online-softmax tile-rescale path.
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 4;
+    cfg.dropout = 0.2;
+    let seq = 40usize;
+    let seqs: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let tgts: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let ids = batched_ids_tensor(&seqs, device);
+    let tgt = batched_ids_tensor(&tgts, device);
+
+    // Both: same init, train mode (dropout active), same step_seed progression →
+    // identical masks; one composed SDPA, one flash SDPA.
+    let off = build(cfg, device).with_training(true);
+    let on = build(cfg, device).with_flash(true).with_training(true);
+
+    let (off_logits, off_loss, off_grads) = fwd_bwd(&off, &ids, &tgt, batch);
+    let (on_logits, on_loss, on_grads) = fwd_bwd(&on, &ids, &tgt, batch);
+
+    assert!(
+        on_logits.iter().all(|v| v.is_finite()) && on_grads.iter().flatten().all(|v| v.is_finite()),
+        "flash+dropout produced non-finite logits/grads"
+    );
+
+    let logit_rel = off_logits
+        .iter()
+        .zip(&on_logits)
+        .map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
+        .fold(0.0f32, f32::max);
+    let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
+    let mut grad_rel = 0.0f32;
+    for (a, b) in off_grads.iter().flatten().zip(on_grads.iter().flatten()) {
+        grad_rel = grad_rel.max((a - b).abs() / a.abs().max(1e-3));
+    }
+    println!(
+        "[F32] flash+dropout vs composed+dropout: loss rel {loss_rel:.2e}, \
+         logits max rel {logit_rel:.2e}, grad max rel {grad_rel:.3e}"
+    );
+    // Same tolerances as the flash-vs-composed gate (flash.rs run_fp32): flash
+    // differs from composed only by reduction order; dropout masks are identical.
+    assert!(
+        logit_rel < 1e-3,
+        "[F32] flash+dropout logits diverged: {logit_rel:.2e}"
+    );
+    assert!(
+        loss_rel < 1e-3,
+        "[F32] flash+dropout loss diverged: {loss_rel:.2e}"
+    );
+    assert!(
+        grad_rel < 2e-2,
+        "[F32] flash+dropout grads diverged: {grad_rel:.3e}"
+    );
+}

crates/xtrain-model/tests/flash.rs

@@ -0,0 +1,209 @@
+// T14 flash-attention correctness gate: the fused flash SDPA core must match the
+// composed T10 path (cublasSgemmStridedBatched×2 + causal-softmax kernel) in
+// forward logits, loss, AND every parameter gradient — flash is the SAME SDPA
+// math (online softmax never materializes the [bh,S,S] scores), so it differs
+// from composed only by reduction order (in-kernel fp32 FMA vs cuBLAS, and the
+// dK/dV atomicAdd order in backward). This test makes that a closed on-GPU loop:
+//
+//   build two identical models (same init), one with `--flash` on, one off, run
+//   the SAME batched loss + backward on both, and assert
+//     1. the forward logits match within tolerance
+//     2. the loss matches
+//     3. EVERY parameter's grad matches within tolerance
+//
+// Parameterised over fp32 AND bf16 (T12). bf16 just adds the bf16 rounding band on
+// top — flash's bf16 path upcasts Q/K/V to fp32 for the kernel exactly like the
+// composed path's fp32 softmax, so the two are still the same softmax numerics.
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
+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, 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: &xtrain_tensor::Tensor) -> Vec<f32> {
+    t.to_dtype(DType::F32)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec()
+}
+
+// fp32: same SDPA math, differs only by reduction order → tight per-element check.
+fn run_fp32(logit_tol: f32, grad_tol: f32) {
+    let (off_logits, off_loss, off_grads, on_logits, on_loss, on_grads) = run_both(DType::F32);
+
+    let logit_rel = off_logits
+        .iter()
+        .zip(&on_logits)
+        .map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
+        .fold(0.0f32, f32::max);
+    let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
+    println!(
+        "[F32] flash on/off: loss {off_loss:.6}/{on_loss:.6} (rel {loss_rel:.2e}), \
+         logits max rel {logit_rel:.2e}"
+    );
+    assert!(
+        logit_rel < logit_tol,
+        "[F32] logits diverged: {logit_rel:.2e}"
+    );
+    assert!(loss_rel < logit_tol, "[F32] loss diverged: {loss_rel:.2e}");
+
+    let mut max_grad_rel = 0.0f32;
+    for (off_g, on_g) in off_grads.iter().zip(&on_grads) {
+        for (a, b) in off_g.iter().zip(on_g) {
+            max_grad_rel = max_grad_rel.max((a - b).abs() / a.abs().max(1e-3));
+        }
+    }
+    println!("[F32] flash on/off: grad max rel err = {max_grad_rel:.3e}");
+    assert!(
+        max_grad_rel < grad_tol,
+        "[F32] flash grads diverged from composed: {max_grad_rel:.3e}"
+    );
+}
+
+// bf16: ~2-3 decimal digits → robust comparison (mean + p99 with abs().max(1.0)
+// for logits, per-tensor scale-relative mean for grads), the same convention as
+// the repo's bf16.rs gate (per-element max-rel blows up on near-zero bf16 logits).
+fn run_bf16() {
+    let (off_logits, off_loss, off_grads, on_logits, on_loss, on_grads) = run_both(DType::BF16);
+
+    let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
+    println!("[BF16] flash on/off: loss {off_loss:.5}/{on_loss:.5} (rel {loss_rel:.3e})");
+    assert!(loss_rel < 2e-2, "[BF16] loss diverged: {loss_rel:.3e}");
+
+    let n = off_logits.len();
+    let mut rels: Vec<f32> = off_logits
+        .iter()
+        .zip(&on_logits)
+        .map(|(f, b)| (b - f).abs() / f.abs().max(1.0))
+        .collect();
+    rels.sort_by(|a, b| a.partial_cmp(b).unwrap());
+    let p99 = rels[(n as f32 * 0.99) as usize];
+    let mean: f32 = rels.iter().sum::<f32>() / n as f32;
+    println!("[BF16] flash on/off logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
+    assert!(mean < 1e-2, "[BF16] logits mean rel too high: {mean:.3e}");
+    assert!(p99 < 5e-2, "[BF16] logits p99 rel too high: {p99:.3e}");
+
+    let mut worst = 0.0f32;
+    for (off_g, on_g) in off_grads.iter().zip(&on_grads) {
+        let scale = off_g
+            .iter()
+            .map(|v| v.abs())
+            .fold(0.0f32, f32::max)
+            .max(1e-6);
+        let mean_err: f32 = off_g
+            .iter()
+            .zip(on_g)
+            .map(|(f, b)| (f - b).abs())
+            .sum::<f32>()
+            / off_g.len() as f32
+            / scale;
+        worst = worst.max(mean_err);
+    }
+    println!("[BF16] flash on/off grads: worst per-tensor scaled-mean err = {worst:.3e}");
+    assert!(worst < 3e-2, "[BF16] flash grads diverged: {worst:.3e}");
+}
+
+#[allow(clippy::type_complexity)]
+fn run_both(dtype: DType) -> (Vec<f32>, f32, Vec<Vec<f32>>, Vec<f32>, f32, Vec<Vec<f32>>) {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    // seq=40 > FA_TILE=32 so the online-softmax tile-rescale path is exercised.
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 4;
+    let batch = 3usize;
+    let seq = 40usize;
+    let seqs: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let tgts: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let ids = batched_ids_tensor(&seqs, device);
+    let tgt = batched_ids_tensor(&tgts, device);
+
+    // --- flash OFF (composed reference) ---
+    let off = build(cfg, device, dtype, false);
+    let off_logits = host(&off.forward_batched(&ids, batch).value());
+    let off_loss = off.loss_batched(&ids, &tgt, batch);
+    let off_loss_val = host(&off_loss.value())[0];
+    off_loss.backward();
+    let off_grads: Vec<Vec<f32>> = off
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("off grad")))
+        .collect();
+
+    // --- flash ON ---
+    let on = build(cfg, device, dtype, true);
+    let on_logits = host(&on.forward_batched(&ids, batch).value());
+    let on_loss = on.loss_batched(&ids, &tgt, batch);
+    let on_loss_val = host(&on_loss.value())[0];
+    on_loss.backward();
+    let on_grads: Vec<Vec<f32>> = on
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("on grad")))
+        .collect();
+
+    (
+        off_logits,
+        off_loss_val,
+        off_grads,
+        on_logits,
+        on_loss_val,
+        on_grads,
+    )
+}
+
+#[test]
+fn flash_matches_composed_fp32() {
+    // fp32: same SDPA math, differs only by reduction order (in-kernel fp32 FMA vs
+    // cuBLAS, dK/dV atomicAdd order). Tight per-element check, not bit-exact.
+    run_fp32(1e-3, 2e-2);
+}
+
+#[test]
+fn flash_matches_composed_bf16() {
+    // bf16 (T12 composition): bf16 rounding band on the fp32-softmax core; robust
+    // (mean/p99/scaled-mean) comparison per the repo's bf16 convention.
+    run_bf16();
+}

crates/xtrain-model/tests/gqa.rs

@@ -0,0 +1,269 @@
+// T15 GQA correctness gate. Real grouped-query attention (num_kv_heads <
+// num_heads): K/V project to num_kv_heads·head_dim and are repeat_kv-broadcast to
+// the full head set before the SDPA. This test pins three things:
+//
+//   1. GQA flash == GQA composed (forward logits + loss + EVERY param grad) — the
+//      repeat_kv broadcast feeds both SDPA paths unchanged, so they must agree; in
+//      particular the wk/wv grads (which flow back through repeat_kv's group-sum)
+//      must match. Parameterised over fp32 (tight) and bf16 (rounding band).
+//   2. group==1 (num_kv_heads == n_heads) is BIT-IDENTICAL to the pre-T15 MHA path
+//      (a model with num_kv_heads explicitly == n_heads vs the default config):
+//      forward logits + every grad |Δ|=0. The regression guard.
+//   3. wk/wv really shrank to [dim, kv_dim] under GQA (shape check).
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
+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, 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: &xtrain_tensor::Tensor) -> Vec<f32> {
+    t.to_dtype(DType::F32)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec()
+}
+
+// A real GQA config: 8 query heads, 2 kv heads → group 4. seq=40 > FA_TILE=32 so
+// the flash online-softmax tile path is exercised too.
+fn gqa_cfg() -> Config {
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 3;
+    // tiny() is 2 heads; rebuild with 8 query / 2 kv heads keeping head_dim=16.
+    Config::from_arch(cfg.vocab, 8, cfg.head_dim, cfg.n_layers, cfg.ffn_hidden).with_kv_heads(2)
+}
+
+fn ids_targets(cfg: &Config, batch: usize, seq: usize) -> (Vec<Vec<i32>>, Vec<Vec<i32>>) {
+    let seqs = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let tgts = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    (seqs, tgts)
+}
+
+#[allow(clippy::type_complexity)]
+fn run_both(
+    cfg: Config,
+    dtype: DType,
+) -> (Vec<f32>, f32, Vec<Vec<f32>>, Vec<f32>, f32, Vec<Vec<f32>>) {
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+    let (batch, seq) = (3usize, 40usize);
+    let (seqs, tgts) = ids_targets(&cfg, batch, seq);
+    let ids = batched_ids_tensor(&seqs, device);
+    let tgt = batched_ids_tensor(&tgts, device);
+
+    let off = build(cfg, device, dtype, false);
+    let off_logits = host(&off.forward_batched(&ids, batch).value());
+    let off_loss = off.loss_batched(&ids, &tgt, batch);
+    let off_loss_val = host(&off_loss.value())[0];
+    off_loss.backward();
+    let off_grads: Vec<Vec<f32>> = off
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("off grad")))
+        .collect();
+
+    let on = build(cfg, device, dtype, true);
+    let on_logits = host(&on.forward_batched(&ids, batch).value());
+    let on_loss = on.loss_batched(&ids, &tgt, batch);
+    let on_loss_val = host(&on_loss.value())[0];
+    on_loss.backward();
+    let on_grads: Vec<Vec<f32>> = on
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("on grad")))
+        .collect();
+
+    (
+        off_logits,
+        off_loss_val,
+        off_grads,
+        on_logits,
+        on_loss_val,
+        on_grads,
+    )
+}
+
+// GQA flash vs composed: same SDPA math on the same repeat_kv-broadcast K/V → fp32
+// agrees to reduction-order, bf16 to its rounding band.
+#[test]
+fn gqa_flash_matches_composed_fp32() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    let cfg = gqa_cfg();
+    assert!(cfg.num_kv_heads < cfg.n_heads, "test must be real GQA");
+    let (off_l, off_loss, off_g, on_l, on_loss, on_g) = run_both(cfg, DType::F32);
+
+    let logit_rel = off_l
+        .iter()
+        .zip(&on_l)
+        .map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
+        .fold(0.0f32, f32::max);
+    let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
+    println!(
+        "[GQA F32] flash on/off: loss {off_loss:.6}/{on_loss:.6} (rel {loss_rel:.2e}), \
+         logits max rel {logit_rel:.2e}"
+    );
+    assert!(
+        logit_rel < 1e-3,
+        "[GQA F32] logits diverged: {logit_rel:.2e}"
+    );
+    assert!(loss_rel < 1e-3, "[GQA F32] loss diverged: {loss_rel:.2e}");
+
+    let mut worst = 0.0f32;
+    for (a_g, b_g) in off_g.iter().zip(&on_g) {
+        for (a, b) in a_g.iter().zip(b_g) {
+            worst = worst.max((a - b).abs() / a.abs().max(1e-3));
+        }
+    }
+    println!("[GQA F32] flash on/off grad max rel = {worst:.3e}");
+    assert!(worst < 2e-2, "[GQA F32] grads diverged: {worst:.3e}");
+}
+
+#[test]
+fn gqa_flash_matches_composed_bf16() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    let (off_l, off_loss, off_g, on_l, on_loss, on_g) = run_both(gqa_cfg(), DType::BF16);
+
+    let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
+    println!("[GQA BF16] flash on/off: loss {off_loss:.5}/{on_loss:.5} (rel {loss_rel:.3e})");
+    assert!(loss_rel < 2e-2, "[GQA BF16] loss diverged: {loss_rel:.3e}");
+
+    let n = off_l.len();
+    let mut rels: Vec<f32> = off_l
+        .iter()
+        .zip(&on_l)
+        .map(|(f, b)| (b - f).abs() / f.abs().max(1.0))
+        .collect();
+    rels.sort_by(|a, b| a.partial_cmp(b).unwrap());
+    let mean: f32 = rels.iter().sum::<f32>() / n as f32;
+    let p99 = rels[(n as f32 * 0.99) as usize];
+    println!("[GQA BF16] logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
+    assert!(
+        mean < 1e-2,
+        "[GQA BF16] logits mean rel too high: {mean:.3e}"
+    );
+    assert!(p99 < 5e-2, "[GQA BF16] logits p99 rel too high: {p99:.3e}");
+
+    let mut worst = 0.0f32;
+    for (a_g, b_g) in off_g.iter().zip(&on_g) {
+        let scale = a_g.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
+        let mean_err: f32 =
+            a_g.iter().zip(b_g).map(|(f, b)| (f - b).abs()).sum::<f32>() / a_g.len() as f32 / scale;
+        worst = worst.max(mean_err);
+    }
+    println!("[GQA BF16] grads: worst per-tensor scaled-mean err = {worst:.3e}");
+    assert!(worst < 3e-2, "[GQA BF16] grads diverged: {worst:.3e}");
+}
+
+// REGRESSION GUARD: num_kv_heads == n_heads (group 1) must be BIT-IDENTICAL to the
+// pre-T15 MHA path. Build one model with the default config (num_kv_heads ==
+// n_heads, the untouched path: repeat_kv not even invoked) and one that explicitly
+// sets num_kv_heads = n_heads, then assert forward logits + every grad match to the
+// bit. (Same composed path, so this is exact equality, not a tolerance.)
+#[test]
+fn gqa_group1_bit_identical_to_mha() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let mut base = Config::tiny();
+    base.vocab = 16;
+    base.n_layers = 3;
+    let base = Config::from_arch(base.vocab, 4, base.head_dim, base.n_layers, base.ffn_hidden);
+    // `explicit` sets num_kv_heads = n_heads (already the default, but exercises the
+    // with_kv_heads path); they are the same config → must produce identical output.
+    let explicit = base.with_kv_heads(base.n_heads);
+    assert_eq!(base.num_kv_heads, explicit.num_kv_heads);
+
+    let (batch, seq) = (2usize, 8usize);
+    let (seqs, tgts) = ids_targets(&base, batch, seq);
+    let ids = batched_ids_tensor(&seqs, device);
+    let tgt = batched_ids_tensor(&tgts, device);
+
+    let run = |cfg: Config| -> (Vec<f32>, f32, Vec<Vec<f32>>) {
+        let m = build(cfg, device, DType::F32, false);
+        let logits = host(&m.forward_batched(&ids, batch).value());
+        let loss = m.loss_batched(&ids, &tgt, batch);
+        let loss_v = host(&loss.value())[0];
+        loss.backward();
+        let grads = m
+            .params()
+            .iter()
+            .map(|p| host(&p.grad().unwrap()))
+            .collect();
+        (logits, loss_v, grads)
+    };
+    let (la, sa, ga) = run(base);
+    let (lb, sb, gb) = run(explicit);
+    assert_eq!(la, lb, "group-1 logits must be bit-identical to MHA");
+    assert_eq!(sa, sb, "group-1 loss must be bit-identical to MHA");
+    for (a, b) in ga.iter().zip(&gb) {
+        assert_eq!(a, b, "group-1 grad must be bit-identical to MHA");
+    }
+    println!("[GQA group1] bit-identical to MHA: logits + loss + all grads |Δ|=0");
+}
+
+// Under GQA, wk/wv must be [dim, kv_dim] (= num_kv_heads·head_dim), wq stays [dim,dim].
+#[test]
+fn gqa_kv_proj_shape() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+    let cfg = gqa_cfg();
+    let m = build(cfg, device, DType::F32, false);
+    let p = m.params();
+    // params order: embed[0], then block 0 = [attn_norm[1], wq[2], wk[3], wv[4],
+    // q_norm[5], k_norm[6], wo[7], ...]
+    let wq = p[2].value().shape().to_vec();
+    let wk = p[3].value().shape().to_vec();
+    let wv = p[4].value().shape().to_vec();
+    assert_eq!(wq, vec![cfg.dim, cfg.dim], "wq must be [dim,dim]");
+    assert_eq!(wk, vec![cfg.dim, cfg.kv_dim()], "wk must be [dim,kv_dim]");
+    assert_eq!(wv, vec![cfg.dim, cfg.kv_dim()], "wv must be [dim,kv_dim]");
+    println!(
+        "[GQA shapes] wq {:?} wk {:?} wv {:?} (kv_dim {})",
+        wq,
+        wk,
+        wv,
+        cfg.kv_dim()
+    );
+}

crates/xtrain-model/tests/parity.py

@@ -52,6 +52,10 @@ cfg = read_cfg()
 DIM = int(cfg["dim"])
 NL = int(cfg["n_layers"])
 NH = int(cfg["n_heads"])
+# GQA (T15): num_kv_heads <= n_heads; each kv head shared by group query heads.
+# Default to NH (MHA) for fixtures dumped before the field existed.
+NKV = int(cfg.get("num_kv_heads", str(NH)))
+GROUP = NH // NKV
 HD = int(cfg["head_dim"])
 EPS = float(cfg["eps"])
 THETA = float(cfg["rope_theta"])
@@ -114,17 +118,23 @@ for L in layers:
     # Attention
     x = rms_norm(h, L["attn_norm"])
     q = (x @ L["wq"]).reshape(B * SEQ, NH, HD)
-    k = (x @ L["wk"]).reshape(B * SEQ, NH, HD)
-    v = (x @ L["wv"]).reshape(B * SEQ, NH, HD)
+    # GQA: K/V project to NKV heads, then repeat each kv head GROUP times to NH.
+    k = (x @ L["wk"]).reshape(B * SEQ, NKV, HD)
+    v = (x @ L["wv"]).reshape(B * SEQ, NKV, HD)
     # Per-head QK-norm (Qwen3-style), before RoPE.
     q = rms_norm(q, L["q_norm"])
     k = rms_norm(k, L["k_norm"])
     q = rope(q)  # [B*SEQ, nh, hd]
-    k = rope(k)
-    # Reshape to [B, NH, SEQ, HD] so attention runs within each sequence.
+    k = rope(k)  # [B*SEQ, nkv, hd]
+    # Reshape to [B, *, SEQ, HD]; broadcast kv heads to NH (repeat_interleave along
+    # the head axis: kv head kvh → query heads [kvh*GROUP, (kvh+1)*GROUP), matching
+    # xtrain repeat_kv + xserv repeat_kv).
     q = q.reshape(B, SEQ, NH, HD).transpose(1, 2)  # [B, nh, seq, hd]
-    k = k.reshape(B, SEQ, NH, HD).transpose(1, 2)
-    v = v.reshape(B, SEQ, NH, HD).transpose(1, 2)
+    k = k.reshape(B, SEQ, NKV, HD).transpose(1, 2)  # [B, nkv, seq, hd]
+    v = v.reshape(B, SEQ, NKV, HD).transpose(1, 2)
+    if GROUP > 1:
+        k = k.repeat_interleave(GROUP, dim=1)  # [B, nh, seq, hd]
+        v = v.repeat_interleave(GROUP, dim=1)
     scale = 1.0 / math.sqrt(HD)
     scores = (q @ k.transpose(-1, -2)) * scale + mask  # [B, nh, seq, seq]
     probs = torch.softmax(scores, dim=-1)

crates/xtrain-model/tests/parity_dump.rs

@@ -58,8 +58,20 @@ fn dump_for_parity() {
     // sequence-major to [B*S]=8 ids. Per-sequence RoPE position (resets at the
     // sequence boundary) + per-sequence causal masking (no cross-sequence
     // attention) are both checked against PyTorch.
+    // Default: tiny MHA (2 heads). With XTRAIN_PARITY_KV_HEADS=k set, dump a real
+    // GQA config (8 query heads / k kv heads) so parity.py checks GQA at B>1 — the
+    // kv-projection shapes + the repeat_kv group-sum backward against PyTorch.
     let mut cfg = Config::tiny();
     cfg.vocab = 12;
+    if let Ok(kv) = std::env::var("XTRAIN_PARITY_KV_HEADS") {
+        let kv: usize = kv.parse().expect("XTRAIN_PARITY_KV_HEADS");
+        cfg = Config::from_arch(cfg.vocab, 8, cfg.head_dim, cfg.n_layers, cfg.ffn_hidden)
+            .with_kv_heads(kv);
+        println!(
+            "parity: GQA config (n_heads {} kv_heads {})",
+            cfg.n_heads, cfg.num_kv_heads
+        );
+    }
     let batch = 2usize;
     let seq = 4usize;
     let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major
@@ -67,7 +79,7 @@ fn dump_for_parity() {
 
     // Same deterministic init as the overfit test.
     let mut seed = 1u64;
-    let model = TinyTransformer::new(cfg, device, |shape| {
+    let mut model = TinyTransformer::new(cfg, device, |shape| {
         seed = seed.wrapping_add(1);
         let n: usize = shape.iter().product();
         if shape.len() == 1 {
@@ -76,6 +88,14 @@ fn dump_for_parity() {
             fill(n, seed, 0.08)
         }
     });
+    // T14: with XTRAIN_PARITY_FLASH set, dump from the fused flash-attention path.
+    // flash is the SAME SDPA math, so the SAME parity.py PyTorch oracle is the
+    // reference for both paths — running this once per path checks flash against
+    // PyTorch at B>1 (forward logits + every parameter grad).
+    if std::env::var("XTRAIN_PARITY_FLASH").is_ok() {
+        model = model.with_flash(true);
+        println!("parity: FLASH attention path");
+    }
 
     // config + ids
     {
@@ -84,6 +104,7 @@ fn dump_for_parity() {
         writeln!(f, "dim {}", cfg.dim).unwrap();
         writeln!(f, "n_layers {}", cfg.n_layers).unwrap();
         writeln!(f, "n_heads {}", cfg.n_heads).unwrap();
+        writeln!(f, "num_kv_heads {}", cfg.num_kv_heads).unwrap();
         writeln!(f, "head_dim {}", cfg.head_dim).unwrap();
         writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
         writeln!(f, "eps {:e}", cfg.eps).unwrap();

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)`.
@@ -1178,6 +1374,210 @@ impl Tensor {
         (dq, dk, dv)
     }
 
+    // --- Fused flash-attention (the T14 op) ---
+
+    /// Fused flash-attention forward (Phase T14). `self`=Q, `k`, `v` each
+    /// `[bh, seq, head_dim]`, contiguous on one GPU. Computes, per batch element,
+    /// `out = softmax(causal(Q·Kᵀ·scale))·V` in a SINGLE kernel that streams over
+    /// KV tiles with an online softmax — the `[bh,seq,seq]` score matrix is NEVER
+    /// materialized. Returns `(out, lse)` where `lse`:[bh,seq] (F32) is the per-row
+    /// logsumexp cached for backward (O(N), vs the composed path's O(N²) probs).
+    ///
+    /// The fused kernel is fp32; for bf16 we upcast Q/K/V → f32 → kernel → downcast
+    /// `out` back to bf16 (same fp32-softmax policy as the composed [`attention`]),
+    /// so flash and composed produce the same softmax numerics. `lse` stays fp32.
+    #[cfg(not(no_cuda))]
+    pub fn flash_attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> (Tensor, Tensor) {
+        assert_eq!(
+            self.ndim(),
+            3,
+            "flash_attention Q must be [bh,seq,head_dim]"
+        );
+        assert_eq!(self.shape(), k.shape(), "Q/K shape mismatch");
+        assert_eq!(self.shape(), v.shape(), "Q/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, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
+        let dev = self.device();
+        let dt = self.dtype;
+
+        let qf = self.to_dtype(DType::F32);
+        let kf = k.to_dtype(DType::F32);
+        let vf = v.to_dtype(DType::F32);
+        let out_f32 = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        let lse = Tensor::zeros(&[bh, seq], DType::F32, dev);
+        unsafe {
+            xtrain_cuda::ffi::launch_flash_attention_fwd_f32(
+                qf.data_ptr() as *const f32,
+                kf.data_ptr() as *const f32,
+                vf.data_ptr() as *const f32,
+                out_f32.data_ptr() as *mut f32,
+                lse.data_ptr() as *mut f32,
+                bh as i32,
+                seq as i32,
+                hd as i32,
+                scale,
+                std::ptr::null_mut(),
+            );
+        }
+        (out_f32.to_dtype(dt), lse)
+    }
+
+    /// Backward of [`flash_attention`](Self::flash_attention). Inputs: forward
+    /// `q`,`k`,`v`, the forward output `out`, the cached `lse`:[bh,seq], the upstream
+    /// `dout`, and the same `scale`. Returns `(dq, dk, dv)`.
+    ///
+    /// flash-style: NO cached probs. Recomputes scores from Q/K/V + `lse`, uses
+    /// `D[i]=Σ dOᵢ·Oᵢ` to collapse the softmax Jacobian, streams KV in tiles. dQ is
+    /// owned per query row; dK/dV are accumulated across rows (atomicAdd). Same
+    /// fp32 kernel; bf16 callers get fp32 grads which the autograd `cast` op casts.
+    #[cfg(not(no_cuda))]
+    pub fn flash_attention_backward(
+        q: &Tensor,
+        k: &Tensor,
+        v: &Tensor,
+        out: &Tensor,
+        lse: &Tensor,
+        dout: &Tensor,
+        scale: f32,
+    ) -> (Tensor, Tensor, Tensor) {
+        let (bh, seq, hd) = (q.shape[0], q.shape[1], q.shape[2]);
+        let dev = q.device();
+        let dt = q.dtype;
+
+        let qf = q.to_dtype(DType::F32);
+        let kf = k.to_dtype(DType::F32);
+        let vf = v.to_dtype(DType::F32);
+        let of = out.to_dtype(DType::F32);
+        let dof = dout.to_dtype(DType::F32);
+        // D[i] = Σ_d dO[i,d]·O[i,d] (one scalar per query row, O(N)).
+        let d = Tensor::zeros(&[bh, seq], DType::F32, dev);
+        unsafe {
+            xtrain_cuda::ffi::launch_flash_attention_rowdot_f32(
+                dof.data_ptr() as *const f32,
+                of.data_ptr() as *const f32,
+                d.data_ptr() as *mut f32,
+                (bh * seq) as i32,
+                hd as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        // dq/dk/dv pre-zeroed (Tensor::zeros memsets); dk/dv accumulate via atomicAdd.
+        let dq = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        let dk = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        let dv = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        unsafe {
+            xtrain_cuda::ffi::launch_flash_attention_bwd_f32(
+                qf.data_ptr() as *const f32,
+                kf.data_ptr() as *const f32,
+                vf.data_ptr() as *const f32,
+                dof.data_ptr() as *const f32,
+                lse.data_ptr() as *const f32,
+                d.data_ptr() as *mut f32,
+                dq.data_ptr() as *mut f32,
+                dk.data_ptr() as *mut f32,
+                dv.data_ptr() as *mut f32,
+                bh as i32,
+                seq as i32,
+                hd as i32,
+                scale,
+                std::ptr::null_mut(),
+            );
+        }
+        (dq.to_dtype(dt), dk.to_dtype(dt), dv.to_dtype(dt))
+    }
+
+    // --- GQA repeat_kv head broadcast (the T15 op) ---
+
+    /// GQA repeat_kv (Phase T15). `self` is a K or V tensor `[batch·num_kv, seq,
+    /// head_dim]`; returns `[batch·nh, seq, head_dim]` with each KV head broadcast
+    /// to its `group = nh/num_kv` query heads (query head `qh` ← kv head `qh/group`,
+    /// query heads contiguous within a group — matching xserv's repeat_kv). When
+    /// `nh == num_kv` (group 1) this is a plain copy (identity → MHA bit-identical).
+    /// bf16 callers are upcast→f32→kernel→downcast (the attention stack's fp32 policy).
+    #[cfg(not(no_cuda))]
+    pub fn repeat_kv(&self, nh: usize, batch: usize) -> Self {
+        assert_eq!(
+            self.ndim(),
+            3,
+            "repeat_kv input must be [batch·num_kv,seq,hd]"
+        );
+        let (bh_kv, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
+        assert_eq!(
+            bh_kv % batch,
+            0,
+            "repeat_kv: bh_kv {bh_kv} not div by batch {batch}"
+        );
+        let num_kv = bh_kv / batch;
+        assert_eq!(
+            nh % num_kv,
+            0,
+            "repeat_kv: nh {nh} not divisible by num_kv {num_kv}"
+        );
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .repeat_kv(nh, batch)
+                .to_dtype(DType::BF16);
+        }
+        assert_eq!(self.dtype, DType::F32, "repeat_kv supports F32/BF16");
+        let out = Tensor::zeros(&[batch * nh, seq, hd], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_repeat_kv_fwd_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                batch as i32,
+                nh as i32,
+                num_kv as i32,
+                seq as i32,
+                hd as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        out
+    }
+
+    /// Backward of [`repeat_kv`](Self::repeat_kv). `dout`:[batch·nh,seq,hd] →
+    /// `din`:[batch·num_kv,seq,hd], summing the `group` query heads that share each
+    /// KV head (the multi-group grad accumulation GQA correctness hinges on).
+    /// Deterministic (no atomics: each kv element serially sums its group rows).
+    #[cfg(not(no_cuda))]
+    pub fn repeat_kv_backward(dout: &Tensor, num_kv: usize, batch: usize) -> Self {
+        assert_eq!(
+            dout.ndim(),
+            3,
+            "repeat_kv_backward dout must be [batch·nh,seq,hd]"
+        );
+        let (bh_q, seq, hd) = (dout.shape[0], dout.shape[1], dout.shape[2]);
+        assert_eq!(bh_q % batch, 0, "repeat_kv_backward: bh_q not div by batch");
+        let nh = bh_q / batch;
+        assert_eq!(
+            nh % num_kv,
+            0,
+            "repeat_kv_backward: nh not divisible by num_kv"
+        );
+        let dt = dout.dtype;
+        if dt == DType::BF16 {
+            return Tensor::repeat_kv_backward(&dout.to_dtype(DType::F32), num_kv, batch)
+                .to_dtype(DType::BF16);
+        }
+        assert_eq!(dt, DType::F32, "repeat_kv_backward supports F32/BF16");
+        let din = Tensor::zeros(&[batch * num_kv, seq, hd], DType::F32, dout.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_repeat_kv_bwd_f32(
+                dout.data_ptr() as *const f32,
+                din.data_ptr() as *mut f32,
+                batch as i32,
+                nh as i32,
+                num_kv as i32,
+                seq as i32,
+                hd as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        din
+    }
+
     /// 4D axis-(1,2) transpose: `self`:[a,b,c,d] → [a,c,b,d],
     /// `out[i,k,j,l]=self[i,j,k,l]`. Lays out batched multi-head attention
     /// (`[B,S,nh,hd] <-> [B,nh,S,hd]`). Its own backward is the same op (swap b,c).

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

@@ -0,0 +1,209 @@
+//! Verifiable-task eval (post-training, M1+). Load a checkpoint, greedily generate an
+//! answer for each held-out arithmetic prompt, parse the `\boxed{}` answer, and report
+//! the exact-match pass-rate against the gold file. Two signals are printed:
+//! **format** (fraction that emitted any boxed integer) and **correctness** (fraction
+//! whose boxed answer matches gold). This is the M1 format-baseline metric and the
+//! reusable verifiable-eval harness for M3 (DPO) / M4 (GRPO).
+//!
+//!   eval_arith <ckpt> <tokenizer.json> --heads 52 --head-dim 32 --kv-heads 13 \
+//!       --layers 22 --ffn 6656 \
+//!       --prompts-file <dir>/arith_eval_prompts.txt \
+//!       --gold-file <dir>/arith_eval_gold.txt --max-tokens 48 --show 8
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("eval_arith: built without CUDA (no_cuda); run on a GPU host (dash5).");
+}
+
+#[cfg(not(no_cuda))]
+use std::path::PathBuf;
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::Device;
+#[cfg(not(no_cuda))]
+use xtrain_train::sample::generate;
+#[cfg(not(no_cuda))]
+use xtrain_train::task::{check_answer, parse_boxed_answer};
+
+// Same deterministic LCG init scheme as bin/train.rs / bin/greedy_sample.rs (the
+// values are overwritten by the loaded checkpoint; init just shapes the tensors).
+#[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 decode_escapes(s: &str) -> String {
+    s.replace("\\n", "\n").replace("\\t", "\t")
+}
+
+/// The model keeps generating past the answer (no EOS stop in the sampler), so keep
+/// only the first answer "turn": cut at the first `<|endoftext|>` and then at the
+/// first newline. The arithmetic answer is a single line, so this isolates it.
+#[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()
+        .map(|s| PathBuf::from(s.as_str()))
+        .expect("usage: eval_arith <ckpt> <tokenizer.json> [flags]");
+    let tok_path = positionals
+        .get(1)
+        .map(|s| PathBuf::from(s.as_str()))
+        .unwrap_or_else(|| PathBuf::from("/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 max_new = flag(&args, "--max-tokens", 48usize);
+    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.
+    let prompts: Vec<String> = std::fs::read_to_string(&prompts_file)
+        .unwrap_or_else(|e| panic!("read prompts {prompts_file}: {e}"))
+        .lines()
+        .map(str::trim)
+        .filter(|l| !l.is_empty() && !l.starts_with('#'))
+        .map(decode_escapes)
+        .collect();
+    let golds: Vec<i64> = std::fs::read_to_string(&gold_file)
+        .unwrap_or_else(|e| panic!("read gold {gold_file}: {e}"))
+        .lines()
+        .map(str::trim)
+        .filter(|l| !l.is_empty())
+        .map(|l| l.parse::<i64>().expect("gold line not an integer"))
+        .collect();
+    assert_eq!(
+        prompts.len(),
+        golds.len(),
+        "prompt/gold count mismatch ({} vs {})",
+        prompts.len(),
+        golds.len()
+    );
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let tok = Tokenizer::from_file(&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 = 1u64;
+    let model = 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)
+        }
+    });
+    xtrain_train::checkpoint::load_into(&ckpt, &model.params()).expect("load checkpoint");
+
+    println!(
+        "eval_arith: ckpt {} | {} prompts | max_new {} | decode={}",
+        ckpt.display(),
+        prompts.len(),
+        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 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() {
+            n_boxed += 1;
+        }
+        let ok = check_answer(seg, gold);
+        if ok {
+            n_correct += 1;
+        }
+        if shown < n_show {
+            let q = prompt.replace('\n', " ");
+            println!("  [{}] gold={gold} got={seg:?} {}", q, if ok { "OK" } else { "x" });
+            shown += 1;
+        }
+    }
+
+    let elapsed = t0.elapsed().as_secs_f64();
+    let n = prompts.len() as f64;
+    println!(
+        "RESULT format(boxed)={}/{} ({:.1}%) | correct={}/{} ({:.1}%)",
+        n_boxed,
+        prompts.len(),
+        100.0 * n_boxed as f64 / n,
+        n_correct,
+        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/export_safetensors.rs

@@ -174,7 +174,7 @@ fn config_json(cfg: &Config) -> String {
         ffn = cfg.ffn_hidden,
         layers = cfg.n_layers,
         heads = cfg.n_heads,
-        kv_heads = cfg.n_heads, // xtrain is MHA → kv heads == query heads
+        kv_heads = cfg.num_kv_heads, // GQA (T15): real num_key_value_heads (= n_heads for MHA)
         head_dim = cfg.head_dim,
         eps = cfg.eps,
         theta = cfg.rope_theta,
@@ -206,6 +206,8 @@ fn main() {
     let head_dim = flag(&args, "--head-dim", 16usize);
     let n_layers = flag(&args, "--layers", 4usize);
     let ffn = flag(&args, "--ffn", 64usize);
+    // GQA (Phase T15): num K/V heads (must match the trained ckpt; default = --heads).
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
 
     assert!(device::device_count().unwrap() > 0, "no CUDA device");
     device::set_device(0).unwrap();
@@ -214,15 +216,16 @@ fn main() {
     // Size the model from the arch flags + gpt2 vocab; must match the checkpoint.
     let tok = Tokenizer::from_file(&tok_path);
     let vocab = tok.vocab_size();
-    let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
+    let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
     println!(
-        "export: ckpt {} → {} (vocab {}, dim {}, layers {}, heads {}, head_dim {})",
+        "export: ckpt {} → {} (vocab {}, dim {}, layers {}, heads {}, kv_heads {}, head_dim {})",
         ckpt.display(),
         out_dir.display(),
         cfg.vocab,
         cfg.dim,
         cfg.n_layers,
         cfg.n_heads,
+        cfg.num_kv_heads,
         cfg.head_dim,
     );
 

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

@@ -0,0 +1,106 @@
+//! Generate the M1 verifiable-arithmetic post-training dataset. Pure host tool (no
+//! CUDA): writes
+//!   <out>/arith_sft.tsv           user<TAB>assistant rows for `train --sft-tsv`
+//!   <out>/arith_eval_prompts.txt  greedy_sample `--prompts-file` format (held out)
+//!   <out>/arith_eval_gold.txt     parallel gold integers for the checker
+//!
+//! Eval problems are deduped against train (no leakage). The SFT rows carry just the
+//! user/assistant content; `data::load_sft_tsv_cached` adds the `User:/Assistant:`
+//! frame + `<|endoftext|>` and masks the prompt, so the eval prompt lines here
+//! reconstruct exactly that frame (`User: <q>\nAssistant:`, literal `\n` decoded by
+//! greedy_sample).
+//!
+//! Example:
+//!   cargo run -p xtrain-train --release --bin gen_arith_task -- \
+//!       --n 20000 --eval 500 --seed 1 --out-dir /dashscope-tmp/wjh/xtrain_post/arith
+
+use std::collections::HashSet;
+use std::fs::{self, File};
+use std::io::{BufWriter, Write};
+use std::path::PathBuf;
+
+use xtrain_train::task::{GenConfig, Op, gen_problem, unique_space};
+
+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(|v| v.parse().ok())
+        .unwrap_or(default)
+}
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    let n_train: usize = flag(&args, "--n", 20000);
+    let n_eval: usize = flag(&args, "--eval", 500);
+    let seed: u64 = flag(&args, "--seed", 1);
+    let max_add: i64 = flag(&args, "--max-add", 999);
+    let max_mul: i64 = flag(&args, "--max-mul", 99);
+    let out_dir: PathBuf = args
+        .iter()
+        .position(|a| a == "--out-dir")
+        .and_then(|i| args.get(i + 1))
+        .map(PathBuf::from)
+        .expect("--out-dir <dir> is required");
+
+    fs::create_dir_all(&out_dir).expect("create out dir");
+    let cfg = GenConfig {
+        max_add,
+        max_mul,
+        ops: vec![Op::Add, Op::Sub, Op::Mul],
+    };
+
+    // Guard: train + eval are deduped (and eval is held out from train), so the
+    // request must fit comfortably inside the unique key space. Cap at 80% to keep
+    // dedup fast and the disjoint-eval loop terminating.
+    let space = unique_space(&cfg);
+    let need = (n_train + n_eval) as u64;
+    assert!(
+        need * 5 <= space * 4,
+        "requested {need} unique problems but the space is only {space} \
+         (max_add={max_add}, max_mul={max_mul}); raise --max-add/--max-mul or lower --n/--eval"
+    );
+
+    let mut rng = seed.max(1);
+
+    // Train: dedup so the same problem is not repeated and so eval can be held out.
+    let mut train_keys = HashSet::new();
+    let mut tsv = BufWriter::new(File::create(out_dir.join("arith_sft.tsv")).expect("create tsv"));
+    while train_keys.len() < n_train {
+        let p = gen_problem(&mut rng, &cfg);
+        if !train_keys.insert(p.key()) {
+            continue;
+        }
+        writeln!(tsv, "{}\t{}", p.question(), p.sft_answer()).expect("write tsv");
+    }
+    tsv.flush().expect("flush tsv");
+
+    // Eval: disjoint from train (skip any key seen in train) and from itself.
+    let mut prompts =
+        BufWriter::new(File::create(out_dir.join("arith_eval_prompts.txt")).expect("create eval"));
+    let mut golds =
+        BufWriter::new(File::create(out_dir.join("arith_eval_gold.txt")).expect("create gold"));
+    writeln!(prompts, "# verifiable arithmetic eval prompts (held out from arith_sft.tsv)")
+        .expect("write header");
+    let mut eval_keys = HashSet::new();
+    while eval_keys.len() < n_eval {
+        let p = gen_problem(&mut rng, &cfg);
+        if train_keys.contains(&p.key()) || !eval_keys.insert(p.key()) {
+            continue;
+        }
+        writeln!(prompts, "User: {}\\nAssistant:", p.question()).expect("write prompt");
+        writeln!(golds, "{}", p.answer()).expect("write gold");
+    }
+    prompts.flush().expect("flush prompts");
+    golds.flush().expect("flush golds");
+
+    println!(
+        "wrote {} train rows + {} eval prompts to {} (ops=+,-,* max_add={} max_mul={} seed={})",
+        train_keys.len(),
+        eval_keys.len(),
+        out_dir.display(),
+        max_add,
+        max_mul,
+        seed
+    );
+}

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/greedy_sample.rs

@@ -7,7 +7,8 @@
 //!   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
 //!   cargo run -p xtrain-train --release --bin greedy_sample -- \
 //!       /tmp/xtrain_v4.ckpt /opt/wjh/models/gpt2/tokenizer.json \
-//!       --heads 24 --head-dim 32 --layers 18 --ffn 2048
+//!       --heads 24 --head-dim 32 --layers 18 --ffn 2048 \
+//!       --prompts-file scripts/chat_alpha_fixed_prompts.txt --max-tokens 120
 
 #[cfg(no_cuda)]
 fn main() {
@@ -52,6 +53,60 @@ fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
         .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 flag_values(args: &[String], name: &str) -> Vec<String> {
+    args.iter()
+        .enumerate()
+        .filter_map(|(i, a)| {
+            if a == name {
+                args.get(i + 1).cloned()
+            } else {
+                None
+            }
+        })
+        .collect()
+}
+
+#[cfg(not(no_cuda))]
+fn decode_prompt_escapes(s: &str) -> String {
+    s.replace("\\n", "\n").replace("\\t", "\t")
+}
+
+#[cfg(not(no_cuda))]
+fn load_prompts(args: &[String]) -> Vec<String> {
+    let mut prompts = Vec::new();
+    if let Some(path) = flag_value(args, "--prompts-file") {
+        let text = std::fs::read_to_string(&path)
+            .unwrap_or_else(|e| panic!("failed to read prompts file {path}: {e}"));
+        prompts.extend(
+            text.lines()
+                .map(str::trim)
+                .filter(|line| !line.is_empty() && !line.starts_with('#'))
+                .map(decode_prompt_escapes),
+        );
+    }
+    prompts.extend(
+        flag_values(args, "--prompt")
+            .into_iter()
+            .map(|p| decode_prompt_escapes(&p)),
+    );
+    if prompts.is_empty() {
+        prompts = ["Once upon a time", "One day", "The little"]
+            .into_iter()
+            .map(String::from)
+            .collect();
+    }
+    prompts
+}
+
 #[cfg(not(no_cuda))]
 fn main() {
     use xserv_tokenizer::Tokenizer;
@@ -72,21 +127,27 @@ fn main() {
     let head_dim = flag(&args, "--head-dim", 16usize);
     let n_layers = flag(&args, "--layers", 4usize);
     let ffn = flag(&args, "--ffn", 64usize);
+    // GQA (Phase T15): num K/V heads (must match the ckpt; default = --heads).
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
     let max_new = flag(&args, "--max-tokens", 40usize);
+    let temperature = flag(&args, "--temperature", 0.0f32);
+    let prompts = load_prompts(&args);
 
     assert!(device::device_count().unwrap() > 0, "no CUDA device");
     device::set_device(0).unwrap();
     let device = Device::Cuda(0);
 
     let tok = Tokenizer::from_file(&tok_path);
-    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn);
+    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn)
+        .with_kv_heads(kv_heads);
     println!(
-        "greedy_sample: ckpt {} (vocab {}, dim {}, layers {}, heads {}, head_dim {})",
+        "greedy_sample: ckpt {} (vocab {}, dim {}, layers {}, heads {}, kv_heads {}, head_dim {})",
         ckpt.display(),
         cfg.vocab,
         cfg.dim,
         cfg.n_layers,
         cfg.n_heads,
+        cfg.num_kv_heads,
         cfg.head_dim,
     );
 
@@ -102,11 +163,16 @@ fn main() {
     });
     xtrain_train::checkpoint::load_into(&ckpt, &model.params()).expect("load checkpoint");
 
-    let prompts = ["Once upon a time", "One day", "The little"];
+    println!(
+        "decode: prompts={} max_new={} temperature={}",
+        prompts.len(),
+        max_new,
+        temperature
+    );
     for p in prompts {
-        let ids: Vec<i32> = tok.encode(p).into_iter().map(|t| t as i32).collect();
+        let ids: Vec<i32> = tok.encode(&p).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 = generate(&model, device, &ids, max_new, temperature, &mut rng);
         let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
         println!("[{p}] → {text}");
     }

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

@@ -88,6 +88,8 @@ fn main() {
     let head_dim = flag(&args, "--head-dim", 16usize);
     let n_layers = flag(&args, "--layers", 4usize);
     let ffn = flag(&args, "--ffn", 64usize);
+    // GQA (Phase T15): num K/V heads (must divide --heads). Default = --heads (MHA).
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
     // `--dim` is informational; dim is always n_heads*head_dim. Warn on mismatch.
     let dim_flag = flag(&args, "--dim", 0usize);
     if dim_flag != 0 && dim_flag != n_heads * head_dim {
@@ -101,6 +103,10 @@ fn main() {
     // Optimization knobs.
     let steps: usize = flag(&args, "--steps", 2000);
     let batch_size: usize = flag(&args, "--batch", 8);
+    // Micro-batch gradient accumulation (Phase T16): effective batch =
+    // accum_steps × batch, at one micro-batch's activation-memory cost. Default 1
+    // = no accumulation (bit-identical to the pre-T16 path).
+    let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
     let seq_len: usize = flag(&args, "--seq", 64);
     let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
     let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
@@ -109,6 +115,7 @@ 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);
+    let sft_tsv = args.iter().any(|a| a == "--sft-tsv");
     // Dropout (Phase T18): residual-path dropout prob, active at training time
     // only (inverted scaling), identity at eval/sampling/export. Default 0 = off
     // (forward graph bit-identical to the no-dropout path).
@@ -120,6 +127,9 @@ fn main() {
     // exact grads, lower peak activation memory (lets dim1024 batch32 fit). Opt-in;
     // default off stores every activation (unchanged numerics).
     let recompute = args.iter().any(|a| a == "--recompute");
+    // Fused flash-attention (Phase T14): single fused SDPA kernel, online softmax,
+    // no materialized [bh,S,S] scores. Opt-in; default off keeps the composed path.
+    let flash = args.iter().any(|a| a == "--flash");
     let ckpt: PathBuf = PathBuf::from(
         args.iter()
             .position(|a| a == "--ckpt")
@@ -127,6 +137,11 @@ fn main() {
             .cloned()
             .unwrap_or_else(|| "/tmp/xtrain_tinystories.ckpt".to_string()),
     );
+    let init_ckpt: Option<PathBuf> = args
+        .iter()
+        .position(|a| a == "--init-ckpt")
+        .and_then(|i| args.get(i + 1))
+        .map(PathBuf::from);
 
     assert!(device::device_count().unwrap() > 0, "no CUDA device");
     device::set_device(0).unwrap();
@@ -137,12 +152,19 @@ fn main() {
         tok_path.display(),
         corpus_path.display()
     );
-    let corpus = Corpus::load_cached(&tok_path, &corpus_path);
+    let corpus = if sft_tsv {
+        Corpus::load_sft_tsv_cached(&tok_path, &corpus_path)
+    } else {
+        Corpus::load_cached(&tok_path, &corpus_path)
+    };
     println!(
         "corpus: {} tokens, vocab {}",
         corpus.len(),
         corpus.vocab_size
     );
+    if sft_tsv {
+        println!("SFT TSV: ON (assistant-only loss via ignore-index labels)");
+    }
     let vocab = corpus.vocab_size;
     // Hold out a tail slice for validation (if requested and the corpus is big).
     let (train_corpus, valid) = if val_tokens > 0 {
@@ -153,14 +175,16 @@ fn main() {
         (corpus, None)
     };
 
-    let mut cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
+    let mut cfg =
+        Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
     cfg.dropout = dropout;
     println!(
-        "model: dim {} layers {} heads {} head_dim {} ffn {} → core {:.3}M params \
+        "model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
          (+ embed/lm {:.2}M = {:.2}M total)",
         cfg.dim,
         cfg.n_layers,
         cfg.n_heads,
+        cfg.num_kv_heads,
         cfg.head_dim,
         cfg.ffn_hidden,
         cfg.core_params() as f32 / 1e6,
@@ -188,9 +212,17 @@ fn main() {
         model = model.with_recompute(true);
         println!("activation recompute: ON (per-block gradient checkpointing)");
     }
+    if flash {
+        model = model.with_flash(true);
+        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)");
     }
+    if let Some(path) = &init_ckpt {
+        xtrain_train::checkpoint::load_into(path, &model.params()).expect("load init checkpoint");
+        println!("init checkpoint: loaded {}", path.display());
+    }
 
     // Eval-only mode: load a checkpoint and score it on the held-out val set, then
     // exit. Used to put an EXISTING model (e.g. v0) and a new one on the same
@@ -209,6 +241,7 @@ fn main() {
     let tcfg = TrainConfig {
         seq_len,
         batch_size,
+        accum_steps,
         steps,
         schedule: LrSchedule {
             max_lr,
@@ -227,10 +260,13 @@ fn main() {
     };
 
     println!(
-        "training: {} steps, seq {}, batch {}, lr {:.1e}→{:.1e}, eval every {}",
+        "training: {} steps, seq {}, batch {} × accum {} = effective batch {}, \
+         lr {:.1e}→{:.1e}, eval every {}",
         tcfg.steps,
         tcfg.seq_len,
         tcfg.batch_size,
+        tcfg.accum_steps,
+        tcfg.batch_size * tcfg.accum_steps,
         tcfg.schedule.max_lr,
         tcfg.schedule.min_lr,
         tcfg.eval_every

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/data.rs

@@ -15,6 +15,7 @@ use xserv_tokenizer::Tokenizer;
 /// A tokenized corpus: one flat stream of token ids, plus the vocab size.
 pub struct Corpus {
     pub tokens: Vec<i32>,
+    pub labels: Option<Vec<i32>>,
     pub vocab_size: usize,
 }
 
@@ -33,6 +34,7 @@ impl Corpus {
         let ids: Vec<i32> = tok.encode(text).into_iter().map(|t| t as i32).collect();
         Self {
             tokens: ids,
+            labels: None,
             vocab_size: tok.vocab_size(),
         }
     }
@@ -52,7 +54,11 @@ impl Corpus {
                 tokens.len(),
                 cache.display()
             );
-            return Self { tokens, vocab_size };
+            return Self {
+                tokens,
+                labels: None,
+                vocab_size,
+            };
         }
         let me = Self::load(tokenizer_path, corpus_path);
         write_u16_cache(&cache, &me.tokens);
@@ -64,22 +70,103 @@ impl Corpus {
         me
     }
 
+    /// Load assistant-only SFT data from a two-column TSV:
+    ///
+    /// ```text
+    /// user<TAB>assistant
+    /// ```
+    ///
+    /// Literal `\n` and `\t` escapes are decoded. Each row is formatted as
+    /// `User: ...\nAssistant:` + answer + `<|endoftext|>`. Labels are `-100`
+    /// for prompt tokens and the token id itself for answer/EOS tokens, so the
+    /// cross-entropy op ignores prompt rows while still training the assistant
+    /// answer and stop token.
+    pub fn load_sft_tsv_cached(tokenizer_path: &Path, corpus_path: &Path) -> Self {
+        let token_cache = cache_path(corpus_path);
+        let label_cache = label_cache_path(corpus_path);
+        let vocab_size = Tokenizer::from_file(tokenizer_path).vocab_size();
+        if token_cache.exists() && label_cache.exists() {
+            let tokens = read_u16_cache(&token_cache);
+            let labels = read_i32_cache(&label_cache);
+            assert_eq!(
+                tokens.len(),
+                labels.len(),
+                "SFT cache token/label length mismatch"
+            );
+            println!(
+                "corpus: read {} cached SFT tokens from {} (+ labels {})",
+                tokens.len(),
+                token_cache.display(),
+                label_cache.display()
+            );
+            return Self {
+                tokens,
+                labels: Some(labels),
+                vocab_size,
+            };
+        }
+
+        let tok = Tokenizer::from_file(tokenizer_path);
+        let text = std::fs::read_to_string(corpus_path)
+            .unwrap_or_else(|e| panic!("failed to read SFT corpus {}: {e}", corpus_path.display()));
+        let mut tokens = Vec::new();
+        let mut labels = Vec::new();
+        for (lineno, line) in text.lines().enumerate() {
+            if line.trim().is_empty() {
+                continue;
+            }
+            let (user, assistant) = line
+                .split_once('\t')
+                .unwrap_or_else(|| panic!("SFT TSV line {} missing tab", lineno + 1));
+            let user = decode_tsv_escapes(user);
+            let assistant = decode_tsv_escapes(assistant);
+            let prompt = format!("User: {user}\nAssistant:");
+            let answer = format!(" {assistant}\n<|endoftext|>");
+            let prompt_ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
+            let answer_ids: Vec<i32> = tok.encode(&answer).into_iter().map(|t| t as i32).collect();
+            let (row_tokens, row_labels) = sft_row(&prompt_ids, &answer_ids);
+            tokens.extend(row_tokens);
+            labels.extend(row_labels);
+        }
+        assert_eq!(tokens.len(), labels.len(), "SFT tokens/labels mismatch");
+        write_u16_cache(&token_cache, &tokens);
+        write_i32_cache(&label_cache, &labels);
+        println!(
+            "corpus: tokenized {} SFT tokens → cached to {} (+ labels {})",
+            tokens.len(),
+            token_cache.display(),
+            label_cache.display()
+        );
+        Self {
+            tokens,
+            labels: Some(labels),
+            vocab_size: tok.vocab_size(),
+        }
+    }
+
     /// Split off the last `n` tokens as a held-out validation corpus, leaving the
     /// rest as the train corpus. Returns `(train, valid)`. Used for periodic val
     /// loss during training without leaking the eval window into training.
     pub fn split_tail(self, n: usize) -> (Self, Self) {
         let n = n.min(self.tokens.len() / 10); // never hand off more than 10%
         let cut = self.tokens.len() - n;
-        let valid = self.tokens[cut..].to_vec();
+        let valid_tokens = self.tokens[cut..].to_vec();
+        let valid_labels = self.labels.as_ref().map(|labels| labels[cut..].to_vec());
         let mut train = self.tokens;
         train.truncate(cut);
+        let train_labels = self.labels.map(|mut labels| {
+            labels.truncate(cut);
+            labels
+        });
         (
             Self {
                 tokens: train,
+                labels: train_labels,
                 vocab_size: self.vocab_size,
             },
             Self {
-                tokens: valid,
+                tokens: valid_tokens,
+                labels: valid_labels,
                 vocab_size: self.vocab_size,
             },
         )
@@ -101,11 +188,27 @@ impl Corpus {
     pub fn sample(&self, seq: usize, rng_state: &mut u64) -> (Vec<i32>, Vec<i32>) {
         assert!(self.tokens.len() > seq + 1, "corpus shorter than a window");
         let max_start = self.tokens.len() - seq - 1;
-        let start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
+        let mut start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
+        if let Some(labels) = &self.labels {
+            for _ in 0..16 {
+                if labels[start + 1..start + seq + 1].iter().any(|&t| t >= 0) {
+                    break;
+                }
+                start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
+            }
+        }
         let input = self.tokens[start..start + seq].to_vec();
-        let target = self.tokens[start + 1..start + seq + 1].to_vec();
+        let target = self.target_window(start, seq);
         (input, target)
     }
+
+    /// Deterministic target labels for an input window starting at `start`.
+    pub fn target_window(&self, start: usize, seq: usize) -> Vec<i32> {
+        match &self.labels {
+            Some(labels) => labels[start + 1..start + seq + 1].to_vec(),
+            None => self.tokens[start + 1..start + seq + 1].to_vec(),
+        }
+    }
 }
 
 /// Drop a leading partial line (before the first newline) and everything after
@@ -127,6 +230,12 @@ fn cache_path(corpus_path: &Path) -> PathBuf {
     PathBuf::from(s)
 }
 
+fn label_cache_path(corpus_path: &Path) -> PathBuf {
+    let mut s = corpus_path.as_os_str().to_os_string();
+    s.push(".labels.i32.bin");
+    PathBuf::from(s)
+}
+
 /// Read a flat little-endian `[u16]` cache into an `i32` id stream.
 fn read_u16_cache(path: &Path) -> Vec<i32> {
     let mut r = BufReader::new(
@@ -140,6 +249,18 @@ fn read_u16_cache(path: &Path) -> Vec<i32> {
         .collect()
 }
 
+fn read_i32_cache(path: &Path) -> Vec<i32> {
+    let mut r = BufReader::new(
+        std::fs::File::open(path).unwrap_or_else(|e| panic!("open cache {}: {e}", path.display())),
+    );
+    let mut buf = Vec::new();
+    r.read_to_end(&mut buf).expect("read cache");
+    assert!(buf.len() % 4 == 0, "corrupt i32 cache (odd byte count)");
+    buf.chunks_exact(4)
+        .map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
+        .collect()
+}
+
 /// Write an id stream as a flat little-endian `[u16]` cache. Ids must fit in u16
 /// (GPT-2 vocab = 50257 < 65536); asserts otherwise.
 fn write_u16_cache(path: &Path, tokens: &[i32]) {
@@ -154,6 +275,35 @@ fn write_u16_cache(path: &Path, tokens: &[i32]) {
     w.flush().expect("flush cache");
 }
 
+fn write_i32_cache(path: &Path, labels: &[i32]) {
+    let mut w = BufWriter::new(
+        std::fs::File::create(path)
+            .unwrap_or_else(|e| panic!("create cache {}: {e}", path.display())),
+    );
+    for &t in labels {
+        w.write_all(&t.to_le_bytes()).expect("write cache");
+    }
+    w.flush().expect("flush cache");
+}
+
+fn decode_tsv_escapes(s: &str) -> String {
+    s.replace("\\n", "\n").replace("\\t", "\t")
+}
+
+/// Build one SFT example's `(tokens, labels)` from already-tokenized prompt/answer
+/// ids: prompt tokens are masked to the ignore-index (`-100`, which `cross_entropy`
+/// skips) so only the answer + EOS tokens contribute to the loss. Pure (no tokenizer
+/// / no CUDA) so the assistant-only masking is unit-testable directly.
+fn sft_row(prompt_ids: &[i32], answer_ids: &[i32]) -> (Vec<i32>, Vec<i32>) {
+    let mut tokens = Vec::with_capacity(prompt_ids.len() + answer_ids.len());
+    tokens.extend_from_slice(prompt_ids);
+    tokens.extend_from_slice(answer_ids);
+    let mut labels = Vec::with_capacity(prompt_ids.len() + answer_ids.len());
+    labels.extend(std::iter::repeat(-100).take(prompt_ids.len()));
+    labels.extend_from_slice(answer_ids);
+    (tokens, labels)
+}
+
 /// Tiny LCG (same constants as the model tests' deterministic fill) so dataset
 /// sampling is reproducible from a single u64 seed.
 fn next_rand(state: &mut u64) -> u64 {
@@ -162,3 +312,27 @@ fn next_rand(state: &mut u64) -> u64 {
         .wrapping_add(1442695040888963407);
     *state >> 16
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn sft_row_masks_prompt_supervises_answer() {
+        let prompt = [5, 6, 7];
+        let answer = [8, 9]; // includes the EOS token in real use
+        let (tokens, labels) = sft_row(&prompt, &answer);
+        // Tokens are prompt then answer, in order.
+        assert_eq!(tokens, vec![5, 6, 7, 8, 9]);
+        // Prompt positions are ignore-index (-100); answer positions are supervised.
+        assert_eq!(labels, vec![-100, -100, -100, 8, 9]);
+        assert_eq!(tokens.len(), labels.len());
+    }
+
+    #[test]
+    fn sft_row_handles_empty_answer() {
+        let (tokens, labels) = sft_row(&[1, 2], &[]);
+        assert_eq!(tokens, vec![1, 2]);
+        assert_eq!(labels, vec![-100, -100]);
+    }
+}

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

@@ -10,10 +10,13 @@
 pub mod clip;
 pub mod data;
 pub mod schedule;
+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/src/task.rs

@@ -0,0 +1,240 @@
+//! Verifiable arithmetic task (post-training, M1). A tiny two-operand integer
+//! arithmetic task with a deterministic, rule-based checker: the assistant must end
+//! its answer with `\boxed{N}`, and the reward is exact-match on `N`.
+//!
+//! This single module is the shared task spec for the whole post-training stack —
+//! M1 SFT-data generation, M3 DPO preference-pair construction, and M4 GRPO reward
+//! scoring all parse/score through here, so the task lives in exactly one place.
+//!
+//! Host-only (no CUDA): generation + parsing + checking are pure, so this compiles
+//! and unit-tests on a GPU-less host.
+
+use std::fmt;
+
+/// The supported binary operations.
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum Op {
+    Add,
+    Sub,
+    Mul,
+}
+
+impl Op {
+    pub fn symbol(self) -> char {
+        match self {
+            Op::Add => '+',
+            Op::Sub => '-',
+            Op::Mul => '*',
+        }
+    }
+}
+
+impl fmt::Display for Op {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.symbol())
+    }
+}
+
+/// A single two-operand arithmetic problem.
+#[derive(Clone, Copy, Debug)]
+pub struct Problem {
+    pub a: i64,
+    pub b: i64,
+    pub op: Op,
+}
+
+impl Problem {
+    /// The exact integer answer (the verifiable gold label).
+    pub fn answer(self) -> i64 {
+        match self.op {
+            Op::Add => self.a + self.b,
+            Op::Sub => self.a - self.b,
+            Op::Mul => self.a * self.b,
+        }
+    }
+
+    /// The user-turn question text. No template wrapping — the SFT loader
+    /// (`data::load_sft_tsv_cached`) adds the `User:/Assistant:` frame.
+    pub fn question(self) -> String {
+        format!("What is {} {} {}?", self.a, self.op, self.b)
+    }
+
+    /// The assistant-turn SFT target: restate the equation and end with the boxed
+    /// answer. This teaches the answer FORMAT (the checker only reads `\boxed{}`);
+    /// arithmetic correctness is what DPO (M3) / GRPO (M4) later improve.
+    pub fn sft_answer(self) -> String {
+        format!("{} {} {} = \\boxed{{{}}}.", self.a, self.op, self.b, self.answer())
+    }
+
+    /// A stable dedup key, so eval problems can be held out from train.
+    pub fn key(self) -> (i64, char, i64) {
+        (self.a, self.op.symbol(), self.b)
+    }
+}
+
+/// Operand-range configuration for problem sampling. Multiplication uses a smaller
+/// range (`max_mul`) so products stay modest; add/sub use `max_add`.
+#[derive(Clone)]
+pub struct GenConfig {
+    pub max_add: i64,
+    pub max_mul: i64,
+    pub ops: Vec<Op>,
+}
+
+impl Default for GenConfig {
+    fn default() -> Self {
+        Self {
+            max_add: 999,
+            max_mul: 99,
+            ops: vec![Op::Add, Op::Sub, Op::Mul],
+        }
+    }
+}
+
+/// Number of distinct problems this config can produce (the key space). Used to
+/// guard the dedup generator against requesting more unique problems than exist —
+/// otherwise train/eval dedup loops near saturation get pathologically slow or, for
+/// a disjoint eval, never terminate.
+pub fn unique_space(cfg: &GenConfig) -> u64 {
+    cfg.ops
+        .iter()
+        .map(|op| {
+            let max = if *op == Op::Mul { cfg.max_mul } else { cfg.max_add };
+            ((max as u64) + 1).pow(2) // ordered (a, b) pairs in [0, max]
+        })
+        .sum()
+}
+
+/// Sample one problem deterministically from the LCG state `rng`. Operands are drawn
+/// in `[0, max]` per the op; subtraction may yield a negative answer (the checker /
+/// parser handle a leading `-`).
+pub fn gen_problem(rng: &mut u64, cfg: &GenConfig) -> Problem {
+    let op = cfg.ops[(next_rand(rng) as usize) % cfg.ops.len()];
+    let max = if op == Op::Mul { cfg.max_mul } else { cfg.max_add };
+    let a = rand_range(rng, max);
+    let b = rand_range(rng, max);
+    Problem { a, b, op }
+}
+
+/// Parse the integer inside the LAST `\boxed{...}` in `text`. Returns `None` if there
+/// is no well-formed boxed integer (no box, empty, or non-integer contents). "Last"
+/// so a model that emits intermediate boxes still scores on its final answer.
+pub fn parse_boxed_answer(text: &str) -> Option<i64> {
+    const TAG: &str = "\\boxed{";
+    let mut found = None;
+    let mut rest = text;
+    while let Some(i) = rest.find(TAG) {
+        let after = &rest[i + TAG.len()..];
+        match after.find('}') {
+            Some(j) => {
+                if let Ok(n) = after[..j].trim().parse::<i64>() {
+                    found = Some(n);
+                }
+                rest = &after[j + 1..];
+            }
+            None => break,
+        }
+    }
+    found
+}
+
+/// Verifiable reward: does the completion's boxed answer exactly match `gold`?
+pub fn check_answer(completion: &str, gold: i64) -> bool {
+    parse_boxed_answer(completion) == Some(gold)
+}
+
+/// `[0, max]` inclusive draw from the LCG.
+fn rand_range(rng: &mut u64, max: i64) -> i64 {
+    debug_assert!(max >= 0);
+    (next_rand(rng) % (max as u64 + 1)) as i64
+}
+
+/// Same LCG constants as the dataset sampler (`data::next_rand`), kept local so the
+/// task module stays dependency-free and host-only.
+fn next_rand(state: &mut u64) -> u64 {
+    *state = state
+        .wrapping_mul(6364136223846793005)
+        .wrapping_add(1442695040888963407);
+    *state >> 1
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn answer_question_and_sft_target() {
+        let p = Problem {
+            a: 12,
+            b: 13,
+            op: Op::Mul,
+        };
+        assert_eq!(p.answer(), 156);
+        assert_eq!(p.question(), "What is 12 * 13?");
+        assert_eq!(p.sft_answer(), "12 * 13 = \\boxed{156}.");
+        let s = Problem {
+            a: 3,
+            b: 8,
+            op: Op::Sub,
+        };
+        assert_eq!(s.answer(), -5);
+    }
+
+    #[test]
+    fn parse_takes_last_boxed_and_handles_edges() {
+        assert_eq!(parse_boxed_answer("\\boxed{3} then \\boxed{156}."), Some(156));
+        assert_eq!(parse_boxed_answer("\\boxed{-7}"), Some(-7));
+        assert_eq!(parse_boxed_answer("\\boxed{ 42 }"), Some(42));
+        assert_eq!(parse_boxed_answer("no box here"), None);
+        assert_eq!(parse_boxed_answer("\\boxed{abc}"), None);
+        assert_eq!(parse_boxed_answer("\\boxed{unterminated"), None);
+    }
+
+    #[test]
+    fn check_is_exact_match() {
+        assert!(check_answer("the result is \\boxed{156}.", 156));
+        assert!(!check_answer("the result is \\boxed{155}.", 156));
+        assert!(!check_answer("no boxed answer at all", 156));
+    }
+
+    #[test]
+    fn sft_target_is_always_self_consistent() {
+        // The SFT target's boxed answer must always check against the problem's own
+        // gold — across all ops/operands. This is the M1 data invariant.
+        let cfg = GenConfig::default();
+        let mut rng = 12345u64;
+        for _ in 0..2000 {
+            let p = gen_problem(&mut rng, &cfg);
+            assert!(
+                check_answer(&p.sft_answer(), p.answer()),
+                "self-inconsistent SFT target for {p:?}"
+            );
+        }
+    }
+
+    #[test]
+    fn unique_space_counts_ordered_pairs_per_op() {
+        // add+sub+mul each contribute (max+1)^2 ordered pairs.
+        let cfg = GenConfig {
+            max_add: 9,
+            max_mul: 4,
+            ops: vec![Op::Add, Op::Sub, Op::Mul],
+        };
+        assert_eq!(unique_space(&cfg), 100 + 100 + 25);
+        // The shipped default is comfortably large (millions), so 20k requests are
+        // a tiny fraction and dedup stays fast.
+        assert!(unique_space(&GenConfig::default()) > 1_000_000);
+    }
+
+    #[test]
+    fn generation_is_deterministic_from_seed() {
+        let cfg = GenConfig::default();
+        let (mut r1, mut r2) = (7u64, 7u64);
+        for _ in 0..200 {
+            assert_eq!(
+                gen_problem(&mut r1, &cfg).key(),
+                gen_problem(&mut r2, &cfg).key()
+            );
+        }
+    }
+}

crates/xtrain-train/src/train_loop.rs

@@ -27,6 +27,12 @@ use crate::schedule::LrSchedule;
 pub struct TrainConfig {
     pub seq_len: usize,
     pub batch_size: usize,
+    /// Micro-batch gradient accumulation (Phase T16): each optimizer step
+    /// accumulates grads over `accum_steps` micro-batches of `batch_size`
+    /// sequences, giving an EFFECTIVE batch of `accum_steps × batch_size` at the
+    /// activation-memory cost of a single micro-batch. `1` = no accumulation
+    /// (bit-identical to the pre-T16 path).
+    pub accum_steps: usize,
     pub steps: usize,
     pub schedule: LrSchedule,
     pub weight_decay: f32,
@@ -74,31 +80,47 @@ pub fn train(
     // Best-val checkpointing only kicks in when we actually evaluate.
     let track_best = valid.is_some() && cfg.eval_every > 0;
 
+    let accum = cfg.accum_steps.max(1);
     for step in 0..cfg.steps {
         let lr = cfg.schedule.lr(step);
 
-        // Sample `batch_size` sequences and run them as ONE batched forward/
-        // backward. The CE mean over all batch*seq rows is the batch-mean loss, so
-        // backward already yields the batch-mean gradient (clip pre-scale = 1.0).
-        let mut inputs = Vec::with_capacity(cfg.batch_size);
-        let mut targets_v = Vec::with_capacity(cfg.batch_size);
-        for _ in 0..cfg.batch_size {
-            let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
-            inputs.push(input);
-            targets_v.push(target);
-        }
-        let ids = batched_ids_tensor(&inputs, device);
-        let targets = batched_ids_tensor(&targets_v, device);
+        // Accumulate grads over `accum` micro-batches of `batch_size` sequences,
+        // then take ONE optimizer step (Phase T16). Each micro-batch is ONE batched
+        // forward/backward; its loss is the CE mean over batch*seq rows, so backward
+        // yields that micro-batch's mean grad. To make the SUM over `accum` micro-
+        // batches equal a single step over an `accum × batch` batch, each micro-loss
+        // is scaled by 1/accum before backward (the tape SUM-accumulates the scaled
+        // grads). `accum == 1` skips the scale entirely → bit-identical to pre-T16.
+        let mut step_loss_sum = 0.0f32;
         // Training mode → dropout active (T18; no-op when cfg.dropout == 0). Set
         // each step so it is restored after a periodic eval flips to eval mode.
+        // Each micro-step's forward bumps the per-step seed → fresh masks.
         model.train();
-        let loss = model.loss_batched(&ids, &targets, cfg.batch_size);
-        let step_loss = read_scalar(&loss);
-        loss.backward();
-        tokens_seen += (cfg.batch_size * cfg.seq_len) as u64;
+        for _ in 0..accum {
+            let mut inputs = Vec::with_capacity(cfg.batch_size);
+            let mut targets_v = Vec::with_capacity(cfg.batch_size);
+            for _ in 0..cfg.batch_size {
+                let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
+                inputs.push(input);
+                targets_v.push(target);
+            }
+            let ids = batched_ids_tensor(&inputs, device);
+            let targets = batched_ids_tensor(&targets_v, device);
+            let loss = model.loss_batched(&ids, &targets, cfg.batch_size);
+            step_loss_sum += read_scalar(&loss);
+            if accum == 1 {
+                loss.backward();
+            } else {
+                xtrain_autodiff::ops::scale(&loss, 1.0 / accum as f32).backward();
+            }
+            tokens_seen += (cfg.batch_size * cfg.seq_len) as u64;
+        }
+        // Reported loss = mean over the effective batch = mean of the raw micro
+        // losses (each is itself a micro-batch mean of equal size).
+        let step_loss = step_loss_sum / accum as f32;
         losses.push(step_loss);
 
-        // Backward already produced the batch-mean gradient — just clip it.
+        // Backward already produced the effective-batch mean gradient — just clip.
         let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
         opt.step(lr, &params);
         for p in &params {
@@ -185,7 +207,7 @@ pub fn eval_loss(
             break;
         }
         let input: Vec<i32> = valid.tokens[s..s + seq].to_vec();
-        let target: Vec<i32> = valid.tokens[s + 1..s + seq + 1].to_vec();
+        let target = valid.target_window(s, seq);
         let ids = ids_tensor(&input, device);
         let targets = ids_tensor(&target, device);
         let loss = model.loss(&ids, &targets);

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
+    );
+}

crates/xtrain-train/tests/grad_accum.rs

@@ -0,0 +1,295 @@
+// T16 gradient-accumulation correctness gates.
+//
+// Gradient accumulation is mathematically EXACT: accumulating the grads of N
+// micro-batches of B sequences (each micro-loss scaled by 1/N before backward,
+// the tape SUM-accumulating) equals a single step over one N·B-sequence batch.
+// This file makes that a closed loop on-GPU, plus the accum_steps=1 bit-identity
+// regression guard.
+//
+//   1. accum_equiv_big_batch: same init, same N·B sequences in the same order.
+//      Path A = ONE batched loss over all N·B (the big-batch baseline). Path B =
+//      N micro-backwards of B each, scale(1/N), tape SUM. Assert loss and EVERY
+//      parameter grad match within fp tolerance (only the summation order differs,
+//      like the T8 DDP-vs-single-GPU and T13 recompute gates).
+//   2. accum1_bit_identical: accum_steps=1 must reproduce the no-accum path
+//      bit-for-bit (the implementation skips the ×1/1 scale entirely) — every
+//      parameter grad max|Δ| == 0.0.
+//   3. accum_train_converges: drive the real `train()` loop with accum and assert
+//      the per-step effective-batch loss trace tracks a big-batch baseline (errors
+//      stay bounded over many AdamW steps, not just one).
+#![cfg(not(no_cuda))]
+
+use xtrain_autodiff::ops;
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
+use xtrain_tensor::Device;
+use xtrain_train::data::Corpus;
+use xtrain_train::schedule::LrSchedule;
+use xtrain_train::{TrainConfig, train};
+
+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)
+        }
+    })
+}
+
+fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
+    t.to_device(Device::Cpu).as_slice::<f32>().to_vec()
+}
+
+// `n` deterministic (seq, target) pairs for the equivalence tests.
+fn make_seqs(n: usize, seq: usize, vocab: usize) -> (Vec<Vec<i32>>, Vec<Vec<i32>>) {
+    let seqs = (0..n)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 7 + i * 3 + 1) % vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let tgts = (0..n)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 5 + i * 2 + 2) % vocab) as i32)
+                .collect()
+        })
+        .collect();
+    (seqs, tgts)
+}
+
+// Run one big-batch forward/backward over all `seqs` and return the grads.
+fn big_batch_grads(
+    model: &TinyTransformer,
+    device: Device,
+    seqs: &[Vec<i32>],
+    tgts: &[Vec<i32>],
+) -> (f32, Vec<Vec<f32>>) {
+    let n = seqs.len();
+    let ids = batched_ids_tensor(seqs, device);
+    let tgt = batched_ids_tensor(tgts, device);
+    let loss = model.loss_batched(&ids, &tgt, n);
+    let loss_val = host(&loss.value())[0];
+    loss.backward();
+    let grads = model
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("grad")))
+        .collect();
+    (loss_val, grads)
+}
+
+// Accumulate over `accum` micro-batches of `b` sequences (drawn in order from the
+// flat `seqs`/`tgts`), scaling each micro-loss by 1/accum before backward; the
+// tape SUM-accumulates. Returns the mean of the raw micro losses + accumulated grads.
+fn accum_grads(
+    model: &TinyTransformer,
+    device: Device,
+    seqs: &[Vec<i32>],
+    tgts: &[Vec<i32>],
+    accum: usize,
+    b: usize,
+    scale: bool,
+) -> (f32, Vec<Vec<f32>>) {
+    let mut loss_sum = 0.0f32;
+    for m in 0..accum {
+        let s = &seqs[m * b..(m + 1) * b];
+        let t = &tgts[m * b..(m + 1) * b];
+        let ids = batched_ids_tensor(s, device);
+        let tgt = batched_ids_tensor(t, device);
+        let loss = model.loss_batched(&ids, &tgt, b);
+        loss_sum += host(&loss.value())[0];
+        if scale {
+            ops::scale(&loss, 1.0 / accum as f32).backward();
+        } else {
+            loss.backward(); // accum==1 bit-identity path
+        }
+    }
+    let grads = model
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("grad")))
+        .collect();
+    (loss_sum / accum as f32, grads)
+}
+
+#[test]
+fn accum_equiv_big_batch() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 3;
+    let b = 2usize; // micro-batch
+    let accum = 4usize; // → effective batch 8
+    let seq = 6usize;
+    let (seqs, tgts) = make_seqs(b * accum, seq, cfg.vocab);
+
+    // Big-batch baseline (accum_steps=1, batch = b·accum).
+    let big = build(cfg, device);
+    let (big_loss, big_grads) = big_batch_grads(&big, device, &seqs, &tgts);
+
+    // Accumulated (accum micro-batches of b, scale 1/accum).
+    let acc = build(cfg, device);
+    let (acc_loss, acc_grads) = accum_grads(&acc, device, &seqs, &tgts, accum, b, true);
+
+    let loss_rel = (big_loss - acc_loss).abs() / big_loss.abs().max(1e-4);
+    let mut max_grad_rel = 0.0f32;
+    for (bg, ag) in big_grads.iter().zip(&acc_grads) {
+        for (x, y) in bg.iter().zip(ag) {
+            max_grad_rel = max_grad_rel.max((x - y).abs() / x.abs().max(1e-3));
+        }
+    }
+    println!(
+        "accum=={accum}×b{b} vs big-batch{}: loss {big_loss:.6}/{acc_loss:.6} (rel {loss_rel:.2e}), \
+         grad max rel {max_grad_rel:.3e}",
+        b * accum
+    );
+    // fp summation order differs (big batch sums b·accum rows once; accum sums per
+    // micro then across micros) → tight fp tol, same convention as T13 recompute.
+    assert!(loss_rel < 1e-5, "loss diverged: {loss_rel:.2e}");
+    assert!(
+        max_grad_rel < 1e-4,
+        "accum grads diverged from big batch: {max_grad_rel:.3e}"
+    );
+}
+
+#[test]
+fn accum1_bit_identical() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 3;
+    let b = 4usize;
+    let seq = 6usize;
+    let (seqs, tgts) = make_seqs(b, seq, cfg.vocab);
+
+    // No-accum reference: one batched loss + backward (the pre-T16 path).
+    let reference = build(cfg, device);
+    let (_, ref_grads) = big_batch_grads(&reference, device, &seqs, &tgts);
+
+    // accum_steps=1 path: the loop runs ONE micro-batch and (by design) skips the
+    // ×1/1 scale → must be byte-for-byte identical to the reference backward.
+    let accum1 = build(cfg, device);
+    let (_, a1_grads) = accum_grads(&accum1, device, &seqs, &tgts, 1, b, false);
+
+    let mut max_abs = 0.0f32;
+    for (r, a) in ref_grads.iter().zip(&a1_grads) {
+        for (x, y) in r.iter().zip(a) {
+            max_abs = max_abs.max((x - y).abs());
+        }
+    }
+    println!("accum_steps=1 vs no-accum: grad max |Δ| = {max_abs:.3e}");
+    assert_eq!(
+        max_abs, 0.0,
+        "accum_steps=1 not bit-identical to no-accum: {max_abs:.3e}"
+    );
+}
+
+// A self-contained synthetic corpus (no tokenizer / data file needed).
+fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
+    Corpus {
+        tokens: (0..n_tokens)
+            .map(|i| (i * 7 + 3) as i32 % vocab as i32)
+            .collect(),
+        labels: None,
+        vocab_size: vocab,
+    }
+}
+
+#[test]
+fn accum_train_converges() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let vocab = 64usize;
+    let mut cfg = Config::tiny();
+    cfg.vocab = vocab;
+    cfg.n_layers = 2;
+    let corpus = synth_corpus(vocab, 4096);
+    let steps = 20usize;
+    let seq = 32usize;
+
+    // Same per-step RNG stream + effective batch 8 either way: the big-batch run
+    // (accum=1, batch=8) and the accumulated run (accum=4, batch=2) draw the SAME
+    // 8 sequences per step in the same order, so the per-step loss/grads — and thus
+    // the whole AdamW trajectory — track within fp tolerance.
+    let sched = LrSchedule {
+        max_lr: 3e-3,
+        min_lr: 3e-4,
+        warmup: 3,
+        total: steps,
+    };
+    let base = |batch, accum| TrainConfig {
+        seq_len: seq,
+        batch_size: batch,
+        accum_steps: accum,
+        steps,
+        schedule: sched.clone(),
+        weight_decay: 0.1,
+        max_grad_norm: 1.0,
+        log_every: 1_000_000,
+        ckpt_path: None,
+        ckpt_every: 0,
+        eval_every: 0,
+        eval_batches: 0,
+        seed: 7,
+    };
+
+    let big_model = build(cfg, device);
+    let big = train(&big_model, device, &corpus, None, &base(8, 1)).train_losses;
+
+    let acc_model = build(cfg, device);
+    let acc = train(&acc_model, device, &corpus, None, &base(2, 4)).train_losses;
+
+    let mut max_rel = 0.0f32;
+    for (x, y) in big.iter().zip(&acc) {
+        max_rel = max_rel.max((x - y).abs() / x.abs().max(1e-6));
+    }
+    // Final params should also stay close (errors don't blow up over the run).
+    let mut max_pdiff = 0.0f32;
+    for (p, q) in big_model.params().iter().zip(&acc_model.params()) {
+        for (x, y) in host(&p.value()).iter().zip(host(&q.value())) {
+            max_pdiff = max_pdiff.max((x - y).abs() / x.abs().max(1e-6));
+        }
+    }
+    println!(
+        "accum(4×2) vs big(8) over {steps} steps: loss[last] {:.6}/{:.6} max_rel {max_rel:.2e}, \
+         final param max rel {max_pdiff:.2e}",
+        big.last().unwrap(),
+        acc.last().unwrap()
+    );
+    assert!(
+        max_rel < 1e-3,
+        "accum loss trajectory diverged: {max_rel:.3e}"
+    );
+    assert!(
+        max_pdiff < 1e-2,
+        "accum final params diverged: {max_pdiff:.3e}"
+    );
+}

crates/xtrain-train/tests/real_training.rs

@@ -84,6 +84,7 @@ fn trains_on_tinystories() {
     let tcfg = TrainConfig {
         seq_len: 64,
         batch_size: 8,
+        accum_steps: 1,
         steps,
         schedule: LrSchedule {
             max_lr: 3e-3,

csrc/ops/flash_attention.cu

@@ -0,0 +1,281 @@
+// Hand-written fused flash-attention (Phase T14).
+//
+// The T10 composed SDPA path is 3 launches that MATERIALIZE the [bh,S,S] score
+// matrix: cublasSgemmStridedBatched (Q·Kᵀ) → causal-softmax kernel (writes the
+// whole probs) → cublasSgemmStridedBatched (P·V), and backward caches that whole
+// probs. flash-attention NEVER materializes N×N: a single fused kernel streams
+// over KV tiles with an ONLINE softmax (running max/sum + rescaled V accumulator),
+// so peak attention activation drops from O(S²) to O(S·hd) (= the output itself).
+//
+// Layout (matches the T10 op): Q/K/V/out are [bh, S, hd] row-major contiguous,
+// bh = batch·n_heads. The query's position within its sequence is the row index
+// within its [S,hd] block (so the flat row's qpos = (row % S) is automatic here —
+// we index per (bh, row)). CAUSAL: a query at position i attends to keys j ≤ i.
+// `scale` (= 1/sqrt(hd)) is folded into the logits before the max/exp.
+//
+// All F32, contiguous. (bf16 callers upcast Q/K/V → f32 on the Rust side and
+// downcast the f32 out, mirroring the composed path's fp32 softmax policy, so the
+// kernel only ever sees fp32.) Reduction helpers are inlined (self-contained file,
+// matching the csrc/ layout).
+//
+// Parallelisation: grid = bh*S, one block per query row; blockDim.x threads
+// cooperate. Forward keeps m (running max), l (running sum), acc[hd] (rescaled
+// V accumulator) in shared memory, streams KV in tiles of BK. Backward recomputes
+// scores from Q/K/V + the saved logsumexp L[bh,S] (NO cached probs), uses
+// D[i]=Σ dOᵢ·Oᵢ to collapse the softmax Jacobian, and atomicAdds dK/dV (which are
+// accumulated across query rows).
+
+#include <math.h>
+
+extern "C" {
+
+__device__ __forceinline__ float fa_warp_sum(float v) {
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+        v += __shfl_down_sync(0xffffffff, v, off);
+    return v;
+}
+__device__ __forceinline__ float fa_warp_max(float v) {
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off));
+    return v;
+}
+__device__ __forceinline__ float fa_block_sum(float v) {
+    __shared__ float sh[32];
+    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
+    int nwarps = (blockDim.x + 31) >> 5;
+    v = fa_warp_sum(v);
+    if (lane == 0) sh[warp] = v;
+    __syncthreads();
+    v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : 0.0f;
+    if (warp == 0) v = fa_warp_sum(v);
+    __shared__ float bc;
+    if (threadIdx.x == 0) bc = v;
+    __syncthreads();
+    return bc;
+}
+__device__ __forceinline__ float fa_block_max(float v) {
+    __shared__ float sh[32];
+    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
+    int nwarps = (blockDim.x + 31) >> 5;
+    v = fa_warp_max(v);
+    if (lane == 0) sh[warp] = v;
+    __syncthreads();
+    v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : -INFINITY;
+    if (warp == 0) v = fa_warp_max(v);
+    __shared__ float bc;
+    if (threadIdx.x == 0) bc = v;
+    __syncthreads();
+    return bc;
+}
+
+#define FA_TILE 32  // KV tile width (columns streamed per step)
+
+// One block per (bh-row, query-position). Computes out[bh, i, :] and L[bh, i] via
+// an online softmax that streams the keys in tiles of FA_TILE — the [S,S] score
+// row is never stored, only the per-tile partials flow through shared memory.
+__global__ void flash_attn_fwd_k(const float* Q, const float* K, const float* V,
+                                 float* O, float* L, int seq, int hd, float scale) {
+    int row = blockIdx.x;          // global query row over bh*S
+    int b = row / seq;             // which (batch,head) block
+    int i = row % seq;             // query position within the sequence (causal limit)
+    int t = threadIdx.x;
+    int nthreads = blockDim.x;
+
+    const float* q = Q + (size_t)row * hd;
+    const float* kb = K + (size_t)b * seq * hd;  // this block's keys  [seq,hd]
+    const float* vb = V + (size_t)b * seq * hd;  // this block's values[seq,hd]
+
+    // Q row in shared memory (reused every tile); acc accumulator over hd.
+    extern __shared__ float smem[];
+    float* sq = smem;          // [hd]
+    float* acc = smem + hd;    // [hd]
+    for (int d = t; d < hd; d += nthreads) {
+        sq[d] = q[d];
+        acc[d] = 0.0f;
+    }
+    __shared__ float m_run, l_run;
+    if (t == 0) { m_run = -INFINITY; l_run = 0.0f; }
+    __syncthreads();
+
+    int valid = i + 1;  // causal: attend to keys [0, i]
+    for (int j0 = 0; j0 < valid; j0 += FA_TILE) {
+        int tile = min(FA_TILE, valid - j0);
+        // Each thread computes whole logits for a strided subset of the tile's
+        // columns: s = scale * (q · k_j). hd is small (≤128) so the per-thread
+        // dot loop is cheap; this avoids a block-reduce per column.
+        __shared__ float s_tile[FA_TILE];
+        for (int c = t; c < tile; c += nthreads) {
+            const float* kj = kb + (size_t)(j0 + c) * hd;
+            float dot = 0.0f;
+            for (int d = 0; d < hd; ++d) dot += sq[d] * kj[d];
+            s_tile[c] = dot * scale;
+        }
+        __syncthreads();
+
+        // Tile max, then online rescale of (m, l, acc).
+        float tmax = -INFINITY;
+        for (int c = t; c < tile; c += nthreads) tmax = fmaxf(tmax, s_tile[c]);
+        tmax = fa_block_max(tmax);
+
+        __shared__ float m_new, corr;
+        if (t == 0) {
+            float mn = fmaxf(m_run, tmax);
+            corr = (m_run == -INFINITY) ? 0.0f : expf(m_run - mn);  // rescale old state
+            m_new = mn;
+        }
+        __syncthreads();
+
+        // Overwrite s_tile with the softmax weights p = exp(s - m_new) ONCE per
+        // column (instead of recomputing expf inside the per-dim V loop, which
+        // would cost hd× the transcendentals). Sum them for l.
+        float lsum = 0.0f;
+        for (int c = t; c < tile; c += nthreads) {
+            float p = expf(s_tile[c] - m_new);
+            s_tile[c] = p;
+            lsum += p;
+        }
+        lsum = fa_block_sum(lsum);
+
+        // Rescale old accumulator + add this tile's p·V (p cached in s_tile).
+        // Each thread owns a strided subset of hd; loops over the tile columns.
+        for (int d = t; d < hd; d += nthreads) {
+            float a = acc[d] * corr;
+            for (int c = 0; c < tile; ++c)
+                a += s_tile[c] * vb[(size_t)(j0 + c) * hd + d];
+            acc[d] = a;
+        }
+        if (t == 0) {
+            l_run = l_run * corr + lsum;
+            m_run = m_new;
+        }
+        __syncthreads();
+    }
+
+    // out = acc / l ; L = m + log(l) (logsumexp, saved for backward).
+    float inv = 1.0f / l_run;
+    for (int d = t; d < hd; d += nthreads) O[(size_t)row * hd + d] = acc[d] * inv;
+    if (t == 0) L[row] = m_run + logf(l_run);
+}
+
+void launch_flash_attention_fwd_f32(const float* q, const float* k, const float* v,
+                                    float* o, float* l, int bh, int seq, int hd,
+                                    float scale, void* s) {
+    int blk = hd < 1024 ? hd : 1024;
+    if (blk < 32) blk = 32;
+    size_t shmem = (size_t)2 * hd * sizeof(float);  // sq[hd] + acc[hd]
+    flash_attn_fwd_k<<<bh * seq, blk, shmem, (cudaStream_t)s>>>(q, k, v, o, l, seq, hd, scale);
+}
+
+// Per-row D[i] = Σ_d dO[i,d] · O[i,d]. One block per row (bh*S rows). Used to
+// collapse the softmax Jacobian in backward (Σ_j P_ij dP_ij = dOᵢ·Oᵢ).
+__global__ void flash_attn_rowdot_k(const float* dO, const float* O, float* D, int hd) {
+    int row = blockIdx.x;
+    int t = threadIdx.x;
+    const float* d = dO + (size_t)row * hd;
+    const float* o = O + (size_t)row * hd;
+    float v = 0.0f;
+    for (int c = t; c < hd; c += blockDim.x) v += d[c] * o[c];
+    v = fa_block_sum(v);
+    if (t == 0) D[row] = v;
+}
+
+// Backward: one block per query row i. Recomputes scores from Q/K/V + the saved
+// logsumexp L (NO cached probs), streams KV in tiles. dQ accumulates locally (this
+// row owns it). dK/dV are accumulated ACROSS query rows so they atomicAdd into the
+// shared global buffers (pre-zeroed by the caller).
+//   p_ij = exp(Qᵢ·Kⱼ·scale - L[i]) ; dp_ij = dOᵢ·Vⱼ ;
+//   ds_ij = p_ij·(dp_ij - D[i])·scale
+//   dQᵢ += Σ_j ds_ij·Kⱼ ; dKⱼ += ds_ij·Qᵢ ; dVⱼ += p_ij·dOᵢ
+__global__ void flash_attn_bwd_k(const float* Q, const float* K, const float* V,
+                                 const float* dO, const float* L, const float* D,
+                                 float* dQ, float* dK, float* dV,
+                                 int seq, int hd, float scale) {
+    int row = blockIdx.x;
+    int b = row / seq;
+    int i = row % seq;
+    int t = threadIdx.x;
+    int nthreads = blockDim.x;
+
+    const float* q = Q + (size_t)row * hd;
+    const float* doi = dO + (size_t)row * hd;
+    const float* kb = K + (size_t)b * seq * hd;
+    const float* vb = V + (size_t)b * seq * hd;
+    float* dkb = dK + (size_t)b * seq * hd;
+    float* dvb = dV + (size_t)b * seq * hd;
+
+    extern __shared__ float smem[];
+    float* sq = smem;          // [hd] Qᵢ
+    float* sdo = smem + hd;    // [hd] dOᵢ
+    float* dqa = smem + 2*hd;  // [hd] dQᵢ accumulator
+    for (int d = t; d < hd; d += nthreads) {
+        sq[d] = q[d];
+        sdo[d] = doi[d];
+        dqa[d] = 0.0f;
+    }
+    __shared__ float Li, Di;
+    if (t == 0) { Li = L[row]; Di = D[row]; }
+    __syncthreads();
+
+    int valid = i + 1;
+    for (int j0 = 0; j0 < valid; j0 += FA_TILE) {
+        int tile = min(FA_TILE, valid - j0);
+        // Phase 1: per-column ds[c] and p[c] (the column owner does the dots).
+        __shared__ float s_ds[FA_TILE];
+        __shared__ float s_p[FA_TILE];
+        for (int c = t; c < tile; c += nthreads) {
+            const float* kj = kb + (size_t)(j0 + c) * hd;
+            const float* vj = vb + (size_t)(j0 + c) * hd;
+            float sdot = 0.0f, dpdot = 0.0f;
+            for (int d = 0; d < hd; ++d) {
+                sdot += sq[d] * kj[d];
+                dpdot += sdo[d] * vj[d];
+            }
+            float p = expf(sdot * scale - Li);
+            s_p[c] = p;
+            s_ds[c] = p * (dpdot - Di) * scale;
+        }
+        __syncthreads();
+        // Phase 2: dV_j += p·dOᵢ ; dK_j += ds·Qᵢ (accumulated across rows → atomic).
+        // Spread the tile×hd atomics over ALL threads (was serial in the column
+        // owner) — flatten (c,d) so every thread issues a balanced share.
+        for (int idx = t; idx < tile * hd; idx += nthreads) {
+            int c = idx / hd, d = idx % hd;
+            size_t off = (size_t)(j0 + c) * hd + d;
+            atomicAdd(&dvb[off], s_p[c] * sdo[d]);
+            atomicAdd(&dkb[off], s_ds[c] * sq[d]);
+        }
+        // dQᵢ += Σ_c ds[c] · K_{j0+c}  (this row owns dQ — no atomic).
+        for (int d = t; d < hd; d += nthreads) {
+            float a = 0.0f;
+            for (int c = 0; c < tile; ++c)
+                a += s_ds[c] * kb[(size_t)(j0 + c) * hd + d];
+            dqa[d] += a;
+        }
+        __syncthreads();
+    }
+    for (int d = t; d < hd; d += nthreads) dQ[(size_t)row * hd + d] = dqa[d];
+}
+
+void launch_flash_attention_bwd_f32(const float* q, const float* k, const float* v,
+                                    const float* d_o, const float* l, float* d_d,
+                                    float* dq, float* dk, float* dv,
+                                    int bh, int seq, int hd, float scale, void* s) {
+    int blk = hd < 1024 ? hd : 1024;
+    if (blk < 32) blk = 32;
+    // d_d is the pre-computed D[i]=Σ dOᵢ·Oᵢ (the Rust wrapper runs rowdot first,
+    // since it holds the forward O). dq/dk/dv are pre-zeroed by the caller.
+    flash_attn_bwd_k<<<bh * seq, blk, (size_t)3 * hd * sizeof(float), (cudaStream_t)s>>>(
+        q, k, v, d_o, l, d_d, dq, dk, dv, seq, hd, scale);
+}
+
+// Standalone D = rowdot(dO, O) launcher (the Rust wrapper calls this before bwd).
+void launch_flash_attention_rowdot_f32(const float* d_o, const float* o, float* d_d,
+                                       int rows, int hd, void* s) {
+    int blk = hd < 1024 ? hd : 1024;
+    if (blk < 32) blk = 32;
+    flash_attn_rowdot_k<<<rows, blk, 0, (cudaStream_t)s>>>(d_o, o, d_d, hd);
+}
+
+} // extern "C"

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;
@@ -338,7 +428,7 @@ __global__ void cross_entropy_fwd_k(const float* x, const int* target,
     for (int c = threadIdx.x; c < cols; c += blockDim.x) pr[c] *= inv;
     if (threadIdx.x == 0) {
         int t = target[r];
-        loss[r] = -logf(pr[t]);
+        loss[r] = t < 0 ? 0.0f : -logf(pr[t]);
     }
 }
 void launch_cross_entropy_fwd_f32(const float* x, const int* target,
@@ -354,8 +444,13 @@ __global__ void cross_entropy_dx_k(const float* probs, const int* target,
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= rows * cols) return;
     int r = i / cols, c = i % cols;
-    float g = probs[i] - (c == target[r] ? 1.0f : 0.0f);
-    dx[i] = g * scale;
+    int t = target[r];
+    if (t < 0) {
+        dx[i] = 0.0f;
+    } else {
+        float g = probs[i] - (c == t ? 1.0f : 0.0f);
+        dx[i] = g * scale;
+    }
 }
 void launch_cross_entropy_dx_f32(const float* probs, const int* target,
                                  float* dx, int rows, int cols, float scale, void* s) {

csrc/ops/repeat_kv.cu

@@ -0,0 +1,84 @@
+// repeat_kv: the grouped-query-attention (GQA) head broadcast (Phase T15).
+//
+// GQA projects K/V to fewer heads than Q (num_kv_heads < num_heads); each KV head
+// is SHARED by a group of `group = num_heads / num_kv_heads` query heads. Before
+// the SDPA (composed or fused-flash, both untouched), we expand the KV tensor from
+// [B·num_kv, S, hd] to the full [B·nh, S, hd] so the existing per-(batch,head)
+// attention sees a full set of heads. GQA is then "free" for both SDPA paths.
+//
+// Layout: K/V are [bh_kv, S, hd] = [B·num_kv, S, hd] row-major, contiguous; the
+// output is [bh_q, S, hd] = [B·nh, S, hd]. The head ordering matches xserv's
+// repeat_kv (crates/xserv-model/src/qwen3.rs): query head qh reads kv head
+// qh/group, query heads CONTIGUOUS within a group (dst = kvh*group + r). So:
+//
+//   out[b·nh + qh, :, :] = in[b·num_kv + qh/group, :, :]
+//
+// Forward is a gather (each output row copies one input row). Backward is its
+// transpose: a kv head receives the SUM of the `group` query heads that share it
+//   din[b·num_kv + kvh, e] = Σ_{r∈[0,group)} dout[b·nh + kvh·group + r, e]
+// — the multi-group-to-one grad accumulation GQA's correctness hinges on. Each
+// input element is owned by exactly one thread that serially sums its `group`
+// source rows: race-free, NO atomics, run-to-run DETERMINISTIC. group==1 makes
+// both directions a plain copy (identity → bit-identical to the MHA path).
+//
+// All F32, contiguous. (bf16 callers upcast → f32 on the Rust side and downcast
+// the f32 result, mirroring the rest of the attention stack's fp32 policy.)
+
+#include <math.h>
+
+extern "C" {
+
+// Forward gather. grid-stride over the bh_q·S·hd output elements; each output
+// element copies from its kv-head source row. b = (out_bh / nh), qh = out_bh % nh,
+// kv source bh = b·num_kv + qh/group.
+__global__ void repeat_kv_fwd_k(const float* in, float* out, int nh, int num_kv,
+                                int group, int seq, int hd) {
+    long row_elems = (long)seq * hd;       // S·hd per head block
+    // One block per (batch, query-head); threads cover the S·hd block.
+    int out_bh = blockIdx.x;               // over B·nh
+    int b = out_bh / nh;
+    int qh = out_bh % nh;
+    int kvh = qh / group;
+    const float* src = in + ((long)b * num_kv + kvh) * row_elems;
+    float* dst = out + (long)out_bh * row_elems;
+    for (long e = threadIdx.x; e < row_elems; e += blockDim.x) dst[e] = src[e];
+}
+
+void launch_repeat_kv_fwd_f32(const float* in, float* out, int batch, int nh,
+                              int num_kv, int seq, int hd, void* s) {
+    int group = nh / num_kv;
+    int blk = (seq * hd) < 256 ? (seq * hd) : 256;
+    if (blk < 32) blk = 32;
+    repeat_kv_fwd_k<<<batch * nh, blk, 0, (cudaStream_t)s>>>(in, out, nh, num_kv,
+                                                             group, seq, hd);
+}
+
+// Backward sum. One block per (batch, kv-head); threads cover the S·hd block.
+// Each owned input element sums the `group` contiguous query-head source rows.
+__global__ void repeat_kv_bwd_k(const float* dout, float* din, int nh, int num_kv,
+                                int group, int seq, int hd) {
+    long row_elems = (long)seq * hd;
+    int in_bh = blockIdx.x;                // over B·num_kv
+    int b = in_bh / num_kv;
+    int kvh = in_bh % num_kv;
+    int qh0 = kvh * group;                 // first query head sharing this kv head
+    float* dst = din + (long)in_bh * row_elems;
+    const float* base = dout + ((long)b * nh + qh0) * row_elems;
+    for (long e = threadIdx.x; e < row_elems; e += blockDim.x) {
+        float acc = 0.0f;
+        for (int r = 0; r < group; ++r) acc += base[(long)r * row_elems + e];
+        dst[e] = acc;
+    }
+}
+
+void launch_repeat_kv_bwd_f32(const float* dout, float* din, int batch, int nh,
+                              int num_kv, int seq, int hd, void* s) {
+    int group = nh / num_kv;
+    int blk = (seq * hd) < 256 ? (seq * hd) : 256;
+    if (blk < 32) blk = 32;
+    repeat_kv_bwd_k<<<batch * num_kv, blk, 0, (cudaStream_t)s>>>(dout, din, nh,
+                                                                 num_kv, group,
+                                                                 seq, hd);
+}
+
+} // extern "C"

docs/13-flash-attention.md

@@ -0,0 +1,183 @@
+# Phase T14: 融合 Flash-Attention Kernel — Design Document
+
+## Goal
+
+T10 把 attention 批量化了，但它的 SDPA 走的是 **「物化 N×N scores」** 的组合路径：
+`cublasSgemmStridedBatched`（Q·Kᵀ）→ 一个 causal-softmax kernel（写出整张 probs）→
+`cublasSgemmStridedBatched`（P·V），**3 次 launch + 一张 `[bh, S, S]` 的 scores/probs 张量**
+常驻显存（反向还要缓存这张 probs）。S 一大，这张 N×N 就成了激活显存与带宽的主导项。
+
+T14 的目标：手写一个**单 kernel 的 fused flash-attention**——streaming / online softmax、**tiled
+over KV**、**绝不物化 N×N**。前向一发 kernel 直接吐出 `out[bh,S,hd]`（外加 `O(N)` 的 logsumexp）；
+反向一发 kernel（flash 式：重算 scores + dQ/dK/dV，同样不物化 N×N）。接进 model + autograd 作
+**opt-in `--flash`**，默认保留 T10 的 composed 路径以便 A/B。
+
+**硬闸门是诚实正确性**：新 kernel 的 dQ/dK/dV finite-diff grad-check 过；fwd/bwd 对现有 composed-SDPA
+路径数值贴合（进 bf16 容差）；PyTorch SDPA 对拍 B>1；峰值显存↓（不物化 scores）+ tok/s before/after 实测；
+全回归套（含 xserv 闭环 md5）开/关 flag 都绿——默认（flag off）图不变 → 不回归。
+
+## 什么是 flash-attention
+
+标准 attention 是 `O = softmax(causal(Q·Kᵀ/√d)) · V`，朴素实现把 `S[i,j] = Qᵢ·Kⱼ/√d` 整张
+`[S,S]` 算出来、softmax、再乘 V——显存 `O(S²)`、HBM 读写 `O(S²)`。
+
+**flash-attention** 的洞察：softmax 可以 **online（streaming）** 地算。把 K/V 切成若干 **tile**，对一个
+query 行 `i`，依次扫过 KV tile，用 **running max `m` + running sum `l`** 维护 softmax 的归一化，并把
+部分加权的 `V` 累加进一个 `[hd]` 的 accumulator `acc`，每来一个新 tile 就用「新旧 max 的差」对旧 `acc`/`l`
+做 rescale。扫完所有 tile，`out = acc / l`。**整张 `[S,S]` 从不落地**——只有 `[hd]` 的 acc 和两个标量
+在寄存器/共享内存里流动。峰值激活从 `O(S²)` 降到 `O(S·hd)`（就是 O 本身）。
+
+online softmax 的核心递推（block `j` 的部分 logits 行 `s_j`，旧状态 `m, l, acc`）：
+
+```text
+m_new = max(m, max_k s_j[k])
+p     = exp(s_j - m_new)                       # 本 tile 的未归一化权重
+l     = l * exp(m - m_new) + sum(p)            # 旧 sum 先 rescale，再加本 tile
+acc   = acc * exp(m - m_new) + p · V_tile      # 旧 acc 同样 rescale，再加本 tile 贡献
+m     = m_new
+# 扫完所有 tile：
+out   = acc / l
+L     = m + log(l)                             # logsumefp，O(N) 存给反向
+```
+
+**因果 mask 内联**：query 全局位置 = `i % S`（沿用 T10 的 per-seq 复位约定），KV 位置 `j` 满足
+`j > i%S` 的列直接当 `-inf`（`p=0`）。tile 整块在对角线之上可**直接 skip**（causal 的天然稀疏，省一半算力）。
+
+**反向（flash 式，[Dao 2022] 的标准做法）**：不缓存 probs，从 Q/K/V + 前向存的 `L[bh,S]` **重算** scores。
+关键预计算 `D[i] = Σ_d dOᵢ[d]·Oᵢ[d]`（每 query 一个标量，`O(N)`），则对每个 `(i,j)`：
+
+```text
+s_ij = Qᵢ·Kⱼ * scale                           # 重算 logit
+p_ij = exp(s_ij - L[i])                         # 重算 softmax 权重（L 是前向存的 logsumexp）
+dp_ij = dOᵢ · Vⱼ                                # 对 P 的梯度
+ds_ij = p_ij * (dp_ij - D[i]) * scale           # softmax 雅可比，化简掉了显式 N×N
+dQᵢ  += ds_ij * Kⱼ ;  dKⱼ += ds_ij * Qᵢ ;  dVⱼ += p_ij * dOᵢ
+```
+
+`ds = P ∘ (dP - D)` 是 softmax 反向用 `Σⱼ Pⱼ·dPⱼ = D`（因为 `D[i]=Σ dOᵢ·Oᵢ = Σⱼ Pᵢⱼ dPᵢⱼ`）化简的结果，
+**不需要 N×N 的 softmax 雅可比矩阵**。同样 tiled、同样不物化 N×N。
+
+## Module Layout（surgical：composed 路径逐字节不动，flash 全程新增并行路径）
+
+```
+csrc/ops/flash_attention.cu              # 新：fwd kernel（online softmax，tiled KV）+ bwd kernel（重算 + dQ/dK/dV）
+crates/xtrain-cuda/
+├── src/ffi.rs                           # +launch_flash_attention_fwd_f32 / _bwd_f32 声明
+└── build.rs                             # +flash_attention.cu
+crates/xtrain-tensor/src/tensor.rs       # +Tensor::flash_attention / flash_attention_backward（fwd 存 logsumexp L；bf16 upcast→f32 kernel→downcast）
+crates/xtrain-autodiff/
+├── src/ops.rs                           # +ops::flash_attention 节点（前向调 fwd，缓存 L，反向调 bwd）
+└── tests/autograd.rs                    # +flash_attention(batched) dQ/dK/dV grad-check
+crates/xtrain-model/
+├── src/model.rs                         # attention() 按 use_flash 选 ops::attention | ops::flash_attention；+with_flash(bool) builder；flash 标志透传 block_forward（recompute 段内也走 flash）
+└── tests/flash.rs                       # 新：flash == composed（fwd logits + 每参数梯度），参数化 fp32/bf16
+crates/xtrain-train/src/bin/train.rs     # +--flash flag → model.with_flash(true)
+crates/xtrain-distributed/src/bin/train_ddp.rs  # +--flash flag（DDP 路径）
+crates/xtrain-model/tests/parity_dump.rs # PyTorch B>1 对拍跑两遍：composed 与 flash（共用 PyTorch oracle）
+```
+
+## Key Design Decisions
+
+### ① 一个 block 负责一行 query（先做对，再谈快）
+
+最直接、最易验证正确的并行划分：**`grid = bh * S`，每个 block 算一整行 query 的 `out[bh, i, :]`**。
+block 内 `hd` 个线程（hd ≤ 128，正好一个 warp 多一点），共享 `m/l` 标量 + `acc[hd]`。block 顺序扫
+KV tile（tile 宽 `BK`，沿 `j` 维），每个 tile：线程并行算 `BK` 个 logit（点积 over hd 用 block-reduce）、
+求 tile max、online-rescale `m/l/acc`、累加 `p·V`。扫完写 `out = acc/l` 与 `L[i] = m + log(l)`。
+
+**为什么先这样而不是 FA2 的 query-tile 划分**：本项目的硬闸门是**正确性 + 不物化 N×N + 显存↓**，不是
+打榜峰值 FLOPs。一行一 block 的版本：(a) online softmax 与 N×N skip 已经完全落地（显存与带宽收益拿到），
+(b) 代码直白、逐 query 行可对拍，正确性风险最低。它**不会**比 cuBLAS 两发 GEMM 更快（cuBLAS tensor-core
+吃满），所以 tok/s 上 flash 在我们这种 `hd=32` 小头维下大概率**持平或略慢**——这正是 flash 的已知权衡
+（flash 的胜场是**显存**，不是小模型的 wall-clock）。把这点诚实写进 perf 表，不掩饰。
+
+### ② 前向只存 `L[bh,S]`（logsumefp），不存 probs
+
+composed 路径反向要缓存整张 `probs[bh,S,S]`（`O(N²)`）。flash 反向**只需要前向的 logsumexp
+`L[i]=m_i+log(l_i)`**（每 query 一个 fp32，`O(N)`）即可重算任意 `p_ij = exp(Qᵢ·Kⱼ·scale - L[i])`。
+所以 fwd kernel 顺手把 `L` 写出来，autograd 节点缓存它（外加 Q/K/V/O parents 本就在）。**这就是显存闸门的来源**：
+attention 的反向缓存从 `[bh,S,S]` 砍到 `[bh,S]`。
+
+### ③ 反向用 `D[i]=Σ dOᵢ·Oᵢ` 化简 softmax 雅可比
+
+softmax 反向通项 `ds_ij = p_ij·(dp_ij - Σ_k p_ik·dp_ik)`。注意 `Σ_k p_ik·dp_ik = Σ_k p_ik (dOᵢ·V_k)
+= dOᵢ·(Σ_k p_ik V_k) = dOᵢ·Oᵢ = D[i]`。所以一趟先算 `D[bh,S]`（每行 `dO·O` 的点积，`O(N)`），反向
+扫 KV tile 时直接 `ds = p·(dp - D)·scale`，**不需要再算或物化整行的 `Σ p·dp`**。
+dQ/dK/dV 三者：dQ 由「该 query 行」累加（block 私有，无竞争）；dK/dV 跨 query 行累加同一个 `(j)`
+→ 用 `atomicAdd` 到全局 dK/dV（fp32 原子加，确定 race-free）。
+
+### ④ bf16：kernel 内 fp32，边界 cast（与 composed 路径一致的数值策略）
+
+T10/T12 的 composed attention 对 bf16 也是 **softmax 用 fp32**（scores 升 f32 → kernel → probs 降回 bf16）。
+flash 沿用同策略，最省心且数值最稳：bf16 模式下 `flash_attention` 把 Q/K/V `to_dtype(F32)` 喂给 fp32 kernel，
+`out` 再 `to_dtype(BF16)`；反向同理。kernel 本身只有一份 fp32 实现。这样 flash 的 bf16 数值与 composed 的
+bf16 数值是**同一套 fp32 softmax 算的**，只差 GEMM rounding（cuBLAS tensor-core vs kernel 内 fp32 FMA）→ 落在
+既有 bf16 容差内。`L` 始终 fp32。
+
+> 备选（不采纳）：bf16 全程 in-kernel half。收益是少两次 cast，但 (a) 引入与 composed 不同的 softmax 累加路径，
+> 威胁 on-vs-off 贴合闸门；(b) 本规模 attention 非瓶颈。escape hatch：先 fp32-core 把正确性钉死，纯 half flash 留 follow-up。
+
+### ⑤ opt-in 透传：`use_flash` 是运行时旗标，不是架构
+
+`use_flash` 不进 `Config`（它不改模型尺寸、不改导出、不该污染 `num_params`），而是 `TinyTransformer` 的一个
+`bool` 字段 + `with_flash(bool)` builder（对齐 `with_recompute` / `with_compute_dtype`）。`block_forward` 已经
+是 `(cfg, cdt, …)` 的自由函数（T13 为 recompute 抽的），给它加一个 `flash: bool` 形参，model 的 `attention()`
+据此选 `ops::attention`（composed）或 `ops::flash_attention`。recompute 闭包捕获 `flash`（`Copy`）→ **重算段内也走
+flash**，flash×recompute 组合天然成立。默认 `false` = composed 路径**逐字节不变**（硬闸门：默认图不变 → 不回归）。
+
+## 验证方法
+
+**硬闸门全绿（dash5 实跑 capture）：**
+
+### 1. 正确性
+
+- **新 kernel dQ/dK/dV finite-diff grad-check**（`xtrain-autodiff/tests/autograd.rs::flash_attention_batched_bwd`）：
+  与既有 `attention_batched_bwd` 同构（`L = sum(W∘out)`，中心差分），断 dQ/dK/dV 在 `cfg_nonlinear`/`cfg_linear` 容差内。
+- **flash == composed**（`xtrain-model/tests/flash.rs`）：同 init 两个模型（flash on/off），同一 batched
+  loss + backward，断**前向 logits / loss / 每参数梯度**在紧容差内一致；参数化 fp32（近逐位）与 bf16（bf16 舍入级）。
+- **PyTorch SDPA 对拍 B>1**（`parity_dump.rs` + `parity.py`）：等价 PyTorch 模型（per-seq RoPE、per-seq causal、
+  QK-norm、SwiGLU）对拍 forward logits + 全部参数梯度——**composed 与 flash 两条都跑**，共用同一 PyTorch oracle。
+- **全回归套开/关 `--flash`**：autograd 15、structural、batched==looped、bf16、recompute（逐位）、overfit 27/27、
+  AdamW（GPU bit-exact + host 对 torch）、DDP loss-match + 跨 rank、**xserv 闭环（导出 safetensors → md5 对 registry →
+  xserv 贪心逐 token 一致）**。flag off 默认图不变 → composed 数值不回归。
+
+### 2. 显存（payoff）—— 不物化 N×N 的直接收益
+
+dash5 1× RTX 5090，同 config，nvidia-smi 峰值，flash off vs on：attention 反向缓存 `[bh,S,S]→[bh,S]`，
+峰值显存应↓（尤其 seq 大时）。capture 实际数字进表。
+
+### 3. 吞吐
+
+同 config steady-state tok/s flash off vs on。预期：本规模 `hd=32` 下 flash kernel **持平或略慢于** cuBLAS 双
+GEMM（小头维喂不满 tensor-core 是 flash 的已知权衡，胜场在显存）——诚实报告，不为绿而调。
+
+## 实测结果（dash5 1× RTX 5090）
+
+**正确性（硬闸门全绿）：**
+
+| 闸门 | 结果 |
+|---|---|
+| ① 新 kernel dQ/dK/dV finite-diff grad-check | **过** — dQ 9.3e-3 / dK 1.7e-2 / dV 5.6e-4（单 tile 干净区；多 tile 由②兜） |
+| flash fwd 对 composed | max rel **6.7e-5** |
+| flash bwd 对（已 grad-check 的）composed bwd | dQ **1.7e-5** / dK 1.2e-5 / dV 4.3e-5 |
+| ② flash==composed（model 级，logits/loss/每参数梯度） | fp32: loss rel **0.0**、logits 1.7e-4、grad 4.4e-5；bf16: loss 1.5e-4、logits mean 1.6e-3/p99 5.9e-3、grad scaled-mean 1.2e-2 |
+| ③ PyTorch SDPA 对拍 B>1（flash 路径，共用 composed oracle） | loss relerr **4.98e-8**、logits **7.92e-6**、25 参数 grad 全进 rtol 0.02 |
+| ⑤ 回归套（flag off 默认 + flash 路径都测）：autograd 18 / structural 5 / batched / bf16 / **flash 3** / overfit 27/27 / recompute 2 / AdamW(GPU+host) / GEMM / DDP 2 / checkpoint-roundtrip | **全绿** |
+| ⑤ xserv 闭环 md5（v3 ckpt 用 T14 代码重导 safetensors） | **逐位一致** `b04fc9f9a0c9af04c47d9ca649aea12e`（与 registry 同）→ 默认 export 零漂移 |
+| ⑤ xserv 闭环（flash 训练 → 导出 → xserv 服务贪心） | flash-训出 coherent TinyStories；xserv(BF16) 对 xtrain(F32) 贪心：3 prompt 中 "One day" 逐 token 一致，其余在 ~0.5% BF16 漂移处晚分叉（与 v1/v2/v3 同款） |
+
+> **finite-diff 的诚实记录**：长 softmax（seq>tile）会产生大量近零梯度元素，中心差分在那些元素上不可靠（出现伪 0.0 / 符号翻转——不是 backward bug）。故 ① 的 finite-diff 跑**单 tile 干净区**（seq=5，对齐既有 composed grad-check 的良态区），**多 tile 的 streaming/online 路径**用「flash bwd 对已 grad-check 的 composed bwd」（seq=40，dQ 1.7e-5）兜——比 finite-diff 更利。dQ/dK 用 eps=2e-3 压低 f32 舍入项（~4e-4 小梯度上舍入项压过截断项）。**没有为凑绿放宽容差**。
+
+**④ 显存 + 吞吐（payoff vs tradeoff，dim768=8L/12h×64/ffn3072, bf16, steady-state）：**
+
+| config | path | 峰值显存 | tok/s |
+|---|---|---|---|
+| batch8 seq1024 | composed (off) | 24670 MiB | **58.6K** |
+| batch8 seq1024 | **flash (on)** | **20736 MiB（−16%）** | 25.0K（−57%, ~2.3× 慢） |
+| batch2 seq2048 | composed (off) | 17264 MiB | 36.7K |
+| batch2 seq2048 | **flash (on)** | **13246 MiB（−23%）** | 13.2K（−64%） |
+
+→ **显存按预期降**（不物化 `[bh,S,S]`），且**收益随 seq 增长**（seq1024 −16% → seq2048 −23%，O(S²) 砍掉）。
+**tok/s 如设计 ① 预测的「持平或略慢」实为 ~2.3–2.8× 慢**：hd=64 的小头维下，手写「一行一 block + 串行扫 KV」kernel 喂不满 SM，干不过 cuBLAS tensor-core 的两发批量 GEMM——这正是 flash 的已知权衡（**胜场在显存，不是小模型 wall-clock**），诚实报告不掩饰。两个落地的优化（softmax 权重缓存进 shared 省 hd× 的 expf；dK/dV 原子加摊到全 block 而非串行在列 owner 内）把 backward 从 6.8× 慢拉到 2.3× 慢——主瓶颈是 backward 的跨行原子累加（FA2 用 K-block 拥有 dK/dV 的独立 pass 解，本版未做，留 follow-up）。
+
+> **escape hatch（follow-up，未做，记给后续）**：① FA2 式 query-tile 划分（一 block 多 query 行，K/V 进 shared 复用）提 SM 占用；② backward 的 dK/dV 改 K-block-owned 独立 pass 消跨行原子；③ 纯 bf16 in-kernel（省两次 cast）。本规模 attention 非训练瓶颈、且会动数值贴合闸门，按 escape hatch 推迟——T14 先把**正确性 + 不物化 N×N + 显存↓**钉死。

docs/14-gqa.md

@@ -0,0 +1,180 @@
+# Phase T15: Grouped-Query Attention (GQA) — Design Document
+
+## Goal
+
+到 T14 为止，xtrain 的 attention 都是 **MHA**（`num_kv_heads = num_heads`）——每个
+query 头有自己独立的 K/V 头。导出 xserv 时 `num_key_value_heads = num_attention_heads`
+（退化 GQA，docs/08）。
+
+T15 做**真正的 grouped-query attention**：`num_kv_heads < num_heads`，K/V 只投影到
+`num_kv_heads · head_dim`，每个 KV 头被一组 `group = num_heads / num_kv_heads` 个 query
+头**共享**（repeat_kv / broadcast）。GQA 是现代 LLM（Llama-2-70B、Qwen2/3、Mistral）的标配
+——它把 KV cache 显存（推理）与 K/V 投影参数（训练）按 `group` 倍压缩，几乎不掉质量。
+
+**硬闸门是诚实正确性**，重点在 **repeat_kv 的反向梯度累加**：一个 KV 头被 `group` 个 query 头
+共享，反向时这 `group` 个 query 头各自对该 KV 头的梯度必须**正确求和**回到那一个共享 KV 头上。
+这条「多组 q 头梯度汇到一个 kv 头」的累加路径是本任务最易错处，单列为首要 grad-check 闸门。
+
+GQA 必须**同时**接进 T14 的 fused flash kernel（优先）与旧 composed/batched SDPA 路径，且
+`num_kv_heads == num_heads`（`group = 1`）时与现有 MHA 路径**逐位一致**（回归保护）。
+
+## 什么是 GQA
+
+MHA：`num_heads` 个 query 头，每个配一个独立 K/V 头。
+MQA（multi-query）：所有 query 头共享**一个** K/V 头（极端）。
+GQA：折中——`num_kv_heads` 个 K/V 头，每个被 `group = num_heads/num_kv_heads` 个相邻 query
+头共享。`num_kv_heads = num_heads` 退化为 MHA，`num_kv_heads = 1` 退化为 MQA。
+
+```
+num_heads = 8, num_kv_heads = 2  →  group = 4
+q heads:  0 1 2 3 4 5 6 7
+kv heads: 0 0 0 0 1 1 1 1        # q head qh 用 kv head qh/group（相邻分组，连续）
+```
+
+**分组约定必须与 xserv repeat_kv 一致**（闭环命门）：xserv 的 `repeat_kv`
+（`crates/xserv-model/src/qwen3.rs`）把 kv 头 `kvh` 复制到目标头 `dst = kvh*group + r`
+（`r∈[0,group)`），即**query 头 `qh` 读 kv 头 `qh/group`，组内 query 头连续**。xtrain 的
+repeat_kv 用同一映射，否则导出的 `q_proj` 行块与 kv 头对不上 → 闭环必崩。
+
+## Module Layout（surgical：复用已验证的两条 SDPA，GQA = 头维 broadcast op）
+
+```
+csrc/ops/repeat_kv.cu                     # 新：repeat_kv fwd（头块 gather）+ bwd（组内 group 行求和，无 atomic，确定性）
+crates/xtrain-cuda/
+├── src/ffi.rs                            # +launch_repeat_kv_fwd_f32 / _bwd_f32 声明（no_cuda 门控）
+└── build.rs                              # +repeat_kv.cu
+crates/xtrain-tensor/src/tensor.rs        # +Tensor::repeat_kv / repeat_kv_backward（[B*kvh,S,hd]→[B*nh,S,hd]；bf16 upcast→f32→downcast）
+crates/xtrain-autodiff/
+├── src/ops.rs                            # +ops::repeat_kv 节点（fwd broadcast，bwd 组内求和）
+└── tests/autograd.rs                     # +repeat_kv grad-check（含 group>1 的多组梯度累加）
+crates/xtrain-model/
+├── src/config.rs                         # +num_kv_heads 字段（默认 = n_heads → MHA）；from_arch 加形参；num_params 计 K/V 投影按 kv_dim
+├── src/model.rs                          # wk/wv 投影到 kv_dim；attention() 在 SDPA 前对 K/V 做 ops::repeat_kv；两条路径都吃到 GQA
+└── tests/gqa.rs                          # 新：GQA(group>1) flash==composed + group=1 与 MHA 逐位一致
+crates/xtrain-train/src/bin/train.rs            # +--kv-heads flag
+crates/xtrain-distributed/src/bin/train_ddp.rs  # +--kv-heads flag（DDP 路径）
+crates/xtrain-train/src/bin/export_safetensors.rs # +--kv-heads；config.json 写真 num_key_value_heads
+crates/xtrain-model/tests/parity{.py,_dump.rs}  # PyTorch 对拍加 GQA（kv 投影 + repeat_kv）
+```
+
+## Key Design Decisions
+
+### ① GQA = K/V 头维 broadcast op，喂给**未改动**的两条 SDPA（不写第三套 attention）
+
+T14 已经有两条**逐位/数值都验证过**的 SDPA：composed（`ops::attention`）与 fused flash
+（`ops::flash_attention`），二者都吃 `[bh, S, hd]`（`bh = batch·heads`）。GQA 的本质只是「K/V
+比 Q 少 `group` 倍头，用前把每个 kv 头复制 `group` 份」。所以**最外科**的做法：
+
+- wk/wv 投影到 `kv_dim = num_kv_heads · head_dim`，按 `[B, num_kv, S, hd] → [B·num_kv, S, hd]`
+  排好（和 Q 的 `[B·nh, S, hd]` 同流水线，只是头数不同）。
+- 在调 SDPA 之前，对 K、V 各做一个新 autograd op `ops::repeat_kv`，把 `[B·num_kv, S, hd]`
+  **broadcast** 成 `[B·nh, S, hd]`（输出行 `b·nh + qh` = 输入行 `b·num_kv + qh/group` 的拷贝）。
+- 之后 `ops::attention` / `ops::flash_attention` **一字不改**——它们看到的就是满头的
+  `[B·nh, S, hd]`，GQA 对两条路径**同时、免费**生效。flash kernel / composed kernel 都不用碰。
+
+**为什么不在 kernel 内做 GQA**：那要给 flash fwd/bwd 两个 kernel 各加 kv-head 索引、给 composed
+的两次 strided GEMM 各算 GQA stride，且两套都要重测——是「第三套 attention 改动」。而 broadcast-op
+方案：(a) 两条 SDPA 路径零改动、其 T14 闸门不回归；(b) repeat_kv 的 fwd/bwd 是独立可 grad-check
+的小 op，正确性风险隔离在一处；(c) 关键的「多组 q 头梯度汇到一个 kv 头」就是 repeat_kv 的**反向**，
+干净地落在一个 op 上单测。代价是 K/V 在显存里被物化成满头 `[B·nh,S,hd]`（多 `group` 倍）——本规模
+（训练、seq 不极端）可接受；真要省这份显存是 follow-up（kernel 内 GQA 读取），记进逃生舱不在 T15 做。
+
+> 备选（不采纳）：flash/composed kernel 内直接按 `kv_head = q_head/group` 索引 K/V。省 broadcast
+> 物化，但动两套已验证 kernel + 重写两套 backward 的 kv 累加，违反「不写第三套 attention」与回归保护。
+> escape hatch：先 broadcast-op 把正确性 + 闭环钉死，kernel-内 GQA（省显存）留 follow-up。
+
+### ② repeat_kv 的反向 = 组内求和（确定性，无 atomic）
+
+`repeat_kv` 前向：`out[b·nh + qh] = in[b·num_kv + qh/group]`（按 `S·hd` 整行拷贝）。
+
+反向是它的**转置**：一个 kv 头收到它那 `group` 个 query 头的梯度之**和**：
+```
+din[b·num_kv + kvh] = Σ_{r=0}^{group-1} dout[b·nh + kvh·group + r]
+```
+这正是闸门要求的「多组 q 头梯度累加到一个 kv 头」。实现上**不用 atomicAdd**：每个输入
+（kv-head, 元素）由唯一一个 block 负责，它**串行累加自己那 group 个连续源行**——天然 race-free
+且**run-to-run 确定**（不像 flash bwd 的跨行 atomic 反向有归约序不确定问题）。`group=1` 时
+反向退化为单行拷贝（identity）。
+
+autograd 层面其实也可以靠引擎的扇出 SUM（把一个 kv Var 喂给 group 个下游），但那样图里要
+显式建 group 份 view、且 flash/composed 的 batched 布局不是按头切的——做成一个专门的
+broadcast op，fwd/bwd 各一发 kernel，最简且能单独 grad-check。
+
+### ③ `num_kv_heads` 进 Config（它改模型尺寸/导出），默认 = n_heads → 退化 MHA
+
+不同于 T14 的 `use_flash`（运行时旗标，不进 Config），`num_kv_heads` **改 K/V 投影的形状、改参数量、
+改导出的 `num_key_value_heads`**——它是**架构**的一部分，必须进 `Config` 并落进 checkpoint/导出。
+
+- `Config` 加 `num_kv_heads: usize`；`from_arch` 加该形参；`Config::tiny()` 默认 `num_kv_heads =
+  n_heads`（MHA）。约束：`num_heads % num_kv_heads == 0`（断言）。
+- `num_params()`：K/V 投影从 `2·dim·dim` 改成 `2·dim·(num_kv_heads·head_dim)`；QK-norm 的
+  `k_norm` 仍是 `[head_dim]`（per-head，作用在单个 head 向量上，与头数无关）→ 不变。
+- **`num_kv_heads == n_heads` 时 `group=1`**：`ops::repeat_kv` 是 identity（fwd 单行拷贝、bwd 单行
+  拷贝），wk/wv 形状回到 `[dim,dim]` → 整条图与 T14 的 MHA 路径**逐位一致**（回归保护闸门）。
+
+> wk/wv 形状从 `[dim,dim]` 变成 `[dim, kv_dim]`：`Block` 里 wk/wv 的 `mk(&[dim, kv_dim])`，
+> `params()`/`block_params()` 顺序不变（还是 attn_norm,wq,wk,wv,q_norm,k_norm,wo,...），只是
+> wk/wv 的 shape 跟着 Config。导出转置照旧按各自 shape 走（`transpose` 读 `v.value().shape()`）。
+
+### ④ bf16 / recompute / dropout / DDP 全部自动兼容
+
+- **bf16**：`Tensor::repeat_kv` 沿用全 repo 一致的 cast 策略——bf16 入则 upcast f32 → kernel →
+  downcast；kernel 只一份 f32。`ops::repeat_kv` 的 fwd/bwd 都在 SDPA 之前/之后，dtype 与 K/V 流一致。
+- **recompute（T13）**：repeat_kv 在 `block_forward` 内、`attention()` 里，重算段重跑 `attention()`
+  自然重跑 repeat_kv（无额外状态，确定性）→ 梯度仍逐位一致。
+- **dropout（T18）**：dropout 接在 attn/mlp 子块**输出**，与 attention 内部的 repeat_kv 正交，不交互。
+- **DDP**：repeat_kv 不引入新参数；wk/wv 变小（kv_dim）只是参数张量小一圈，`params()` 泛化迭代
+  + all-reduce 照旧；跨 rank 一致性不受影响。
+
+### ⑤ 导出 xserv：写真 `num_key_value_heads`，分组约定对齐 repeat_kv
+
+`export_safetensors.rs` 的 `config.json` 把 `num_key_value_heads` 从「= num_attention_heads」改成
+**真 `cfg.num_kv_heads`**；`--kv-heads` flag 传入（须与训练 ckpt 一致）。q/k/v_proj 各自按其 shape
+转置导出（k/v_proj 现在是 `[kv_dim, dim]`，xserv loader 期望的 GQA 形状）。xserv 的 `repeat_kv`
+用 `dst = kvh·group + r` 分组，与 ① 的 xtrain 约定**逐头对齐** → 同一份权重在两侧前向数学一致，
+闭环（贪心逐 token 一致）成立。
+
+## 验证方法
+
+全部 `#![cfg(not(no_cuda))]` 门控，本地 `cargo check`/`fmt`，构建+实跑在 dash5（8× RTX 5090）。
+
+### 1. 正确性（硬闸门全绿，dash5 实跑 capture）
+
+- **repeat_kv finite-diff grad-check**（`autograd.rs::repeat_kv_grad`）：**核心闸门**——`group>1`
+  （如 bh: 2 kv 头 → 6 q 头）下 grad-check `din`，验证「多组 q 头梯度求和到一个 kv 头」的反向。
+  外加 `group=1` identity 自检。
+- **GQA flash==composed**（`gqa.rs`）：真 GQA 配置（`num_kv_heads < n_heads`，如 8 头/2 kv 头）下，
+  flash on/off 两个同 init 模型，断 forward logits / loss / **每参数梯度**一致（fp32 紧容差 + bf16
+  舍入带）——尤其 wk/wv 的梯度（它们经过 repeat_kv 反向的组内求和）。
+- **group=1 与 MHA 逐位一致**（`gqa.rs`）：`num_kv_heads = n_heads` 的模型对 T14 的 MHA 模型，
+  forward + 每参数梯度 `|Δ|=0`（回归保护）。
+- **PyTorch GQA 对拍 B>1**（`parity_dump.rs` + `parity.py`）：等价 PyTorch 模型加 GQA（k/v 投影到
+  kv_dim + `repeat_interleave(group)` 分组，与 xserv/xtrain 约定一致），对拍 forward logits + 全部
+  参数梯度（composed 与 flash 两条都跑，共用同一 oracle）。
+- **小 GQA 配置短训收敛**：一个真 GQA 小模型短训，loss 单调降、无 NaN、采样连贯。
+- **全回归套开/关**：autograd / structural / batched==looped / bf16 / recompute（逐位）/ overfit 27/27 /
+  AdamW（GPU bit-exact + host 对 torch）/ DDP loss-match + 跨 rank（`--test-threads=1`）/ flash /
+  grad_accum / dropout / **xserv 闭环 md5**。MHA 默认（kv=heads）图不变 → 不回归。
+
+### 2. 闭环（payoff）—— 真 GQA 端到端
+
+导出一个 `num_key_value_heads < num_attention_heads` 的 GQA checkpoint → xserv 加载 → 贪心生成
+**对 xtrain 自身逐 token 一致**（BF16 推理 vs f32 训练，与 v1–v8 同款判据）。这是 GQA 真正落地的证明：
+训练侧的分组、导出的分组、推理侧 xserv 的 repeat_kv 分组三方对齐。
+
+## 实测结果（dash5 1× / 2× RTX 5090）
+
+**硬闸门全绿：**
+
+| 闸门 | 结果 |
+|---|---|
+| ① repeat_kv grad-check（**多组 q 头梯度求和到一个 kv 头**，group=3） | **过** — din max_rel **2.05e-4**；group=1 identity 双向**逐位**（fwd/bwd |Δ|=0） |
+| GQA flash==composed（model 级 8h/2kv，logits/loss/每参数梯度） | fp32: loss rel **0.0**、logits 3.0e-4、grad **4.1e-5**；bf16: loss 9.0e-5、logits mean 2.9e-3/p99 1.0e-2、grad scaled-mean 8.9e-3 |
+| group=1 对 MHA**逐位一致**（回归保护） | **过** — logits + loss + 全部梯度 |Δ|=0 |
+| ② PyTorch GQA 对拍 B>1（composed & flash，repeat_interleave 分组对齐） | composed: loss **1.74e-8**/logits 2.04e-5/25 grad 进 rtol；flash: loss 1.74e-8/logits 2.28e-5/25 grad 进 rtol |
+| ③ 小 GQA 配置短训收敛（8h/2kv/hd32/4L/ffn1024，600 步） | train **10.90→3.15** 无 NaN、gnorm 稳 ~1.2、采样连贯英文（~200K tok/s） |
+| ④ **xserv 闭环真 GQA**（导出 `num_key_value_heads=2 < num_attention_heads=8`，xserv 加载 `heads=8/2 kv`，贪心） | "One day"/"The little" 两 prompt **逐 token 一致**；"Once upon a time" 在 `...Lily's mommy ` 处 BF16 漂移晚分叉（said vs came）——与 v1/v2/v3/T14 同款判据 |
+| ⑤ 回归套：autograd 23（含 repeat_kv 2）/ structural 5 / batched / bf16 / flash 2 / **gqa 4** / overfit 27/27 / recompute 2 / dropout 6 / grad_accum 3 / checkpoint-roundtrip / AdamW(host 对 torch 4.8e-6) / DDP 3(`--test-threads=1`, loss 5.67e-7+跨 rank 一致) / GEMM / tensor | **全绿** |
+| ⑤ MHA 默认 export md5（v3 ckpt 用 T15 代码重导 safetensors） | **逐位一致** `b04fc9f9a0c9af04c47d9ca649aea12e`（与 registry/T14 同）→ 默认（kv=heads）export 零漂移 |
+
+> **诚实记录**：闭环 2/3 prompt 完全 token-identical、1/3 在 BF16 漂移点晚分叉——这恰证明 GQA 分组**正确**：若 kv→q 头映射错，attention 会从第一个生成 token 起就崩（不会是深处近-tie 的晚分叉）。GQA 把 K/V 在显存里物化成满头 `[B·nh,S,hd]`（broadcast-op 方案的代价）——本规模可接受，kernel-内 GQA（省这份显存）留 follow-up。未为凑绿放宽任何容差。

docs/15-grad-accum.md

@@ -0,0 +1,165 @@
+# Phase T16: Gradient Accumulation — Design Document
+
+## Goal
+
+在已有的训练 loop（T6/T10）与 DDP（T8）之上，加 **micro-batch 梯度累积**：把 `accum_steps=N`
+个 **micro-step** 的梯度在 tape 里累加起来，再做**一次** `AdamW.step` + `zero_grad`——得到
+**有效 batch = N × micro_batch** 的更新，而显存只占**一个 micro-batch** 的激活峰值（不随 N 增长）。
+
+两条硬约束：
+
+1. **数值等效**：`accum_steps=N`（N 个 micro-step 后一次 step）必须对住「一个 N× 大 batch
+   的单 step」——梯度/loss 在仓内既有容差内**逐位贴合**。这是核心等效性证明。
+2. **DDP 只在累积边界通信**：`world>1` 下，N 个 micro-step 里**只在最后一个**做 all-reduce
+   （中间 micro-step **跳过跨卡通信**），最终喂给优化器的仍是 global 有效 batch 的均值梯度，
+   loss 对单卡。
+
+并暴露 train 入口的 `--accum-steps` flag。`accum_steps=1` 必须对当前无累积路径**逐位一致**
+（回归保护）。
+
+**不做**：micro-batch 间变 LR / 变 batch（恒定 micro_batch）；累积里换 dropout RNG（T18 才有
+dropout）；ZeRO（T17）。本 Phase 只动**优化器 step 的节奏**与 **DDP 通信门控**，复用 tape 既有
+的 SUM 累加。
+
+## Module Layout
+
+```
+crates/xtrain-train/src/
+├── train_loop.rs        # TrainConfig += accum_steps；inner micro-loop（缩放 loss + tape SUM）
+└── bin/train.rs         # 新 --accum-steps flag；打印有效 batch
+
+crates/xtrain-distributed/src/
+└── ddp.rs               # DdpConfig += accum_steps；all-reduce 门控到累积边界
+
+crates/xtrain-train/tests/
+└── grad_accum.rs        # 等效性硬闸门 + accum_steps=1 逐位回归（单卡）
+
+crates/xtrain-distributed/tests/
+└── ddp_correctness.rs   # += DDP+accum 对单卡（复用既有 ddp_matches… 框架）
+
+docs/15-grad-accum.md     # 本文
+```
+
+无新 crate、无新 kernel、无新 autograd op——梯度累积是**纯调度**：tape 早已 SUM 累加，
+缩放用既有 `ops::scale`，DDP 通信用既有 `all_reduce_average_grads`，只是改**调用节奏与门控**。
+
+## Key Design Decisions
+
+### ① 等效性的数学：缩放每个 micro-loss 为 `1/N`
+
+模型的 `loss_batched` 是 **CE-mean over `batch*seq` 行**（见 `model.rs`）。设一个 micro-batch 有
+`B` 序列、seq 长 `S`，记某 micro-step 那批 `B*S` 行的 per-row 梯度之和为 `Σ_micro`：
+
+- **大 batch 基线**（有效 batch `N·B`）：一次 `loss_batched(N·B 序列)` = CE-mean over `N·B·S` 行
+  → backward 给 `G_big = Σ_all / (N·B·S)`，其中 `Σ_all = Σ_n Σ_micro_n`。
+- **累积**（N 个 micro-step，每个 `B`）：micro-step n 的 `loss_batched(B)` = CE-mean over `B·S` 行
+  → 若直接 backward 得 `Σ_micro_n / (B·S)`；**N 个 backward 之间不 `zero_grad`**，tape SUM 累加 →
+  `Σ_n Σ_micro_n / (B·S) = Σ_all / (B·S) = N · G_big`。
+
+差一个因子 N。修正：**每个 micro-loss 先 `ops::scale(loss, 1/N)` 再 backward**——`scale` 的
+backward 把上游梯度乘 `1/N`（见 `ops.rs`），于是每个 micro 贡献 `Σ_micro_n / (N·B·S)`，
+累积后 `Σ_all / (N·B·S) = G_big`，**与大 batch 逐位等效**（仅 fp 求和顺序不同 → 进容差，和
+T8 DDP-vs-单卡同性质）。
+
+> 为什么不在 clip 里用 `pre_scale=1/N`？clip 的 `pre_scale` 已被 batch-mean 占用（=1.0）。
+> 在 loss 上 `scale(1/N)` 更内聚：缩放穿过既有 autograd，不碰 clip/optimizer，且 `N=1` 时
+> `scale(1.0)` 的 backward 是恒等乘 1 —— 这正是 `accum_steps=1` 逐位回归的保证（见 ④）。
+
+报告的 step-loss = N 个 micro 的**原始** loss（未缩放值）之和 / N = 有效 batch 的 mean loss，
+和大 batch 的单一 mean loss 一致（同样仅求和顺序差）。
+
+### ② 单卡 train loop：inner micro-loop
+
+每个 optimizer step：
+
+```text
+for micro in 0..N:
+    抽 B 序列 → loss = loss_batched(B)
+    step_loss_acc += raw_loss(loss)              # 累报告用的原始 loss
+    scale(loss, 1/N).backward()                  # tape SUM 累加缩放后的梯度
+# —— 累积边界 ——
+clip_grad_norm_gpu(params, max_norm, 1.0)        # 梯度已是有效 batch 均值
+opt.step(lr); zero_grad()
+losses.push(step_loss_acc / N)
+tokens_seen += N * B * S                         # 有效 batch tok
+```
+
+`accum_steps` 默认 1 → micro-loop 跑一次、`scale(loss,1.0)`、不在 micro 间 zero_grad（本就如此）
+→ 与现路径完全等价。**每个 micro-step 的计算图在它自己的 backward 后即可释放**（Rust `Rc` 在
+循环变量出作用域时 drop），所以**显存峰值 = 单个 micro-batch 的激活**，不随 N 增长（③ 实测）。
+
+抽样次序保持：单卡仍是连续从 RNG 抽 `N·B` 序列；与「大 batch 抽 `N·B`」逐序列对齐，只是分 N 组
+forward——并集同序，所以 `Σ_all` 的项一致。
+
+### ③ 显存平 + 有效 batch 实测
+
+「显存不随 N 增长」是 grad-accum 的卖点，要**实测**而非断言：固定有效 batch `E = N·B`，跑
+`(N=1,B=E)`（大 batch）vs `(N=E,B=1)`（极端累积），用 `nvidia-smi`/`cudaMemGetInfo` 量峰值显存——
+后者应**显著低**（少 N× 激活）。train 入口打印 `effective batch = accum_steps × batch`。
+
+### ④ `accum_steps=1` 逐位回归
+
+`N=1` 时 inner loop 跑一次、`scale(loss, 1.0)`。`ops::scale(_, 1.0)` 的 fwd 是
+`value.scale(1.0)`、bwd 是 `grad.scale(1.0)`——数学恒等。为**绝对**逐位（连一次 `×1.0` kernel
+都不引入），实现里 `N==1` 直接 `loss.backward()`（跳过 scale），与现路径**字节一致**。测试
+`accum1_bit_identical_to_no_accum` 锁这条。
+
+### ⑤ DDP：all-reduce 门控到累积边界
+
+T8 的 `all_reduce_average_grads(params)` 每 step 调一次。grad-accum 下**只在最后一个
+micro-step 之后调一次**——中间 micro-step 的 backward 只在本卡 tape 里 SUM，**不发 NCCL**。
+
+均值的账（沿用 T8 的「通信里 /world，clip 里 /b_local」拆分，再叠加 ① 的 /N）：
+
+```text
+每卡每 micro:  scale(loss, 1/N).backward()  → 本卡 tape SUM 该 micro 的 (Σ_micro / N)/...
+N 个 micro 后, 本卡 grad = Σ_{micro∈本卡所有micro} ... = 本卡 N·B_local 行的 (1/N) 缩放和
+all-reduce(sum)+/world (累积边界一次):  跨卡求和后 /world
+→ 每卡持有 Σ_global,(N·B) / (N · world · ?)   # 见下：用 1/N·scale 替代单卡的 1/b
+clip pre_scale = 1.0
+```
+
+精确推导：每卡每 micro 的 `loss_batched(B_local)` 是 **本卡 mean over `B_local·S` 行**。
+`scale(1/N)` 后 backward = `Σ_local_micro / (N · B_local · S)`。N 个 micro tape SUM →
+`Σ_local_all / (N · B_local · S)`，其中 `Σ_local_all` = 本卡 `N·B_local` 行之和。
+`all_reduce(sum)` 跨 world 卡 → `Σ_global_all / (N · B_local · S)`（`Σ_global_all` = 全
+`world·N·B_local = N·B_global` 行之和）；`/world` → `Σ_global_all / (N · B_local · S · world)`
+`= Σ_global_all / (N · B_global · S)`（因 `B_global = world·B_local`）。这正是**有效 batch
+`N·B_global` 的 mean 梯度**——与单卡「有效 batch `N·B_global` 的大 batch 单 step」逐位等效
+（求和顺序差进容差）。
+
+> 关键正确性点：`all_reduce_average_grads` 里的 `/world` 是按 **world** 缩放（与 N 无关）；N 的
+> 那个 `1/N` 已由 ① 的 `scale` 在每个 micro 的 backward 里完成。两者正交，不会互相污染。
+> 单卡（`world=1`）退化：all-reduce 是 no-op，`/world=1`，只剩 ① 的 `1/N` → 与 ② 一致。
+
+DDP 报告 loss = N 个 micro 的本卡原始 loss·B_local 之和、跨卡 all-reduce(sum)、/(N·B_global)。
+
+### ⑥ 不变量小结
+
+| | 单卡基线（大 batch E） | 单卡 accum（N×B=E） | DDP accum（world, N×B_local·world=E） |
+|---|---|---|---|
+| loss 缩放 | 无（CE-mean） | 每 micro `×1/N` | 每 micro `×1/N` |
+| grad 累加 | tape SUM 一批 | tape SUM N 批 | tape SUM N 批/卡 |
+| 跨卡通信 | — | — | **仅累积边界 1 次** all-reduce + /world |
+| clip pre_scale | 1.0 | 1.0 | 1.0 |
+| 显存峰值 | E 的激活 | **B 的激活** | **B_local 的激活** |
+
+## 验证方法（验收，全部 dash5 实跑 capture）
+
+GPU 测试 `#[cfg(not(no_cuda))]` 门控。
+
+1. **等效性（核心硬闸门）** `grad_accum.rs::accum_equiv_big_batch`：同 init、同数据同序，
+   跑「`accum_steps=N`, micro_batch=B」与「`accum_steps=1`, batch=N·B」各一 step，断言
+   ①loss、②**每个参数的 grad** rel-err 进 fp 容差（求和顺序差，~1e-4 量级，对齐 recompute/DDP
+   闸门约定）。多步版（跑 K 个 optimizer step）再断言**终参**贴合（误差不发散）。
+2. **`accum_steps=1` 逐位回归** `grad_accum.rs::accum1_bit_identical`：`accum_steps=1` 与现
+   no-accum 路径同 init/同数据 → 每参数 grad `max|Δ| == 0.0`（④ 跳过 scale，字节一致）。
+3. **DDP+accum 对单卡** `ddp_correctness.rs`（扩既有 `ddp_matches_single_gpu…`）：单卡
+   有效 batch `E` 的大 batch baseline vs `world=2 + accum_steps=N`（每卡每 micro `B_local`，
+   `world·N·B_local=E`）→ loss 轨迹 `max_rel<1e-3`、跨 rank 参数一致、且 only-at-boundary 通信
+   （micro 间不发 NCCL，由实现保证 + 不变量推导）。
+4. **显存平 + 有效 batch** ：固定有效 batch，量 `(N=1,大batch)` vs `(N=大,micro=1)` 峰值显存
+   （后者显著低），train 入口打印 effective batch。capture nvidia-smi。
+5. **全回归套**：autograd grad-check / structural / batched==looped / bf16 / recompute（逐位）/
+   overfit 27/27 / AdamW（GPU bit-exact + host vs torch）/ DDP loss-match + 跨 rank / **xserv
+   闭环 md5**——`accum_steps=1` 默认值保证全部不回归。

docs/16-process-per-gpu.md

@@ -0,0 +1,266 @@
+# Phase T17: Process-per-GPU DDP（torchrun 式独立 CUDA context）— Design Document
+
+## Goal
+
+T8 的 DDP 是**单进程 thread-per-GPU**：一个进程开 N 个 OS 线程，每线程 `cudaSetDevice` 绑一张卡、
+在**同一个 CUDA primary context** 里跑自己 rank 的训练。T11 修掉 per-op `cudaMalloc` 串行后，8 卡
+scaling 从 ~1.3× 恢复到 **~5×@8**，但**残留 5×@8 而非 ~8× 的非线性**——根因在 T11 doc / KI-5 已点明：
+**N 个 rank 线程共享同一个 CUDA context，driver 层很多调用（kernel launch、cuBLAS handle、stream
+排队）在单 context 内进程级串行**，pool allocator 只消掉了其中最大的一笔（malloc），剩下的 launch /
+cuBLAS 串行仍在。
+
+T17 的目标 = **torchrun 式 process-per-GPU**：每个 rank 是一个**独立 OS 进程**，各自持有**独立的
+CUDA context**，彼此的 driver 调用不再在同一 context 排队 → 移除 thread-per-GPU 的残留串行，把 8 卡
+scaling 推向更接近线性。这是 Phase 2 里**改动最大**的一项（launcher 结构性重写 + 跨进程 NCCL
+bootstrap），所以本 doc 先行。
+
+**Scope（用户已拍板）：process-per-GPU ONLY。ZeRO-1 / sharded optimizer 明确 drop**——本尺度
+optimizer state 小、收益薄。本任务只换**启动模型与 NCCL bootstrap**，**训练 step（grad all-reduce →
+本地 AdamW）原样复用、零改动**。
+
+**保留 thread-per-GPU 路径**：T8 的 `launch()` + `train_ddp` bin 不删（回归保护 + 闸门 ①要求新旧路径
+loss 对齐）。process-per-GPU 作为**并列的新 launcher** 加上去。
+
+验收（硬闸门全绿，诚实正确性，不放宽容差）：
+1. 多进程（world=2 / world=4）训练 loss **对单卡贴合**（进既有 DDP 容差 `<1e-3`），且对住旧 thread-per-GPU 路径；
+2. 跨 rank 参数一致（repo 既有 `<1e-6` 约定）；
+3. **8 卡线性度 before→after 实测**：thread-per-GPU baseline（~5×@8）vs process-per-GPU @ {1,2,4,8}，给数字；
+4. 全回归套绿（含 xserv 闭环 md5 / token-identical）；单卡与旧 thread-per-GPU 路径不回归。
+
+## 什么变、什么不变
+
+```
+                       thread-per-GPU (T8, 保留)          process-per-GPU (T17, 新增)
+启动                   1 进程 × N 线程                     1 launcher 进程 → fork/exec N 个 worker 进程
+CUDA context           N 线程共享 1 个 primary context     每 worker 进程 1 个独立 context
+rank/world/device      闭包捕获 + thread::scope            env: RANK / WORLD_SIZE / LOCAL_RANK
+模型构建               每线程闭包内 build_model（!Send）   每进程 main 内 build_model（天然隔离）
+NCCL UniqueId 分发     move 一个 Copy struct 进线程闭包    launcher 生成 → hex 编码进子进程 env
+NCCL comm init         DdpContext::init（不变）            DdpContext::init（不变）
+─────────────────────────────────────────────────────────────────────────────────────
+grad all-reduce        all_reduce_average_grads（不变）    ← 同一份代码，零改动
+本地 AdamW step        train_rank（不变）                  ← 同一份代码，零改动
+batch sharding         i % world == rank（不变）           ← 同一份代码，零改动
+参数一致性证明         同 init+同 grad+同 opt（不变）       ← 同一论证
+```
+
+**核心洞察**：T8 早把训练 step 写成「**per-rank**、接受 `&DdpContext`」的形状（`train_rank`）。
+thread-per-GPU 与 process-per-GPU **唯一的区别只在「怎么把 rank 跑起来 + 怎么把 UniqueId 递给每个
+rank」**——前者跨线程 move，后者跨进程 env。`train_rank` / `all_reduce_average_grads` / sharding /
+一致性论证**全部原样复用**。这正是把启动模型与训练逻辑解耦的回报。
+
+## Module Layout
+
+```
+crates/xtrain-distributed/src/
+├── lib.rs        ← 加 pub mod proc; re-export hex_encode/decode_unique_id + run_worker entry
+├── proc.rs       ← 新增：① launcher（spawn N worker 进程，env 注入 rank/world/local_rank/uid）
+│                    ② worker entry（读 env → DdpContext::init → build_model → train_rank）
+│                    ③ UniqueId hex 编解码（跨进程 env 传 128 字节）
+├── ddp.rs        ← 不变（train_rank / build_model / DdpConfig 复用）
+├── lib.rs::DdpContext / all_reduce_average_grads / get_unique_id  ← 不变
+└── bin/
+    ├── train_ddp.rs       ← 不变（thread-per-GPU，保留）
+    └── train_ddp_mp.rs    ← 新增：multi-process launcher / worker 二合一入口
+crates/xtrain-distributed/tests/
+└── ddp_proc.rs            ← 新增：spawn 多进程跑几步 → loss 对单卡 + 跨 rank 参数一致 +（顺手）before/after 吞吐
+docs/16-process-per-gpu.md ← 本文
+```
+
+`proc.rs` 全程 `#[cfg(not(no_cuda))]` 门控（同 crate 既有约定）；本地无 nvcc 时 crate 编空，`cargo
+check` 过；dash5 上全量编译链 NCCL。
+
+## Key Design Decisions
+
+### ① Launch model：同一 binary 双模（launcher / worker）
+
+`train_ddp_mp` 一个可执行文件，靠**环境变量是否存在**自判角色（torchrun 的 `LOCAL_RANK` 注入 worker
+是同一思路）：
+
+- **launcher 模式**（直接被用户 / 测试调用，env 里没有 `XTRAIN_RANK`）：
+  1. 看 `CUDA_VISIBLE_DEVICES` / `device_count()` 定 `world`；
+  2. 调 `get_unique_id()` 生成一个 `ncclUniqueId`（128 字节），**hex 编码**成字符串；
+  3. `for rank in 0..world`：`Command::new(current_exe())`，**复制自己全部 argv**（超参/路径透传），
+     额外设 env `XTRAIN_RANK=rank`、`XTRAIN_WORLD=world`、`XTRAIN_LOCAL_RANK=rank`、
+     `XTRAIN_NCCL_ID=<hex>`，spawn 为子进程；
+  4. `wait()` 所有子进程，任一非零退出码 → launcher 以非零退出（CI / 闸门可感知）。
+- **worker 模式**（被 launcher spawn，env 里有 `XTRAIN_RANK`）：
+  1. 从 env 读 `rank / world / local_rank / uid_hex`；
+  2. `device::set_device(local_rank)` 绑卡（**每进程独立 primary context** 在此首次 CUDA 调用时建立）；
+  3. hex 解码出 `NcclUniqueId`；`DdpContext::init(rank, world, id, local_rank)`（**复用 T8 的 init**）；
+  4. `build_model(cfg, device)`（**复用 T8 的确定性 init** → 同种子 → 跨进程逐位同起点）；
+  5. `train_rank(&ctx, &model, …, &cfg)`（**复用 T8 的训练 step，零改动**）；
+  6. 退出码 0（成功）/ 非零（panic → 进程崩，launcher 感知）。
+
+**单机 `CUDA_VISIBLE_DEVICES` 处理**：launcher 看到的 visible 设备集就是 `0..world`；每个 worker
+继承同一个 `CUDA_VISIBLE_DEVICES`（env 默认透传），`local_rank` 直接当作 visible 集内的 device
+ordinal → `set_device(local_rank)`。这与 thread-per-GPU 的 `devices = 0..count` 语义一致，单节点足够。
+（真·多节点要把 `LOCAL_RANK` 与全局 `RANK` 分离 + 每节点 `CUDA_VISIBLE_DEVICES` 切片，单节点不需要，
+记为 follow-up。）
+
+### ② 跨进程 NCCL UniqueId 分发：launcher 生成 + hex-env 注入（**最简、无竞态**）
+
+这是 T17 最该想清楚的一处。候选机制（任务列了文件 / TCP / 共享 FS）逐一权衡：
+
+| 机制 | 怎么做 | 单节点取舍 |
+|---|---|---|
+| **共享文件** | rank0-worker 写 `/tmp/xtrain.id`，其余 worker **轮询读** | 要处理「文件还没写好」的 race（轮询 + 重试 + 超时），还要 worker 间约定谁是 rank0、何时清理 |
+| **TCP rendezvous** | 起一个 c10d-store 式小 server 派发 id | 最贴 torchrun，但要写 socket server/client、端口选择、握手协议——单节点 overkill |
+| **launcher 生成 + env 注入** ✅ | **launcher**（而非 rank0-worker）调 `ncclGetUniqueId`，hex 编码后**在 spawn 时就写进每个子进程的 env** | 无文件 race、无轮询、无 TCP server、无清理——env 在子进程出生前就备好。子进程读 env 即得 id |
+
+**选 env 注入**，诚实理由：单节点下 launcher 是所有 worker 的**共同父进程**，env 是父→子最朴素的带外
+通道，且**在子进程创建那一刻就原子地确定**——天然没有「id 还没就绪」的竞态，比文件轮询 / TCP 握手都
+简单且更鲁棒。代价是 launcher 进程要链 NCCL 调 `ncclGetUniqueId`（它本就在 distributed crate 里，已链
+NCCL），可接受。
+
+> **与「rank 0 生成」的关系**：torchrun 概念上是 rank 0 把 id 放进 c10d store、别的 rank 取。这里
+> **launcher 充当协调者**替 rank 0 生成——功能等价（id 只是个一次性握手 token，谁生成不影响正确性，
+> 只要全 rank 拿到同一个），但单节点下省掉了「worker 间再来一轮带外同步」。**128 字节 → hex = 256
+> 字符**，远低于环境变量长度上限，env 传输安全。
+
+`hex_encode`/`decode_unique_id` 是 `proc.rs` 里两个纯函数（`[c_char;128] ↔ 256-char hex`），单测可在
+host 侧验 roundtrip（不需 GPU）。
+
+### ③ 独立 CUDA context = 移除残留串行（这才是 T17 的 payoff）
+
+thread-per-GPU 的残留非线性（KI-5 / T11 doc）来自：**N rank 线程共享同一 CUDA primary context**，driver
+对该 context 的很多操作（kernel launch 队列、cuBLAS handle、内部锁）是进程级 / context 级串行的——
+pool allocator 消掉了 malloc 这一最大笔，但 launch / cuBLAS 串行仍在，表现为 8 卡 ~5× 而非 ~8×。
+
+process-per-GPU 下**每个 rank 是独立进程 → 独立 CUDA context → 独立 driver 状态**：各进程的 kernel
+launch / cuBLAS 调用**互不在同一 context 排队**，残留串行（按此假设）应被结构性移除。这正是闸门 ③
+（before→after 线性度）要量出来的东西——若 process-per-GPU 把 8 卡从 ~5× 推到明显更高，即验证此假设。
+**诚实原则**：若提升有限，如实报告（说明残留瓶颈在 NCCL all-reduce / PCIe 拓扑，那是另一层，非本任务 scope）。
+
+> ⚠️ **此假设被实测证伪**——见下方「实测结果 · 闸门 ③」：process-per-GPU 与 thread-per-GPU 吞吐统计上一致
+> （~5.3×@8 都一样），且 8 卡全 95–99% util。残留非线性是通信/PCIe 墙，不是单 context 串行。结论钉死、留档。
+
+### ④ 训练 step / 一致性论证：原样复用 T8，零改动
+
+process-per-GPU 不碰任何训练数学：
+
+- **grad all-reduce**：`all_reduce_average_grads(params)` 一字不改——NCCL collective 跨**进程**和跨
+  **线程**对调用方完全一样（comm 是 rank 维度的，与进程/线程无关）。
+- **batch sharding**：`i % world == rank` 不变——每进程推进**同一个 seed 的 RNG**抽出整批 `B_global`
+  序列、只算自己那片。各进程的并集 == 单卡同序批 → all-reduce 后的 grad 和与单卡逐序列一致。
+- **参数一致性**：同 ③个充分条件（T8 doc ④）——(a) 同确定性 `build_model`（同 LCG 种子，跨进程同样
+  成立）；(b) NCCL all-reduce 跨 rank 返回逐位相同的归约（PCIe-only run-to-run 几 ULP 抖动，故闸门
+  ②用 `<1e-6` 而非 `==0`，与 T11 既有约定一致）；(c) 同 optimizer 超参/状态演化。
+- **对单卡**：与单卡只在 **fp 求和顺序**上差（单卡 tape SUM B 个；DDP 各 rank 先 SUM 分片再 NCCL SUM）
+  → `<1e-3` rel，不逐位。与 thread-per-GPU 路径则应**数值同量级**（同一 sharding + 同一 all-reduce）。
+
+### ⑤ 进程生命周期 / 失败传播 / 资源清理
+
+- **失败传播**：worker panic → 进程非零退出；launcher `wait()` 收集所有退出码，任一非零 → launcher
+  非零退出并打印哪个 rank 挂了。NCCL comm 在进程退出时由 OS 回收 context（`DdpContext::Drop` 调
+  `ncclCommDestroy`，正常退出路径走到；崩溃时 OS 兜底回收）。
+- **不需要跨进程 barrier**：每个 worker 独立跑完 `cfg.steps` 自然退出；NCCL collective 本身是同步点
+  （所有 rank 必须到齐才返回），训练循环天然对齐。
+- **资源清理**：无临时文件（env 注入，无 `/tmp` id 文件）；ckpt 由 rank0-worker 写到 `--ckpt` 指定路径，
+  与 thread-per-GPU 一致；测试用的 ckpt / 进程在测试结束清理。
+
+## 验证方法（硬闸门全绿，dash5 实跑）
+
+### 闸门 ①②：正确性 —— `tests/ddp_proc.rs`（`#[cfg(not(no_cuda))]`，<2 卡 skip）
+
+测试本身是 launcher：用 `Command` spawn N 个 worker 进程（worker = 同测试 binary 的一个特殊模式，或复用
+`train_ddp_mp`），跑固定步数，worker 把最终 loss / 参数 dump 到各自的 stdout / 临时文件，测试父进程读回：
+
+- **(a) loss 对单卡**：单卡 baseline（既有 `run_single_gpu`）vs 2-进程 / 4-进程 DDP，整条 loss 轨迹
+  `max_rel < 1e-3`（与 thread-per-GPU 测试同容差）。
+- **(b) 跨 rank 参数一致**：`max|p_i - p_j| < 1e-6`（KI-5 既有约定）。
+- **(c) 对住 thread-per-GPU 路径**：同 config 同 seed，process-per-GPU 的 loss 轨迹 vs thread-per-GPU
+  的 loss 轨迹应在 `<1e-3`（两者只差进程/线程，sharding+all-reduce 同）。
+
+> **harness 注意**：分布式测试在共享 GPU 上并行会争卡 deadlock → 一律 `--test-threads=1`（已知 harness
+> 属性，capstone/known-issues 记过）。
+
+### 闸门 ③：线性度 before→after —— `train_ddp`(thread) vs `train_ddp_mp`(process) @ {1,2,4,8}
+
+固定**每卡 batch 32 / seq 256 / dim384**（与 T11 KI-5 表同口径，便于直接对比），各跑 steady-state tok/s：
+
+```
+                  thread-per-GPU (T11 baseline)   process-per-GPU (T17)
+  world   tok/s(global)   speedup                 tok/s(global)   speedup
+     1          ~92K        1.00×                       ?            1.00×
+     2         ~147K        1.59×                       ?              ?
+     4         ~270K        2.92×                       ?              ?
+     8         ~461K        4.99×    ← 残留非线性          ?            ? ← 目标更接近 8×
+```
+
+8 卡跑时 `nvidia-smi` 抽样确认 8 卡 util。**资源纪律**：线性度 bench 合法地短用 8 卡，但**短跑**（每个
+world 几十~一两百步够测 steady-state），跑完清 ckpt / 中间物。
+
+### 闸门 ④：全回归套（标准 `--test-threads=1`）
+
+autograd / structural / batched / bf16 / recompute / overfit / AdamW / 既有 DDP loss-match + 跨 rank /
+flash / gqa / grad_accum / dropout，**+ xserv 闭环**（导出 → md5 对 registry → token-identical）。单卡与
+旧 thread-per-GPU 路径不得回归（process-per-GPU 是**新增**路径，旧路径代码未动 → 天然不回归，测试确认）。
+
+### dash5 实跑
+
+```bash
+export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
+# 正确性（多进程）：
+CUDA_VISIBLE_DEVICES=0,1 cargo test -p xtrain-distributed --release --test ddp_proc -- --nocapture --test-threads=1
+# 多进程训练 / 线性度 driver（process-per-GPU launcher）：
+CUDA_VISIBLE_DEVICES=0,1,2,3 cargo run -p xtrain-distributed --release --bin train_ddp_mp -- \
+  /opt/wjh/models/gpt2/tokenizer.json data/tinystories-valid-3mb.txt \
+  --dim 384 --heads 12 --head-dim 32 --layers 12 --ffn 1536 --steps 200 --batch 128 --seq 256
+```
+
+实测数字见下方「实测结果」。
+
+## 实测结果（dash5, 8× RTX 5090, sm_120）
+
+### 正确性（闸门 ①②④ 全绿）
+
+- **闸门 ① loss 对单卡 / 对 thread 路径**（`ddp_proc`, world=2，合成语料 20 步）：
+  proc-per-GPU vs single-GPU `max_rel = 5.67e-7`；**proc-per-GPU vs thread-per-GPU `max_rel = 1.5e-7`**
+  （两条路径数值同量级，符合预期——只差进程/线程，sharding+all-reduce 同）。
+- **闸门 ② 跨 rank 参数**：`max|p0−p1| = 1.19e-7`（< 1e-6，KI-5 既有 ULP 容差，PCIe NCCL run-to-run 抖动）。
+- **闸门 ④ 全回归**：全 workspace `--test-threads=1` 全绿（autograd/structural/batched/bf16/recompute/
+  overfit/AdamW/既有 DDP/flash/gqa/grad_accum/dropout）+ **xserv 闭环**：v3 ckpt 用 T17 代码重导
+  safetensors 与 registry **md5 逐位一致 `b04fc9f9a0c9af04c47d9ca649aea12e`**（T17 不碰任何数值路径 → 必然一致）。
+
+### 闸门 ③ 线性度 before→after —— **本任务的关键发现：process-per-GPU 在本尺度对吞吐中性**
+
+固定每卡 batch 32 / dim384 / seq256 / 150 步（与 T11 KI-5 表同口径），steady-state tok/s：
+
+| world | thread-per-GPU (`train_ddp`) | speedup | process-per-GPU (`train_ddp_mp`) | speedup |
+|---|---|---|---|---|
+| 1 | 93257 | 1.00× | 92952 | 1.00× |
+| 2 | 149747 | 1.61× | 148809 | 1.60× |
+| 4 | 278276 | 2.98× | 273308 | 2.94× |
+| 8 | **491360** | **5.27×** | **493128** | **5.31×** |
+
+（world=8 重复 2 次确认非噪声：thread 493671/493292，proc 491102/494123——**两路差异 < 1%，落在 run-to-run 噪声内**。）
+
+→ **process-per-GPU 与 thread-per-GPU 吞吐统计上一致（~5.3×@8 都一样）**。本 doc 设计假设 ③
+（「残留 5×@8 来自单 CUDA context 的 kernel-launch/cuBLAS 串行，process-per-GPU 给独立 context 即可移除」）
+**被实测证伪**——这正是 ③ 里预留的「诚实原则」分支。
+
+**根因重定位（实测佐证）**：proc-per-GPU world=8 跑时 `nvidia-smi` 抽样 **8 卡全部 95–99% util**
+（每卡 ~23GB）——GPU **已 compute-bound 喂满、并非串行空转**（KI-5 当年「1–2/8 在忙」的串行病在 T11 的
+caching allocator 就已治好）。8 卡已满载却仍只 5.3×，缺的 ~35% 吞吐只能去向**每步 grad all-reduce +
+本机 PCIe-only 拓扑在 8 rank 下的通信开销**——即 T11 早已点明的「~7% all-reduce + 8 卡 PCIe 余量」那一层，
+在 8 卡下被放大。换独立 context 不动这一层，故吞吐不变。
+
+**这与 T11 自身的方法论一致**：T11 实测证伪了「分桶 all-reduce」；T17 实测证伪了「process-per-GPU 解残留
+串行」。两次都靠 profile/measure 推翻假设而非硬上。**结论**：本尺度（dim384–1024、单机 8× PCIe RTX 5090）
+残留非线性是**通信/拓扑墙**，不是 launch 模型；要再逼近线性得动 all-reduce overlap / 更快互联（NVLink），
+那是另一条线，**非 T17 scope**。
+
+**T17 的净价值（诚实记账）**：① 学到 / 落地了 torchrun 式 process-per-GPU 这条训练栈标准链路（独立进程 +
+独立 CUDA context + 跨进程 NCCL bootstrap）——**项目本职「学训练全栈」的目标达成**；② **实测把「process-per-GPU
+是残留非线性的解」这个长期挂在 KI-5/T11 doc 里的猜想钉死为「在本尺度无吞吐收益」**，移除一个误导性 backlog
+项；③ 正确性零回归、与 thread 路径数值对齐。**吞吐上它与 thread-per-GPU 等价**——故默认训练路径**不变**
+（thread-per-GPU 仍是 v1–v8 用的那条），process-per-GPU 作为并列可选路径 + 这条诊断结论留档。
+
+## 不做（本任务范围外，记 follow-up）
+
+- **ZeRO-1 / sharded optimizer**：用户已 drop（本尺度 optimizer state 小、收益薄）。
+- **真·多节点 bootstrap**：本任务单节点（env 注入足够）；跨节点要 TCP rendezvous（c10d-store 式）+
+  `LOCAL_RANK`/`RANK` 分离 + 每节点 `CUDA_VISIBLE_DEVICES` 切片 → 留 follow-up。
+- **NCCL 通信压缩 / overlap with backward**：与 T8/T11 同理由，all-reduce 当前非主瓶颈。
+- **删除 thread-per-GPU 路径**：保留（回归 baseline + 闸门 ①要求对齐）。

docs/17-dropout.md

@@ -131,6 +131,11 @@ forward 里保持不变**——本设计天然满足：mask 只由 `(seed, i)`
 - **DDP（T8）**：每 rank 独立跑自己的 forward/backward，各自的 mask 由各 rank 的 `base_seed` 决定。
   本任务的 DDP 闸门是「loss 对单卡 / 跨 rank 参数一致」，在 **dropout 关（默认 p=0）** 的回归配置下跑，
   不引入跨 rank mask 同步需求（p>0 时各 rank mask 本就该不同，属正常 DDP 语义）。
+  - **⚠️ T18 的 launcher wiring gap → FIXED in T21**：T18 只把 dropout 接进**单卡** `train.rs`，
+    `train_ddp` bin/`train_rank` loop **没接**（无 `--dropout` flag、从不调 `model.train()`），
+    所以 DDP 路径下 dropout 被静默忽略——V9-PILOT 全栈实跑才暴露（op + 单卡测试覆盖不到 launcher 级）。
+    **T21** 补齐：`train_ddp` 加 `--dropout`、`train_rank` 每步 `model.train()`（eval 后 restore），
+    并加 DDP-dropout 回归测试（p>0 下 dropout live + p=0 逐位一致）。见 known-issues「DDP-dropout wiring」。
 - **梯度累积（T16）/ flash（T14）**：本分支独立于二者，不依赖其未合并改动。
 
 ## 验证方法

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

@@ -0,0 +1,629 @@
+# Phase: Post-Training Infra — SFT / DPO / Reward Model / GRPO — Design Document
+
+> Status: **DESIGN — decisions locked, pending go-ahead to implement.** Nothing
+> implemented yet. This doc proposes the scope, the staged build, the new infra pieces,
+> and the correctness gates for a standard post-training stack on top of the xtrain
+> training framework. Decisions D1–D4 are resolved (see "Resolved decisions"):
+> **DPO → GRPO (reward model optional) · rule-based/verifiable reward · KV-cache decode
+> engine built up front · a verifiable task as the optimization/eval target.**
+
+## Goal
+
+Build a **standard, from-scratch post-training infrastructure** — the systems layer that
+turns a pretrained base LM into an aligned chat model — and use it to run chat
+alignment. The deliverable that matters here is the **infra and the lessons**, not the
+end-to-end chat quality (see the project's learning-axis framing). Each stage should
+teach exactly one new post-training systems concept and ship with a hard correctness
+gate, matching the Phase-1/Phase-2 culture (grad-checks, PyTorch parity, bit-identical
+default paths, profile-first).
+
+Concretely we want to be able to answer, with our own code:
+
+- How does **offline preference optimization (DPO)** differ from SFT in the training
+  loop — what is the reference model, why two forwards, what is the loss?
+- How does a **reward model** turn preferences into a scalar signal?
+- How does **online RL (GRPO)** actually run — the rollout engine, reward scoring,
+  group-relative advantage, the clipped policy-gradient update, the KL leash?
+- Where are the **memory and throughput** pressure points that make post-training infra
+  different from pretraining infra (multiple models resident, generation in the loop)?
+
+## Baseline: what already exists vs. what is missing
+
+What the framework already gives us (verified in code, reused as-is):
+
+| capability | where | reuse for post-training |
+|---|---|---|
+| batched forward → logits `[B*S, vocab]` | `model.rs::forward_batched` | logprob extraction for DPO/RM/GRPO |
+| cross-entropy with **ignore-index −100** | `ops.rs::cross_entropy`, `nn.cu` | assistant-only / completion-only masking |
+| assistant-only **SFT** (TSV, masked labels) | `data.rs::load_sft_tsv_cached` (commit `fbf4ac2`) | SFT chat baseline = DPO init + reference |
+| bf16 mixed precision, fp32 master | `with_compute_dtype` | policy + frozen reference both bf16 compute |
+| recompute / flash / grad-accum | `with_recompute` / `with_flash` / `--accum-steps` | bound activation memory with 2–3 models resident |
+| DDP (thread + process-per-GPU) | `xtrain-distributed` | data-parallel post-training |
+| AdamW + clip + LR sched + checkpoint | `xtrain-optim`, `checkpoint.rs`, `schedule.rs` | unchanged optimizer path |
+| single-seq greedy/temperature sampling | `sample.rs::generate` | **slow** rollout fallback (no KV cache) |
+
+What is **missing** and must be built (these are the actual lessons):
+
+1. **Per-sequence completion logprob** — a way to read `Σ log πθ(y_t | x, y_<t)` over the
+   completion tokens of a sequence. CE gives a *mean* scalar; DPO/GRPO need a *per-sequence
+   masked sum*. New op or thin wrapper over the CE per-row machinery.
+2. **Frozen reference model** held in memory alongside the trainable policy (no grad, no
+   optimizer), or its logprobs precomputed and cached.
+3. **Pairwise preference loss** (DPO) and **Bradley-Terry ranking loss** (RM).
+4. **Reward head** — a `[dim,1]` scalar head reading the last non-pad position (RM only).
+5. **Rollout / generation engine** — batched autoregressive sampling. Current `generate`
+   is single-sequence and re-runs the full forward each step (no KV cache). Online RL needs
+   batched rollouts; a real **KV-cache incremental-decode engine** is the centerpiece infra
+   build.
+6. **GRPO machinery** — group sampling, group-relative advantage, clipped PG loss, KL
+   penalty, the actor-learner loop.
+
+## The post-training landscape — where the infra lives
+
+```
+                 data                 models in memory        new systems concept
+  SFT      (prompt, answer)           policy                  loss masking (have it)
+  DPO      (prompt, chosen, reject)   policy + ref(frozen)    dual forward, pairwise logσ loss
+  RM       (prompt, chosen, reject)   reward model            scalar head, ranking loss
+  PPO      prompts + reward source    policy+ref+RM+critic    rollout + GAE + clipped PG  (4 models)
+  GRPO     prompts + reward source    policy+ref(+RM)         rollout + group baseline + clipped PG
+```
+
+The pedagogical ladder is **SFT → DPO → (RM) → GRPO**. DPO is the cheapest "real" alignment
+method (no generation, no reward model, reuses the training loop almost verbatim) and is the
+right first rung. GRPO is chosen over PPO as the online-RL rung because it **drops the value
+critic** (group-relative advantage replaces the learned baseline) — that removes a whole
+model and the GAE machinery while still teaching the complete online-RL loop. PPO is noted
+as an optional later extension, not a primary target.
+
+## Proposed scope & sequencing (recommended path)
+
+> ✅ **DECISION D1 (scope/sequencing) — LOCKED: P0 → P1(DPO) → P3(GRPO), P2(reward
+> model) optional.** With D3 locked to "KV-cache engine up front", the engine becomes a
+> foundational milestone that both DPO pair-generation and GRPO rollouts sit on. Effective
+> build order: **P0 → KV-cache decode engine → P1(DPO) → P3(GRPO) → P2(optional)** (see
+> "Milestones").
+
+### Stage P0 — SFT chat baseline (light; mostly reuse)
+
+Goal: a clean SFT checkpoint to serve as **both the DPO/GRPO init and the frozen
+reference**. With D4 = verifiable task, P0 SFT teaches the **task format** (e.g. arithmetic
+prompts → a parseable answer such as `\boxed{N}`) so the model emits checker-readable
+completions; the same template is reused by rollout and eval. The current SFT (commit
+`fbf4ac2`) already does single-turn assistant-only masking; P0 only adds what alignment
+needs:
+
+- a fixed **chat template** (the `User:/Assistant:` + `<|endoftext|>` format already used,
+  promoted to a documented constant shared by SFT data prep, rollout, and eval),
+- optional **multi-turn masking** (supervise every assistant turn, mask user turns),
+- optional **sequence packing** (concatenate examples to fill `seq`, reset attention/RoPE
+  per example — note `forward_batched` already isolates sequences, so packing = careful
+  index bookkeeping, not new attention code).
+
+Gate: masking unit test (only assistant tokens contribute to loss); packing does not leak
+loss across example boundaries. **Hypothesis:** a documented chat template + multi-turn mask
+gives a reproducible SFT reference without changing the training numerics for single-turn data
+(bit-identical to `fbf4ac2` on single-turn input).
+
+### Stage P1 — DPO (offline preference optimization) ⭐ first real method
+
+New infra:
+
+1. **Preference data — constructed from the verifiable checker (D4).** On a verifiable task
+   there is no off-the-shelf preference set, so we build pairs: sample several completions
+   per prompt from the P0 SFT model (using the KV-cache engine built in the prior milestone),
+   score each with the rule-based checker, take a **correct** completion as `chosen` and an
+   **incorrect** one as `rejected`. This is a one-time offline data-prep step; DPO training
+   itself is then static. Tokenize each as `template(prompt) + completion + EOS`; build a
+   completion mask (prompt = masked).
+2. **`seq_logprob(logits, target_ids, mask) → [B]`**: per-sequence sum of
+   `log softmax(logits)[target]` over masked positions. Implement by reusing the CE per-row
+   path (CE per-row = `−log πθ(target)`), summing `−per_row` over the mask. Add a grad-checked
+   op so the backward is exact.
+3. **Frozen reference** `πref`: load the SFT checkpoint into a second model in **eval/no-grad**
+   bf16. Its logprobs are **constants** in the loss. Optimization to teach: **precompute and
+   cache reference logprobs** once over the dataset → the reference model need not stay
+   resident during training (one model in memory, like SFT).
+4. **DPO loss** (Rafailov et al.): with
+   `Δ = β[(logπθ(yw|x) − logπref(yw|x)) − (logπθ(yl|x) − logπref(yl|x))]`,
+   `L = −log σ(Δ)`. Only `πθ` terms carry gradient.
+
+Memory: policy (fp32 master + Adam m/v + bf16 + grads) + reference (bf16 only, or cached
+logprobs → zero). Recompute + accum keep activations bounded; 1B fits 32 GB comfortably.
+
+Correctness gates:
+- `seq_logprob` finite-difference grad-check (tiny model).
+- DPO-loss + grad **PyTorch parity** (the project's standard gate).
+- **Degenerate checks**: `πθ == πref` at init ⇒ `Δ = 0`, `L = log 2`, implicit reward 0;
+  `β → 0` ⇒ gradient → 0.
+- **Health metric**: chosen−rejected **reward margin** rises over training; accuracy
+  (margin > 0) increases. Reported, not just loss (the doc-13 lesson: val/loss alone is not a
+  sufficient signal).
+
+Application: chat alignment via DPO on English preference pairs. This is the **offline
+chat-alignment deliverable**.
+
+### Stage P2 — Reward model (Bradley-Terry) — OPTIONAL
+
+> ✅ **DECISION D2 (reward source) — LOCKED: rule-based / verifiable reward first.** GRPO
+> brings up on the deterministic checker; a learned reward model is **deferred/optional** (only
+> if we later want general-chat GRPO). So this whole stage is optional and not on the critical
+> path.
+
+New infra: a **scalar reward head** (`[dim,1]`) reading the hidden state at the last
+non-pad position; **ranking loss** `−log σ(r(x,yw) − r(x,yl))`. Reuses the preference data
+and the dual-sequence forward from P1.
+
+Gates: ranking-loss grad-check; held-out **pairwise accuracy** (`r_w > r_l`); a frozen RM
+loads/serves the scalar correctly.
+
+### Stage P3 — GRPO (online RL, critic-free) ⭐ the deep infra lesson
+
+This is the centerpiece. It introduces **generation inside the training loop**.
+
+**(a) Rollout / generation engine — built up front (its own milestone).**
+
+> ✅ **DECISION D3 (rollout depth) — LOCKED: build the KV-cache incremental-decode engine
+> up front**, as a foundational milestone *before* DPO/GRPO, rather than starting naive. It is
+> then the shared substrate for DPO pair-generation and GRPO rollouts. Tradeoff accepted:
+> front-loads the single hardest build and delays the first alignment result, in exchange for
+> a real generation engine and a clean, isolated infra lesson.
+
+The engine: per-layer **K/V cache**, **single-token incremental forward** (process the prompt
+once to fill the cache, then decode one token at a time), **batched ragged decode** (B prompts
+× G samples; sequences hit EOS at different lengths → finished-mask / left-padding /
+compaction). The current attention assumes a full causal window over `seq`; incremental decode
+needs a **decode-time attention path** — query length 1 against cached K/V of length `t`, with
+RoPE position = `t`. This reuses the composed SDPA shapes (one-row query), so it can land as a
+distinct code path without disturbing the training attention (flash/GQA/composed unchanged).
+
+Hard gate (the centerpiece correctness lesson): **KV-cache decode == full-recompute decode,
+token-identical** greedy output — the same byte-/token-identical discipline the project uses
+for the xserv export closed loop. A throughput baseline (decode tokens/s, cache-fill vs.
+per-token decode) is recorded here, before any rollout optimization (profile-first).
+
+**(b) Reward scoring.** Rule-based verifiable reward first (e.g., exact-match on a synthetic
+arithmetic/format task) or RM from P2. Returns a scalar per completion.
+
+**(c) Group-relative advantage.** Sample `G` completions per prompt; advantage
+`A_i = (r_i − mean(r_group)) / (std(r_group) + ε)`. No critic, no GAE.
+
+**(d) Clipped policy-gradient loss with KL leash.** Per completion token,
+`ρ_t = exp(logπθ_t − logπθ_old_t)` (old = policy at rollout time), token loss
+`−min(ρ_t A, clip(ρ_t, 1±ε) A) + βKL(πθ‖πref)`, masked to completion tokens. KL via the k3
+estimator.
+
+**(e) Actor-learner loop.** sample prompt batch → rollout G each → score → advantage →
+capture `πθ_old` logprobs → K inner epochs of clipped PG updates → repeat. Reference `πref`
+fixed throughout.
+
+Memory: policy + reference (+ RM if learned). Each 1B; recompute + accum bound activations.
+Throughput note: rollout (generation) will dominate wall-clock — a baseline must be recorded
+(tokens/s of generation vs. update) **before** any rollout optimization, per the project's
+profile-first rule.
+
+Correctness gates:
+- PG-loss finite-diff grad-check.
+- **Degenerate checks**: `G = 1` ⇒ advantage 0 ⇒ no PG signal, only KL; `ε → ∞` ⇒ vanilla PG;
+  `β = 0` ⇒ no KL term.
+- (KV-cache decode token-identical to full-recompute is gated in the engine milestone, a
+  prerequisite of GRPO.)
+- **Synthetic RL overfit**: on a tiny verifiable task with a known optimum, mean reward must
+  rise to the optimum (the RL analogue of T5's "overfit 27/27" — a hard, falsifiable signal
+  that the loop is correct, independent of fuzzy chat quality).
+
+## Evaluation
+
+- **Offline (DPO/RM)**: reward margin, preference accuracy, KL drift from reference, plus the
+  fixed chat-prompt generation suite (`scripts/chat_alpha_fixed_prompts.txt`) judged before/
+  after — reusing and extending the doc-13 recommendation for a generation-based eval harness
+  (exact-match math, code syntax, stop-token, refusal appropriateness, corruption).
+- **Online (GRPO)**: mean reward curve, KL-to-reference, response length, the verifiable-task
+  pass rate, and the same fixed-prompt suite.
+- **Selection by generation eval, not loss** — the recurring doc-13/v11 lesson: lower
+  post-training loss did not mean better generations.
+
+## Memory & throughput budget (8× RTX 5090, 1.05B model, indicative)
+
+- Params (bf16) ~2.1 GB; fp32 master ~4.2 GB; AdamW m/v ~8.4 GB; grads ~2.1 GB → policy
+  optimizer state alone ~17 GB before activations. Recompute + grad-accum keep activations
+  small; this is why post-training reuses the Phase-1/2 memory levers unchanged.
+- DPO: + reference (bf16 ~2.1 GB, or 0 if logprobs cached). Fits.
+- GRPO: + reference (~2.1 GB) (+ RM ~2.1 GB if learned). Fits; rollout activations are the new
+  variable. **Generation, not the update, is expected to be the throughput bottleneck** — to be
+  measured, not assumed.
+
+## Correctness-gate philosophy (unchanged from Phase 1/2)
+
+Every stage ships: (1) a finite-difference grad-check on the new loss/op, (2) PyTorch parity
+on loss + grads where applicable, (3) explicit degenerate-case bit/again checks (β→0, G=1,
+ε→∞, ref==policy), (4) a falsifiable "it actually learns" signal (reward margin up / synthetic
+RL overfit), and (5) **no change to the default training path** when post-training flags are
+off. New CUDA kernels (if any, e.g. decode-time attention) get the same fwd/bwd-vs-reference
+gates as flash/GQA.
+
+## Risks & tradeoffs
+
+- **Rollout engine is the long pole.** A correct KV-cache incremental-decode path is a real
+  build (decode-time attention, ragged batch). Mitigation: naive rollout first; KV-cache as an
+  isolated, separately-gated sub-phase.
+- **RL is finicky.** KL leash, advantage normalization, clip range, reward hacking. Mitigation:
+  synthetic verifiable task with a known optimum as the bring-up gate before any real chat reward.
+- **Reward-model noise** can mislead GRPO. Mitigation: rule-based reward first.
+- **Tokenizer (KI-4)** — gpt2 50257 vocab is kept for the xserv closed loop; unchanged here.
+- **Two/three resident models** raise memory; bounded by recompute/accum and (for DPO) reference
+  logprob caching.
+
+## Resolved decisions (aligned 2026-06-29)
+
+- **D1 — Scope & sequencing → DPO → GRPO, reward model optional.**
+- **D2 — Online-RL reward source → rule-based / verifiable reward first** (RM deferred/optional).
+- **D3 — Rollout engine depth → build the KV-cache incremental-decode engine up front** (not
+  naive-first), as a foundational milestone before DPO/GRPO.
+- **D4 — Alignment task / eval target → a verifiable task** (arithmetic/format/GSM8K-style) with
+  a deterministic exact-match reward, for a clean, falsifiable RL signal.
+
+## Milestones (locked order)
+
+1. **M1 — P0 SFT task baseline.** Chat template + assistant-only masking on the verifiable
+   task; produces the reference + init checkpoint. Gate: masking unit test; single-turn
+   bit-identical to `fbf4ac2`.
+2. **M2 — KV-cache decode engine** (D3, up front). Per-layer K/V cache + incremental
+   decode-time attention + batched ragged decode. Gate: **token-identical to full-recompute
+   greedy**; record decode throughput baseline.
+3. **M3 — P1 DPO.** Verifiable-checker pair construction (via M2) → `seq_logprob` op
+   (grad-check) → DPO loss (PyTorch parity; ref==policy and β→0 degenerate checks) → DPO
+   training loop → run + reward-margin / preference-accuracy curve.
+4. **M4 — P3 GRPO.** Group rollout (M2) + rule-based reward + group-relative advantage +
+   clipped PG with KL leash. Gate: PG grad-check; G=1/ε→∞/β=0 degenerate checks; **synthetic
+   verifiable-task RL-overfit** (mean reward → known optimum) → verifiable-task GRPO run.
+5. **M5 (optional) — P2 reward model.** Scalar head + ranking loss + pairwise-accuracy gate;
+   enables GRPO-with-RM for general chat.
+
+> Each milestone is one design+gate cycle; results get appended here (like the run docs) and a
+> row in `docs/evolution.md` (algorithm/infra dimensions) when it lands.
+
+## Implementation log
+
+### M1 — SFT task baseline (landed)
+
+The verifiable task and its data pipeline are implemented and verified host-side (no CUDA
+needed); the SFT run + eval ran on dash5 (1×5090). **Result: SFT moves answer-format
+adherence 0% → 100%, with arithmetic correctness 8% — exactly the intended split (SFT buys
+the format; correctness is M3/M4's job).**
+
+**Verifiable task (the spec, in one Rust module — `crates/xtrain-train/src/task.rs`):**
+
+- Two-operand integer arithmetic, ops `+ − ×`; operands `[0,999]` for `+/−`, `[0,99]` for `×`
+  (modest products); subtraction may be negative. (Ranges enlarged from the first cut to keep
+  the unique-key space ≫ requested rows — see the saturation guard below.)
+- User turn: `What is A op B?`. SFT target: `A op B = \boxed{N}.` — teaches the answer FORMAT;
+  the checker reads only `\boxed{}`, so arithmetic *correctness* is what M3/M4 improve.
+- Rule-based reward: `parse_boxed_answer` (takes the LAST `\boxed{int}`) + `check_answer`
+  (exact match vs. gold). This is the single shared checker reused by M3 (pair construction)
+  and M4 (GRPO reward).
+- Why this task: trivial deterministic checker, freely scalable difficulty, and it directly
+  probes the base model's known arithmetic weakness (v12 SFT failed `12 * 13`).
+
+**Data generator (`crates/xtrain-train/src/bin/gen_arith_task.rs`, pure host bin):**
+writes `arith_sft.tsv` (`user<TAB>assistant` for `--sft-tsv`), `arith_eval_prompts.txt`
+(`greedy_sample --prompts-file` format), and `arith_eval_gold.txt` (parallel gold ints).
+Train rows are deduped; eval is held out from train (no leakage). A **saturation guard**
+(`unique_space()` + `assert need·5 ≤ space·4`) rejects requests that approach the unique-key
+space, since deduped train + disjoint eval near saturation get pathologically slow (or, for
+the disjoint-eval loop, never terminate). With the shipped defaults the space is ~2.01M keys,
+so a 20 000 + 500 request is a tiny fraction (gen runs in ~0.2 s).
+
+**Scorer (`crates/xtrain-train/src/bin/eval_arith.rs`):** loads a checkpoint, greedily
+generates a continuation per held-out prompt, isolates the first answer segment (cut at the
+first `<|endoftext|>` then first newline), and reports two signals via the shared checker —
+**format** (fraction emitting any `\boxed{int}`) and **correctness** (exact-match vs. gold).
+This is the reusable verifiable-eval harness for M3 (DPO) / M4 (GRPO). It uses the *naive*
+no-KV-cache sampler (full forward per token), so even 100 prompts is slow — concrete
+motivation for M2 (the KV-cache decode engine).
+
+**Masking made testable:** the assistant-only label masking in `load_sft_tsv_cached` was
+extracted into a pure `sft_row(prompt_ids, answer_ids)` helper (behavior-preserving — the
+single-turn path is bit-identical to `fbf4ac2`).
+
+**Gate (verified locally in `no_cuda` mode):** `cargo test -p xtrain-train --lib` → 9/9 pass,
+including `sft_row` masks prompt→`-100` / supervises answer, the SFT-target self-consistency
+invariant (always checker-correct over 2000 samples), parser edge cases, and seed determinism.
+A 200/50 generation run confirmed clean 2-column TSV, correct gold (incl. negatives), and 0
+train/eval leakage.
+
+**Run (dash5, 1×5090, from the v12 1.05B base):**
+1. dataset: `gen_arith_task --n 20000 --eval 500 --seed 1 --out-dir <dir>` → 20 000 train +
+   500 held-out eval, 0 leakage.
+2. SFT: `train <tok> <dir>/arith_sft.tsv --sft-tsv --init-ckpt <v12-base.ckpt> --heads 52
+   --head-dim 32 --kv-heads 13 --layers 22 --ffn 6656 --bf16 --recompute --flash --seq 256
+   --batch 16 --steps 250 --max-lr 1e-4 --min-lr 1e-5 --ckpt arith_sft_v12.ckpt` → the P0
+   reference/init checkpoint. Train loss 4.68 → ~0.34, best val 0.386, no OOM, ~4.3K tok/s.
+3. eval: `eval_arith <ckpt> <tok> <arch> --prompts-file <dir>/arith_eval_prompts.txt
+   --gold-file <dir>/arith_eval_gold.txt --max-tokens 32`, base vs. SFT, on 100 held-out prompts.
+
+**M1 result (100 held-out prompts, greedy, max_new 32):**
+
+| checkpoint          | format (`\boxed{}`)  | correct (exact-match) |
+|---------------------|----------------------|-----------------------|
+| v12 base (pre-SFT)  | 0 / 100 (0%)         | 0 / 100 (0%)          |
+| arith SFT           | **100 / 100 (100%)** | 8 / 100 (8%)          |
+
+The base model never emits the format — it answers `"I don't know."` / restates the question
+and stops. SFT moves format **0% → 100%**: every completion cleanly restates the equation and
+boxes an integer (`46 * 80 = \boxed{3380}.`). Correctness is only **8%**: the format is fully
+learned but the *arithmetic* is the base model's own weak capability — e.g. it boxes 3380 for
+gold 3680, −10 for gold 5; it does get some right (`895 − 353 = \boxed{542}.` ✓). That residual
+gap is exactly what the verifiable reward in M3 (DPO) / M4 (GRPO) is built to close.
+
+**Gate met:** format 0% → 100% confirms the assistant-only SFT path is wired end-to-end; the
+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

@@ -7,7 +7,7 @@
 
 ---
 
-## 一、基建 phase（T1–T13）—— 主要动「算法」与「Infra」
+## 一、基建 phase（T1–T13 + Phase 2 systems-depth）—— 主要动「算法」与「Infra」
 
 | Phase | 维度 | 变化 | 结果 / 验证 |
 |---|---|---|---|
@@ -24,13 +24,18 @@
 | T11 | Infra | **device caching/pool allocator**（复用 op 输出显存，消 per-step cudaMalloc） | 单卡 2.3×；**8卡 461K tok/s** 近线性（修 KI-5） |
 | T12 | 算法/Infra | **bf16 混合精度**（fp32 master，cuBLAS GemmEx，norm/softmax/CE 保 fp32） | dim768 OOM 解除，−29% 显存/+13% tok/s（修 KI-2） |
 | T13 | 算法/Infra | **激活重计算**（per-block gradient checkpointing：前向 no-tape + 反向重算，`backward_seeded`） | 梯度对非重计算版**逐位一致**(0.00)；dim768 31.1→14.6GB；**dim1024 batch32 OOM→16.6GB 装下**（修 KI-3，解锁 v8） |
+| T14 | 算法/Infra | **融合 flash-attention kernel**（手写单 kernel：online softmax、tiled over KV、**不物化 N×N scores**；flash 式 bwd：重算 scores + `D=ΣdO·O` 化简雅可比 + dQ/dK/dV）；opt-in `--flash`，默认保 composed（Phase 2） | fwd 对 composed 6.7e-5、bwd 对 composed dQ 1.7e-5、PyTorch B>1 7.9e-6、flash==composed loss rel 0.0；**峰值显存 −16%@seq1024 / −23%@seq2048**（不物化 N×N，收益随 seq 增长）；tok/s ~2.3–2.8× 慢（hd=64 小头维干不过 cuBLAS tensor-core，flash 已知权衡=胜场在显存）；md5 闭环逐位一致 |
+| T15 | 模型架构 | **真 GQA**（`num_kv_heads<num_heads`：wk/wv 投影到 `kv_dim`，新 `repeat_kv` broadcast 算子把 K/V 复制 `group=nh/num_kv` 份喂给**未改动**的 composed/flash 两条 SDPA；分组约定对齐 xserv repeat_kv `dst=kvh·group+r`）；`repeat_kv` 反向=组内 group 行**确定性求和**（无 atomic）→ 多组 q 头梯度汇一个 kv 头；`num_kv_heads` 进 Config(默认=nh→MHA)、`--kv-heads` flag、导出写真 `num_key_value_heads`（Phase 2） | repeat_kv grad-check 2.1e-4(group3)+group1 identity 逐位；GQA flash==composed fp32 grad 4.1e-5/bf16 在带；**group1 对 MHA 逐位一致**(回归保护)；PyTorch GQA B>1 对拍 composed/flash 各 loss 1.7e-8/logits 2.3e-5/25 grad 进 rtol；小 GQA(8h/2kv) 训 600 步 10.9→3.15 连贯；**xserv 闭环真 GQA**(num_kv 2<8)：2/3 prompt token-identical、1 在 BF16 漂移处晚分叉；MHA 默认 export md5 逐位一致(b04fc9f9) |
+| T16 | 算法/Infra | **梯度累积**（N 个 micro-step：每个 micro-loss `×1/N` 再 backward，tape SUM 累加 → 一次 AdamW step+zero；`--accum-steps`）；**DDP 只在累积边界 all-reduce**（中间 micro-step 不发 NCCL，`/world` 与 `1/N` 正交）；显存随 micro 不随有效 batch | 等效大 batch**逐位贴合**（loss rel 8.5e-8、grad rel 3.8e-5）；`accum=1` 逐位回归(0.00)；DDP+accum 对单卡 loss 5.7e-7/跨 rank 一致；**显存平**：同有效 batch 64，big-batch 27.7GB→accum(4×16) **7.2GB(−74%)**（big-batch OOM 而 accum 装下）；全回归+xserv 闭环 md5 一致 |
 | T18 | 算法 | **dropout**（手写 counter-based 设备 RNG → Bernoulli mask，训练 inverted 1/(1-p) scaling、eval 恒等）；新 autodiff `dropout` 算子（fwd 生成+施加 mask，bwd 用同 mask），接 residual/ffn 两处；`--dropout` flag 默认 0 | 固定 seed grad-check 过；E[out]≈input + keep≈1-p；**p=0 与无 dropout 逐位一致**；recompute(T13) 组合下梯度仍逐位一致（counter-based seed 重算复现同 mask）；全回归 + xserv 闭环绿（导出/推理 dropout 关） |
+| T17 | Infra | **process-per-GPU**（torchrun 式：`launch_processes` 每卡 spawn 一个 worker 进程=独立 CUDA context；launcher 一次性铸 `ncclUniqueId` 后 **hex 编码注入子进程 env**——无共享 FS/TCP、无竞态；worker 读 env→bind device→`DdpContext::init`+`build_model`+`train_rank` **全复用 T8 零改动**；新 `train_ddp_mp` bin/`ddp_proc` test，**保留 thread-per-GPU 旧路径**）；scope=process-per-GPU only（ZeRO-1 用户 drop）（Phase 2） | 正确性全绿：proc vs 单卡 loss 5.67e-7、**proc vs thread-per-GPU 1.5e-7**、跨 rank 1.19e-7(<1e-6)、全回归+xserv 闭环 md5 逐位一致 `b04fc9f9`。**⚠️关键发现（实测证伪原假设）：本尺度 process-per-GPU 对吞吐中性**——thread vs proc @ {1,2,4,8} = {1.00/1.61/2.98/**5.27**}× vs {1.00/1.60/2.94/**5.31**}×（差<1% 噪声内）；8 卡全 95–99% util ⇒ 残留 ~5.3×@8 非线性是 **NCCL all-reduce + 本机 PCIe 拓扑墙**，**非**单 CUDA context 串行（KI-5/T11 doc 的猜想被钉死推翻，方法论同 T11 证伪「分桶 all-reduce」）。净价值=落地 torchrun 式标准链路 + 把误导性 backlog 项实测关闭；默认训练路径不变 |
+| T21 | Infra | **DDP-dropout wiring fix**（V9-PILOT 暴露：T18 只把 dropout 接进单卡 `train.rs`，`train_ddp` bin 无 `--dropout` flag、`train_rank` 从不调 `model.train()` → DDP 下 dropout 被静默忽略。补：`--dropout` flag + `train_rank` 每步 `model.train()`，镜像单卡 train/eval 纪律——`eval_loss` 翻 eval 后由每步 `train()` restore）；加 DDP-dropout 回归测试堵缺口 | DDP-dropout 回归测试绿：p>0 下 dropout **live**（loss 轨迹对 p=0 有可观差异，pre-T21 会逐位相同）、p=0 对无 dropout 路径**逐位一致**、run 后 `is_training()==true`；既有 DDP loss-match/跨 rank 测试不变。**元教训：op/单卡单元测试漏掉 launcher 级 integration gap，只有真实启动器端到端跑（pilot）才暴露** |
 
 ---
 
-## 二、Scaling runs（v0–v8）—— 主要动「模型架构」与「数据集」
+## 二、Scaling runs（v0–v10）—— 主要动「模型架构」与「数据集」
 
-架构始终是 **Qwen3-style**（RoPE + RMSNorm + QK-norm + SwiGLU，gpt2 50257 词表），逐版放大 dim/层/头（v8 起首次拨容量轴到 dim1024）；其余维度逐版变化如下：
+架构始终是 **Qwen3-style**（RoPE + RMSNorm + QK-norm + SwiGLU，gpt2 50257 词表），逐版放大 dim/层/头（v8 起首次拨容量轴到 dim1024，v9 进入 dim1280+真 GQA 双轴点，v10 固定架构只补数据轴）；其余维度逐版变化如下：
 
 | ver | 模型架构（dim/层/头·hd · 核心/总参） | 数据集（语料 · 实训 token · epoch） | 算法/精度 | Infra（GPU · 吞吐） | 结果（val · 备注） |
 |---|---|---|---|---|---|
@@ -43,25 +48,70 @@
 | v6 | dim768/18L（同 v4/v5） | **FineWeb-edu** 真实网页 · 2.29B · 1.02ep | bf16 | 8 GPU · 218K | val **3.07**：⚠️**FineWeb 留出集，与 v0–v5 不可比**（真实网页熵高，~3.0 是预期）；判据=采样质量+transfer。第一版脱离 TinyStories，**语言种类质变**（小故事→真实说明文）；transfer→TinyStories val 2.75(v5 native 1.11)，纯通用数据对窄分布有代价；val 末步仍单调降=未饱和 |
 | v7 | dim768/18L（同 v4/v5/v6） | **同 v6 的 FineWeb-edu 子集**（非新数据）· 3.28B · **1.45ep** | bf16 | 8 GPU · 218K | val **3.01**（与 v6 可比）：⚠️**同子集多 epoch 近天花板**——唯一变量=epoch(1.02→1.45)，多喂 ~1B token val 仅 ↓0.05 且 ~step44000 后走平、采样无质变。与 v5 的 TinyStories 数据量饱和同类（重复老数据边际薄）；真·更多数据要**新 shards** |
 | v8 | **dim1024**/18L/**32h** · **226M/329M**（+78% 容量，ffn 2730） | **同 v6/v7 的 FineWeb-edu 子集**（非新数据）· 2.36B · **1.05ep** | bf16 **+ 激活重计算(T13)** | 8 GPU · 129K（重算税） | val **2.98**（与 v6/v7 可比）：⭐**容量轴 A/B——容量有用**：唯一变量=dim768→dim1024，同 ~1ep v6 3.07→**2.98**（↓0.085），且 v8(1.05ep) < v7(1.45ep 更多老数据) 3.01 ⇒ 放大容量 > 重复老数据 ⇒ v6/v7 部分 capacity-limited。⚠️但增益仅 ~3%、val 末步**仍在降未饱和** ⇒ **单轴(数据/容量)单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale(Chinchilla)** |
+| v9 | **dim1280**/18L/**40h/10kv GQA** · **357M/486M**（ffn 4096） | **FineWeb-edu 扩展 shards 000-009** · **6.01B** · **~1.00ep** | bf16 + recompute + **flash + grad-accum + true GQA** | 8 GPU · **78.6K**（21.25h） | val **2.8854**（与 v6-v8 可比）：✅**双轴 Chinchilla 点有效**——容量从 v8 226M→357M，同时数据从 2.255B 子集→6.013B token，best val 比 v8 再降 **0.0947 (~3.2%)**。采样写真实说明文更稳一些，但 greedy 重复仍明显；收益仍是稳健增量而非质变 |
+| v10 | **同 v9** | **FineWeb-edu 扩展 shards 000-010** · **6.765B** · **~1.00ep** | bf16 + recompute + flash + grad-accum + true GQA | 8 GPU · **79.0K**（23.86h） | moving-tail val **2.8816**；固定 eval v1 上 v9 **2.9278**→v10 **2.8814**。结论：补 shard010 对新分布有效，但只补数据轴不解决 greedy 重复；后续应固定 eval set，并优先试更大模型+长 context |
 
-> 实训 token = steps×batch×seq（非数据集大小）。val 同一 1M-token TinyStories 留出集（v0–v5 可比；v6 起换 FineWeb-edu 留出集，分布不同、与 v0–v5 不可比；v6/v7/v8 同一 FineWeb 留出集、三版彼此可比 3.07/3.01/2.98）。
+> 实训 token = steps×batch×seq（非数据集大小）。v0–v5 的 val 是同一 1M-token TinyStories 留出集。v6 起换 FineWeb-edu，
+> 且 v9/v10 追加新 shards 会移动默认 tail-heldout；严格横比改用 fixed eval v1（shard010 tail 1M）：
+> v6/v7/v8/v9/v10 = **3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814**。
 
 ---
 
 ## 三、各维度的累积演进（轴向看一条线怎么走的）
 
-- **算法**：手写 autograd(tape)+扇出累加 → AdamW/LR-sched/grad-clip → +QK-norm(Qwen3) → batched forward → bf16 混合精度(fp32 master) → 激活重计算(T13) → dropout(T18，counter-based 设备 RNG + inverted scaling，train/eval 切换)。
-- **模型架构**：固定 Qwen3-style；dim **32→256→384→512→768→1024**（v8 首拨容量轴，头数 24→32）；核心参数 **41K→226M**（总 3.26M→329M）。
-- **Infra**：单卡 fp32 → cuBLAS/GPU-optim(T7) → NCCL DDP(T8) → batched forward(T10) → caching allocator(T11) → bf16(T12) → 激活重计算(T13，解锁 dim1024)。吞吐 **3.3K→217K tok/s**（dim768 bf16），dim1024+重算 ~129K（重算税）；MFU **0.4%→17%**（每次提升都对应一块 perf 基建，详见 known-issues + MFU 分析）。
-- **数据集**：TinyStories 3MB 切片 → 全量 TinyStories（epoch 0.01→5.33，**至饱和**）→ **v6 毕业到 FineWeb-edu 真实网页**（2.255B 语料，1.02ep）→ **v7 同子集多 epoch（1.45ep，近顶）→ v8 同子集换大模型**（dim1024，1.05ep）。tokenizer 全程 gpt2 BPE（复用 xserv-tokenizer；v6 刻意不换 tokenizer 以隔离「数据来源」变量，KI-4 留后续版本）。
+- **算法**：手写 autograd(tape)+扇出累加 → AdamW/LR-sched/grad-clip → +QK-norm(Qwen3) → batched forward → bf16 混合精度(fp32 master) → 激活重计算(T13) → 融合 flash-attention(T14，online softmax + flash 式 bwd) → 梯度累积(T16，复用 tape SUM，等效大 batch 而显存随 micro) → dropout(T18，counter-based 设备 RNG + inverted scaling，train/eval 切换)。
+- **模型架构**：固定 Qwen3-style；dim **32→256→384→512→768→1024→1280**（v8 首拨容量轴，v9 进入 dim1280）；核心参数 **41K→357M**（总 3.26M→486M）。+QK-norm(T9，Qwen3 兼容) → **真 GQA(T15，`num_kv_heads<num_heads`，repeat_kv broadcast + 组内梯度求和；默认=nh→MHA 逐位回归；v9 用 40 query / 10 kv）**——架构补齐到现代 LLM 标配（MHA/GQA/MQA 一条 `num_kv_heads` 轴），两条 SDPA(composed/flash) 共用同一 broadcast，导出真 `num_key_value_heads` 且 xserv 闭环。
+- **Infra**：单卡 fp32 → cuBLAS/GPU-optim(T7) → NCCL DDP(T8) → batched forward(T10) → caching allocator(T11) → bf16(T12) → 激活重计算(T13，解锁 dim1024) → flash-attention(T14，不物化 N×N，attention 显存收益随 seq 增长) → 梯度累积(T16，DDP 只在累积边界通信，显存随 micro 不随有效 batch) → process-per-GPU(T17，torchrun 式独立进程/CUDA context，复用 T8 train_rank 零改动)。吞吐 **3.3K→217K tok/s**（dim768 bf16），dim1024+重算 ~129K（重算税）；MFU **0.4%→17%**（每次提升都对应一块 perf 基建，详见 known-issues + MFU 分析）。T13/T14/T16 是三条**显存杠杆**（重计算压激活峰值、flash 不物化 N×N attention scores、梯度累积解耦有效 batch 与激活显存），可叠加放大有效 batch。**T17 实测=负结果记账**：process-per-GPU 在本尺度对吞吐**中性**（thread ~5.27× vs proc ~5.31×@8，差<1% 噪声），8 卡全 95–99% util ⇒ 残留非线性是 NCCL/PCIe 通信墙、**非**单 context 串行——把 KI-5/T11 doc 长挂的「process-per-GPU 是残留串行的解」猜想实测钉死推翻（方法论同 T11 证伪「分桶 all-reduce」）。
+- **数据集**：TinyStories 3MB 切片 → 全量 TinyStories（epoch 0.01→5.33，**至饱和**）→ **v6 毕业到 FineWeb-edu 真实网页**（2.255B 语料，1.02ep）→ **v7 同子集多 epoch（1.45ep，近顶）→ v8 同子集换大模型**（dim1024，1.05ep）→ **v9 扩新 FineWeb shards 到 6.013B token 并同步放大模型** → **v10 补 shard010 到 6.765B token（只拨数据轴）**。tokenizer 全程 gpt2 BPE（复用 xserv-tokenizer；保闭环优先，KI-4 接受）。
   - **v5→v6 数据轴的质变**：v0–v5 都吃合成幼儿故事（TinyStories，低熵、词汇受控），v5 证明同尺寸模型在它上面已饱和；v6 第一版换成**真实教育类网页文本**（FineWeb-edu），语言种类发生质变——采样从「只会写小故事」变成「能写历史/科学/说明文」。
   - ⚠️ **同子集多 epoch 也有天花板（v6→v7）**：v6 的 FineWeb val 才训 1.02ep、末步仍单调降，曾被读作「还没喂够」；v7 把**同一 2.255B 子集**喂到 1.45ep（多 ~1B token），FineWeb val 仅 ↓0.05（3.07→3.01）且 ~step44000 后走平、采样无质变 ⇒ **该子集在 dim768 已近天花板**。这与 v5 的 TinyStories 数据量饱和是**同一类现象**：**「重复喂老数据」边际都薄，无论是 v5 的同语料多 epoch 还是 v7 的同子集多 epoch**。真正抬天花板的是 v6「换更广的新语料」那一步——**杠杆在「更多样的新 token」，不在「同数据多读几遍」**。后续要继续降 val，必须补**新 FineWeb shards**（更多样、不重复），不是同子集加 epoch。
-  - ⚠️ **val 可比性**：v0–v5 的 val 是同一 TinyStories 1M 留出集（彼此可比）；**v6 起换 FineWeb-edu 留出集，分布不同、val 不能和 v0–v5（~1.1）比大小**——真实网页熵高，~3.0 是预期而非回退；**v6/v7/v8 同一 FineWeb 留出集、三版彼此可比**（3.07→3.01→2.98）。v6 的判据还有采样质量 + **transfer eval**（v6→TinyStories val 2.75 vs v5 native 1.11，量化「纯通用数据对窄分布的代价」）。
+  - ⚠️ **val 可比性**：v0–v5 的 val 是同一 TinyStories 1M 留出集（彼此可比）；**v6 起换 FineWeb-edu 留出集，分布不同、val 不能和 v0–v5（~1.1）比大小**——真实网页熵高，~3.0 是预期而非回退。v9/v10 追加 shards 后默认 tail-heldout 会移动，不能再只看 moving-tail best。为后续建立 fixed eval v1（shard010 tail 1M）：v6/v7/v8/v9/v10 = **3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814**。
   - ⭐ **容量轴有用，但也只有 ~3%（v8）**：v6/v7 在 dim768 上「吃不动更多数据」，v8 用最干净的 A/B 回答了「是数据见够还是容量不够」——**冻结数据子集、纯把 dim768→dim1024（core 127M→226M，+78%）**，同 ~1 epoch 下 FineWeb val **3.07→2.98（↓0.085）**，且 v8（1.05ep）还低于 v7（1.45ep 更多老数据）的 3.01。⇒ **容量有用，v6/v7 部分是 capacity-limited（不全是数据见够）**；放大容量比「给小模型多喂老数据」更值。**但增益只有 ~3%**，与数据轴单步杠杆同量级。
-  - 🧭 **元结论：单轴单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale（Chinchilla 小尺度复现）**：把三条轴并起来看——数据量轴（v5/v7 同子集多 epoch，饱和，~1.6–5%/步）、数据广度轴（v6 换语料，是一次性换分布红利）、容量轴（v8，有用但 ~3%）——**到 v8，任何单轴的单步杠杆都收敛到 ~3%/lever**。而 v8 容量 +78% 却只配同样的 2.36B token、val 末步仍在降 ⇒ 数据立刻成新瓶颈。⇒ **要继续进步，容量与数据必须匹配地一起 scale，而不是单独猛拨一根轴**——这正是 Chinchilla 在这个 toy 尺度上的复现。
+  - ✅ **双轴一起 scale 有效（v9）**：v9 把 v8 的提案落地：模型 core 226M→357M，数据 2.255B 子集→6.013B token（实训 6.012B），best FineWeb val **2.9801→2.8854**，再降 **0.0947 (~3.2%)**。这确认 Chinchilla 式双轴方向正确；但收益仍是 ~3% 级稳健增量，greedy 重复仍在，说明小尺度下“更好 val”尚未完全转化成肉眼质变。
+  - 📌 **只补数据轴边际有限（v10）**：v10 保持 v9 架构，仅补 shard010 到 6.765B token。fixed eval v1 上 v9 2.9278→v10 2.8814，说明新 shard 分布被学到；但 moving-tail best 只从 2.8854→2.8816，且 greedy 复读不变。下一步更值得改模型/context，而不是继续一片片补数据。
+
+## 三·五、Phase 2 系统栈深度综合（T14–T18 五条特性按四维收束）
+
+scaling 科学线（v0–v8）收官后，项目重启回到本职「学训练全栈」，把此前**显式延后**的五条训练栈特性补齐。区别于 Phase 1 的「修真实瓶颈」（T10–T13 每条都治一个 KI），Phase 2 是**补齐标配 + 一次诚实的负结果**。五条按四维落点：
+
+- **算法**三条 = **flash-attention(T14)** + **梯度累积(T16)** + **dropout(T18)**。
+  - 三条里 **T14/T16 与 Phase 1 的 T13 一起构成可叠加的「显存三杠杆」**：T13 压激活峰值、T14 不物化 N×N attention scores（收益随 seq 增长）、T16 解耦有效 batch 与激活显存（显存随 micro 不随 N×）——三者正交叠加可放大有效 batch / seq。
+  - **T18 dropout 的设计点是 stateless counter-based RNG**：mask 由 `(seed, 元素下标)` 无状态产出，所以**与 T13 激活重计算天然 bit-exact 组合**——反向重算时同 seed 重生同一张 mask，梯度逐位一致。这是两条 Phase-2/Phase-1 特性的正交性被正确性闸门钉死的一个例子。
+  - 诚实账：flash-attention **赢在显存不赢墙钟**（hd=64 小头维手写 kernel ~2.3× 慢于 cuBLAS tensor-core），opt-in 默认关、不回归。
+- **模型架构**一条 = **真 GQA(T15)**：架构补齐到现代 LLM 标配（MHA/GQA/MQA 一条 `num_kv_heads` 轴）。实现关键 = `repeat_kv` broadcast 算子的**反向组内确定性求和（无 atomic）**，让 K/V 零改动喂进 composed + flash 两条 SDPA；`group=1` 对 MHA **逐位一致**作回归保护，导出真 `num_key_value_heads` 且 xserv 闭环真 GQA。
+- **Infra**一条 = **process-per-GPU(T17)**，但它是**实测负结果**而非性能提升：落地 torchrun 式独立进程/CUDA context 标准链路（复用 T8 train_rank 零改动），却实测本尺度**吞吐中性**（thread ~5.27× vs proc ~5.31×@8，差<1%，8 卡全 95–99% util），把 KI-5/T11 doc 长挂的「共享单 context 致残留 ~5×@8」猜想**钉死推翻**——残留是 NCCL all-reduce + PCIe 拓扑墙，非 context 串行。**方法论与 Phase 1 的 T11（证伪「分桶 all-reduce」）一脉相承：profile/measure-first。**
+- **数据集**零条：Phase 2 不动数据轴（KI-4 小词表用户拍板 **drop 以保 xserv gpt2-tokenizer 闭环**，转记为接受的建模权衡，见 known-issues）。
+
+**Phase 2 的统一闸门 = 诚实正确性，全程未为凑绿放宽容差**：flash==composed（grad/PyTorch）、GQA group=1 == MHA 逐位、accum=1 逐位、dropout p=0 逐位 + dropout×重算 bit-exact、每条特性默认路径不变、**xserv 闭环 md5 `b04fc9f9` 两阶段全程逐位一致**。
+
+> 📌 两条 integration 发现（非回归，pre-existing，记账）：① **DDP 三个测试并行会争 2 卡 deadlock** → 文档/测试用 `--test-threads=1`（或标 serial）跑。② **fresh-train md5 run-to-run 不定**——反向 atomicAdd 归约序非确定 → 有效的确定性闸门是**导出（export）重确定性**（同 ckpt 重导 safetensors md5 逐位一致），**不是** fresh-train 复现。
+
+## 三·六、Phase 3 后训练栈（SFT → KV-cache → DPO → GRPO，详见 [18-post-training-rl-sft.md](18-post-training-rl-sft.md)）
+
+Phase 1/2 把**预训练全栈**学完后，Phase 3 转向**后训练 infra**（对齐方向）。锁定路线 DPO→GRPO（reward model 可选）、**rule-based 可验证 reward 优先**、**KV-cache 增量解码引擎前置自建**、任务取**可验证算术**（确定性 exact-match，给 RL 干净可证伪信号）。里程碑 M1（SFT baseline）→ M2（KV-cache 解码引擎，token-identical 闸门）→ M3（DPO）→ M4（GRPO）→ M5（可选 RM）。按维度落点：
+
+- **算法**：后训练损失族——SFT（assistant-only masking，已有）→ DPO（`seq_logprob` 算子 + Bradley-Terry/σ(Δ) 偏好损失，frozen reference）→ GRPO（group-relative advantage，无 critic + clipped PG + KL leash）。每条沿用 Phase 1/2 闸门规矩：新损失/算子有限差分 grad-check + PyTorch parity + 退化检查（β→0 / G=1 / ε→∞ / ref==policy）+ 一条可证伪「真在学」信号（reward margin↑ / 合成 RL overfit）。
+- **Infra**：**KV-cache 增量解码引擎（M2，前置）**是这一阶段的硬核——per-layer K/V cache + 单 token 增量 forward（prompt 灌一次 cache 后逐 token 解码）+ ragged 批量解码。硬闸门 = **解码逐 token 等价于全重算 greedy**（同 xserv 导出闭环的逐位纪律），并先记解码吞吐 baseline（profile-first）。它是 DPO 造对 + GRPO rollout 的共享底座。
+- **数据集**：可验证任务自带数据生成器——两操作数整数算术（`+ − ×`），rule-based checker 读 `\boxed{}` 做 exact-match，是 M1 SFT 数据 + M3 造对 + M4 GRPO reward 的单一共享 spec。
+- **模型架构**：复用 v12 1.05B 基座，不动架构。
+
+**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 用上）。
-- **待办**：KI-4 大词表小 vocab · process-per-GPU（要更高多卡线性时）。
-- 两次「先 profile 再动手」证伪了错误的拟修复（KI-1「加大batch」、KI-5「分桶all-reduce」），避免了无效大改——profile-first。
+- **实测关闭（负结果）**：process-per-GPU（T17）——曾挂在 KI-5/T11 doc 作残留非线性的拟修复方向，T17 实测**吞吐中性**（thread ~5.27× vs proc ~5.31×@8，8 卡全满载），残留是 NCCL/PCIe 通信墙非 context 串行 → 不再是 perf 待办，链路本身已落地留作可选路径。
+- **待办**：KI-4 大词表小 vocab（接受的建模权衡）· 要更高多卡线性 → all-reduce overlap / NVLink 互联（非本尺度优先）。
+- **三次「先 profile/measure 再动手」证伪了错误的拟修复**（KI-1「加大batch」、KI-5「分桶all-reduce」、T17「process-per-GPU 解残留串行」），避免了无效大改——profile/measure-first。

docs/known-issues.md

@@ -13,6 +13,35 @@ _(KI-1 fixed in T10. KI-5 fixed in T11. KI-2 fixed in T12. **KI-3（激活重计
 
 ## Fixed
 
+### DDP-dropout wiring（launcher 漏接 dropout）— `FIXED` (T21)
+- **背景（V9-PILOT 暴露）**：T18 dropout 在**单卡** `train.rs` 完整接好（`--dropout` flag → `cfg.dropout`，每步 `model.train()`，eval 用 `model.eval()`），op 级 + 单卡都测过。但 V9-PILOT 全栈端到端跑（DDP 8 卡 + dropout0.1 + flash + GQA + accum + bf16）时发现 **DDP 路径根本没接 dropout**：`train_ddp` bin **无 `--dropout` flag、从不设 `cfg.dropout`**，且 `ddp.rs::train_rank` **从不调 `model.train()`** → 模型停在默认 eval 模式，`ops::dropout` 恒等 → DDP 下 dropout **被静默忽略，无论 config 怎么设**。模型 + autodiff 完全支持 dropout（T18），漏的纯是 **DDP launcher / 训练 loop 的 wiring**。
+- **为何 op/单卡测试没抓到**：dropout 的测试只覆盖**单卡训练循环 + op 级 grad-check**，从没在 **DDP 路径下**跑过 dropout。`train_rank` 是独立于单卡 `train()` 的另一条 loop，二者共享 model/autodiff 但**各自布线 train/eval 纪律** —— 单卡那条对了不代表 DDP 那条对。**元教训：op 级 + 单 GPU 单元测试能漏掉 launcher 级 integration gap**；只有把特性放进**真实启动器路径**端到端跑（pilot 做的事）才暴露。修复随手补了 DDP 路径的回归测试堵这个缺口。
+- **修复（[docs/17-dropout.md](17-dropout.md)）**：① `train_ddp.rs` 加 `--dropout <p>` flag（默认 0 = 关，对齐旧行为）并设 `cfg.dropout`；② `ddp.rs::train_rank` 每步 micro-batch 循环前调 `model.train()`，镜像单卡 loop 的 train/eval 纪律——**关键**：`eval_loss()` 内部 `model.eval()` 翻成 eval 模式且**不还原**，所以每步重新 assert `model.train()`，dropout 才能跨 eval 边界保持活跃。
+- **正确性（新增 DDP-dropout 回归测试 `ddp_dropout_is_live_and_p0_bit_identical`，跑真实 `train_rank` 启动器路径）**：① **GATE A**——`p=0` 下 DDP loss 轨迹 + 末态参数对无 dropout 路径**逐位一致**（`ops::dropout(p=0)` 是 clone no-op，回归保护）；② **GATE B**——`p=0.2` 的 loss 轨迹对 `p=0` **有可观差异**（>1e-3），证 dropout mask 真在训练 forward 应用（pre-T21 代码停在 eval 模式 → 二者会逐位相同，此 gate 会 FAIL）；③ **GATE C**——run 后 `model.is_training()==true`（直接证 `model.train()` 被调用且跨末步 eval 存活）；④ p>0 run 无 NaN/Inf。测试故意启用 `eval_every < steps` 让 eval 中途翻 eval 模式，验证每步 `model.train()` 的 restore 纪律。默认 `--dropout 0` 下既有 DDP loss-match + 跨 rank 测试**不变**（回归保护）。
+- **commit**：见 T21 提交链（`distributed: --dropout flag + model.train() per step in train_rank` / `test: DDP-dropout regression (live under DDP + p=0 bit-identical)` / 文档更新）。
+
+---
+
+### process-per-GPU（torchrun 式独立 CUDA context）— `CLOSED / 实测负结果` (T17)
+- **背景**：KI-5（T11）修掉 per-op `cudaMalloc` 串行后，8 卡 scaling 从 ~1.3× 恢复到 **~5×@8**，但残留 ~5×@8 非完美线性。T11 doc / KI-5「残留」推测下一步是 **process-per-GPU**（每 rank 独立进程 + 独立 CUDA context，torchrun 式）——理由是「N rank 线程共享单 CUDA primary context，kernel-launch/cuBLAS 仍在 context 级串行」。**T17 把这条 torchrun 式链路落地并实测，证伪了该推测。**
+- **实现（[docs/16-process-per-gpu.md](16-process-per-gpu.md)）**：`xtrain-distributed` 加 `proc.rs`——`launch_processes` 每卡 spawn 一个 worker 进程（re-exec current_exe + `XTRAIN_{RANK,WORLD,LOCAL_RANK,NCCL_ID}` env）；**launcher 一次性铸 `ncclUniqueId` 后 hex 编码注入子进程 env**（无共享 FS/TCP、无轮询、无竞态——id 在子进程出生前就原子就绪）；worker 读 env → bind device（独立 CUDA context）→ `DdpContext::init` + `build_model` + `train_rank` **全部复用 T8 零改动**。新 `train_ddp_mp` bin + `ddp_proc` test；**保留 thread-per-GPU 旧路径**（回归 baseline）。scope=process-per-GPU only（ZeRO-1 用户 drop）。
+- **正确性（全绿，无回归）**：proc vs 单卡 loss `5.67e-7`、**proc vs thread-per-GPU `1.5e-7`**（两路数值同量级）、跨 rank `1.19e-7`（<1e-6）；全回归套 `--test-threads=1` 全绿 + **xserv 闭环 v3 重导 md5 逐位一致 `b04fc9f9`**（T17 不碰任何数值路径）。
+- **实测结果（关键，dash5 8× RTX 5090, dim384 per-rank batch32 seq256, steady-state）**：
+
+  | world | thread-per-GPU (`train_ddp`) | speedup | process-per-GPU (`train_ddp_mp`) | speedup |
+  |---|---|---|---|---|
+  | 1 | 93257 | 1.00× | 92952 | 1.00× |
+  | 2 | 149747 | 1.61× | 148809 | 1.60× |
+  | 4 | 278276 | 2.98× | 273308 | 2.94× |
+  | 8 | **491360** | **5.27×** | **493128** | **5.31×** |
+
+  （world=8 各重复 2 次：thread 493671/493292、proc 491102/494123——**两路差异 <1%，落在噪声内**。）
+- **诊断（证伪原推测）**：process-per-GPU world=8 跑时 `nvidia-smi` 抽样 **8 卡全部 95–99% util**（每卡 ~23GB）——GPU **已 compute-bound 喂满、非串行空转**（KI-5「1–2/8 在忙」的串行病 T11 allocator 已治好）。8 卡满载却仍只 5.3× ⇒ 缺的 ~35% 吞吐去向**每步 grad all-reduce + 本机 PCIe-only 拓扑在 8 rank 下的通信开销**（T11 早点明的「~7% all-reduce + PCIe 余量」那一层，8 卡放大），换独立 context 不动这一层。**结论：本尺度（dim384–1024、单机 8× PCIe RTX 5090）残留非线性是通信/拓扑墙，不是 launch 模型**——要再逼近线性须动 all-reduce overlap / NVLink 互联（非本尺度优先）。
+- **方法论一致**：T11 实测证伪「分桶 all-reduce」、T17 实测证伪「process-per-GPU 解残留串行」——两次都靠 measure 推翻假设而非硬上（profile/measure-first）。**净价值**：落地 torchrun 式 process-per-GPU 标准链路（项目本职「学训练全栈」）+ 把这个误导性 backlog 项**实测钉死关闭**。**默认训练路径不变**（thread-per-GPU），process-per-GPU 作并列可选路径留档。
+- **commit**：见 T17 提交链（`distributed: process-per-GPU launcher + worker` / `distributed: train_ddp_mp bin` / `test: process-per-GPU DDP correctness` / 设计文档 `docs: Phase T17 — process-per-GPU DDP design`）。
+
+---
+
 ### KI-3 · 激活重计算（gradient checkpointing）— `FIXED` (T13)
 - **触发点（v8 surfaced）**：容量轴放大到 dim1024（核心 ~210M+）测是否 capacity-limited。autograd tape 为反向保存所有中间激活，激活显存随 dim 线性增长——dim768 bf16 batch32 已 31.1GB（T12 甜点区），**dim1024 batch32 再次 OOM**（实测撞 32100/32607MiB → `OutOfMemory`）。
 - **设计（per-block gradient checkpointing，opt-in，[docs/12-activation-recompute.md](12-activation-recompute.md)）**：新增 `xtrain_autodiff::checkpoint(segment_fn, input, params)` 高阶原语（类比 `torch.utils.checkpoint`）。**前向**：把 input/params detach 成局部 leaf 跑 `segment_fn`，只取输出值，局部 tape 立即 drop → 段内激活释放（不留在外层 tape）；checkpoint 节点 parents=[input, ..params]。**反向**：从保存的 input + 未变的 param 值重跑 `segment_fn` 重建局部 tape，用上游 grad seed（`Var::backward_seeded`，新增——段输出非标量）回传，恢复的 input/param 梯度 push 给真 parents，局部 tape drop → 重算激活释放。模型每个 transformer block 前向用它包裹（`--recompute` flag，默认关）。切粒度 = 每 block。
@@ -68,7 +97,7 @@ _(KI-1 fixed in T10. KI-5 fixed in T11. KI-2 fixed in T12. **KI-3（激活重计
   → **单卡 40226→92638 tok/s (~2.3×)**；**8 卡 49K→461K tok/s (9.4×)**，scaling 从 ~1.3× 封顶恢复到 **~5×@8**；8 卡 `nvidia-smi` 抽样 **全 8 卡 95–99% util**（KI-5 时只 1–2/8 忙）。loss 轨迹逐位对住（单卡 10.9026→4.8453 before/after 一致）。
 - **正确性（全绿，无回归）**：15 算子 grad-check、5 结构、GEMM 对 cuBLAS、batched==looped、overfit 27/27、AdamW GPU bit-exact + host 对 torch、checkpoint 逐位、DDP loss 对单卡 **5.67e-7** + 跨 rank diff 0.0（loosened `<1e-6`）、**xserv 闭环**（v3 ckpt 重导 safetensors 与 registry md5 逐位一致 + xserv 加载服务贪心 "Once upon a time," 对住）。
 - **顺手**：DDP `ddp_correctness` 的 cross-rank `==0.0` → `<1e-6`（本机 PCIe-only NCCL run-to-run 跨 rank 非逐位可复现，diff≤1.2e-7 几 ULP 无害，承重闸门是 loss-match 5.67e-7）；`ddp_throughput_scaling` 扩到 world=8。
-- **残留**：~5×@8 非完美线性（grad all-reduce ~7% + 8 卡 PCIe/launch 余量），但弱扩展悬崖已消。v4 若要更高线性度，下一步是 **process-per-GPU**（每 rank 独立 CUDA context，torchrun 式）。
+- **残留**：~5×@8 非完美线性（grad all-reduce ~7% + 8 卡 PCIe/launch 余量），但弱扩展悬崖已消。曾以为下一步是 **process-per-GPU**（每 rank 独立 CUDA context，torchrun 式）——**T17 实测证伪该方向**（见下方「process-per-GPU（T17）」）：残留是**通信/PCIe 墙**，不是单 CUDA context 的 launch/cuBLAS 串行。
 - **commit**：见 T11 提交链（`cuda: device caching allocator` / `perf: KI-5 …` 那条带 before/after）。
 - **历史诊断保留如下**（证伪「分桶 all-reduce」的过程）：
 
@@ -136,6 +165,20 @@ _(KI-3 已在 T13 FIXED，见上方 Fixed。本节暂无待启项。)_
 
 ## Modeling notes
 
-### KI-4 · 大词表 embedding 占比过高
+### KI-4 · 大词表 embedding 占比过高 — `接受的建模权衡（用户拍板，T19 DROPPED）`
 - gpt2 `vocab=50257` 在 dim 小时让 embed+lm_head 主导参数：v1 25.7M/34M、v2 38.6M/66.9M；core transformer 才是学习主体。
-- 后续可考虑更贴合 TinyStories 的小 vocab（会牺牲 xserv gpt2-tokenizer 复用）；或在更大 dim 下让 core 自然成为主体（继续 scaling 即可缓解占比）。
+- **决定（不是 open，是一个被接受的权衡）**：曾计划 T19 训一个更贴语料的小 vocab 来压 embed 占比，**用户拍板 DROP**——保住 **xserv gpt2-tokenizer 闭环**（项目皇冠：导出权重回流 xserv 逐 token 一致）比清理 embedding 占比更值。小 dim 下 embed 占比高=复用 gpt2 大词表的**已知、接受的代价**。
+- 缓解路径仍在：更大 dim 时 core 自然成为主体（继续 scaling 即可摊薄占比，v8 dim1024 core 226M 已主导）；若日后愿意放弃该版闭环再重启小词表（见 [`~/toc/projects/xtrain.md`](../../toc/projects/xtrain.md) T19）。
+
+---
+
+## 集成 / 测试注记（pre-existing，非回归，记账）
+
+### DDP 三测并行会争卡 deadlock → `--test-threads=1`
+- `xtrain-distributed` 的三个 DDP 测试（thread-per-GPU correctness / scaling、process-per-GPU `ddp_proc`）若被 cargo **并行**跑，会在共享的 2 卡上互相争 GPU/NCCL 资源 **deadlock**（不是数值 bug，是测试调度）。
+- **跑法**：`cargo test ... -- --test-threads=1`（或把 DDP 测试标 serial）串行跑即全绿。Phase-2 全回归 capture 均在 `--test-threads=1` 下取得。
+
+### fresh-train md5 run-to-run 不定 → 有效确定性闸门是「导出重确定性」而非「fresh-train 复现」
+- **现象**：从随机初始化全新训练（fresh-train）产出的 ckpt，其 md5 **run-to-run 不逐位复现**。
+- **根因**：反向里多处 `atomicAdd` 归约（如跨行 dK/dV、扇出累加）的浮点**加法序非确定**（GPU 原子操作完成序不固定）→ 末位 ULP 抖动逐步累积 → ckpt 字节不同。属本机/本版的已知浮点非确定性，**不是正确性回归**（loss 轨迹仍同量级收敛）。
+- **因此项目的确定性硬闸门定义为「导出（export）重确定性」**：拿**同一个已固定的 ckpt** 重新导出 HF-safetensors，md5 与 registry **逐位一致**（`b04fc9f9`，两阶段全程）——这条是确定性的、承重的；**不要求** fresh-train 字节复现。

docs/runs/09-v9-fineweb-edu-dim1280-gqa.md

@@ -0,0 +1,152 @@
+# Scaling Run v9: Chinchilla 双轴 — dim1280/18L true GQA(core 356.9M) + FineWeb-edu 6.01B token + Phase-2 stack — Design Document
+
+## Goal
+
+v8 给出的元结论是：单独拨容量轴有用，但只有约 3% 的边际；单独重复旧数据也只有约 1.6% 的边际。要继续明显超过
+v8，必须把 **模型容量 + 新 token** 一起放大，而不是只拨一根轴。
+
+v9 就是这个双轴点：
+
+1. **模型轴**：dim1024/core 226M -> **dim1280/core 356.9M**，同时启用真 GQA（40 query heads / 10 kv heads）。
+2. **数据轴**：v6-v8 的 2.255B FineWeb 子集 -> **6.013B token**，追加了新 FineWeb-edu shards 003-009。
+3. **系统栈**：使用 Phase-2 现代路径：`--flash + --accum-steps + bf16 + recompute + DDP`。dropout 设为 0，按标准预训练。
+
+> v9 的 val 仍是 FineWeb-edu 分布，不能和 v0-v5 的 TinyStories val 直接比。注意：v9 扩展 cache 后默认
+> tail-heldout 已经从 v6-v8 的旧 tail 移到新 shards 末尾；严格横比后续以 fixed eval v1 为准。
+
+## Data
+
+| 项 | 值 |
+|----|----|
+| 来源 | FineWeb-edu `sample/10BT`，原 shards 000-002 + 新 shards 003-009 |
+| token cache | `data/fineweb-edu.txt.u16.bin` |
+| 总 token | **6,013,639,492** |
+| held-out val | 末尾 **1,000,000** token |
+| train corpus | 6,012,639,492 token |
+| 训练消费 token | **6,012,600,320** = 91745 steps x effective batch 256 x seq 256 |
+| epoch | ~1.00 |
+
+P3-DATA 目标本来是约 7B token；shard 010 下载 `curl rc=18` 中断，所以最终停在 6.01B。对 core 356.9M 来说，
+D/N 约 **16.8 token/param**，低于理想 Chinchilla 20，但已经远高于 v8 的约 10.4，是一个干净的双轴 scale 点。
+
+## Architecture
+
+| 项 | v8 | **v9** |
+|----|----|----|
+| dim | 1024 | **1280** |
+| layers | 18 | 18 |
+| query heads x head_dim | 32 x 32 | **40 x 32** |
+| kv heads | 32 (MHA) | **10 (true GQA, group=4)** |
+| ffn | 2730 | **4096** |
+| core params | 226.50M | **356.89M** |
+| total params | 329.42M | **485.55M** |
+| export tensors | 201 | **201** |
+
+`config.json` writes real `num_key_value_heads = 10`, so xserv loads v9 as true GQA rather than MHA.
+
+## Training
+
+| 项 | 值 |
+|----|----|
+| optimizer | hand-written AdamW, wd=0.1 |
+| schedule | warmup -> cosine, max_lr 6e-4 -> min_lr 6e-5 |
+| grad clip | global norm 1.0 |
+| steps | **91745** |
+| effective global batch | **256** (`--batch 128 --accum-steps 2`) |
+| seq_len | 256 |
+| precision | bf16 mixed precision, fp32 master |
+| memory stack | activation recompute + flash-attention + gradient accumulation |
+| world size | 8 x RTX 5090 |
+| wall clock | **21h15m** |
+| steady throughput | **~78.6K tok/s** |
+| peak observed memory | ~17GB / GPU |
+
+Command:
+
+```sh
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 cargo run -p xtrain-distributed --release --bin train_ddp -- \
+  /opt/wjh/models/gpt2/tokenizer.json data/fineweb-edu.txt \
+  --heads 40 --head-dim 32 --kv-heads 10 --layers 18 --ffn 4096 \
+  --steps 91745 --batch 128 --accum-steps 2 --seq 256 \
+  --max-lr 6e-4 --min-lr 6e-5 --val-tokens 1000000 --eval-every 1000 \
+  --eval-batches 64 --bf16 --recompute --flash --dropout 0.0 \
+  --ckpt /dashscope-tmp/wjh/xtrain_v9.ckpt
+```
+
+## Results
+
+- train loss: **11.1550 -> 2.9340**
+- first val: step 1000 = **5.1517**
+- best val: step 91000 = **2.8854**
+- final val: step 91745 = **2.8873**
+- exit code: **0**
+
+FineWeb val curve milestones:
+
+| step | 1000 | 10000 | 20000 | 30000 | 40000 | 50000 | 60000 | 70000 | 80000 | 90000 | 91000 | final |
+|------|------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
+| val | 5.1517 | 3.4820 | 3.2953 | 3.2026 | 3.1422 | 3.0844 | 3.0148 | 2.9616 | 2.9160 | 2.8915 | **2.8854** | 2.8873 |
+
+The curve kept improving into the last 1K-step window, then the final eval bounced slightly from 2.8854 to 2.8873. This is close to
+the floor for this run, but not a clear overfit failure.
+
+## Comparison
+
+| | v6 | v7 | v8 | **v9** |
+|---|---|---|---|---|
+| model | dim768/core127M | dim768/core127M | dim1024/core226M | **dim1280/core357M + GQA** |
+| data | 2.29B | 3.28B same subset | 2.36B same subset | **6.01B expanded shards** |
+| best val | 3.0652 | 3.0149 | 2.9801 | **2.8854** |
+
+On the run-local moving tail, v9 beats v8 by **0.0947** val loss (~3.2% relative), essentially the same size as the
+v6->v8 capacity gain but now on top of it. A later fixed eval v1 check still supports the same direction
+(v8 3.1515 -> v9 2.9278 on shard010-tail holdout), while making the moving-tail caveat explicit. This confirms
+the v8 prediction: **双轴 scale 有效**. It is still an incremental gain, not a qualitative jump.
+
+## Samples
+
+xserv greedy samples (`--max-tokens 60`) are more coherent than the v8 examples on some prompts, but repetition remains:
+
+```text
+[The history of] the United States is the story of the people, the places, and the events that have shaped the nation...
+[In science,] the term "scientific method" is used to describe the process of gathering information and testing it...
+[The most important] thing is to be aware of the symptoms and to seek medical attention...
+[Water is] a natural resource that is essential for human life...
+```
+
+The model writes real explanatory English and the domain mix is FineWeb-like. Greedy decoding still falls into repeated clauses on
+some prompts (`scientific method`, symptoms, and earlier fixed prompts), so the val gain is more visible in the metric than in a
+dramatic sample-quality leap.
+
+## xserv validation
+
+Registry path:
+
+```text
+/opt/wjh/projects/tiny-models/v9-fineweb-edu-dim1280-gqa
+```
+
+Files:
+
+- `config.json`
+- `model.safetensors` (BF16, 201 tensors, 927MB)
+- `tokenizer.json`
+- `xtrain.ckpt` (fp32 master checkpoint, 1.9GB)
+- `RUN.md`
+
+xserv loads v9 as:
+
+```text
+Model: qwen3, layers=18, hidden=1280, heads=40/10 kv, vocab=50257
+Loaded 201 tensors
+Ready (KV cache, dtype=bf16).
+```
+
+Token-match check against xtrain greedy (`max-tokens 40`):
+
+- `Once upon a time`: xtrain and xserv matched through the checked continuation.
+- `One day`: diverged after "large, dark," (`very tall man` vs `metallic object`) from BF16 greedy tie sensitivity.
+- `The little`: same repetitive pattern, with a short BF16 path divergence.
+
+This is the same class of BF16-vs-f32 greedy drift seen in v8; the important integration result is that xserv successfully loads
+true GQA (`kv_heads=10 < heads=40`) and generates from the exported weights.

docs/runs/10-v10-fineweb-edu-dim1280-gqa-data6765.md

@@ -0,0 +1,200 @@
+# Scaling Run v10: Data-axis follow-up — dim1280/18L true GQA + FineWeb-edu 6.765B token — Design Document
+
+## Goal
+
+v9 证明了双轴 scale（更大模型 + 更多新 token）有效：best val 从 v8 的 2.9801 降到 2.8854。
+但 v9 的数据量只有 6.013B token，D/N 约 16.8，低于 Chinchilla 经验里的 20。v10 的目标很窄：
+
+1. **只补数据轴**：补上 v9 中断的 FineWeb-edu shard010，把 cache 从 6.013B 推到 6.765B。
+2. **架构不变**：完全复用 v9 dim1280 / 18L / 40q-10kv GQA / ffn4096。
+3. **验证边际**：看 D/N 从 16.8 到 18.95 是否还能显著降低 val。
+
+## Data
+
+| 项 | 值 |
+|----|----|
+| 来源 | FineWeb-edu `sample/10BT`，shards 000-010 |
+| token cache | `data/fineweb-edu.txt.u16.bin` |
+| 总 token | **6,765,333,808** |
+| held-out val | 末尾 **1,000,000** token |
+| train corpus | 6,764,333,808 token |
+| 训练消费 token | **6,764,298,240** = 103215 steps x effective batch 256 x seq 256 |
+| epoch | ~1.00 |
+
+Important caveat: xtrain 当前训练入口用“全 cache 的末尾 1M token”做 held-out。追加 shard010 后，v10 的 val tail
+和 v9 的 val tail 不再是同一个切片。因此 v9 原报告的 2.8854 与 v10 原报告的 2.8816 不能被当作严格同一
+验证集上的横比。
+
+为了解决这个问题，本轮创建了固定 eval set：
+
+```text
+/dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt.u16.bin
+```
+
+它包含 shard010 末尾 11M token；前 10M token 只是为了复用现有 `split_tail(val_tokens=1M)`，真正 eval 的是最后
+1M token。该 fixed eval v1 对 v6-v9 都是未见数据；对 v10 也是训练时 held-out。
+
+## Architecture
+
+v10 与 v9 完全相同：
+
+| 项 | 值 |
+|----|----|
+| dim | 1280 |
+| layers | 18 |
+| query heads x head_dim | 40 x 32 |
+| kv heads | 10 (true GQA, group=4) |
+| ffn | 4096 |
+| core params | 356.89M |
+| total params | 485.55M |
+| export tensors | 201 |
+
+## Training
+
+| 项 | 值 |
+|----|----|
+| optimizer | hand-written AdamW, wd=0.1 |
+| schedule | warmup -> cosine, max_lr 6e-4 -> min_lr 6e-5 |
+| grad clip | global norm 1.0 |
+| steps | **103215** |
+| effective global batch | **256** (`--batch 128 --accum-steps 2`) |
+| seq_len | 256 |
+| precision | bf16 mixed precision, fp32 master |
+| memory stack | activation recompute + flash-attention + gradient accumulation |
+| world size | 8 x RTX 5090 |
+| wall clock | **23h51m** |
+| steady throughput | **~79.0K tok/s** |
+| peak observed memory | ~17GB / GPU |
+
+Command:
+
+```sh
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 cargo run -p xtrain-distributed --release --bin train_ddp -- \
+  /opt/wjh/models/gpt2/tokenizer.json data/fineweb-edu.txt \
+  --heads 40 --head-dim 32 --kv-heads 10 --layers 18 --ffn 4096 \
+  --steps 103215 --batch 128 --accum-steps 2 --seq 256 \
+  --max-lr 6e-4 --min-lr 6e-5 --val-tokens 1000000 --eval-every 1000 \
+  --eval-batches 64 --bf16 --recompute --flash --dropout 0.0 \
+  --ckpt /dashscope-tmp/wjh/xtrain_v10.ckpt
+```
+
+## Results
+
+- train loss: **11.1575 -> 2.9000**
+- first val: step 999 = **5.3048**
+- best val: step 103214 = **2.8816**
+- final val: step 103214 = **2.8816**
+- exit code: **0**
+
+FineWeb moving-tail val milestones:
+
+| step | 999 | 9999 | 19999 | 29999 | 39999 | 49999 | 59999 | 69999 | 79999 | 89999 | 99999 | final |
+|------|-----|------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
+| val | 5.3048 | 3.5622 | 3.3282 | 3.2450 | 3.1886 | 3.1342 | 3.0714 | 3.0202 | 2.9724 | 2.9236 | 2.8950 | **2.8816** |
+
+The curve still improved at the final eval. There is no overfit signal in this run.
+
+## Fixed Eval V1
+
+Fixed eval v1 (`shard010 tail 1M`, seq256, 64 eval batches):
+
+| version | fixed eval v1 |
+|---------|---------------|
+| v6 | 3.2328 |
+| v7 | 3.1850 |
+| v8 | 3.1515 |
+| v9 | 2.9278 |
+| **v10** | **2.8814** |
+
+This is the cleanest cross-version result in the v10 round. It says:
+
+- v9's double-axis gain transfers to a shard010 holdout: v8 3.1515 -> v9 2.9278.
+- v10 further improves on the new shard010 distribution: v9 2.9278 -> v10 2.8814.
+- The apparent v9 moving-tail 2.8854 -> v10 moving-tail 2.8816 delta is tiny and not strict apples-to-apples.
+
+## Decoding
+
+Greedy decoding still repeats. Fixed prompts from xtrain:
+
+```text
+[Once upon a time] there was a king who had a daughter. She was beautiful and beautiful...
+[The little] The little boy was a little boy. The little boy was a little boy...
+[One day] I was walking down the street and I saw a man with a dog...
+```
+
+Temperature 0.8 is more varied and less immediately looped, but coherence remains weak:
+
+```text
+[Once upon a time] I was a kid who did not go to the beach to swim...
+[The little] ones are not as loud as the adults...
+[One day] I was on the edge of the water, and I saw something I had never seen before...
+```
+
+xserv loads the exported v10 true-GQA weights and generates FineWeb-like explanatory prose, but repeated sentence frames remain:
+
+```text
+[The history of] the city of San Francisco is a story of the growth of the city...
+[In science,] the term "observation" is used to describe the act of observing something...
+[Water is] the most important element in the human body...
+```
+
+Conclusion: decoding remains a separate bottleneck. The current xtrain sampler only supports greedy and temperature sampling; top-p and
+repetition penalty exist in xserv's chat path, but not in the raw xtrain sampler or `xserv-cli` path used for weight validation. A clean
+next step is to add a raw generation tool with `temperature/top-p/repetition-penalty` so decoding experiments do not depend on chat
+templates.
+
+## xserv Validation
+
+Registry path:
+
+```text
+/opt/wjh/projects/tiny-models/v10-fineweb-edu-dim1280-gqa-data6765
+```
+
+Files:
+
+- `config.json`
+- `model.safetensors` (BF16, 201 tensors, 927MB)
+- `tokenizer.json`
+- `xtrain.ckpt` (fp32 master checkpoint, 1.9GB)
+
+xserv loads v10 as true GQA:
+
+```text
+Model: qwen3, layers=18, hidden=1280, heads=40/10 kv, vocab=50257
+Loaded 201 tensors
+Ready (KV cache, dtype=bf16).
+```
+
+## v11 Feasibility: Bigger Model + Longer Context
+
+A v11 smoke test prioritized the user's chosen direction: larger model plus longer context.
+
+Candidate:
+
+| item | value |
+|------|-------|
+| dim / layers | 1536 / 20 |
+| heads / kv_heads | 48 / 12 |
+| ffn | 6144 |
+| core / total params | 684.26M / 838.65M |
+| stack | bf16 + recompute + flash + accum + 8 GPU DDP |
+
+Smoke results:
+
+| seq | batch / accum | effective batch | peak mem | tok/s | result |
+|-----|---------------|-----------------|----------|-------|--------|
+| 512 | 64 / 4 | 256 | **30530 MiB** | **44.7K** | 50 steps OK |
+| 1024 | 32 / 8 | 256 | **30530 MiB** | **31.0K** | 20 steps OK |
+
+Both fit, but the memory margin is thin on 32GB RTX 5090. Expected one-epoch wall clock on 6.76B tokens:
+
+- seq512: roughly **42h**
+- seq1024: roughly **61h**
+
+Recommendation: make v11 a controlled run, not a blind launch. Use fixed eval v1, keep data fixed, and choose either:
+
+1. **v11a practical**: dim1536/20L, seq512, batch64/accum4. Faster, still doubles context over v10.
+2. **v11b long-context**: dim1536/20L, seq1024, batch32/accum8. More aligned with "long context", but ~2.5 days and tight memory.
+
+For scientific clarity, v11 should not append more data before training; use the current 6.765B train cache while preserving fixed eval v1.

docs/runs/12-v12-1b-longctx-chat-alpha.md

@@ -0,0 +1,251 @@
+# Scaling Run v12: 1B-class long-context base → chat-alpha-v2 — Design Document
+
+## Goal
+
+v11 proved that a larger `dim1536/20L` model can train at `seq1024` on dash5, and it improved the fixed-eval-data-v1, long-context (`seq1024`) score to **2.7467**. It also proved the current bottleneck: greedy generation still repeats, and broad real-data SFT on top of v11 regressed chat quality despite lower SFT validation loss.
+
+v12 therefore separates the next phase into two gates:
+
+1. **Base gate**: train a stronger English base model around 1B total params with `seq1024`, using the existing FineWeb-edu 6.765B-token cache and fixed eval data v1.
+2. **Chat gate**: only after the base gate is healthy, run assistant-only English SFT and judge it with fixed prompt generation, not SFT loss alone.
+
+Success means the model is serviceable enough for a small chat-alpha: stable base loss, lower fixed eval than v11, less repetitive fixed generation, and SFT that improves instruction behavior without destroying arithmetic/refusal/debug prompts.
+
+## Baseline: What v11 Taught Us
+
+| item | v11 |
+|------|-----|
+| arch | dim1536 / 20L / 48q-12kv GQA / ffn6144 |
+| params | 684.26M core / 838.65M total |
+| data | FineWeb-edu 6.765B token, 1 epoch |
+| context | seq1024 |
+| throughput | ~30.96K tok/s on 8 x RTX 5090 |
+| fixed eval data v1, seq1024 | **2.7467** |
+| issue | greedy repetition remains; direct real SFT regressed generation quality |
+
+SFT result from v11:
+
+| model | train result | generation result |
+|-------|--------------|-------------------|
+| `v11-chat-alpha-sft-v2-anchor` | synthetic assistant-only anchor | current best narrow chat-alpha |
+| `v11-chat-alpha-real-sft-v1` | SFT val 1.4272 | bad hallucination, math failure |
+| `v11-chat-alpha-real-mix-v1` | SFT val 2.0543 | better than direct real-SFT, still worse than anchor |
+
+Conclusion: SFT data quality matters, but v11's base is still too weak for broad real SFT to become a general chat model.
+
+## Architecture
+
+v12 target: slightly above 1B total params while staying close to the proven v11 shape and keeping GQA group size 4.
+
+| item | value |
+|------|-------|
+| dim | **1664** |
+| layers | **22** |
+| query heads x head_dim | **52 x 32** |
+| kv heads | **13** |
+| GQA group | 4 |
+| ffn | **6656** |
+| core params | **883.4M** |
+| embed + lm_head | **167.3M** |
+| total params | **1.0506B** |
+
+Why this shape:
+
+- It is a controlled step from v11 rather than a new architecture family.
+- `52/13` preserves true GQA with group 4.
+- Total params are near the requested 1B target.
+- `dim1664` is less aggressive than `dim1792/22L` and has a better chance to fit `seq1024` on 32GB 5090s.
+
+## Data
+
+Base pretraining stays English-oriented and uses the current token cache. Pass the `.txt` stem to xtrain; `Corpus::load_cached` appends `.u16.bin` internally.
+
+```text
+/opt/wjh/projects/xtrain/data/fineweb-edu.txt
+cache = /opt/wjh/projects/xtrain/data/fineweb-edu.txt.u16.bin
+tokens = 6,765,333,808
+```
+
+Training uses the last 1M tokens as moving-tail validation. Every cross-version v12 claim must also run fixed eval data v1 with the long-context `seq1024` setting, matching the v11 `eval_v11_seq1024.log` score of **2.7467**. This is distinct from the older v10 table that used the same fixed eval data with `seq256`.
+
+```text
+/dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt
+cache = /dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt.u16.bin
+```
+
+No new FineWeb shards are added in this phase. The experiment is model/context scale, not another data-axis change.
+
+## Training Plan
+
+Primary v12 run:
+
+| item | value |
+|------|-------|
+| world | 8 x RTX 5090 on dash5 |
+| precision | bf16 mixed precision, fp32 master |
+| memory stack | recompute + flash + grad accumulation |
+| seq | **1024** |
+| micro global batch | **16** (2 sequences/rank) |
+| accum | **15** |
+| effective global batch | **240** |
+| tokens/step | **245,760** |
+| full steps | **27,524** |
+| max_lr → min_lr | **4e-4 → 4e-5** |
+| eval | moving-tail 1M every 500 steps; fixed eval data v1 at seq1024 after checkpoints |
+| smoke throughput | **~24.5K tok/s** |
+| estimated full wall clock | **~76-78h** |
+
+The reduced micro-batch is intentional: v11 `seq1024` with global batch 32 already sat near the 5090 memory limit. v12 has larger weights; an initial `batch24/accum10` smoke OOMed after step 0, while `batch16/accum15` passed a 10-step smoke at ~29.4GB/GPU and preserved the same 245,760 tokens/step.
+
+Command wrapper:
+
+```sh
+scripts/run_v12_phase.sh start-pilot
+scripts/run_v12_phase.sh start-full
+scripts/run_v12_phase.sh status
+scripts/run_v12_phase.sh eval-fixed
+scripts/run_v12_phase.sh export
+scripts/run_v12_phase.sh sample
+```
+
+## Gates
+
+### Gate 0: build and smoke
+
+Run:
+
+```sh
+scripts/run_v12_phase.sh smoke
+```
+
+Pass criteria:
+
+- no CUDA OOM
+- no NaN loss
+- first 30 steps decrease from initialization
+- peak memory leaves enough margin for eval
+
+### Gate 1: pilot
+
+Run:
+
+```sh
+scripts/run_v12_phase.sh start-pilot
+```
+
+Default pilot is 300 steps with held-out eval every 100 steps.
+
+Pass criteria:
+
+- train loss decreases smoothly
+- grad norm does not spike persistently
+- moving-tail eval is finite and improving
+- checkpoint can be reloaded by `eval-fixed`
+
+### Gate 2: full base
+
+Run only after the pilot passes:
+
+```sh
+scripts/run_v12_phase.sh start-full
+```
+
+Pass criteria:
+
+- fixed eval data v1 at `seq1024` beats v11's **2.7467**
+- generation samples improve or at least do not regress on repetition
+- checkpoint exports and xserv loads the true GQA config
+
+### Gate 3: chat-alpha SFT
+
+After a healthy v12 base:
+
+1. Use assistant-only SFT (`--sft-tsv`) with English-only data.
+2. Start from narrow anchors first, then mix in Smol-SmolTalk.
+3. Judge with fixed generation prompts before calling it useful.
+
+The primary high-quality source remains `HuggingFaceTB/smol-smoltalk` filtered to English single-turn examples, with local anchors preserved to keep deterministic behavior.
+
+## Evaluation
+
+Base metrics:
+
+- moving-tail val during training
+- fixed eval data v1 at `seq1024`
+- xtrain fixed prompt samples from `scripts/chat_alpha_fixed_prompts.txt`
+- xserv exported-model smoke
+
+Chat metrics:
+
+- fixed prompt answers for SFT explanation, SFT data provenance, arithmetic, refusal, repetition-debug checklist, summary, and simple code generation
+- compare against `v11-chat-alpha-sft-v2-anchor`
+- reject models that lower SFT validation loss but hallucinate more in fixed prompts
+
+## Artifacts
+
+Expected paths:
+
+```text
+/dashscope-tmp/wjh/xtrain_v12/
+/dashscope-tmp/wjh/xtrain_v12/xtrain_v12_pilot.ckpt
+/dashscope-tmp/wjh/xtrain_v12/xtrain_v12.ckpt
+/opt/wjh/projects/tiny-models/v12-fineweb-edu-1b-longctx
+```
+
+## Results
+
+### Gate 0/1: smoke + pilot
+
+- `batch24/accum10` smoke OOMed after step 0.
+- `batch16/accum15` smoke passed 10 steps: train loss **11.2347 -> 7.9459**, ~24.5K tok/s, ~29.4GB/GPU.
+- 300-step pilot passed: train loss **11.2296 -> 5.4832**, val **6.5810 -> 5.9642 -> 5.5888**, exit code 0.
+- Pilot checkpoint reload matched final val: fixed eval data v1 at seq1024 = **5.5891**.
+- Fixed chat prompts still repeat heavily, as expected for a 300-step base; use them as a regression baseline, not as chat quality.
+
+### Gate 2: full base
+
+Full run completed on dash5:
+
+| item | result |
+|------|--------|
+| wall clock | **81h01m** |
+| throughput | **~24.55K tok/s** |
+| train loss | **11.2294 -> 2.6696** |
+| moving-tail best val | **2.7411** |
+| moving-tail final val | **2.7412** |
+| fixed eval data v1, seq1024 reload | **2.7410** |
+| exit code | **0** |
+
+Validation milestones:
+
+| step | 499 | 999 | 1499 | 1999 | 2499 | 21999 | 23999 | 25999 | 26999 | 27499 | final |
+|------|-----|-----|------|------|------|-------|-------|-------|-------|-------|-------|
+| val | 5.3029 | 4.4079 | 3.9287 | 3.6964 | 3.5555 | 2.7805 | 2.7637 | 2.7468 | 2.7443 | **2.7411** | 2.7412 |
+
+Compared with v11's fixed eval data v1 at seq1024 (**2.7467**), v12 reaches **2.7410** after reload. This is a real but very small gain
+(~0.006 absolute), despite the parameter increase from 838.65M to 1.0506B total and the slower 24.55K tok/s throughput. The result says the
+larger 1B-class base is viable and marginally better, but this scale step did not produce a qualitative base-model jump.
+
+Generation:
+
+- Raw FineWeb-style prompts are better than the pilot checkpoint and can produce plausible explanatory prose.
+- Greedy repetition remains visible, especially on story-like prompts.
+- Chat prompts are not reliable without SFT: SFT data provenance is hallucinated, arithmetic still fails, and the model repeats template-like text.
+- xserv loads the export correctly as true GQA: `layers=22, hidden=1664, heads=52/13 kv`.
+
+Exported model:
+
+```text
+/opt/wjh/projects/tiny-models/v12-fineweb-edu-1b-longctx
+```
+
+Files:
+
+- `config.json`
+- `model.safetensors` (2.0GB)
+- `tokenizer.json`
+- `xtrain.ckpt` (4.0GB)
+
+Conclusion: v12 passes the base gate and is a better SFT starting point than v11 by metric, but the gain is narrow. The next step should be
+assistant-only chat SFT from v12 with conservative anchors first, then a small Smol-SmolTalk mix. Do not expect the base checkpoint itself to
+serve as a usable chat model.

docs/runs/13-v12-chat-sft-quality.md

@@ -0,0 +1,180 @@
+# v12 Chat SFT Quality Check
+
+Date: 2026-06-29
+
+## Goal
+
+Turn the completed v12 1.05B base checkpoint into a usable chat-alpha model with
+SFT, then judge whether it is stable enough to call a high-quality chat model.
+
+Base checkpoint:
+
+```text
+/dashscope-tmp/wjh/xtrain_v12/xtrain_v12.ckpt
+```
+
+Architecture:
+
+```text
+dim=1664 layers=22 heads=52 kv_heads=13 head_dim=32 ffn=6656
+total params=1.0506B
+```
+
+## Stage A: Synthetic SFT
+
+Data:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_alpha_v2/chat_alpha_v2_sft.tsv
+211,257 examples, about 14.96M SFT tokens
+```
+
+Run:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_v12_alpha_v2/chat_alpha_v12_v2.ckpt
+```
+
+Metrics:
+
+```text
+train loss: 3.5730 -> 0.0426
+eval: step39 0.1078, step79 0.0582, step119 0.0466, step159 0.0423,
+      step199 0.0403, step239 0.0390, step279 0.0389, step319 0.0378
+best/final val loss: 0.0378
+```
+
+Export:
+
+```text
+/opt/wjh/projects/tiny-models/v12-chat-alpha-sft-v2
+```
+
+Quality notes:
+
+- Learns the User/Assistant format and usually stops correctly.
+- Too narrow and template-heavy.
+- Fails basic math and code prompts in fixed greedy evaluation.
+
+## Stage B: Anchor SFT
+
+Data:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_alpha_v2_anchor/chat_alpha_v2_anchor.tsv
+32,020 examples, about 1.73M SFT tokens
+```
+
+Run:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_v12_anchor/chat_alpha_v12_anchor.ckpt
+```
+
+Metrics:
+
+```text
+train loss: 1.7777 -> 0.1165
+eval: step19 0.3447, step39 0.1449, step59 0.1217, step79 0.1158
+best/final val loss: 0.1158
+```
+
+Export:
+
+```text
+/opt/wjh/projects/tiny-models/v12-chat-alpha-sft-v2-anchor
+```
+
+Generation artifacts:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_v12_anchor/generation/anchor_xserv_greedy.txt
+/dashscope-tmp/wjh/xtrain_sft_v12_anchor/generation/anchor_diagnostic_greedy.txt
+```
+
+Quality notes:
+
+- Better project-context answers and summaries than synthetic-only.
+- Still unreliable on basic multiplication, yes/no facts, translation, and code.
+- Overuses "cannot verify" style answers outside appropriate uncertainty cases.
+
+## Stage C: Real-Mix Repair
+
+Data:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_real_mix_v1/smol_smoltalk_real_mix.tsv
+96,287 examples, about 25.3M SFT tokens
+```
+
+Run:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/chat_alpha_v12_real_mix_repair.ckpt
+```
+
+Training setup:
+
+```text
+init=/dashscope-tmp/wjh/xtrain_sft_v12_anchor/chat_alpha_v12_anchor.ckpt
+steps=200
+seq=512
+batch=32
+accum=8
+effective batch=256
+lr=1e-6 -> 2e-7
+```
+
+Metrics:
+
+```text
+train loss: 2.7391 -> 2.0384
+eval: step49 2.1964, step99 2.0383, step149 1.9801, step199 1.9570
+best/final val loss: 1.9570
+```
+
+Export:
+
+```text
+/opt/wjh/projects/tiny-models/v12-chat-alpha-real-mix-repair
+```
+
+Generation artifacts:
+
+```text
+/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/generation/real_mix_repair_xserv_greedy.txt
+/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/generation/real_mix_repair_diagnostic_greedy.txt
+/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/generation/real_mix_repair_diagnostic_greedy_reppenalty1.txt
+```
+
+Quality notes:
+
+- Loss improved cleanly and the model kept chat formatting.
+- Fixed prompt math `17% of 240` improved in the standard suite.
+- General diagnostic math still fails, e.g. `12 * 13`.
+- Code generation remains unusable for simple Python function prompts.
+- Some outputs contain corrupted or off-topic fragments.
+- Reducing repeat penalty from 1.15 to 1.0 did not fix the failures.
+
+## Verdict
+
+The SFT pipeline works, and v12 can be turned into a chat-shaped model that follows
+the prompt format and stops correctly. However, none of the three SFT variants is a
+stable high-quality chat model yet.
+
+The limiting issue is no longer infrastructure. It is data and objective quality:
+the current synthetic/anchor data is too narrow, while the current real-mix data
+adds breadth but also noisy or low-quality behavior. Validation loss alone is not a
+sufficient selection signal for chat quality.
+
+## Recommended Next Step
+
+Build a smaller, higher-precision SFT curriculum before another large run:
+
+1. Keep the anchor data, but reduce over-refusal templates.
+2. Add verified small instruction sets for math, code, translation, summarization,
+   and closed-book common facts.
+3. Add an automatic fixed-prompt eval harness that scores exact-match math, simple
+   code syntax, refusal appropriateness, stop-token behavior, and corruption.
+4. Train a short curriculum from the v12 base or v12 anchor checkpoint, then pick
+   by generation eval rather than SFT loss alone.

docs/runs/README.md

@@ -22,6 +22,10 @@ val loss 一栏给的是各版**各自训练 run 报告的 best val**（held-out
 **v8 改测容量轴**：同 v6/v7 子集、纯把 dim768→dim1024（core 127M→226M），FineWeb val 3.07/3.01→**2.98** ⇒
 **容量有用**（v6/v7 部分 capacity-limited）；但增益仅 ~3%、val 末步仍在降未饱和 ⇒ **到 v8，数据轴与容量轴的
 单步杠杆都收敛到 ~3%/lever = 全面边际递减，要双轴一起 scale**（Chinchilla，详见 [08-v8](08-v8-fineweb-edu-dim1024.md)）。
+**v9 兑现双轴**：dim1024→dim1280（core 226M→357M）并把 FineWeb token 从 2.255B 子集扩到 6.013B，
+best val **2.8854**，相比 v8 再降 0.0947（~3.2%）。结论：双轴 scale 有效，但仍是稳健增量而非质变。
+**v10 只补数据轴**：同 v9 架构，只补 shard010 到 6.765B token，moving-tail best/final val **2.8816**。
+注意追加 shard 会移动 held-out tail；固定 eval v1 上 v6→v10 为 **3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814**。
 
 ⚠️ **v6 起换了保留集（语料）**：v0–v5 的 val 都是 **TinyStories** 1M 留出集（彼此可比）；v6 换成纯
 **FineWeb-edu**（真实网页文本），它的 val（3.07）是**另一把尺子上的另一个分布**，**不能**和 v0–v5 的
@@ -39,18 +43,15 @@ val loss 一栏给的是各版**各自训练 run 报告的 best val**（held-out
 | [v6-fineweb-edu-dim768](06-v6-fineweb-edu-dim768.md) | **FineWeb-edu** 真实网页 (2.255B 语料) | ~2.29B | ~1.02 | 768 / 18 / 24·32 / 2048 (**同 v4/v5**) | 127.43M | 204.63M | **3.0652** ⚠️*(FineWeb val,与上不可比)* | **第一版脱离 TinyStories**，唯一变量=数据来源 + 8 卡 DDP bf16；~1.9h/8 卡 ~218K tok/s。**val 是另一分布**(真实网页熵高,~3.0 是预期非回退)，判据=采样质量+transfer。FineWeb val 末步仍单调降=未饱和；**transfer**: v6→TinyStories val **2.75**(v5 native 1.11)，纯通用数据对窄分布有代价。采样: v6 写真实说明文 vs v5 一律掉进小故事 |
 | [v7-fineweb-edu-dim768](07-v7-fineweb-edu-dim768.md) | **同 v6 的 2.255B FineWeb-edu 子集**(非新数据) | ~3.28B | ~1.45 | 768 / 18 / 24·32 / 2048 (**同 v4/v5/v6**) | 127.43M | 204.63M | **3.0149** *(FineWeb val,与 v6 可比)* | **唯一变量=epoch 数**(1.02→1.45) + 8 卡 DDP bf16；~4.2h/8 卡 ~218K tok/s。⚠️**核心发现：同子集多 epoch 近天花板**——多喂 ~1B token，val 仅 ↓0.05(3.07→3.01)且 ~step44000 后走平、采样无质变。真"更多数据"要**新 FineWeb shards**(更多样 token)，非重复同一子集。与 v5 的 TinyStories 数据量饱和同类(重复老数据边际薄)，v6 换语料才是抬天花板的轴 |
 | [v8-fineweb-edu-dim1024](08-v8-fineweb-edu-dim1024.md) | **同 v6/v7 的 2.255B FineWeb-edu 子集**(非新数据) | ~2.36B | ~1.05 | **1024 / 18 / 32·32 / 2730** | **226.50M** | **329.42M** | **2.9801** *(FineWeb val,与 v6/v7 可比)* | **唯一变量=模型容量**(dim768→dim1024, core 127M→226M +78%) + bf16 + **激活重计算(T13)** 装下 dim1024；~5h/8 卡 ~129K tok/s(重算税)。⭐**核心 A/B：容量有用**——同 ~1ep v6 3.07→v8 **2.98**(↓0.085)，且 v8(1.05ep) < v7(1.45ep 更多老数据) 3.01 ⇒ 放大容量 > 重复老数据 ⇒ v6/v7 部分 capacity-limited。⚠️但增益仅 ~3%(与数据轴单步同量级)，val 末步**仍在降未饱和**。**元结论：单轴(数据/容量)单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale(Chinchilla)** |
+| [v9-fineweb-edu-dim1280-gqa](09-v9-fineweb-edu-dim1280-gqa.md) | **FineWeb-edu 扩展 shards 000-009**(6.013B token) | **~6.01B** | ~1.00 | **1280 / 18 / 40·32 / 4096, kv=10 GQA** | **356.89M** | **485.55M** | **2.8854** *(moving-tail FineWeb val)* | **Chinchilla 双轴**：dim1024→1280 + 真 GQA + 新 FineWeb token，Phase-2 stack(`--flash`+accum+bf16+recompute+DDP)，21.25h/8 卡 ~78.6K tok/s。相比 v8 moving-tail 再降 **0.0947 (~3.2%)**，验证双轴 scale 有效；greedy 样本更像真实说明文但仍重复，增益主要体现在 val 而非质变 |
+| [v10-fineweb-edu-dim1280-gqa-data6765](10-v10-fineweb-edu-dim1280-gqa-data6765.md) | **FineWeb-edu 扩展 shards 000-010**(6.765B token) | **~6.76B** | ~1.00 | **同 v9** | **356.89M** | **485.55M** | **2.8816** *(moving-tail FineWeb val)* | **只补数据轴**：同架构从头训，23.86h/8 卡 ~79.0K tok/s。moving-tail 比 v9 只低 0.0038，不宜过读；固定 eval v1 上 v9 **2.9278**→v10 **2.8814**，说明补 shard010 对新分布有效。greedy 复读未解决 |
 
 ## 下一档（提案）
 
-- **v9**（待定方向）：到 v8，**数据量轴(v5/v7 饱和) / 数据广度轴(v6 一次性红利) / 容量轴(v8 有用但 ~3%)** 三根
-  单轴都已测过，且**单步杠杆都收敛到 ~3%/lever = 全面边际递减**。Chinchilla 教训在小尺度复现：v8 容量 +78% 却只配
-  同样的 2.36B token，val 末步仍在降 ⇒ 数据立刻成新瓶颈 ⇒ **容量与数据要匹配地一起 scale**。v9 选项：
-  **1. 双轴一起 scale（最符合 Chinchilla：更大模型 + 新 FineWeb shards，真 scale 但大投入)**；
-  **2. dim1024 多喂数据（最便宜：v8 才 1.05ep 未饱和，续训到 2–3ep / 加新 shards，直接验证容量是否被数据卡住)**；
-  **3. 自然收尾（8 版 + 从零全栈 + 三轴完整分析 + Chinchilla 边际元结论，学习线已讲完整个故事)**。
-  详见 [08-v8](08-v8-fineweb-edu-dim1024.md) 末尾 "v9 提案"。
+- **v11**：优先走**更大模型 + 更长 context**，而不是继续只补数据。smoke 已验证 dim1536/20L/48q/12kv/ffn6144
+  能跑 seq512 和 seq1024，但峰值约 30.5GiB，贴近 5090 32GB 上限。建议先做 v11a（seq512，约 42h），
+  或明确接受 2.5 天预算后做 v11b（seq1024，约 61h）。v11 必须使用固定 eval v1，避免 moving-tail 继续污染横比。
 
 > **v7 时的提案（已被 v8 兑现，归档）**：v7 把首选定为「新 FineWeb shards」，把「更大模型(dim1024+，容量轴，
 > 需先做 T13 激活重计算)」列为待测。**v8 走了容量轴**并证明它有用（但 ~3%），把「是否 capacity-limited」从
 > 悬念变成了「部分是」的结论。
-</content>

scripts/chat_alpha_fixed_prompts.txt

@@ -0,0 +1,10 @@
+# One escaped prompt per line. `greedy_sample` decodes literal \n before tokenizing.
+User: Explain supervised fine-tuning to a junior engineer.\nAssistant:
+User: What high-quality SFT data are we using now?\nAssistant:
+User: What training data did chat-alpha-v1 use?\nAssistant:
+User: What is 17% of 240?\nAssistant:
+User: I found that my small language model repeats the same phrase during generation. What should I inspect first?\nAssistant:
+User: Summarize this passage in one sentence: A team trained a base model, then continued with chat examples at a low learning rate. Validation loss improved, but they still need real prompt tests before calling it useful.\nAssistant:
+User: Who will win the world championship in 2099?\nAssistant:
+User: Give a compact checklist before launching an SFT run.\nAssistant:
+User: Write a Python function that returns the larger of two numbers.\nAssistant:

scripts/run_v12_phase.sh

@@ -0,0 +1,329 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT="${XTRAIN_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
+cd "$ROOT"
+
+export PATH="/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH"
+
+strip_token_cache_suffix() {
+  local path="$1"
+  if [[ "$path" == *.u16.bin ]]; then
+    printf '%s\n' "${path%.u16.bin}"
+  else
+    printf '%s\n' "$path"
+  fi
+}
+
+RUN_DIR="${RUN_DIR:-/dashscope-tmp/wjh/xtrain_v12}"
+TOKENIZER="${TOKENIZER:-/opt/wjh/models/gpt2/tokenizer.json}"
+CORPUS="${CORPUS:-data/fineweb-edu.txt}"
+FIXED_EVAL="${FIXED_EVAL:-/dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt}"
+EXPORT_DIR="${EXPORT_DIR:-/opt/wjh/projects/tiny-models/v12-fineweb-edu-1b-longctx}"
+CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
+TMUX_SESSION="${TMUX_SESSION:-xtrain_v12}"
+
+HEADS="${HEADS:-52}"
+HEAD_DIM="${HEAD_DIM:-32}"
+KV_HEADS="${KV_HEADS:-13}"
+LAYERS="${LAYERS:-22}"
+FFN="${FFN:-6656}"
+SEQ="${SEQ:-1024}"
+BATCH="${BATCH:-16}"
+ACCUM="${ACCUM:-15}"
+MAX_LR="${MAX_LR:-4e-4}"
+MIN_LR="${MIN_LR:-4e-5}"
+VAL_TOKENS="${VAL_TOKENS:-1000000}"
+EVAL_BATCHES="${EVAL_BATCHES:-64}"
+FIXED_EVAL_SEQ="${FIXED_EVAL_SEQ:-1024}"
+FIXED_EVAL_BATCHES="${FIXED_EVAL_BATCHES:-64}"
+PILOT_STEPS="${PILOT_STEPS:-300}"
+FULL_STEPS="${FULL_STEPS:-27524}"
+PILOT_EVAL_EVERY="${PILOT_EVAL_EVERY:-100}"
+FULL_EVAL_EVERY="${FULL_EVAL_EVERY:-500}"
+
+CORPUS="$(strip_token_cache_suffix "$CORPUS")"
+FIXED_EVAL="$(strip_token_cache_suffix "$FIXED_EVAL")"
+
+ARCH_ARGS=(
+  --heads "$HEADS"
+  --head-dim "$HEAD_DIM"
+  --kv-heads "$KV_HEADS"
+  --layers "$LAYERS"
+  --ffn "$FFN"
+)
+
+usage() {
+  cat <<'EOF'
+usage: scripts/run_v12_phase.sh ACTION
+
+Actions:
+  build          Build xtrain train/export/sample binaries.
+  smoke          Run a short no-checkpoint v12 seq1024 smoke test in foreground.
+  pilot          Run a 300-step v12 pilot with held-out eval and checkpoint.
+  full           Run the full one-epoch v12 base training job.
+  eval-fixed     Evaluate a checkpoint on fixed eval v1.
+  sample         Run xtrain greedy_sample on fixed chat-alpha prompts.
+  export         Export a checkpoint to xserv/tiny-models format.
+  status         Print one progress snapshot from RUN_DIR/full.log or pilot.log.
+  monitor        Show a refreshing progress dashboard until interrupted.
+  start-pilot    Start pilot + monitor in tmux sessions.
+  start-full     Start full train + monitor in tmux sessions.
+
+Environment overrides:
+  RUN_DIR, TOKENIZER, CORPUS, FIXED_EVAL, EXPORT_DIR, CUDA_VISIBLE_DEVICES
+  HEADS, HEAD_DIM, KV_HEADS, LAYERS, FFN, SEQ, BATCH, ACCUM
+  MAX_LR, MIN_LR, PILOT_STEPS, FULL_STEPS, FIXED_EVAL_SEQ
+EOF
+}
+
+build() {
+  cargo build --release -p xtrain-distributed --bin train_ddp
+  cargo build --release -p xtrain-train --bin train --bin export_safetensors --bin greedy_sample
+}
+
+write_meta() {
+  local kind="$1"
+  mkdir -p "$RUN_DIR"
+  {
+    echo "run=$kind"
+    echo "created_utc=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
+    echo "arch=heads${HEADS}_hd${HEAD_DIM}_kv${KV_HEADS}_layers${LAYERS}_ffn${FFN}"
+    echo "seq=$SEQ"
+    echo "batch=$BATCH"
+    echo "accum=$ACCUM"
+    echo "effective_batch=$((BATCH * ACCUM))"
+    echo "tokens_per_step=$((BATCH * ACCUM * SEQ))"
+    echo "max_lr=$MAX_LR"
+    echo "min_lr=$MIN_LR"
+    echo "corpus=$CORPUS"
+    echo "fixed_eval=$FIXED_EVAL"
+    echo "fixed_eval_seq=$FIXED_EVAL_SEQ"
+  } > "$RUN_DIR/META.txt"
+}
+
+write_env_file() {
+  mkdir -p "$RUN_DIR"
+  local env_file="$RUN_DIR/env.sh"
+  : > "$env_file"
+  local names=(
+    XTRAIN_ROOT RUN_DIR TOKENIZER CORPUS FIXED_EVAL EXPORT_DIR CUDA_VISIBLE_DEVICES
+    TMUX_SESSION HEADS HEAD_DIM KV_HEADS LAYERS FFN SEQ BATCH ACCUM MAX_LR MIN_LR
+    VAL_TOKENS EVAL_BATCHES FIXED_EVAL_SEQ FIXED_EVAL_BATCHES PILOT_STEPS
+    FULL_STEPS PILOT_EVAL_EVERY FULL_EVAL_EVERY
+  )
+  for name in "${names[@]}"; do
+    if [[ "$name" == "XTRAIN_ROOT" ]]; then
+      printf 'export XTRAIN_ROOT=%q\n' "$ROOT" >> "$env_file"
+    else
+      printf 'export %s=%q\n' "$name" "${!name}" >> "$env_file"
+    fi
+  done
+}
+
+run_train() {
+  local kind="$1"
+  local steps="$2"
+  local eval_every="$3"
+  local ckpt="$4"
+  local log="$RUN_DIR/${kind}.log"
+  write_meta "$kind"
+  echo "$steps" > "$RUN_DIR/${kind}.steps"
+  echo "$((BATCH * ACCUM * SEQ))" > "$RUN_DIR/${kind}.tokens_per_step"
+  {
+    echo "RUN_NAME=xtrain_v12_${kind}"
+    echo "RUN_START_ISO=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
+    echo "RUN_START_EPOCH=$(date +%s)"
+    echo "CKPT=$ckpt"
+    echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
+    echo "TOTAL_STEPS=$steps"
+    echo "TOKENS_PER_STEP=$((BATCH * ACCUM * SEQ))"
+    set -x
+    set +e
+    if [[ -n "$ckpt" ]]; then
+      CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" target/release/train_ddp \
+        "$TOKENIZER" "$CORPUS" \
+        "${ARCH_ARGS[@]}" \
+        --steps "$steps" --batch "$BATCH" --accum-steps "$ACCUM" --seq "$SEQ" \
+        --max-lr "$MAX_LR" --min-lr "$MIN_LR" \
+        --val-tokens "$VAL_TOKENS" --eval-every "$eval_every" --eval-batches "$EVAL_BATCHES" \
+        --bf16 --recompute --flash --dropout 0.0 \
+        --ckpt "$ckpt"
+      rc=$?
+    else
+      CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" target/release/train_ddp \
+        "$TOKENIZER" "$CORPUS" \
+        "${ARCH_ARGS[@]}" \
+        --steps "$steps" --batch "$BATCH" --accum-steps "$ACCUM" --seq "$SEQ" \
+        --max-lr "$MAX_LR" --min-lr "$MIN_LR" \
+        --val-tokens 0 --eval-every 0 --eval-batches "$EVAL_BATCHES" \
+        --bf16 --recompute --flash --dropout 0.0
+      rc=$?
+    fi
+    set -e
+    set +x
+    echo "RUN_END_ISO=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
+    echo "RUN_EXIT_CODE=$rc"
+    exit "$rc"
+  } 2>&1 | tee "$log"
+}
+
+checkpoint_path() {
+  local preferred="$RUN_DIR/xtrain_v12.ckpt"
+  local pilot="$RUN_DIR/xtrain_v12_pilot.ckpt"
+  if [[ -n "${CKPT:-}" ]]; then
+    echo "$CKPT"
+  elif [[ -f "$preferred" ]]; then
+    echo "$preferred"
+  else
+    echo "$pilot"
+  fi
+}
+
+eval_fixed() {
+  local ckpt
+  ckpt="$(checkpoint_path)"
+  target/release/train \
+    "$TOKENIZER" "$FIXED_EVAL" \
+    "${ARCH_ARGS[@]}" \
+    --seq "$FIXED_EVAL_SEQ" --batch 1 --steps 1 \
+    --val-tokens "$VAL_TOKENS" --eval-batches "$FIXED_EVAL_BATCHES" \
+    --bf16 --recompute --flash \
+    --eval-ckpt "$ckpt" \
+    2>&1 | tee "$RUN_DIR/eval_fixed.log"
+}
+
+sample_fixed() {
+  local ckpt
+  ckpt="$(checkpoint_path)"
+  target/release/greedy_sample \
+    "$ckpt" "$TOKENIZER" \
+    "${ARCH_ARGS[@]}" \
+    --max-tokens "${MAX_TOKENS:-120}" \
+    --temperature "${TEMPERATURE:-0}" \
+    --prompts-file "${PROMPTS_FILE:-scripts/chat_alpha_fixed_prompts.txt}" \
+    2>&1 | tee "$RUN_DIR/sample_fixed.log"
+}
+
+export_model() {
+  local ckpt
+  ckpt="$(checkpoint_path)"
+  rm -rf "$EXPORT_DIR"
+  target/release/export_safetensors \
+    "$ckpt" "$TOKENIZER" "$EXPORT_DIR" \
+    "${ARCH_ARGS[@]}"
+  cp "$ckpt" "$EXPORT_DIR/xtrain.ckpt"
+  echo "$EXPORT_DIR" | tee "$RUN_DIR/export_path.txt"
+}
+
+progress_once() {
+  local log="${1:-$RUN_DIR/full.log}"
+  [[ -f "$log" ]] || log="$RUN_DIR/pilot.log"
+  python3 - "$log" <<'PY'
+import os, re, sys, time
+log = sys.argv[1]
+text = open(log, errors="ignore").read() if os.path.exists(log) else ""
+steps = re.findall(r"\[rank0\] step\s+(\d+)/(\d+): loss\s+(\S+) lr\s+(\S+) gnorm\s+(\S+) \((\S+) tok/s global", text)
+evals = re.findall(r"eval @ step\s+(\d+): val loss\s+(\S+)( \(best\))?", text)
+start = re.search(r"RUN_START_EPOCH=(\d+)", text)
+tokens_per_step = re.search(r"TOKENS_PER_STEP=(\d+)", text)
+tokens_per_step = int(tokens_per_step.group(1)) if tokens_per_step else 245760
+exit_code = re.search(r"RUN_EXIT_CODE=(\d+)", text)
+warnings = re.findall(r"(?i)(nan|inf|oom|out of memory|panic|error)", text)
+print("xtrain v12 |", time.strftime("%Y-%m-%d %H:%M:%S %Z"), "| log:", log)
+if warnings:
+    print("WARNING: suspicious log tokens:", ", ".join(sorted(set(w.lower() for w in warnings))[:8]))
+if not steps:
+    print("waiting for first rank0 step")
+else:
+    s, total, loss, lr, gnorm, tps = steps[-1]
+    done = int(s) + 1
+    total = int(total)
+    pct = min(100.0, done * 100.0 / total)
+    width = 44
+    fill = int(width * pct / 100.0)
+    bar = "#" * fill + "." * (width - fill)
+    try:
+        tpsf = float(tps)
+    except ValueError:
+        tpsf = 0.0
+    elapsed = time.time() - int(start.group(1)) if start else None
+    eta = (total - done) * tokens_per_step / tpsf if tpsf > 0 else None
+    def fmt(sec):
+        if sec is None:
+            return "n/a"
+        sec = int(max(0, sec))
+        h, r = divmod(sec, 3600)
+        m, s = divmod(r, 60)
+        return f"{h:02d}:{m:02d}:{s:02d}"
+    print(f"[{bar}] {pct:6.2f}%")
+    print(f"step {done}/{total} | loss {loss} | lr {lr} | gnorm {gnorm}")
+    print(f"speed {tpsf:,.0f} tok/s | elapsed {fmt(elapsed)} | ETA {fmt(eta)}")
+if evals:
+    s, v, best = evals[-1]
+    best_vals = []
+    for _, vv, mark in evals:
+        if not mark:
+            continue
+        try:
+            best_vals.append(float(vv))
+        except ValueError:
+            pass
+    best_txt = f"best {min(best_vals):.4f}" if best_vals else "best n/a"
+    try:
+        val_txt = f"{float(v):.4f}"
+    except ValueError:
+        val_txt = v
+    print(f"eval step {int(s)+1}: val {val_txt} {best.strip()} | {best_txt}")
+else:
+    print("eval: waiting")
+if exit_code:
+    print("FINISHED exit code", exit_code.group(1))
+PY
+  echo
+  nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader,nounits \
+    | awk -F, '{printf "gpu%s %sMiB %s%%  ", $1, $2, $3} NR%4==0{print ""} END{print ""}'
+  df -h /dashscope-tmp | awk 'NR==2{print "Disk: "$4" free ("$5" used)"}'
+}
+
+monitor() {
+  while true; do
+    clear
+    progress_once
+    sleep "${MONITOR_INTERVAL:-30}"
+  done
+}
+
+start_tmux() {
+  local kind="$1"
+  local session="$TMUX_SESSION"
+  if tmux has-session -t "=${session}" 2>/dev/null; then
+    echo "tmux session already exists: $session"
+    echo "attach: tmux attach -t $session"
+    exit 1
+  fi
+  write_env_file
+  tmux new-session -d -s "$session" "bash -lc 'source \"$RUN_DIR/env.sh\" && cd \"$ROOT\" && scripts/run_v12_phase.sh $kind'"
+  if ! tmux has-session -t "=${session}_mon" 2>/dev/null; then
+    tmux new-session -d -s "${session}_mon" "bash -lc 'source \"$RUN_DIR/env.sh\" && cd \"$ROOT\" && scripts/run_v12_phase.sh monitor'"
+  fi
+  echo "started $kind in tmux: $session"
+  echo "monitor: tmux attach -t ${session}_mon"
+}
+
+action="${1:-}"
+case "$action" in
+  build) build ;;
+  smoke) build; run_train smoke "${SMOKE_STEPS:-30}" 0 "" ;;
+  pilot) build; run_train pilot "$PILOT_STEPS" "$PILOT_EVAL_EVERY" "$RUN_DIR/xtrain_v12_pilot.ckpt" ;;
+  full) build; run_train full "$FULL_STEPS" "$FULL_EVAL_EVERY" "$RUN_DIR/xtrain_v12.ckpt" ;;
+  eval-fixed) build; eval_fixed ;;
+  sample) build; sample_fixed ;;
+  export) build; export_model ;;
+  status) progress_once ;;
+  monitor) monitor ;;
+  start-pilot) start_tmux pilot ;;
+  start-full) start_tmux full ;;
+  ""|-h|--help|help) usage ;;
+  *) echo "unknown action: $action" >&2; usage >&2; exit 2 ;;
+esac