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

.gitignore

@@ -13,5 +13,7 @@
 # Large scaling-run corpora + tokenized id caches live on dash5 only, never in
 # git (the small data/tinystories-valid-3mb.txt is committed as a fixture).
 /data/tinystories-train.txt
+/data/fineweb-edu.txt
+/data/*.parquet
 *.u16.bin
 *.ckpt

README.md

@@ -1,50 +1,210 @@
 # xtrain
 
-A from-scratch **Rust + CUDA** LLM **training** engine — the sibling of
+A from-scratch **Rust + CUDA** LLM **training** engine — the sibling of **xserv** (the
-[xserv](https://github.com/) (the inference side). GPU-first.
+inference side). A learning project: hand-write the entire training-systems stack
+(autograd → backward → optimizer → training loop → distributed → mixed precision →
+gradient checkpointing), then use it to run a multi-version **scaling study** that maps
+the data-vs-capacity frontier for a tiny model.
 
-The goal is to learn the full training-systems stack by hand: autograd / backward
+> **Status: complete — three phases.**
-passes / optimizers (AdamW) / the training loop / distributed logic. Heavy lifting
+> **Phase 1** = the from-scratch full stack (T1–T13) + an 8-version scaling study (v0–v8):
-is borrowed where it makes sense (GEMM → cuBLAS after a hand-written version,
+> hand-write the whole training-systems stack, then map the data-vs-capacity frontier.
-multi-GPU comms → NCCL, tokenizer → reused from xserv), but the core is written
+> **Phase 2** = systems-stack depth (T14–T18): hand-write the five deferred training-stack
-from scratch. The target architecture is a tiny modern transformer
+> features — fused flash-attention, real GQA, gradient accumulation, process-per-GPU DDP,
-(RoPE + RMSNorm + SwiGLU, ~1–30M params) whose forward aligns with xserv's Qwen3,
+> dropout. **Phase 3** = one Chinchilla-style double-axis run (v9): dim1280 true-GQA +
-so the backward passes map one-to-one onto xserv's existing forward kernels and
+> 6.01B FineWeb tokens, validating the v8 conclusion that data and capacity must scale
-trained weights can flow back into xserv.
+> 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/).
 
-## Status
+---
 
-Bootstrapping (P0). This repo currently contains only the project skeleton and a
+## What got built (from scratch, by hand)
-working Rust↔CUDA build chain, verified by a trivial vector-add CUDA kernel.
 
-## Layout
+7 crates, no ML framework — only cuBLAS / NCCL / safetensors as deliberate "heavy-lifting"
+borrows, the rest hand-written CUDA + Rust:
 
-```
+| crate | what's hand-written |
-xtrain/
+|---|---|
-├── Cargo.toml              # workspace
+| `xtrain-cuda` | CUDA Runtime FFI, RAII `GpuBuffer`, **caching/pool allocator**, cuBLAS (sgemm + bf16 GemmEx) bindings |
-├── csrc/                   # CUDA sources (.cu)
+| `xtrain-tensor` | tensor (dtype/shape/strides/storage), elementwise + transpose + embedding kernels |
-│   └── test/vecadd.cu      # trivial element-wise vector-add (smoke test)
+| `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) |
-└── crates/
+| `xtrain-model` | tiny **Qwen3-style** transformer (RoPE + RMSNorm + QK-norm + SwiGLU), batched forward, **GQA** (`num_kv_heads<num_heads`), residual/MLP **dropout** |
-    └── xtrain-cuda/        # CUDA Runtime FFI + build.rs (nvcc → sm_120)
+| `xtrain-optim` | hand-written **AdamW** (host + GPU kernels) |
-        ├── build.rs        # compiles csrc/*.cu via the `cc` crate, links cudart
+| `xtrain-train` | training loop, LR schedule, grad clip, **gradient accumulation**, checkpoint, BPE corpus + cache, samplers, safetensors export |
-        ├── src/            # ffi / error / device / memory
+| `xtrain-distributed` | **NCCL DDP** (thread-per-GPU + torchrun-style process-per-GPU launcher / cross-process `ncclUniqueId`, all-reduce) |
-        └── tests/          # vecadd smoke test
+
+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); deterministic fixtures produce token-identical greedy output,
+and large checkpoints are validated end-to-end in the serving path.
+
+## 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). **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 |
+|---|---|---|
+| T1–T2 | Rust↔CUDA build chain · tensor abstraction | vector-add verified · roundtrip |
+| T3–T4 | hand GEMM fwd/bwd + finite-diff · **tape autograd** + 11 op backwards | grads vs cuBLAS 1e-7 / finite-diff |
+| T5 | tiny transformer (RoPE+RMSNorm+SwiGLU) | overfit + **PyTorch parity** |
+| T6 | AdamW + training loop + checkpoint · GPT-2 BPE + TinyStories | first **coherent English** |
+| T7 | cuBLAS + GPU optimizer + drop syncs | ~3× (2.7K→8.5K tok/s) |
+| T8 | NCCL DDP | multi-GPU (weak scaling, then) |
+| T9 | + per-head **QK-norm** (Qwen3-compat) + safetensors export | **xserv closed loop, token-identical** |
+| **T10** | **batched multi-sequence forward** (fixes KI-1) | **single-GPU 15–24×**; MFU 0.4%→14% |
+| **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) — which is where **Phase 1** closed.
+
+## 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/)).
+
+| ver | data (trained tok / epoch) | dim / core params | val loss | axis explored |
+|---|---|---|---|---|
+| v0–v3 | TinyStories (↑) | 32→512 / 41K→67M | 3.80 → 1.30 | bring-up |
+| v4 | TinyStories 1.54ep | 768 / 127M | 1.17 | — |
+| v5 | TinyStories 5.33ep | 768 / 127M | **1.11** | **data volume → saturates** |
+| 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.
+
+### 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.
+2. **Corpus > more-of-the-same.** Swapping TinyStories → FineWeb-edu (v5→v6) was a
+   *qualitative* jump: the model went from only-writes-kid-stories to writing genuine
+   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 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
+
+The throughput story is the perf-infra report card (RTX 5090, bf16/fp32):
+
+| | v1 | v2 | v3 | v4 | v5 |
+|---|---|---|---|---|---|
+| tok/s | 3.3K (1 GPU) | 3.6K (4 GPU) | 26K (1 GPU) | 145K (8 GPU) | 217K (8 GPU) |
+| **MFU** | 0.4% | 0.2% | 14% | 17% | 13% |
+| enabled by | — | DDP (weak) | **batched (T10)** | **alloc (T11)** | **bf16 (T12)** |
+
+v1/v2 ran at **<0.5% MFU** — the single-sequence design left the GPU idle (launch-bound).
+**Batched forward (T10) was the single biggest unlock** (~35× MFU jump). 6ND is an accurate
+FLOPs count, but predicting *time* needs the *realized* MFU, which varied ~40× across
+versions — a fixed-MFU estimate is off by up to ~100× for the early launch-bound runs.
+
+## Engineering lessons
+
+- **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…v10`). Serve any trained version in xserv:
+
+```bash
+# on the GPU box
+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,
 ```
 
-The build mirrors xserv's approach: `build.rs` invokes `nvcc` (via the `cc` crate)
+Build/test the engine itself (CUDA compiles + runs on the GPU box; host-side `cargo check`
-to compile `csrc/*.cu` targeting `sm_120` (RTX 5090) and links them into the Rust
+works anywhere via the `no_cuda` cfg):
-crate over hand-written `extern "C"` FFI.
-
-## Building & testing
-
-CUDA compilation and execution happen on a GPU box (dash5, 8× RTX 5090, sm_120):
 
 ```sh
 export PATH=/usr/local/cuda/bin:$HOME/.cargo/bin:$PATH
-cargo build
+cargo test --workspace            # autograd grad-checks, PyTorch parity, DDP, etc.
-cargo test -p xtrain-cuda -- --nocapture   # runs the vecadd smoke test
 ```
 
-On a machine without `nvcc`/GPU, `build.rs` detects the missing toolchain, skips
+## Doc index
-CUDA compilation, and sets a `no_cuda` cfg — so host-side `cargo check` still
+
-works (the GPU smoke test is compiled out).
+- [`docs/evolution.md`](docs/evolution.md) — per-milestone changes across algorithm / architecture / infra / dataset.
+- [`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/checkpoint.rs

@@ -0,0 +1,103 @@
+//! Activation recomputation / gradient checkpointing (Phase T13, KI-3).
+//!
+//! A higher-order autograd primitive — the analogue of `torch.utils.checkpoint`.
+//! It runs a *segment* of the model (a transformer block, here) WITHOUT recording
+//! the segment's internal ops on the surrounding tape, so the segment's
+//! intermediate activations are freed right after the forward instead of being
+//! kept alive until backward. When the segment's output-grad arrives in backward,
+//! the segment forward is **re-run** from the saved input (into a throwaway local
+//! tape), the recomputed output is seeded with the upstream grad, and the gradient
+//! is backpropagated through the local tape to recover the input-grad and the
+//! parameter-grads — which are then pushed to the real tape's parents. The local
+//! tape is dropped at the end of the closure, freeing the recomputed activations.
+//!
+//! ## Why it is exact (the hard gate)
+//! The recompute runs the *same* `segment_fn` from the *same* input value and the
+//! *same* parameter values (parameters are leaves that persist across forward and
+//! backward; only their grad slot changes). The forward kernels are deterministic,
+//! so the recomputed output equals the original output bit-for-bit, and the local
+//! backward is the ordinary analytic backward of that segment. Therefore the input-
+//! and parameter-grads are identical to those a non-checkpointed forward would
+//! produce — checkpointing trades compute (one extra forward per segment) for
+//! memory, never correctness.
+//!
+//! ## Composition
+//! - **bf16 (T12):** `segment_fn` is the unchanged block forward, so the recompute
+//!   runs the same bf16 path; the `cast` op's grad upcast still bridges bf16→fp32.
+//! - **DDP (T8):** each rank checkpoints its own forward/backward independently;
+//!   the param-grads recovered here feed the same per-rank `.grad()` slots that the
+//!   all-reduce averages — no change to the distributed path.
+//! - **batched (T10):** the segment input/output carry the `[batch*seq, …]` batch
+//!   dim transparently; `checkpoint` is shape-agnostic.
+
+#![cfg(not(no_cuda))]
+
+use crate::tape::Var;
+use std::rc::Rc;
+
+/// Run `segment_fn(input, params)` with activation recomputation.
+///
+/// `segment_fn(x, p)` must build the segment's forward graph from a single input
+/// `x` and the parameter slice `p`, returning the single segment output. It is
+/// called once now (forward, result detached from the outer tape) and once per
+/// backward (recompute). It MUST be deterministic and depend only on `x` and `p`
+/// (this is what makes the recompute exact).
+///
+/// `params` are the segment's learnable leaves; their grads are accumulated into
+/// the SAME leaves the optimizer reads (so DDP / AdamW are unchanged).
+///
+/// Returns the segment output as a `Var` on the outer tape whose backward triggers
+/// the recompute. Equivalent — grad-for-grad — to calling `segment_fn(input,
+/// params)` directly, but without keeping the segment's internal activations alive.
+pub fn checkpoint<F>(segment_fn: F, input: &Var, params: &[Var]) -> Var
+where
+    F: Fn(&Var, &[Var]) -> Var + 'static,
+{
+    let segment_fn = Rc::new(segment_fn);
+
+    // --- Forward (no taping of internals) ---
+    // Detach the input and params into fresh leaves so `segment_fn` builds a LOCAL
+    // tape disconnected from the outer graph. We only keep the output's value; the
+    // local `Var`s (and thus the segment's intermediate activations) are dropped
+    // when this scope ends.
+    let out_value = {
+        let x_det = Var::leaf(input.value());
+        let params_det: Vec<Var> = params.iter().map(|p| Var::leaf(p.value())).collect();
+        let out_local = segment_fn(&x_det, &params_det);
+        out_local.value()
+    };
+
+    // Parents on the OUTER tape: the segment input, then the params (so their grads
+    // land in the leaves the optimizer reads).
+    let mut parents = Vec::with_capacity(1 + params.len());
+    parents.push(input.clone());
+    parents.extend(params.iter().cloned());
+
+    let segment_fn = segment_fn.clone();
+    Var::from_op(
+        out_value,
+        parents,
+        Box::new(move |dout, parents| {
+            // --- Backward (recompute) ---
+            // Rebuild fresh leaves from the CURRENT input/param values (params are
+            // unchanged since forward; input is the saved segment input), re-run the
+            // forward to rebuild the local tape, seed the recomputed output with the
+            // upstream grad, and backprop through the local tape.
+            let x_det = Var::leaf(parents[0].value());
+            let params_det: Vec<Var> = parents[1..].iter().map(|p| Var::leaf(p.value())).collect();
+            let out_local = segment_fn(&x_det, &params_det);
+            out_local.backward_seeded(dout.clone());
+
+            // Push the recovered grads to the real parents (engine SUMs on fan-out).
+            if let Some(dx) = x_det.grad() {
+                Var::push_grad(&parents[0], dx);
+            }
+            for (det, parent) in params_det.iter().zip(&parents[1..]) {
+                if let Some(dp) = det.grad() {
+                    Var::push_grad(parent, dp);
+                }
+            }
+            // `out_local` / the local tape drop here → recomputed activations freed.
+        }),
+    )
+}

crates/xtrain-autodiff/src/lib.rs

@@ -18,6 +18,8 @@ pub use finite_diff::{GradCheckConfig, GradCheckResult, ParamFn, grad_check};
 // kernels via xtrain-tensor, so they are gated behind `not(no_cuda)` (the
 // per-crate convention); the grad_check harness above stays host-only.
 #[cfg(not(no_cuda))]
+pub mod checkpoint;
+#[cfg(not(no_cuda))]
 pub mod ops;
 #[cfg(not(no_cuda))]
 pub mod tape;

crates/xtrain-autodiff/src/ops.rs

@@ -13,7 +13,27 @@
 #![cfg(not(no_cuda))]
 
 use crate::tape::Var;
-use xtrain_tensor::Tensor;
+use xtrain_tensor::{DType, Tensor};
+
+/// dtype cast as an autograd node (Phase T12 — the AMP bridge between fp32 master
+/// weights / fp32 reductions and the bf16 compute stream). Forward casts `x` to
+/// `target`; **backward casts the upstream grad back to `x`'s dtype**. So a fp32
+/// master-weight leaf fed through `cast(w, BF16)` into a bf16 matmul accumulates
+/// an **fp32** grad — AdamW / clip / DDP all-reduce stay fp32, untouched.
+pub fn cast(x: &Var, target: DType) -> Var {
+    let src = x.value().dtype();
+    if src == target {
+        return x.clone();
+    }
+    let out = x.value().to_dtype(target);
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            Var::push_grad(&parents[0], d.to_dtype(src));
+        }),
+    )
+}
 
 /// `C = A @ B` (2D). Backward: `dA = dC @ Bᵀ`, `dB = Aᵀ @ dC`.
 pub fn matmul(a: &Var, b: &Var) -> Var {
@@ -120,6 +140,31 @@ pub fn swiglu(gate: &Var, up: &Var) -> Var {
     mul(&silu(gate), up)
 }
 
+/// Dropout (Phase T18). With probability `p` zero each element, scale the kept
+/// ones by `1/(1-p)` (inverted dropout — `E[out] == x`). The keep/drop mask is
+/// drawn by a counter-based RNG from `(seed, element index)`, so it is fully
+/// determined by `seed` (same `seed` ⇒ same mask: stable across the T13 recompute
+/// re-run, and held fixed across the ± perturbation of a finite-diff grad-check).
+/// Forward caches the per-element scale `mask`; **backward applies the same mask**
+/// (`dx = d ⊙ mask`), making dropout a fixed elementwise linear map of `x`.
+///
+/// `p == 0` is a no-op: returns `x.clone()` (no node added) so the default graph
+/// is bit-identical to the no-dropout path. eval-time identity is handled by the
+/// caller simply not invoking dropout (the model's train/eval switch).
+pub fn dropout(x: &Var, p: f32, seed: u64) -> Var {
+    if p == 0.0 {
+        return x.clone();
+    }
+    let (out, mask) = x.value().dropout(p, seed);
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            Var::push_grad(&parents[0], Tensor::dropout_backward(d, &mask));
+        }),
+    )
+}
+
 /// RoPE (rotate_half) over `x:[tokens,heads,head_dim]` with per-sequence position
 /// `row % period` (`period` = sequence length; `period == tokens` for a single
 /// sequence). Orthogonal map, so the backward is the inverse rotation of `dy` — no
@@ -305,15 +350,80 @@ 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
+    // probs). The grad must match the logits' dtype so it chains into a bf16
+    // lm_head matmul backward — cast dx back. Keeping logits bf16 (no persistent
+    // 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();
@@ -323,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);
+            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/src/tape.rs

@@ -108,14 +108,24 @@ impl Var {
             "backward() expects a scalar loss; got shape {:?}",
             self.value().shape()
         );
+        self.backward_seeded(ones_like(&self.value()));
+    }
 
+    /// Reverse-mode backward from this node seeded with an explicit upstream grad
+    /// `seed` (same shape as this node's value), instead of the scalar `dL/dL = 1`.
+    ///
+    /// This is the entry point for **activation recomputation** (Phase T13): a
+    /// checkpointed segment re-runs its forward into a fresh local tape, then
+    /// backprops the upstream output-grad through it via this method (the segment
+    /// output is generally NOT a scalar). For a scalar root, [`backward`] is the
+    /// thin wrapper that seeds ones.
+    pub fn backward_seeded(&self, seed: Tensor) {
         // 1. Topological order (post-order DFS), parents before children.
         let mut topo: Vec<Var> = Vec::new();
         let mut visited: Vec<*const RefCell<VarNode>> = Vec::new();
         build_topo(self, &mut topo, &mut visited);
 
-        // 2. Seed the loss gradient with ones.
+        // 2. Seed this node's gradient with the supplied upstream grad.
-        let seed = ones_like(&self.value());
         self.accumulate(seed);
 
         // 3. Walk in reverse: each node hands its grad to its parents' closures.

crates/xtrain-autodiff/tests/autograd.rs

@@ -625,6 +625,343 @@ 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
+// (it depends only on (seed, index), NOT on x), so dropout is a fixed elementwise
+// linear map `out_i = c_i·x_i` and the central difference of L is differentiable:
+// the ± perturbation of each x_i sees the SAME mask. The forward function in the
+// closure calls `ops::dropout(x, p, SEED)` with the same SEED, so it reproduces
+// the same mask both times.
+#[test]
+fn dropout_bwd() {
+    require_gpu();
+    const SEED: u64 = 0xD120_FE5E;
+    let p = 0.3f32;
+    let (m, n) = (16, 12);
+    let x_h = fill(m * n, 71);
+    let w = fill(m * n, 72);
+
+    let x = Var::leaf(cuda(&x_h, &[m, n]));
+    let out = ops::dropout(&x, p, SEED);
+    scalar_loss(&out, &w).backward();
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+
+    let wf = w.clone();
+    let lx = move |v: &[f32], s: &[usize]| {
+        let o = ops::dropout(&Var::leaf(cuda(v, s)), p, SEED);
+        weighted_sum(&o.value(), &wf)
+    };
+    report(
+        "dropout dX",
+        &grad_check(&x_h, &[m, n], &lx, dx.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// Inverted-dropout expectation + keep-rate check. Over a large tensor and a sweep
+// of seeds, the mean of dropout(x) tracks the mean of x (E[out] ≈ x, the inverted
+// 1/(1-p) scaling), and the kept fraction tracks 1-p (the RNG is ~Bernoulli).
+#[test]
+fn dropout_expectation_and_keep_rate() {
+    require_gpu();
+    let p = 0.25f32;
+    let n = 200_000usize;
+    let x_h = vec![1.0f32; n]; // mean(x) = 1 → mean(out) should ≈ 1
+    let x = cuda(&x_h, &[n]);
+
+    let trials = 8;
+    let mut mean_out_acc = 0.0f64;
+    let mut keep_acc = 0.0f64;
+    for t in 0..trials {
+        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();
+        mean_out_acc += mean_out;
+        keep_acc += kept as f64 / n as f64;
+    }
+    let mean_out = mean_out_acc / trials as f64;
+    let keep_rate = keep_acc / trials as f64;
+    println!(
+        "dropout p={p}: E[out]={mean_out:.5} (input mean 1.0), keep_rate={keep_rate:.5} (1-p={:.3})",
+        1.0 - p
+    );
+    assert!(
+        (mean_out - 1.0).abs() < 0.01,
+        "E[out] {mean_out} not ≈ input mean 1.0 (inverted scaling broken)"
+    );
+    assert!(
+        (keep_rate - (1.0 - p) as f64).abs() < 0.01,
+        "keep_rate {keep_rate} not ≈ 1-p {}",
+        1.0 - p
+    );
+}
+
+// p=0 is a no-op (the op returns x.clone(), no node) → output is bit-identical to
+// x and its grad flows straight through (the default-graph regression guard at the
+// op level; the model-level bit-identity is in xtrain-model/tests/dropout.rs).
+#[test]
+fn dropout_p0_is_identity() {
+    require_gpu();
+    let (m, n) = (8, 5);
+    let x_h = fill(m * n, 91);
+    let x = cuda(&x_h, &[m, n]);
+    let (out, _mask) = x.dropout(0.0, 12345);
+    let out_h = out.to_device(Device::Cpu);
+    for (a, b) in x_h.iter().zip(out_h.as_slice::<f32>()) {
+        assert_eq!(*a, *b, "p=0 dropout must be identity");
+    }
+}
+
 // --- test helpers ---
 
 // Scalar loss node L = sum(W ∘ out): wraps a fixed-weight Var and reduces. We
@@ -668,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,10 @@ 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/cublas.rs

@@ -19,6 +19,7 @@
 
 use crate::ffi::{self, CublasHandle};
 use std::cell::RefCell;
+use std::ffi::c_void;
 
 thread_local! {
     static HANDLE: RefCell<Option<CublasHandle>> = const { RefCell::new(None) };
@@ -159,3 +160,131 @@ pub fn sgemm_strided_batched(
         assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
     });
 }
+
+/// bf16 row-major GEMM `C[m,n] = opA(A)·opB(B)` via `cublasGemmEx`: bf16 in/out,
+/// **fp32 accumulation** (`CUBLAS_COMPUTE_32F`) — the standard AMP matmul (Phase
+/// T12). `a`/`b`/`c` are device pointers to row-major **bf16** matrices; the
+/// row-major⟺col-major transpose algebra is identical to [`sgemm`] (we compute
+/// the col-major `Cᵀ`). `alpha`/`beta` are fp32 host scalars (compute is fp32).
+#[allow(clippy::too_many_arguments)]
+pub fn gemm_ex(
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: f32,
+    a: *const c_void,
+    b: *const c_void,
+    beta: f32,
+    c: *mut c_void,
+) {
+    let lda = if trans_a { m } else { k };
+    let ldb = if trans_b { k } else { n };
+    let ldc = n;
+    let op_a = if trans_a {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let op_b = if trans_b {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let bf16 = ffi::CUDA_R_16BF;
+
+    with_handle(|handle| {
+        let status = unsafe {
+            ffi::cublasGemmEx(
+                handle,
+                op_b,
+                op_a,
+                n as i32,
+                m as i32,
+                k as i32,
+                &alpha as *const f32 as *const c_void,
+                b,
+                bf16,
+                ldb as i32,
+                a,
+                bf16,
+                lda as i32,
+                &beta as *const f32 as *const c_void,
+                c,
+                bf16,
+                ldc as i32,
+                ffi::CUBLAS_COMPUTE_32F,
+                ffi::CUBLAS_GEMM_DEFAULT,
+            )
+        };
+        assert_eq!(status, 0, "cublasGemmEx failed: {status}");
+    });
+}
+
+/// Strided-batched bf16 GEMM (Phase T12) — the [`gemm_ex`] analogue of
+/// [`sgemm_strided_batched`] for the batched attention GEMMs. bf16 in/out, fp32
+/// accumulation; strides are in ELEMENTS.
+#[allow(clippy::too_many_arguments)]
+pub fn gemm_ex_strided_batched(
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: f32,
+    a: *const c_void,
+    stride_a: usize,
+    b: *const c_void,
+    stride_b: usize,
+    beta: f32,
+    c: *mut c_void,
+    stride_c: usize,
+    batch: usize,
+) {
+    let lda = if trans_a { m } else { k };
+    let ldb = if trans_b { k } else { n };
+    let ldc = n;
+    let op_a = if trans_a {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let op_b = if trans_b {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let bf16 = ffi::CUDA_R_16BF;
+
+    with_handle(|handle| {
+        let status = unsafe {
+            ffi::cublasGemmStridedBatchedEx(
+                handle,
+                op_b,
+                op_a,
+                n as i32,
+                m as i32,
+                k as i32,
+                &alpha as *const f32 as *const c_void,
+                b,
+                bf16,
+                ldb as i32,
+                stride_b as i64,
+                a,
+                bf16,
+                lda as i32,
+                stride_a as i64,
+                &beta as *const f32 as *const c_void,
+                c,
+                bf16,
+                ldc as i32,
+                stride_c as i64,
+                batch as i32,
+                ffi::CUBLAS_COMPUTE_32F,
+                ffi::CUBLAS_GEMM_DEFAULT,
+            )
+        };
+        assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
+    });
+}

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))]
@@ -324,3 +453,171 @@ unsafe extern "C" {
 pub const CUBLAS_OP_N: i32 = 0;
 #[cfg(not(no_cuda))]
 pub const CUBLAS_OP_T: i32 = 1;
+
+// --- bf16 mixed precision (Phase T12) ---
+//
+// cudaDataType / cublasComputeType enum values (same as xserv's gemm.rs). The
+// bf16 GEMM uses bf16 in/out with fp32 accumulation (CUBLAS_COMPUTE_32F).
+#[cfg(not(no_cuda))]
+pub const CUDA_R_32F: i32 = 0;
+#[cfg(not(no_cuda))]
+pub const CUDA_R_16BF: i32 = 14;
+#[cfg(not(no_cuda))]
+pub const CUBLAS_COMPUTE_32F: i32 = 68;
+/// CUBLAS_GEMM_DEFAULT — let cuBLAS pick the algorithm.
+#[cfg(not(no_cuda))]
+pub const CUBLAS_GEMM_DEFAULT: i32 = -1;
+
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    // General GEMM with explicit in/out + compute types (bf16 path). `alpha`/
+    // `beta` are fp32 host scalars (compute type is fp32). Pointers are void* so
+    // the same FFI serves bf16 / fp32.
+    #[allow(clippy::too_many_arguments)]
+    pub fn cublasGemmEx(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const std::ffi::c_void,
+        a: *const std::ffi::c_void,
+        a_type: i32,
+        lda: i32,
+        b: *const std::ffi::c_void,
+        b_type: i32,
+        ldb: i32,
+        beta: *const std::ffi::c_void,
+        c: *mut std::ffi::c_void,
+        c_type: i32,
+        ldc: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+    #[allow(clippy::too_many_arguments)]
+    pub fn cublasGemmStridedBatchedEx(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const std::ffi::c_void,
+        a: *const std::ffi::c_void,
+        a_type: i32,
+        lda: i32,
+        stride_a: i64,
+        b: *const std::ffi::c_void,
+        b_type: i32,
+        ldb: i32,
+        stride_b: i64,
+        beta: *const std::ffi::c_void,
+        c: *mut std::ffi::c_void,
+        c_type: i32,
+        ldc: i32,
+        stride_c: i64,
+        batch_count: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+}
+
+// bf16 cast + elementwise kernels (csrc/ops/cast.cu). Pointers are void* (bf16
+// buffers); f32 sides are typed. The activation stream flows bf16; the math
+// accumulates in fp32 inside each kernel.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    pub fn launch_cast_f32_to_bf16(input: *const f32, out: *mut c_void, n: i32, s: CudaStream);
+    pub fn launch_cast_bf16_to_f32(input: *const c_void, out: *mut f32, n: i32, s: CudaStream);
+
+    pub fn launch_add_bf16(
+        a: *const c_void,
+        b: *const c_void,
+        out: *mut c_void,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_mul_bf16(
+        a: *const c_void,
+        b: *const c_void,
+        out: *mut c_void,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_scale_bf16(
+        input: *const c_void,
+        out: *mut c_void,
+        alpha: f32,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_silu_bf16(x: *const c_void, y: *mut c_void, n: i32, s: CudaStream);
+    pub fn launch_silu_dx_bf16(
+        x: *const c_void,
+        dy: *const c_void,
+        dx: *mut c_void,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_add_bias_bf16(
+        x: *const c_void,
+        bias: *const c_void,
+        out: *mut c_void,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+    pub fn launch_sum_rows_bf16(
+        dout: *const c_void,
+        dbias: *mut c_void,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+}
+
+// Dropout (Phase T18, csrc/ops/dropout.cu). A counter-based (stateless) RNG: the
+// keep/drop decision for element `i` is `hash(seed, i)` — no global state, so a
+// re-run with the same `seed` reproduces the same mask (compatible with T13
+// activation recomputation). Forward writes `out = x ⊙ mask` and the fp32 `mask`
+// buffer (mask[i] = (1/(1-p)) if kept else 0, the inverted-dropout scale);
+// backward applies the SAME mask: dx = d ⊙ mask. fp32 + bf16 activation variants
+// (mask is fp32 in both; the uniform is computed in fp32, dtype-independent).
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    pub fn launch_dropout_fwd_f32(
+        x: *const f32,
+        out: *mut f32,
+        mask: *mut f32,
+        p: f32,
+        scale: f32,
+        seed: u64,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_dropout_bwd_f32(
+        d: *const f32,
+        mask: *const f32,
+        dx: *mut f32,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_dropout_fwd_bf16(
+        x: *const c_void,
+        out: *mut c_void,
+        mask: *mut f32,
+        p: f32,
+        scale: f32,
+        seed: u64,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_dropout_bwd_bf16(
+        d: *const c_void,
+        mask: *const f32,
+        dx: *mut c_void,
+        n: i32,
+        s: CudaStream,
+    );
+}

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,11 +88,33 @@ 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");
+    // Activation recomputation (Phase T13): per-block gradient checkpointing — each
+    // 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).
@@ -107,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 {
@@ -123,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,
@@ -140,6 +178,7 @@ fn main() {
     let dcfg = DdpConfig {
         seq_len,
         batch_size: batch,
+        accum_steps,
         steps,
         schedule: LrSchedule {
             max_lr,
@@ -157,16 +196,49 @@ fn main() {
     };
 
     println!(
-        "training: {steps} steps, seq {seq_len}, global batch {batch}, lr {max_lr:.1e}→{min_lr:.1e}, \
+        "training: {steps} steps, seq {seq_len}, global batch {batch} × accum {accum_steps} = \
-         eval every {eval_every}"
+         effective global batch {}, lr {max_lr:.1e}→{min_lr:.1e}, eval every {eval_every}",
+        batch * accum_steps
     );
 
+    if bf16 {
+        println!("bf16 mixed precision: ON (fp32 master weights)");
+    }
+    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,
         valid.as_ref(),
         &dcfg,
-        move |device| build_model(cfg, device),
+        move |device| {
+            let mut m = build_model(cfg, device);
+            if bf16 {
+                m = m.with_compute_dtype(xtrain_tensor::DType::BF16);
+            }
+            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
+        },
     );
     let r0 = &results[0];
     let start = r0.losses.first().copied().unwrap_or(0.0);

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);
+        // Accumulate grads over `accum` micro-batches, then ONE optimizer step
-        // collect only this rank's shard (global index % world == rank) and run it
+        // (Phase T16). Per micro-batch: draw the whole micro global batch from the
-        // as ONE batched forward/backward. The union of shards == the single-GPU
+        // shared RNG (same on every rank), keep only this rank's shard (global index
-        // batch; each rank's backward yields its local mean (Σ_local / b_local).
+        // % world == rank), run it as ONE batched forward/backward. Each micro-loss
-        let mut inputs = Vec::with_capacity(batch_local);
+        // is scaled by 1/accum before backward (the tape SUM-accumulates the scaled
-        let mut targets_v = Vec::with_capacity(batch_local);
+        // grads across the `accum` micro-backwards) so the boundary grad equals a
-        for i in 0..cfg.batch_size {
+        // single step over an `accum × batch_size` global batch. `accum == 1` skips
-            let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
+        // the scale → bit-identical to the pre-T16 DDP path. The cross-rank
-            if i % ctx.world == ctx.rank {
+        // all-reduce fires ONLY after the last micro-step (intermediate micro-steps
-                inputs.push(input);
+        // are local-only, no NCCL).
-                targets_v.push(target);
+        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
+        // Accumulation boundary: ONE AllReduce(sum) + /world over the accumulated
-        // (local means summed over ranks, /world = global mean). See note above.
+        // 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)
+        // Reported loss = effective-batch mean: AllReduce(sum) the per-rank local
-        // across ranks, /B_global. With equal b_local this is mean over ranks.
+        // sums across ranks, /(accum·B_global).
-        let step_loss =
+        let step_loss = all_reduce_loss(ctx, local_sum) / (accum * cfg.batch_size) as f32;
-            all_reduce_loss(ctx, local_mean * batch_local as f32) / 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).
@@ -20,6 +27,11 @@ pub struct Config {
     pub eps: f32,
     /// RoPE base frequency (theta).
     pub rope_theta: f32,
+    /// Dropout probability `p` (Phase T18). Applied at the attention/MLP sub-block
+    /// outputs (before each residual add) at TRAINING time, with inverted scaling
+    /// `1/(1-p)`; disabled (identity) at eval. Default `0.0` = no dropout, and the
+    /// forward graph is then bit-identical to the pre-T18 path.
+    pub dropout: f32,
 }
 
 impl Config {
@@ -32,10 +44,12 @@ 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,
             rope_theta: 10000.0,
+            dropout: 0.0,
         }
     }
 
@@ -56,13 +70,36 @@ 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,
             rope_theta: 10000.0,
+            dropout: 0.0,
         }
     }
 
+    /// 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
@@ -75,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

@@ -2,17 +2,19 @@
 
 #![cfg(not(no_cuda))]
 
+use std::cell::Cell;
+
 use crate::config::Config;
 use xtrain_autodiff::ops;
 use xtrain_autodiff::tape::Var;
-use xtrain_tensor::{Device, Tensor};
+use xtrain_tensor::{DType, Device, Tensor};
 
 /// One decoder block's learnable tensors.
 struct Block {
     attn_norm: Var, // [dim]
     wq: Var,        // [dim, dim]
-    wk: Var,        // [dim, dim]
+    wk: Var,        // [dim, kv_dim] — kv_dim = num_kv_heads·head_dim (GQA; = dim for MHA)
-    wv: Var,        // [dim, dim]
+    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]
@@ -30,6 +32,44 @@ pub struct TinyTransformer {
     blocks: Vec<Block>,
     final_norm: Var, // [dim]
     lm_head: Var,    // [dim, vocab]
+    /// Compute dtype for the forward graph (Phase T12). `F32` (default) = the
+    /// original path, bit-identical to T10/T11. `BF16` = mixed precision: the
+    /// parameter leaves stay fp32 (master), but each linear's weight is cast to
+    /// bf16 on the fly and the activation stream flows bf16 (see
+    /// `docs/11-bf16-mixed-precision.md`). The cast op's backward upcasts the bf16
+    /// weight grad back to fp32, so AdamW/clip/DDP stay fp32 and unchanged.
+    compute_dtype: DType,
+    /// Activation recomputation / gradient checkpointing (Phase T13, KI-3). When
+    /// `true`, each transformer block's forward runs through
+    /// [`xtrain_autodiff::checkpoint`]: the block's internal activations are NOT
+    /// kept on the tape during forward (only the block input is), and the block
+    /// forward is re-run during backward to recover them. Trades ~one extra forward
+    /// per block for a large drop in peak activation memory → lets dim1024 batch32
+    /// fit. Default `false` = the unchanged path (every activation stored), so
+    /// 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/
+    /// export). `Cell` so `train()`/`eval()` flip it through `&self` (the forward
+    /// takes `&self`). When `cfg.dropout == 0` this flag is irrelevant — the graph
+    /// is bit-identical to the no-dropout path either way.
+    training: Cell<bool>,
+    /// Per-step dropout RNG seed (Phase T18). Bumped once at the start of each
+    /// TRAINING forward so every step draws fresh masks; combined with the layer
+    /// index + a per-site constant to give each dropout site its own seed. The RNG
+    /// is counter-based, so re-running a checkpointed block's forward in backward
+    /// (T13) reproduces the same seed → the same mask (recompute stays exact).
+    step_seed: Cell<u64>,
 }
 
 impl TinyTransformer {
@@ -51,8 +91,9 @@ impl TinyTransformer {
             .map(|_| Block {
                 attn_norm: mk(&[cfg.dim]),
                 wq: mk(&[cfg.dim, cfg.dim]),
-                wk: mk(&[cfg.dim, cfg.dim]),
+                // GQA (T15): K/V project to num_kv_heads·head_dim (= dim when MHA).
-                wv: mk(&[cfg.dim, cfg.dim]),
+                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]),
@@ -71,6 +112,11 @@ impl TinyTransformer {
             blocks,
             final_norm,
             lm_head,
+            compute_dtype: DType::F32,
+            recompute: false,
+            use_flash: false,
+            training: Cell::new(false),
+            step_seed: Cell::new(0),
         }
     }
 
@@ -78,6 +124,73 @@ impl TinyTransformer {
         &self.cfg
     }
 
+    /// Set the forward compute dtype (Phase T12). `BF16` enables mixed precision
+    /// (fp32 master weights, bf16 linears + activations); `F32` (the default) is
+    /// the unchanged full-precision path. Builder-style so existing call sites
+    /// that don't opt in keep the fp32 numerics bit-for-bit.
+    pub fn with_compute_dtype(mut self, dtype: DType) -> Self {
+        assert!(
+            matches!(dtype, DType::F32 | DType::BF16),
+            "compute_dtype must be F32 or BF16"
+        );
+        self.compute_dtype = dtype;
+        self
+    }
+
+    pub fn compute_dtype(&self) -> DType {
+        self.compute_dtype
+    }
+
+    /// Enable per-block activation recomputation / gradient checkpointing (Phase
+    /// T13). Builder-style and opt-in; default off keeps the unchanged tape (every
+    /// activation stored). On, each block's forward is wrapped in
+    /// [`xtrain_autodiff::checkpoint`] — exact grads, lower peak activation memory.
+    pub fn with_recompute(mut self, recompute: bool) -> Self {
+        self.recompute = recompute;
+        self
+    }
+
+    pub fn recompute(&self) -> bool {
+        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) {
+        self.training.set(true);
+    }
+
+    /// Switch to eval mode (Phase T18): dropout is identity. Held-out eval,
+    /// autoregressive sampling, and weight export all run in this mode (default).
+    pub fn eval(&self) {
+        self.training.set(false);
+    }
+
+    pub fn is_training(&self) -> bool {
+        self.training.get()
+    }
+
+    /// Builder-style train/eval toggle (Phase T18) — handy for tests that want a
+    /// model fixed in one mode. Equivalent to [`train`](Self::train) /
+    /// [`eval`](Self::eval) but chains off `new(..)`.
+    pub fn with_training(self, training: bool) -> Self {
+        self.training.set(training);
+        self
+    }
+
     /// All learnable parameters, in a stable order. The optimizer (a hand-written
     /// GD step in T5, AdamW in T6) iterates this; each holds its `.grad()` after
     /// `backward()`.
@@ -127,21 +240,72 @@ impl TinyTransformer {
         );
         let seq = total / batch;
 
-        let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim]
+        // Dropout (T18) is active only in training mode with p>0; otherwise it is
-        for b in &self.blocks {
+        // identity (`ops::dropout` no-ops at p==0). Bump the per-step seed ONCE per
-            // --- Attention sub-block (pre-norm + residual) ---
+        // training forward so each step draws fresh masks (counter-based RNG, so a
-            let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
+        // checkpointed block's recompute reproduces the same seed → same mask).
-            let attn = self.attention(b, &normed, batch, seq);
+        let dropout_p = if self.training.get() {
-            h = ops::add(&h, &attn);
+            self.cfg.dropout
+        } else {
+            0.0
+        };
+        if dropout_p > 0.0 {
+            self.step_seed.set(self.step_seed.get().wrapping_add(1));
+        }
+        let base_seed = self.step_seed.get();
 
-            // --- MLP sub-block (pre-norm + residual) ---
+        // Embedding gathers from the fp32 master table; in bf16 mode cast the
-            let normed = ops::rms_norm(&h, &b.ffn_norm, self.cfg.eps);
+        // activation stream to bf16 here (norms are cast to bf16 gammas too).
-            let mlp = self.swiglu_mlp(b, &normed);
+        let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim], fp32
-            h = ops::add(&h, &mlp);
+        if self.compute_dtype == DType::BF16 {
+            h = ops::cast(&h, DType::BF16);
+        }
+        for (li, b) in self.blocks.iter().enumerate() {
+            // Per-layer dropout seed: a deterministic function of (base_seed,
+            // layer index) — NOT a mutable counter — so the checkpoint recompute
+            // (which re-derives it from the captured base_seed/li) gets the same
+            // masks. The block derives its two per-site seeds from this.
+            let block_seed = base_seed
+                .wrapping_mul(0x100000001B3)
+                .wrapping_add(li as u64);
+            h = if self.recompute {
+                // Activation recomputation (T13): run the whole block forward inside
+                // `checkpoint` so its internal activations aren't kept on the tape;
+                // 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.
+                // `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, 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,
+                    block_seed,
+                    &h,
+                    &b.block_params(),
+                )
+            };
         }
 
-        let h = ops::rms_norm(&h, &self.final_norm, self.cfg.eps);
+        let h = ops::rms_norm(
-        ops::matmul(&h, &self.lm_head) // [batch*seq, vocab]
+            &h,
+            &norm_gamma(self.compute_dtype, &self.final_norm),
+            self.cfg.eps,
+        );
+        // lm_head matmul in compute dtype. Logits stay bf16 in bf16 mode — the
+        // cross_entropy op upcasts to fp32 internally (no persistent fp32 logits
+        // buffer, a real saving at vocab 50257), and its backward casts dx back.
+        linear(self.compute_dtype, &h, &self.lm_head) // [batch*seq, vocab]
     }
 
     /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
@@ -158,68 +322,185 @@ impl TinyTransformer {
         let logits = self.forward_batched(ids, batch);
         ops::cross_entropy(&logits, targets)
     }
+}
 
-    /// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
+impl Block {
-    /// (already normed), laid out sequence-major. The Q/K/V/O projections are big
+    /// The block's learnable leaves, in the fixed order the segment forward
-    /// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
+    /// (`block_forward`) indexes them — matches the per-block slice in
-    /// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each
+    /// [`TinyTransformer::params`]. This is the param order `checkpoint` passes to
-    /// attends within its own `[seq,seq]` causal window (NO cross-sequence
+    /// the recompute closure.
-    /// attention), with RoPE positions reset per sequence (`period = seq`). Causal
+    fn block_params(&self) -> Vec<Var> {
-    /// masking is applied inside the fused op's softmax kernel (no additive
+        vec![
-    /// `[seq,seq]` mask tensor).
+            self.attn_norm.clone(),
-    fn attention(&self, b: &Block, x: &Var, batch: usize, seq: usize) -> Var {
+            self.wq.clone(),
-        let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
+            self.wk.clone(),
-        let total = batch * seq;
+            self.wv.clone(),
-        let bh = batch * nh;
+            self.q_norm.clone(),
-        let scale = 1.0 / (hd as f32).sqrt();
+            self.k_norm.clone(),
+            self.wo.clone(),
+            self.ffn_norm.clone(),
+            self.w_gate.clone(),
+            self.w_up.clone(),
+            self.w_down.clone(),
+        ]
+    }
+}
 
-        // Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
+/// Project `x` (activation, in the compute dtype) by weight `w` (an fp32 master
-        //   [B*S,dim] @ [dim,dim] = [B*S,dim]
+/// leaf). In bf16 mode the weight is cast to bf16 via the autograd `cast` op (whose
-        //   reshape   [B*S, nh, hd]
+/// backward upcasts the grad to fp32); in fp32 mode this is just `matmul(x, w)`.
-        //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
+fn linear(cdt: DType, x: &Var, w: &Var) -> Var {
-        //   rope      [B*S, nh, hd] with per-sequence position (period = seq)
+    match cdt {
-        //   reshape   [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
+        DType::F32 => ops::matmul(x, w),
-        let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
+        DType::BF16 => ops::matmul(x, &ops::cast(w, DType::BF16)),
-            let r = ops::reshape(&proj, &[total, nh, hd]);
+        _ => unreachable!(),
-            let r = match norm {
+    }
-                // Per-head RMSNorm: flatten the (B*S,nh) head rows, norm over hd,
+}
-                // restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
+
-                Some(gamma) => {
+/// A norm/QK-norm gamma in the compute dtype. fp32 master leaf → bf16 (cast op,
-                    let flat = ops::reshape(&r, &[total * nh, hd]);
+/// grad upcast) in bf16 mode; identity in fp32 mode.
-                    let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
+fn norm_gamma(cdt: DType, gamma: &Var) -> Var {
-                    let r = ops::reshape(&normed, &[total, nh, hd]);
+    match cdt {
-                    ops::rope(&r, self.cfg.rope_theta, seq)
+        DType::F32 => gamma.clone(),
-                }
+        DType::BF16 => ops::cast(gamma, DType::BF16),
-                None => r,
+        _ => unreachable!(),
-            };
+    }
-            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]
+/// One transformer block's forward: pre-norm + multi-head causal attention +
+/// (T18) dropout + residual, then pre-norm + SwiGLU MLP + dropout + residual.
+/// 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
+/// reproduced bit-for-bit. `dropout_p == 0` makes `ops::dropout` a no-op (the
+/// graph is then identical to the pre-T18 path). `params` is the block's leaves in
+/// [`Block::block_params`] order.
+#[allow(clippy::too_many_arguments)]
+fn block_forward(
+    cfg: Config,
+    cdt: DType,
+    flash: bool,
+    batch: usize,
+    seq: usize,
+    dropout_p: f32,
+    block_seed: u64,
+    h: &Var,
+    p: &[Var],
+) -> Var {
+    let (attn_norm, wq, wk, wv) = (&p[0], &p[1], &p[2], &p[3]);
+    let (q_norm, k_norm, wo) = (&p[4], &p[5], &p[6]);
+    let (ffn_norm, w_gate, w_up, w_down) = (&p[7], &p[8], &p[9], &p[10]);
+
+    // Per-site dropout seeds (XOR a site constant into the block seed) so the two
+    // residual-path dropouts draw independent masks within the same step/layer.
+    let attn_seed = block_seed ^ 0x0A7700;
+    let ffn_seed = block_seed ^ 0x0FF700;
+
+    // --- 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, 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);
+
+    // --- MLP sub-block (pre-norm + dropout + residual) ---
+    let normed = ops::rms_norm(&h, &norm_gamma(cdt, ffn_norm), cfg.eps);
+    let mlp = swiglu_mlp(cdt, &normed, w_gate, w_up, w_down);
+    let mlp = ops::dropout(&mlp, dropout_p, ffn_seed);
+    ops::add(&h, &mlp)
+}
+
+/// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
+/// (already normed), laid out sequence-major. The Q/K/V/O projections are big
+/// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
+/// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each attends
+/// within its own `[seq,seq]` causal window (NO cross-sequence attention), with
+/// RoPE positions reset per sequence (`period = seq`). Causal masking is applied
+/// inside the fused op's softmax kernel (no additive `[seq,seq]` mask tensor).
+#[allow(clippy::too_many_arguments)]
+fn attention(
+    cfg: Config,
+    cdt: DType,
+    flash: bool,
+    batch: usize,
+    seq: usize,
+    x: &Var,
+    wq: &Var,
+    wk: &Var,
+    wv: &Var,
+    q_norm: &Var,
+    k_norm: &Var,
+    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 scale = 1.0 / (hd as f32).sqrt();
+
+    // 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, 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,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 * heads, hd]);
+                let normed = ops::rms_norm(&flat, &norm_gamma(cdt, gamma), cfg.eps);
+                let r = ops::reshape(&normed, &[total, heads, hd]);
+                ops::rope(&r, cfg.rope_theta, seq)
+            }
+            None => r,
         };
+        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(ops::matmul(x, &b.wq), Some(&b.q_norm));
+    let q = to_bh(linear(cdt, x, wq), nh, Some(q_norm));
-        let k = to_bh(ops::matmul(x, &b.wk), Some(&b.k_norm));
+    // K/V are laid out with num_kv heads, then repeat_kv-broadcast to nh heads so
-        let v = to_bh(ops::matmul(x, &b.wv), None);
+    // 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
+    // Causal SDPA over all B*nh (sequence,head) blocks. `flash` (T14) picks the
-        // (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
+    // single fused flash kernel (online softmax, no materialized [bh,S,S] scores);
-        let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]
+    // 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) →
+    // Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
-        // [B,S,nh,hd] → [B*S, dim].
+    // [B,S,nh,hd] → [B*S, dim].
-        let out = ops::reshape(&out, &[batch, nh, seq, hd]);
+    let out = ops::reshape(&out, &[batch, nh, seq, hd]);
-        let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
+    let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
-        let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
+    let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
-        ops::matmul(&concat, &b.wo) // out projection
+    linear(cdt, &concat, wo) // out projection
-    }
+}
 
-    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*seq,dim].
+/// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*seq,dim].
-    fn swiglu_mlp(&self, b: &Block, x: &Var) -> Var {
+fn swiglu_mlp(cdt: DType, x: &Var, w_gate: &Var, w_up: &Var, w_down: &Var) -> Var {
-        let gate = ops::matmul(x, &b.w_gate); // [seq, ffn_hidden]
+    let gate = linear(cdt, x, w_gate); // [seq, ffn_hidden]
-        let up = ops::matmul(x, &b.w_up); // [seq, ffn_hidden]
+    let up = linear(cdt, x, w_up); // [seq, ffn_hidden]
-        let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
+    let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
-        ops::matmul(&act, &b.w_down) // [seq, dim]
+    linear(cdt, &act, w_down) // [seq, dim]
-    }
 }
 
 /// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step

crates/xtrain-model/tests/bf16.rs

@@ -0,0 +1,151 @@
+// T12 bf16 mixed-precision correctness gate (on-GPU, no PyTorch).
+//
+// The SAME model (identical fp32 master weights) run in fp32 vs bf16 compute
+// mode must agree within a LOOSE bf16 tolerance (bf16 = 7-bit mantissa ≈ 2-3
+// decimal digits → ~1e-2 relative error is expected and acceptable), both for
+// the forward loss/logits AND every parameter's gradient. We also assert no
+// NaN/Inf leaks and that the fp32 grads are fp32 (the cast op upcast the bf16
+// weight grad back to the fp32 master, so AdamW/clip/DDP stay fp32).
+//
+// This is the "bf16 within looser tol vs fp32 reference" gate; the short-run
+// convergence comparison is the train_loop-level bench on dash5.
+#![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) -> 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()
+}
+
+#[test]
+fn bf16_matches_fp32_within_loose_tol() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    // A few layers / heads so the bf16 rounding accumulates through the depth
+    // the real model has (not just a single matmul).
+    let mut cfg = Config::tiny();
+    cfg.vocab = 32;
+    cfg.n_layers = 3;
+    let batch = 2usize;
+    let seq = 8usize;
+
+    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);
+
+    // fp32 reference.
+    let fp32 = build(cfg, device);
+    let f_logits = host(&fp32.forward_batched(&ids, batch).value());
+    let f_loss = fp32.loss_batched(&ids, &tgt, batch);
+    let f_loss_val = host(&f_loss.value())[0];
+    f_loss.backward();
+    let f_params = fp32.params();
+
+    // bf16 — SAME init (build re-runs the same deterministic fill). The forward
+    // now returns bf16 logits (CE upcasts internally); cast to f32 to read.
+    let bf16 = build(cfg, device).with_compute_dtype(DType::BF16);
+    let b_logits = host(
+        &bf16
+            .forward_batched(&ids, batch)
+            .value()
+            .to_dtype(DType::F32),
+    );
+    let b_loss = bf16.loss_batched(&ids, &tgt, batch);
+    let b_loss_val = host(&b_loss.value())[0];
+    b_loss.backward();
+    let b_params = bf16.params();
+
+    // No NaN/Inf in the bf16 forward.
+    assert!(
+        b_logits.iter().all(|v| v.is_finite()) && b_loss_val.is_finite(),
+        "bf16 forward produced non-finite values"
+    );
+
+    // Forward loss within loose bf16 tol.
+    let loss_rel = (b_loss_val - f_loss_val).abs() / f_loss_val.abs().max(1e-4);
+    println!("bf16 vs fp32: loss {b_loss_val:.5} vs {f_loss_val:.5} (rel {loss_rel:.3e})");
+    assert!(
+        loss_rel < 2e-2,
+        "bf16 loss too far from fp32: {loss_rel:.3e}"
+    );
+
+    // Logits: bf16 has ~2-3 decimal digits → compare on a robust (median-style)
+    // basis, requiring the bulk to be within ~3e-2 and the mean error small.
+    let n = f_logits.len();
+    let mut rels: Vec<f32> = f_logits
+        .iter()
+        .zip(&b_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 vs fp32 logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
+    assert!(mean < 1e-2, "bf16 logits mean rel err too high: {mean:.3e}");
+    assert!(p99 < 5e-2, "bf16 logits p99 rel err too high: {p99:.3e}");
+
+    // Gradients: fp32 master grads must be fp32 (cast op upcast), finite, and
+    // within loose bf16 tol of the fp32 reference (mean over each param tensor).
+    let mut worst_param_mean = 0.0f32;
+    for (fp, bp) in f_params.iter().zip(&b_params) {
+        let bg = bp.grad().expect("bf16 grad");
+        assert_eq!(bg.dtype(), DType::F32, "bf16-mode grad must be fp32 master");
+        let fg = host(&fp.grad().expect("fp32 grad"));
+        let bg = host(&bg);
+        assert!(bg.iter().all(|v| v.is_finite()), "bf16 grad has non-finite");
+        // Scale-relative mean error over the tensor (robust to a few small entries).
+        let scale = fg.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
+        let mean_err: f32 =
+            fg.iter().zip(&bg).map(|(f, b)| (f - b).abs()).sum::<f32>() / fg.len() as f32 / scale;
+        worst_param_mean = worst_param_mean.max(mean_err);
+    }
+    println!("bf16 vs fp32 grads: worst per-tensor scaled-mean err = {worst_param_mean:.3e}");
+    assert!(
+        worst_param_mean < 3e-2,
+        "bf16 grads too far from fp32: {worst_param_mean:.3e}"
+    );
+}

crates/xtrain-model/tests/dropout.rs

@@ -0,0 +1,322 @@
+// T18 dropout model-level gates.
+//
+//   1. p=0 bit-identical: a model built with cfg.dropout=0 (in either train or
+//      eval mode) produces logits/loss/grads bit-for-bit identical to the same
+//      model with no dropout field touched — the default forward graph is
+//      unchanged (the regression guard).
+//   2. eval identity: with p>0 but eval mode, the forward equals the p=0 forward
+//      bit-for-bit (dropout is OFF at eval).
+//   3. train vs eval differ: with p>0 and train mode, the forward differs from
+//      eval (dropout actually does something) and grads are still finite.
+//   4. recompute compatibility: with p>0 + train + recompute, grads match the
+//      non-recompute path (the counter-based seed reproduces the same mask on the
+//      backward re-run — T13 stays exact even with dropout in the block).
+//
+// (The fixed-seed grad-check of the dropout op and the E[out]≈x / keep-rate check
+// live in xtrain-autodiff/tests/autograd.rs; p>0 training convergence is the
+// dash5 short run noted in docs/17-dropout.md.)
+#![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) -> 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_dtype(DType::F32)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec()
+}
+
+fn tiny_cfg(dropout: f32) -> Config {
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 4;
+    cfg.dropout = dropout;
+    cfg
+}
+
+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()
+        })
+        .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();
+    (
+        batched_ids_tensor(&seqs, device),
+        batched_ids_tensor(&tgts, device),
+    )
+}
+
+fn require_gpu() -> Device {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    Device::Cuda(0)
+}
+
+// Run forward+backward, return (logits, loss, per-param grads).
+fn fwd_bwd(
+    m: &TinyTransformer,
+    ids: &xtrain_tensor::Tensor,
+    tgt: &xtrain_tensor::Tensor,
+    batch: usize,
+) -> (Vec<f32>, f32, Vec<Vec<f32>>) {
+    let logits = host(&m.forward_batched(ids, batch).value());
+    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();
+    (logits, loss_val, grads)
+}
+
+// --- Gate 3: p=0 is bit-identical to the no-dropout path (default graph). ---
+#[test]
+fn dropout_p0_bit_identical() {
+    let device = require_gpu();
+    let batch = 3;
+
+    // Reference: cfg.dropout default (0.0), never touched train/eval.
+    let cfg0 = tiny_cfg(0.0);
+    let (ids, tgt) = batch_data(&cfg0, device);
+    let ref_m = build(cfg0, device);
+    let (ref_logits, ref_loss, ref_grads) = fwd_bwd(&ref_m, &ids, &tgt, batch);
+
+    // p=0 in TRAINING mode: the seed bump is gated on p>0, the op no-ops at p==0,
+    // so the graph must be byte-identical.
+    let p0_train = build(tiny_cfg(0.0), device);
+    p0_train.train();
+    let (lt, lst, gt) = fwd_bwd(&p0_train, &ids, &tgt, batch);
+
+    assert_eq!(ref_logits, lt, "p=0 train logits not bit-identical");
+    assert_eq!(ref_loss, lst, "p=0 train loss not bit-identical");
+    for (i, (a, b)) in ref_grads.iter().zip(&gt).enumerate() {
+        assert_eq!(a, b, "p=0 train grad[{i}] not bit-identical");
+    }
+    println!("p=0 (train) vs no-dropout: logits/loss/grads bit-identical ✅");
+}
+
+// --- Gate 2: eval is exact identity (p>0 but eval mode == p=0). ---
+#[test]
+fn dropout_eval_is_identity() {
+    let device = require_gpu();
+    let batch = 3;
+    let cfg = tiny_cfg(0.2);
+    let (ids, tgt) = batch_data(&cfg, device);
+
+    // p=0 reference and a p=0.2 model held in eval — outputs must match bit-for-bit.
+    let ref_m = build(tiny_cfg(0.0), device);
+    let (ref_logits, ref_loss, ref_grads) = fwd_bwd(&ref_m, &ids, &tgt, batch);
+
+    let eval_m = build(cfg, device);
+    eval_m.eval(); // explicit; also the default
+    let (el, els, eg) = fwd_bwd(&eval_m, &ids, &tgt, batch);
+
+    assert_eq!(ref_logits, el, "eval (p>0) logits not identity");
+    assert_eq!(ref_loss, els, "eval (p>0) loss not identity");
+    for (i, (a, b)) in ref_grads.iter().zip(&eg).enumerate() {
+        assert_eq!(a, b, "eval (p>0) grad[{i}] not identity");
+    }
+    println!("eval (p=0.2) == no-dropout: bit-identical (eval is identity) ✅");
+}
+
+// --- Gate (train vs eval differ): with p>0 + train, dropout actually fires. ---
+#[test]
+fn dropout_train_differs_from_eval() {
+    let device = require_gpu();
+    let batch = 3;
+    let cfg = tiny_cfg(0.3);
+    let (ids, _tgt) = batch_data(&cfg, device);
+
+    let m = build(cfg, device);
+    m.eval();
+    let eval_logits = host(&m.forward_batched(&ids, batch).value());
+    m.train();
+    let train_logits = host(&m.forward_batched(&ids, batch).value());
+
+    let max_diff = eval_logits
+        .iter()
+        .zip(&train_logits)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0.0f32, f32::max);
+    assert!(
+        max_diff > 1e-4 && train_logits.iter().all(|v| v.is_finite()),
+        "train logits should differ from eval (dropout active) and be finite; max_diff={max_diff}"
+    );
+    println!("train vs eval logits max diff {max_diff:.4e} (dropout active in train) ✅");
+}
+
+// --- Gate 4: p>0 + recompute grads match non-recompute (T13 stays exact). ---
+// The counter-based seed is a pure function of (step_seed, layer, site); the
+// checkpoint backward re-runs block_forward and re-derives the SAME seeds, so the
+// recomputed dropout masks match the forward — grads stay bit-identical.
+fn recompute_with_dropout(dtype: DType, grad_tol: f32) {
+    let device = require_gpu();
+    let batch = 3;
+    let cfg = tiny_cfg(0.2);
+    let (ids, tgt) = batch_data(&cfg, device);
+
+    // 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 on = build(cfg, device)
+        .with_compute_dtype(dtype)
+        .with_recompute(true)
+        .with_training(true);
+
+    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 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 mut max_rel = 0.0f32;
+    for (a, b) in off_grads.iter().flatten().zip(on_grads.iter().flatten()) {
+        max_rel = max_rel.max((a - b).abs() / a.abs().max(1e-3));
+    }
+    println!("[{dtype:?}] dropout p=0.2 recompute on/off grad max rel = {max_rel:.3e}");
+    assert!(
+        max_rel < grad_tol,
+        "[{dtype:?}] recompute grads diverged with dropout: {max_rel:.3e}"
+    );
+}
+
+#[test]
+fn dropout_recompute_matches_fp32() {
+    recompute_with_dropout(DType::F32, 1e-4);
+}
+
+#[test]
+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)
+    # GQA: K/V project to NKV heads, then repeat each kv head GROUP times to NH.
-    v = (x @ L["wv"]).reshape(B * SEQ, NH, HD)
+    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)
+    k = rope(k)  # [B*SEQ, nkv, hd]
-    # Reshape to [B, NH, SEQ, HD] so attention runs within each sequence.
+    # 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)
+    k = k.reshape(B, SEQ, NKV, HD).transpose(1, 2)  # [B, nkv, seq, hd]
-    v = v.reshape(B, SEQ, NH, HD).transpose(1, 2)
+    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-model/tests/recompute.rs

@@ -0,0 +1,161 @@
+// T13 activation-recomputation correctness gate (the HARD gate).
+//
+// Gradient checkpointing is mathematically EXACT: the backward re-runs the same
+// `segment_fn` from the same saved input and the same (unchanged) parameter
+// values, so the recomputed activations equal the originals and the recovered
+// grads equal the non-checkpointed grads — checkpointing trades compute for
+// memory, never correctness. This test makes that a closed loop on-GPU:
+//
+//   build two identical models (same init), one with `--recompute` on, one off,
+//   run the SAME batched loss + backward on both, and assert
+//     1. the forward logits match (recompute doesn't touch forward output)
+//     2. the loss matches
+//     3. EVERY parameter's grad matches within a tight fp tolerance.
+//
+// Composition is covered by parameterising over fp32 AND bf16 (T12): the
+// recompute path is the unchanged block forward, so it runs the same dtype path.
+#![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, 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.08)
+        }
+    });
+    m.with_compute_dtype(dtype).with_recompute(recompute)
+}
+
+/// Upcast to fp32 then read to host — logits are bf16 in bf16 mode (grads are
+/// always fp32 master, but this is uniform and harmless for fp32 tensors).
+fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
+    t.to_dtype(DType::F32)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec()
+}
+
+fn run(dtype: DType, logit_tol: f32, grad_tol: f32) {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    // A few layers so checkpointing actually wraps multiple blocks.
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 4;
+    let batch = 3usize;
+    let seq = 6usize;
+    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);
+
+    // --- recompute OFF (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();
+
+    // --- recompute 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();
+
+    // 1. Forward logits — recompute must not change the forward output.
+    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);
+    // 2. Loss.
+    let loss_rel = (off_loss_val - on_loss_val).abs() / off_loss_val.abs().max(1e-4);
+    println!(
+        "[{dtype:?}] recompute on/off: loss {off_loss_val:.6}/{on_loss_val:.6} (rel {loss_rel:.2e}), \
+         logits max rel {logit_rel:.2e}"
+    );
+    assert!(
+        logit_rel < logit_tol,
+        "[{dtype:?}] logits diverged: {logit_rel:.2e}"
+    );
+    assert!(
+        loss_rel < logit_tol,
+        "[{dtype:?}] loss diverged: {loss_rel:.2e}"
+    );
+
+    // 3. Every parameter grad — the load-bearing gate.
+    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) {
+            let rel = (a - b).abs() / a.abs().max(1e-3);
+            max_grad_rel = max_grad_rel.max(rel);
+        }
+    }
+    println!("[{dtype:?}] recompute on/off: grad max rel err = {max_grad_rel:.3e}");
+    assert!(
+        max_grad_rel < grad_tol,
+        "[{dtype:?}] recompute grads diverged from non-recompute: {max_grad_rel:.3e}"
+    );
+}
+
+#[test]
+fn recompute_matches_non_recompute_fp32() {
+    // fp32: recompute runs the identical deterministic kernels → grads match to
+    // (near) bit-exact; allow a hair for any nondeterministic GPU reduction.
+    run(DType::F32, 1e-5, 1e-4);
+}
+
+#[test]
+fn recompute_matches_non_recompute_bf16() {
+    // bf16 (T12 composition): same bf16 path on recompute. The recompute is still
+    // exact w.r.t. the bf16 forward, so on/off match tightly (looser tol only for
+    // bf16 rounding, not for any recompute discrepancy).
+    run(DType::BF16, 5e-3, 5e-3);
+}

crates/xtrain-tensor/src/dtype.rs

@@ -1,12 +1,16 @@
 //! Tensor data types.
 //!
-//! T2 only needs `F32`, but the enum + `TensorDType` trait are structured so
+//! T2 only needs `F32`; `BF16` was added in T12 for mixed-precision training
-//! half-precision types (F16/BF16) can be added later (T7 mixed precision)
+//! (bf16 linears / activations, fp32 master weights — see
-//! without touching call sites.
+//! `docs/11-bf16-mixed-precision.md`). The enum + `TensorDType` trait keep call
+//! sites dtype-polymorphic.
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum DType {
     F32,
+    /// bfloat16: 1 sign / 8 exponent / 7 mantissa. Same exponent range as f32
+    /// (so no loss scaling needed), ~2-3 decimal digits. The T12 AMP compute type.
+    BF16,
     /// 32-bit signed integers. Used for cross-entropy targets (token ids).
     I32,
 }
@@ -15,6 +19,7 @@ impl DType {
     pub fn size_bytes(self) -> usize {
         match self {
             DType::F32 => 4,
+            DType::BF16 => 2,
             DType::I32 => 4,
         }
     }
@@ -22,6 +27,7 @@ impl DType {
     pub fn name(self) -> &'static str {
         match self {
             DType::F32 => "f32",
+            DType::BF16 => "bf16",
             DType::I32 => "i32",
         }
     }
@@ -50,6 +56,16 @@ impl TensorDType for f32 {
     }
 }
 
+impl TensorDType for half::bf16 {
+    const DTYPE: DType = DType::BF16;
+    fn to_f64(self) -> f64 {
+        self.to_f64()
+    }
+    fn from_f64(v: f64) -> Self {
+        half::bf16::from_f64(v)
+    }
+}
+
 impl TensorDType for i32 {
     const DTYPE: DType = DType::I32;
     fn to_f64(self) -> f64 {

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

@@ -0,0 +1,179 @@
+//! Greedy-generation helper for run verification — load a trained checkpoint with
+//! its arch flags and print xtrain's OWN greedy continuations for the fixed run
+//! prompts, so they can be diffed against xserv's greedy on the exported weights
+//! (the token-match check each scaling run reports). f32 forward, same
+//! model/config/ckpt + init scheme as bin/train.rs and bin/export_safetensors.rs.
+//!
+//!   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 \
+//!       --prompts-file scripts/chat_alpha_fixed_prompts.txt --max-tokens 120
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("greedy_sample: 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;
+
+// Same deterministic LCG init scheme as bin/train.rs / bin/export_safetensors.rs.
+#[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()
+}
+
+// A flag like `--layers 18`: 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 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;
+
+    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()))
+        .unwrap_or_else(|| PathBuf::from("/tmp/xtrain_tinystories.ckpt"));
+    let tok_path = positionals
+        .get(1)
+        .map(|s| PathBuf::from(s.as_str()))
+        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
+
+    // Architecture must match the checkpoint. 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);
+    // 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)
+        .with_kv_heads(kv_heads);
+    println!(
+        "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,
+    );
+
+    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!(
+        "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 mut rng = 7u64;
+        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

@@ -31,6 +31,8 @@ use xtrain_cuda::device;
 #[cfg(not(no_cuda))]
 use xtrain_model::{Config, TinyTransformer};
 #[cfg(not(no_cuda))]
+use xtrain_tensor::DType;
+#[cfg(not(no_cuda))]
 use xtrain_tensor::Device;
 #[cfg(not(no_cuda))]
 use xtrain_train::data::Corpus;
@@ -86,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 {
@@ -99,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);
@@ -107,6 +115,21 @@ 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).
+    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");
+    // Activation recomputation (Phase T13): per-block gradient checkpointing —
+    // 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")
@@ -114,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();
@@ -124,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 {
@@ -140,13 +175,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,
@@ -155,7 +193,7 @@ fn main() {
     );
 
     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 {
@@ -166,6 +204,25 @@ fn main() {
             fill(n, seed, 0.04)
         }
     });
+    if bf16 {
+        model = model.with_compute_dtype(DType::BF16);
+        println!("bf16 mixed precision: ON (fp32 master weights)");
+    }
+    if recompute {
+        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
@@ -184,6 +241,7 @@ fn main() {
     let tcfg = TrainConfig {
         seq_len,
         batch_size,
+        accum_steps,
         steps,
         schedule: LrSchedule {
             max_lr,
@@ -202,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/sample.rs

@@ -26,7 +26,11 @@ pub fn generate(
 
     for _ in 0..max_new {
         let ids_t = ids_tensor(&ids, device);
-        let logits = model.forward(&ids_t).value().to_device(Device::Cpu);
+        // In bf16 mode the logits are bf16; cast to f32 (on device) before reading.
+        let logits = model.forward(&ids_t).value();
+        let logits = logits
+            .to_dtype(xtrain_tensor::DType::F32)
+            .to_device(Device::Cpu);
         let lg = logits.as_slice::<f32>();
         // Last row = next-token distribution for the current prefix.
         let last = &lg[(ids.len() - 1) * vocab..ids.len() * vocab];

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,28 +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/
+        // Accumulate grads over `accum` micro-batches of `batch_size` sequences,
-        // backward. The CE mean over all batch*seq rows is the batch-mean loss, so
+        // then take ONE optimizer step (Phase T16). Each micro-batch is ONE batched
-        // backward already yields the batch-mean gradient (clip pre-scale = 1.0).
+        // forward/backward; its loss is the CE mean over batch*seq rows, so backward
-        let mut inputs = Vec::with_capacity(cfg.batch_size);
+        // yields that micro-batch's mean grad. To make the SUM over `accum` micro-
-        let mut targets_v = Vec::with_capacity(cfg.batch_size);
+        // batches equal a single step over an `accum × batch` batch, each micro-loss
-        for _ in 0..cfg.batch_size {
+        // is scaled by 1/accum before backward (the tape SUM-accumulates the scaled
-            let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
+        // grads). `accum == 1` skips the scale entirely → bit-identical to pre-T16.
-            inputs.push(input);
+        let mut step_loss_sum = 0.0f32;
-            targets_v.push(target);
+        // 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();
+        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;
         }
-        let ids = batched_ids_tensor(&inputs, device);
+        // Reported loss = mean over the effective batch = mean of the raw micro
-        let targets = batched_ids_tensor(&targets_v, device);
+        // losses (each is itself a micro-batch mean of equal size).
-        let loss = model.loss_batched(&ids, &targets, cfg.batch_size);
+        let step_loss = step_loss_sum / accum as f32;
-        let step_loss = read_scalar(&loss);
-        loss.backward();
-        tokens_seen += (cfg.batch_size * cfg.seq_len) as u64;
         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 {
@@ -169,6 +194,8 @@ pub fn eval_loss(
     if valid.len() <= seq + 1 {
         return f32::NAN;
     }
+    // Eval mode → dropout is identity (T18).
+    model.eval();
     let n_win = (valid.len() - 1) / seq; // disjoint windows that fit
     let batches = batches.max(1).min(n_win.max(1));
     let stride = (n_win / batches).max(1);
@@ -180,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/cast.cu

@@ -0,0 +1,141 @@
+// bf16 mixed-precision kernels (Phase T12, KI-2).
+//
+// Two groups:
+//   1. f32 <-> bf16 cast — the bridge between fp32 master weights / fp32
+//      reductions and the bf16 compute/activation stream.
+//   2. bf16 elementwise ops (add / mul / silu / scale + their backwards) — the
+//      residual-stream ops that flow bf16 activations. Each loads bf16 -> float,
+//      computes in fp32, stores bf16 (so the math accumulates in fp32 while the
+//      stored activation is half-size). Matmuls go through cuBLAS GemmEx
+//      (cublas.rs); norm / softmax / rope / cross-entropy stay fp32 (the Rust
+//      wrappers upcast around the existing fp32 kernels).
+//
+// bf16 is __nv_bfloat16; __float2bfloat16 / __bfloat162float round-trip via fp32.
+
+#include <cuda_bf16.h>
+
+extern "C" {
+
+// --- f32 <-> bf16 cast ---
+
+__global__ void cast_f32_to_bf16_k(const float* in, __nv_bfloat16* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __float2bfloat16(in[i]);
+}
+void launch_cast_f32_to_bf16(const float* in, void* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    cast_f32_to_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        in, (__nv_bfloat16*)out, n);
+}
+
+__global__ void cast_bf16_to_f32_k(const __nv_bfloat16* in, float* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __bfloat162float(in[i]);
+}
+void launch_cast_bf16_to_f32(const void* in, float* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    cast_bf16_to_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)in, out, n);
+}
+
+// --- bf16 elementwise (load->fp32->compute->store bf16) ---
+
+__global__ void add_bf16_k(const __nv_bfloat16* a, const __nv_bfloat16* b,
+                           __nv_bfloat16* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n)
+        out[i] = __float2bfloat16(__bfloat162float(a[i]) + __bfloat162float(b[i]));
+}
+void launch_add_bf16(const void* a, const void* b, void* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    add_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+}
+
+__global__ void mul_bf16_k(const __nv_bfloat16* a, const __nv_bfloat16* b,
+                           __nv_bfloat16* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n)
+        out[i] = __float2bfloat16(__bfloat162float(a[i]) * __bfloat162float(b[i]));
+}
+void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    mul_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+}
+
+__global__ void scale_bf16_k(const __nv_bfloat16* in, __nv_bfloat16* out,
+                             float alpha, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __float2bfloat16(__bfloat162float(in[i]) * alpha);
+}
+void launch_scale_bf16(const void* in, void* out, float alpha, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    scale_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, alpha, n);
+}
+
+// SiLU: y = x*sigmoid(x). Backward: dx = dy * (sig + x*sig*(1-sig)).
+__global__ void silu_bf16_k(const __nv_bfloat16* x, __nv_bfloat16* y, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float v = __bfloat162float(x[i]);
+        float sig = 1.0f / (1.0f + expf(-v));
+        y[i] = __float2bfloat16(v * sig);
+    }
+}
+void launch_silu_bf16(const void* x, void* y, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    silu_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)y, n);
+}
+
+__global__ void silu_dx_bf16_k(const __nv_bfloat16* x, const __nv_bfloat16* dy,
+                               __nv_bfloat16* dx, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float v = __bfloat162float(x[i]);
+        float sig = 1.0f / (1.0f + expf(-v));
+        float g = sig + v * sig * (1.0f - sig);
+        dx[i] = __float2bfloat16(__bfloat162float(dy[i]) * g);
+    }
+}
+void launch_silu_dx_bf16(const void* x, const void* dy, void* dx, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    silu_dx_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)dy, (__nv_bfloat16*)dx, n);
+}
+
+// Broadcast bias add: out[r,c] = x[r,c] + bias[c]. x:[rows,cols], bias:[cols].
+__global__ void add_bias_bf16_k(const __nv_bfloat16* x, const __nv_bfloat16* bias,
+                                __nv_bfloat16* out, int rows, int cols) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < rows * cols)
+        out[i] = __float2bfloat16(__bfloat162float(x[i]) +
+                                  __bfloat162float(bias[i % cols]));
+}
+void launch_add_bias_bf16(const void* x, const void* bias, void* out, int rows,
+                          int cols, void* s) {
+    int blk = 256, grid = (rows * cols + blk - 1) / blk;
+    add_bias_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)bias, (__nv_bfloat16*)out,
+        rows, cols);
+}
+
+// Column-sum over rows: dbias[c] = sum_r dout[r,c] (bias backward), fp32 accum.
+__global__ void sum_rows_bf16_k(const __nv_bfloat16* dout, __nv_bfloat16* dbias,
+                                int rows, int cols) {
+    int c = blockIdx.x * blockDim.x + threadIdx.x;
+    if (c < cols) {
+        float acc = 0.0f;
+        for (int r = 0; r < rows; ++r) acc += __bfloat162float(dout[r * cols + c]);
+        dbias[c] = __float2bfloat16(acc);
+    }
+}
+void launch_sum_rows_bf16(const void* dout, void* dbias, int rows, int cols, void* s) {
+    int blk = 256, grid = (cols + blk - 1) / blk;
+    sum_rows_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)dout, (__nv_bfloat16*)dbias, rows, cols);
+}
+
+}  // extern "C"

csrc/ops/dropout.cu

@@ -0,0 +1,109 @@
+// Dropout kernels (Phase T18).
+//
+// A counter-based (stateless) RNG: the keep/drop decision for element `i` is a
+// pure function of (seed, i) — no global RNG state is advanced. This is what
+// makes dropout compatible with activation recomputation (T13): when a
+// checkpointed block re-runs its forward in backward, the SAME seed regenerates
+// the SAME mask, so the recomputed activations / grads stay bit-identical to the
+// forward (no mask drift).
+//
+// Inverted dropout: at training time kept elements are scaled by 1/(1-p) so the
+// expectation E[out] == x (no inference-time rescale needed; eval is identity,
+// handled in Rust by simply not calling dropout).
+//
+//   key  = seed ^ (i * GOLDEN)
+//   h    = splitmix64(key)             // a few rounds of xorshift/multiply
+//   u    = (h >> 40) / 2^24  in [0,1)  // 24-bit uniform
+//   keep = u >= p                      // Bernoulli(keep = 1-p)
+//   out  = keep ? x * scale : 0        // scale = 1/(1-p)
+//   mask = keep ? scale : 0            // cached for backward (dx = d * mask)
+//
+// fp32 + bf16 variants: bf16 loads/stores half-size activations but the uniform
+// `u` is always computed in fp32, so the mask distribution is identical across
+// dtypes (drop decisions don't depend on bf16 rounding). The mask buffer is fp32
+// in both cases (it stores `scale` or 0 — exactly representable, tiny relative to
+// the activation, reused only elementwise in backward).
+
+#include <cuda_bf16.h>
+#include <stdint.h>
+
+extern "C" {
+
+// splitmix64: cheap, well-mixed counter hash. Maps a 64-bit counter to a 64-bit
+// pseudo-random output; we only need the high bits for a uniform.
+__device__ __forceinline__ uint64_t splitmix64(uint64_t x) {
+    x += 0x9E3779B97F4A7C15ULL;
+    x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
+    x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
+    return x ^ (x >> 31);
+}
+
+// Uniform [0,1) for element i under `seed`, computed in fp32 (dtype-independent).
+__device__ __forceinline__ float dropout_uniform(uint64_t seed, int i) {
+    uint64_t key = seed ^ ((uint64_t)i * 0x9E3779B97F4A7C15ULL);
+    uint64_t h = splitmix64(key);
+    // Top 24 bits → [0,1) with 2^-24 resolution.
+    return (float)(h >> 40) * (1.0f / 16777216.0f);  // 1/2^24
+}
+
+// fp32 forward: out[i] = keep ? x[i]*scale : 0 ; mask[i] = keep ? scale : 0.
+__global__ void dropout_fwd_f32_k(const float* x, float* out, float* mask,
+                                  float p, float scale, uint64_t seed, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float keep = (dropout_uniform(seed, i) >= p) ? scale : 0.0f;
+        mask[i] = keep;
+        out[i] = x[i] * keep;
+    }
+}
+void launch_dropout_fwd_f32(const float* x, float* out, float* mask, float p,
+                            float scale, uint64_t seed, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    dropout_fwd_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, out, mask, p, scale,
+                                                         seed, n);
+}
+
+// Backward applies the SAME cached mask elementwise: dx[i] = d[i] * mask[i].
+__global__ void dropout_bwd_f32_k(const float* d, const float* mask, float* dx,
+                                  int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) dx[i] = d[i] * mask[i];
+}
+void launch_dropout_bwd_f32(const float* d, const float* mask, float* dx, int n,
+                            void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    dropout_bwd_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(d, mask, dx, n);
+}
+
+// bf16 forward: activation is bf16; mask is fp32 (stores `scale` or 0). Uniform
+// is fp32, so the mask matches the fp32 path bit-for-bit (same drop decisions).
+__global__ void dropout_fwd_bf16_k(const __nv_bfloat16* x, __nv_bfloat16* out,
+                                   float* mask, float p, float scale,
+                                   uint64_t seed, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float keep = (dropout_uniform(seed, i) >= p) ? scale : 0.0f;
+        mask[i] = keep;
+        out[i] = __float2bfloat16(__bfloat162float(x[i]) * keep);
+    }
+}
+void launch_dropout_fwd_bf16(const void* x, void* out, float* mask, float p,
+                             float scale, uint64_t seed, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    dropout_fwd_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, mask, p, scale, seed, n);
+}
+
+__global__ void dropout_bwd_bf16_k(const __nv_bfloat16* d, const float* mask,
+                                   __nv_bfloat16* dx, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) dx[i] = __float2bfloat16(__bfloat162float(d[i]) * mask[i]);
+}
+void launch_dropout_bwd_bf16(const void* d, const float* mask, void* dx, int n,
+                             void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    dropout_bwd_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)d, mask, (__nv_bfloat16*)dx, n);
+}
+
+}  // extern "C"

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);
+    int t = target[r];
-    dx[i] = g * scale;
+    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/11-bf16-mixed-precision.md

@@ -0,0 +1,98 @@
+# Phase T12: bf16 混合精度（fp32 master）— Design Document
+
+> KI-2 的具体落地。v4（dim768, fp32）在单卡 32GB 下 per-rank batch 32（global 256）**OOM**，被迫降到 batch 16 训练。bf16 把激活显存减半（找回 batch-256 甜点区），并在 dim768 这个已 compute-bound 的规模上加速 tensor-core GEMM。附带收益：xserv 推理是 **BF16-only**，bf16 训练让闭环更贴。
+
+## Goal
+
+在**不动 fp32 路径任何数值**的前提下，新增一个 **opt-in 的 bf16 混合精度模式**（标准 AMP，fp32 master weights）：
+
+1. **正确性硬闸门**：fp32 全套回归（T3 GEMM / T4 12 算子 grad-check / T5 结构+overfit+PyTorch 对拍 / T6 AdamW+checkpoint / T8 DDP / T10 batched / xserv 闭环）在**同样紧容差**下保持绿。bf16 是**加法、可选**的，绝不扰动 fp32。
+2. **bf16 正确性**：bf16 前向/梯度在**更松的 bf16 容差**（≈2–3 位十进制有效数字 → ~1e-2 相对误差）内对住 fp32 参考；一段**短 bf16 训练收敛对住 fp32**（loss 曲线接近、无 NaN/发散）。
+3. **显存+吞吐（payoff）**：dim768 bf16 能跑 per-rank batch 32（解 OOM）；测 dim768 bf16 vs fp32 的显存+tok/s。
+
+## 什么是 bf16、什么是 fp32（标准 AMP split）
+
+| 组件 | 精度 | 理由 |
+|---|---|---|
+| **master weights** + AdamW state(m/v) + 优化器更新 | **fp32** | 小步长更新需要 fp32 精度，否则被 bf16 的 8-bit 尾数吃掉 |
+| **linear GEMM**（q/k/v/o、gate/up/down、lm_head） | **bf16 in/out + fp32 accum** | compute+memory 主体；tensor-core 走 `cublasGemmEx`（`CUDA_R_16BF` in/out，`CUBLAS_COMPUTE_32F` 累加） |
+| **激活流**（残差流、attention Q/K/V/probs/out、MLP 中间） | **bf16** | 激活显存减半——这是解 OOM 的关键，不只是 GEMM 提速 |
+| **RMSNorm / QK-norm** | **fp32**（bf16→fp32 算 reduction→bf16） | 求和/rsqrt 数值敏感 |
+| **softmax（attention）/ RoPE / cross-entropy** | **fp32** | softmax 的 exp/求和、CE 的 log、RoPE 的 sin/cos 都数值敏感 |
+| **梯度 → AdamW** | **fp32** | 见下「cast 算子」——grad 在 fp32 master leaf 上累加，AdamW/clip/DDP all-reduce 全程 fp32、**完全不改** |
+
+**无 loss scaling**：bf16 是 8-bit 指数（与 fp32 同动态范围），不像 fp16（5-bit 指数易下溢）。所以梯度不会下溢到 0，**不需要** loss scaling。
+
+## Module Layout（surgical：fp32 路径逐字节不动）
+
+核心思路：**所有 op 按 `self.dtype()` 分派**。fp32 分支跑原 kernel（一字不改）；bf16 分支是新增代码。
+
+### 1. `xtrain-tensor::dtype` — 加 `BF16`
+- `DType::BF16`，`size_bytes()=2`，`half::bf16` 实现 `TensorDType`（`half` crate 已是依赖）。
+
+### 2. `xtrain-cuda` — bf16 GEMM + cast kernel
+- `ffi.rs`：声明 `cublasGemmEx` / `cublasGemmStridedBatchedEx`（void* 指针、`a_type/b_type/c_type/compute_type`），常量 `CUDA_R_16BF=14`、`CUDA_R_32F=0`、`CUBLAS_COMPUTE_32F=68`（数值同 xserv `gemm.rs`）。
+- `cublas.rs`：`gemm_ex(...)` / `gemm_ex_strided_batched(...)`——和 `sgemm` 同样的 row-major⟺col-major 转置代数，只是走 `GemmEx`、in/out=bf16、accum=fp32。
+- `csrc/ops/cast.cu` + ffi：`launch_cast_f32_to_bf16` / `launch_cast_bf16_to_f32`（逐元素 `__float2bfloat16` / `__bfloat162float`）。
+
+### 3. `xtrain-tensor::tensor` — dtype-polymorphic ops
+- `to_dtype(target)`：f32↔bf16 cast（CUDA），同 dtype 直接 clone。
+- `matmul` / `matmul_backward` / `attention` / `attention_backward`：按 dtype 分派——fp32 走原 `sgemm`（**不动**），bf16 走 `gemm_ex`，输出同 dtype。
+- 逐元素 op（`add`/`mul`/`silu`/`scale`…）+ `embedding`：允许 bf16 输入。逐元素 kernel 对 bf16 走「load→fp32→算→store bf16」（新增 bf16 kernel）或对 norm/softmax/CE 在 wrapper 里 upcast→fp32 kernel→downcast。**fp32 调用走原 f32 kernel 不变。**
+
+### 4. `xtrain-autodiff::ops` — `cast` 算子 + bf16 透传
+- **`cast(x, target_dtype)`**：前向 `x.to_dtype(target)`；**反向把 grad cast 回 `x` 的 dtype**。这是 AMP 的关键钩子：
+  - fp32 master weight leaf → `cast(w, BF16)` 喂给 matmul；matmul 的 bf16 grad 经 cast 反向**升回 fp32**，累加在 fp32 leaf 上。
+  - ⟹ `.grad()` 是 **fp32**，AdamW / `clip_grad_norm_gpu` / DDP `all_reduce_average_grads` **一行不改**，全程 fp32 master。
+- 其它 op 的 backward 自然按张量 dtype 流转（softmax/rms_norm wrapper 内部 upcast→fp32→downcast，对外是 bf16）。
+
+### 5. `xtrain-model::TinyTransformer` — `compute_dtype` 开关
+- `new_amp(cfg, device, dtype, init)` 或 `forward_batched` 接 `compute_dtype: DType`（默认 `F32` = 原路径，逐字节同）。
+- bf16 模式：embedding 输出 `cast→bf16` 进入残差流；每个 weight matmul 前 `cast(w, BF16)`；norm/softmax/rope 对 bf16 激活自动 fp32 内算；最后 logits `cast→fp32` 给 cross_entropy。**fp32 模式 `compute_dtype==F32` 时跳过所有 cast，graph 与 T10/T11 完全一致。**
+
+### 6. `xtrain-train` / `xtrain-distributed` — `--bf16` flag
+- `TrainConfig`/`DdpConfig` 加 `compute_dtype: DType`；`train.rs`/`train_ddp.rs` 加 `--bf16` flag。
+- AdamW / clip / checkpoint / DDP all-reduce **不改**（master 永远 fp32，grad 永远 fp32）。
+
+## Key Design Decisions
+
+- **cast 算子承载 fp32 master ↔ bf16 compute 的桥**：不需要在优化器里维护一份独立的 bf16 weight 副本——fp32 leaf 即 master，前向临时 cast 出 bf16，反向 grad 自动升回 fp32。最小改动、零优化器侵入。
+- **按 dtype 分派而非新类型**：fp32 路径走的还是同一个函数的 `F32` 分支 → 原 kernel、原 cuBLAS 调用、原 launch 顺序，数值逐字节不变（满足硬闸门）。
+- **norm/softmax/CE 不写 bf16 reduction kernel**：wrapper 里 `to_dtype(F32)` → 复用现有 fp32 kernel → `to_dtype(BF16)`。多两个 cast launch，但**复用已验证的 fp32 数值**，且这些不是显存/算力大头。
+- **无 loss scaling**：bf16 8-bit 指数，省掉 fp16 那套 scale/unscale/inf-check。
+
+## 验证方法（双闸门）
+
+### 闸门 ① fp32 不回归（hard gate）
+全套现有测试在原紧容差下保持绿（bf16 是 opt-in，默认 dtype=F32）：
+- `cargo test` 全 crate：grad-check（rel ≤2e-2）、structural、GEMM 对 cuBLAS（~1e-7）、batched==looped、overfit 27/27、AdamW GPU bit-exact + host 对 torch、checkpoint 逐位、DDP loss 对单卡 <1e-6、**PyTorch 对拍**（loss/logits/grad）。
+- **xserv 闭环**：v4 ckpt（fp32 训）重导 safetensors md5 一致 + xserv 贪心逐 token 对住。
+
+### 闸门 ② bf16 正确性 + 收敛
+- **bf16 looser-tol 数值**：同一组随机权重/输入，bf16 forward logits 与 bf16 grad 对 fp32 参考在 **rel ~1e-2**（bf16 2–3 位有效数字）内。
+- **短训练收敛**：dim768（或缩小代理）bf16 跑数百步，loss 曲线对住 fp32，无 NaN/发散，end loss 接近。
+
+### 闸门 ③ 显存 + 吞吐（payoff）
+- **dim768 bf16 能跑 per-rank batch 32**（v4 OOM 的触发点）。
+- 测 dim768 **bf16 vs fp32** 的峰值显存 + steady-state tok/s（预期：显存↓、tok/s↑）。
+
+## 实测结果（dash5, 1× RTX 5090 32GB, sm_120）
+
+**闸门 ① fp32 不回归**：全套测试在原紧容差绿（autograd 15 / structural 5 / GEMM 5 / batched==looped / overfit 27/27 / AdamW GPU bit-exact + host 对 torch / checkpoint 逐位 / DDP 2）。**xserv 闭环**：v3 ckpt 用 T12 代码重导 `model.safetensors` 与 registry **md5 逐位一致**（`b04fc9f9a0c9af04c47d9ca649aea12e`）——export/fp32 数值零漂移。
+
+**闸门 ② bf16 正确性 + 收敛**：
+- **looser-tol（tests/bf16.rs）**：同 fp32 master 跑 fp32 vs bf16——loss rel `1.2e-4`、logits mean rel `2.0e-3` / p99 `6.8e-3`、grad worst scaled-mean `1.0e-2`，无 NaN，grad 仍 fp32（master 未动）。
+- **收敛**：dim768 短训 150 步，bf16-b16 loss 轨迹对住 fp32-b16（step50 `4.40` vs `4.40`、step149 `3.984` vs `3.988`），单调降、无发散。
+
+**闸门 ③ 显存 + 吞吐（dim768/18L/24h×32 ffn2048 seq256, steady-state）**：
+
+| config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
+|---|---|---|---|---|
+| fp32 | 16 (v4 fallback) | 27.2 GB | 31.5K | ✅ |
+| **bf16** | 16 | **19.3 GB（−29%）** | **35.5K（+13%）** | ✅ |
+| fp32 | 32 | — | — | ❌ **OOM** |
+| **bf16** | **32（甜点区）** | **31.1 GB** | **40.8K** | ✅ **解 OOM** |
+
+→ 同 batch：bf16 显存 −29% / tok/s +13%；**bf16 解 fp32-batch32 OOM**，batch32 达 40.8K tok/s（+29% vs fp32-b16）。KI-2 标 **FIXED**。
+</content>
+</invoke>

docs/12-activation-recompute.md

@@ -0,0 +1,90 @@
+# Phase T13: 激活重计算（gradient checkpointing）— Design Document
+
+> KI-3 的具体落地。autograd tape 为反向保存了所有中间激活；dim768/bf16 在单卡 32GB 能跑 batch32（T12 解 OOM），但**容量轴放大到 dim1024 会再次 OOM**——激活显存随 dim 线性增长。激活重计算用「多一次前向」换显存：段内激活不在前向保存，反向时**重算该段前向**重建局部 tape 再回传。峰值激活从「所有 block 同时在显存」降到「~一个 block + 每 block 的输入」→ dim1024 batch32 装得下。
+
+## Goal
+
+在**不动非重计算路径任何数值**的前提下，新增一个 **opt-in 的 per-block 激活重计算**：
+
+1. **正确性硬闸门（exact）**：重计算是**数学精确**的——同一段前向、同一输入、同一（未变的）参数值、确定性 kernel ⟹ 重算出的激活与原激活逐位相同，回传的梯度与非重计算版一致。直接的 **on-vs-off 梯度对拍**（紧容差）+ 全套回归（T4 grad-check、T5 overfit+PyTorch 对拍、T6 AdamW、T8 DDP loss-match+跨 rank、xserv 闭环）开 `--recompute` 全绿。**绝不提交一个改变梯度的重计算。**
+2. **显存（payoff）**：测 dim768 batch32 峰值显存 on vs off（应降）；确认 **dim1024 batch32 现在装得下**（不开重计算时 OOM，开了 fit）。
+3. **吞吐**：测 tok/s on vs off（多一次前向，预计慢 ~20–35%）——报告 compute/memory 权衡。
+
+## 什么是激活重计算
+
+反向传播需要前向的中间激活（如 SwiGLU 的 `gate`、attention 的 `probs`）来算梯度。define-by-run 的 tape 默认把它们全部留在显存，直到对应 op 的 backward 跑完。模型越深、激活越多，峰值显存越高。
+
+**梯度检查点**把模型切成若干**段（segment）**。前向时，段内 op **不记到 tape**（detached / no-grad），只保留**段的输入**（参数作为 leaf 本就常驻，不算激活）。反向时，当段的 output-grad 到达，**从保存的输入重跑该段前向**（输入作为 require-grad 的 leaf），用上游 grad seed 重算出的 output，在局部 tape 上回传，得到输入梯度（并累加参数梯度），然后释放局部 tape。
+
+代价：每段多一次前向（约 +1/3 的总 FLOPs，因为反向本就 ~2× 前向）。收益：峰值激活从「所有段」降到「~一段 + 每段输入」。
+
+**切粒度 = 每个 transformer block**：一个 block（attention 子块 + MLP 子块 + 两个残差）是天然的段边界，输入/输出都是 `[batch*seq, dim]` 的残差流张量，接口最干净。
+
+## Module Layout（surgical：非重计算路径逐字节不动）
+
+### 1. `xtrain-autodiff::checkpoint` — `checkpoint` 高阶原语
+
+新增 `checkpoint(segment_fn, input, params) -> Var`，类比 `torch.utils.checkpoint`：
+
+- `segment_fn: Fn(&Var, &[Var]) -> Var`——从单个输入 `x` 和参数 slice `p` 构建段前向、返回段输出。必须**确定性**、只依赖 `x` 和 `p`（这是重算精确的前提）。
+- **前向（不 tape 内部）**：把 `input`/`params` detach 成新 leaf（`Var::leaf(v.value())`），跑 `segment_fn` 得到 `out_local`，**只取 `out_local.value()`**。局部 `Var` 出作用域即 drop → 段内激活立即释放。checkpoint 节点的 parents = `[input, ..params]`（参数梯度落进优化器读的同一批 leaf）。
+- **反向（重算）**：闭包捕获 `Rc<segment_fn>`。给定 `dout`，从 `parents` 当前值重建 detached leaf，重跑 `segment_fn` 重建局部 tape，调 `out_local.backward_seeded(dout)`，再把 `x_det.grad()` push 给 `parents[0]`、各 `param_det.grad()` push 给对应参数 parent。闭包结束 → 局部 tape drop → 重算激活释放。
+
+### 2. `xtrain-autodiff::tape` — `backward_seeded`
+
+引擎原 `backward()` 只能从标量 root 出发（seed `ones_like` + 断言 numel==1）。段输出一般**非标量**，故新增 `Var::backward_seeded(seed)`：同样的 topo + 反向遍历，但用显式上游 grad seed（不断言标量）。`backward()` 退化为「seed ones」的薄包装——标量 loss 路径逐字节不变。
+
+### 3. `xtrain-model::TinyTransformer` — per-block 包裹 + `recompute` 开关
+
+- 把 block 前向体抽成**自由函数** `block_forward(cfg, compute_dtype, batch, seq, input, params)`（不借 `&self`，才能在反向闭包里重跑）；同时把 `linear / norm_gamma / attention / swiglu_mlp` 改成参数化 `(cfg, compute_dtype)` 的自由函数。`Block::block_params()` 给出 11 个 leaf 的固定序（与 `params()` 每 block 段一致）。
+- `forward_batched` 的 block 循环：`recompute` 开 → `checkpoint(seg, &h, &b.block_params())`；关 → 直接 `block_forward(...)`（**与之前完全同图**）。
+- `with_recompute(bool)` builder（opt-in，默认关 = 原 tape，数值逐字节同）。
+
+### 4. `xtrain-train` / `xtrain-distributed` — `--recompute` flag
+
+- `bin/train` / `bin/train_ddp` 加 `--recompute`，调 `model.with_recompute(true)`。AdamW / clip / checkpoint / DDP all-reduce **不改**（梯度语义与非重计算一致）。
+
+## Key Design Decisions
+
+- **切粒度 = 每个 block**：接口最干净（输入输出都是残差流 `[B*S, dim]`），且峰值激活降到 ~1 个 block。比「整模型一段」省得多，比「逐 op」简单得多。
+- **参数作为 checkpoint 节点的 parents**：参数 leaf 跨前向/反向不变（只 grad 槽变），重算用**当前**参数值即原值。把它们列为 parents，重算恢复的参数梯度直接 push 到优化器读的同一批 leaf → **DDP/AdamW 零改动**。
+- **detached leaf 隔离局部 tape**：前向/反向都把输入和参数 detach 成新 leaf，使 `segment_fn` 构建的图与外层 tape 断开。前向丢弃局部图（释放激活）；反向局部图回传完即 drop（释放重算激活）。
+- **`backward_seeded` 而非改 `backward`**：段输出非标量，需要用上游 output-grad 作 seed 回传局部 tape。新增方法、原标量 `backward()` 不动。
+- **重算精确 → 梯度逐位一致（硬闸门）**：同一 `segment_fn`、同一输入值、同一参数值、确定性前向 kernel ⟹ 重算 output 与原 output 相同，局部反向就是该段的普通解析反向。故输入/参数梯度与非重计算版一致——这是绝不能违反的闸门。
+
+## 与 bf16 / DDP / batched 的组合
+
+- **bf16（T12）**：`segment_fn` 就是不变的 block 前向，重算跑**同一条 bf16 路径**；`cast` 算子的 grad 升精度桥（bf16→fp32）照常。重计算节点的参数 parents 是 fp32 master leaf，恢复的是 fp32 梯度。on-vs-off 对拍同时跑 fp32 和 bf16 两路。
+- **DDP（T8）**：每个 rank 独立 checkpoint 自己的前向/反向；恢复的参数梯度落进各 rank 的 `.grad()` 槽，再被 all-reduce 取均值——分布式路径不感知重计算。
+- **batched（T10）**：段输入/输出透明带 `[batch*seq, …]` batch 维；`checkpoint` 与形状无关。
+
+## 验证方法
+
+### 1. 正确性（exact，硬闸门）
+
+- **on-vs-off 梯度对拍**（`crates/xtrain-model/tests/recompute.rs`）：同 init 建两个模型（recompute on/off），跑同一 batched loss+backward，断言**前向 logits、loss、每个参数梯度**在紧容差内一致——参数化跑 **fp32 和 bf16** 两路。fp32 期望近逐位（容差 1e-4），bf16 仅放松到 bf16 舍入级（非重计算误差）。
+- **全套回归开 `--recompute`**：T4 15 算子 grad-check、T5 overfit 27/27 + PyTorch 对拍、T6 AdamW、T8 DDP loss-match + 跨 rank、**xserv 闭环 md5**——全绿。
+
+### 2. 显存（payoff）
+
+- dash5 1× RTX 5090 32GB，dim768/18L batch32 seq256，bf16：测峰值显存 recompute on vs off（应降）。
+- **dim1024 batch32**：先验证不开重计算 **OOM**，再验证开了 **fit**——capture 实际 `nvidia-smi` 峰值。
+
+### 3. 吞吐
+
+- 同 config 测 steady-state tok/s recompute on vs off，报告慢多少（预计 ~20–35%，多一次前向）。
+
+## 实测结果（dash5 1× RTX 5090 32GB, bf16, batch32 seq256, steady-state）
+
+**正确性（exact，硬闸门）**：on-vs-off 梯度对拍 —— **fp32 与 bf16 双路都逐位一致**：loss rel `0.00e0`、logits max rel `0.00e0`、**每个参数梯度 max rel `0.00e0`**（不是「在容差内」，是逐位相同——证实重算确实精确）。全套回归开/关重计算全绿：T4 15 算子 grad-check、5 结构、batched、bf16、overfit、AdamW（GPU+host）、GEMM、checkpoint roundtrip、**T8 DDP loss 对单卡 5.67e-7 + 跨 rank 0.0**；DDP+recompute 2 卡短训 loss 单调降（11.079→6.010）。
+
+**显存 + 吞吐**（dim768 = 18L/24h×32/ffn2048 core 127M；dim1024 = 18L/32h×32/ffn2730 core 226M）：
+
+| config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
+|---|---|---|---|---|
+| dim768 recompute **off** | 32 | 31144 MiB | 39.7K | ✅ |
+| **dim768 recompute on** | 32 | **14562 MiB（−53%）** | **31.5K（−20%）** | ✅ |
+| **dim1024** recompute **off** | 32 | 32100 → **OOM** | — | ❌ **OOM** |
+| **dim1024 recompute on** | 32 | **16596 MiB** | 23.1K | ✅ **解 OOM** |
+
+→ dim768：重计算把峰值显存 **31.1→14.6GB（−53%，~砍半激活）**，代价 tok/s **−20%**（多一次前向，落在预测 20–35% 区间）。dim1024 batch32：**不开 OOM（撞 32100/32607MiB 上限）→ 开了 16.6GB 稳稳装下**，~23K tok/s 训练正常收敛 —— **KI-3 的目标达成，dim1024 解锁**。

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

@@ -0,0 +1,160 @@
+# Phase T18: Dropout（device RNG + mask）— Design Document
+
+## Goal
+
+在已有的 tape autograd 引擎（T4）+ tiny transformer（T5）之上，**手写一个 dropout 算子**：
+训练时按 Bernoulli(keep = 1−p) 生成一个 0/1 mask，丢弃的元素置 0、保留的元素按
+**inverted dropout** 乘 `1/(1−p)`（让训练期望与推理一致）；推理（eval）时 dropout 是**恒等**。
+新增一个 autodiff `dropout` 节点：**前向生成并施加 mask，反向施加同一个 mask**。
+接到模型的标准位置（residual 之前的 attention / MLP 子块输出；attention-probs dropout 不做，见下）。
+通过 `Config.dropout` / `--dropout` 暴露 `p`，**默认 `p=0`**。
+
+明确范围（T18 只做这些）：
+
+1. 一个 device 端 **counter-based RNG**（Philox 风格的 bit-mix），按 `(seed, 元素下标)` 无状态地产出
+   每元素的 Bernoulli 抽样 → 0/1 mask（保留=1，丢弃=0），同 seed **逐位可复现**。
+2. 一个 `dropout` autodiff 节点（fwd 生成 mask + 施加 inverted scaling；bwd 用**缓存的同一 mask**）。
+3. 模型里加 **training / eval 开关**：train 走 dropout、eval/采样/导出走恒等。
+4. `p` 经 `Config.dropout` 落地，`bin/train` 加 `--dropout` flag。
+
+明确**不做**：attention-probs（softmax 后）dropout——本项目 attention 是**一个 fused batched SDPA 算子**
+（`ops::attention`，softmax 在 kernel 内部不物化 probs 给外部施加 mask），在其上插 dropout 要么改 fused kernel、
+要么退回组合路径，**不值当**且偏离「标准 residual/ffn dropout」这条主线。文档明确记下「只做 residual-path dropout」。
+
+## Module Layout
+
+```
+csrc/ops/dropout.cu                    # 新：counter-based RNG mask 生成 + 施加 (fwd) / 反向施加同 mask
+                                       #     fp32 + bf16 两条（activation 流可能是 bf16，对齐 cast.cu 风格）
+
+crates/xtrain-cuda/
+├── build.rs                           # 新增 dropout.cu
+└── src/ffi.rs                         # 新增 launch_dropout_{f32,bf16} 声明（no_cuda 门控）
+
+crates/xtrain-tensor/
+└── src/tensor.rs                      # 新增 Tensor::dropout_mask_apply(p, seed) -> (out, mask)
+                                       #         Tensor::dropout_apply_mask(&mask) -> out（bwd 用）
+
+crates/xtrain-autodiff/
+├── src/ops.rs                         # 新增节点 dropout(x, p, seed)（p==0 提前返回 x.clone()，零节点）
+└── tests/autograd.rs                  # 新增：固定 seed grad-check（mask 跨 ± 扰动固定）+ 期望保持数值检查
+
+crates/xtrain-model/
+├── src/config.rs                      # Config 加 dropout: f32（默认 0）
+├── src/model.rs                       # train/eval 开关（Cell<bool>）+ 在 attn/ffn 子块输出接 dropout；
+│                                       #   per-site 确定性 seed（与 checkpoint recompute 兼容）
+└── tests/dropout.rs                   # 新增：p=0 逐位一致 / eval 恒等 / 期望保持 / p>0 小训练收敛
+
+crates/xtrain-train/src/bin/train.rs   # --dropout flag → Config.dropout；训练 model.train()，sample 前 model.eval()
+```
+
+为什么 RNG/mask 落在 `tensor.rs`（而非引擎）：和 `scale`/`silu` 一样是一个 device kernel 的薄封装；
+autodiff 层只负责把它包成带 backward 的 `Var` 节点（对齐 T4 既有分层）。
+
+## Key Design Decisions
+
+### RNG：counter-based（Philox 风格），无状态、可复现、与重计算兼容
+
+mask[i] 只由 `(seed, i)` 决定，**不读取任何可变 RNG 状态**：
+
+```
+key  = seed XOR (i * 0x9E3779B97F4A7C15)        // golden-ratio 常数打散下标
+h    = splitmix64(key)                            // 几轮 bit-mix（xorshift+乘法）
+u    = (h >> 40) as f32 / 2^24                    // [0,1) 均匀
+keep = u >= p                                     // Bernoulli(keep = 1−p)
+out[i] = keep ? x[i] * (1/(1−p)) : 0
+```
+
+选 counter-based 而非「per-step 推进一个全局 LCG 状态」的关键原因 = **激活重计算（T13）**：
+checkpoint 的 segment 在 backward 时会**重跑一遍 forward**（`segment_fn` 再执行）。
+若 dropout 用「调用时推进的可变状态」，重跑会拿到**不同的 mask** → 梯度与前向用的 mask 不一致 → 错。
+counter-based + **每个 dropout 站点一个确定性 seed**（见下）保证：重跑同 seed → **同 mask**，
+重计算依旧逐位一致（T13 的硬闸门不被 dropout 破坏）。
+
+> 复现性：同一 `(seed, p, shape)` 下 mask 逐位确定；fp32/bf16 mask 判定都在 fp32 里算 `u`（bf16 仅存/取
+> activation），所以两精度的 mask **同分布**（drop 与否由 fp32 `u` 决定，不受 bf16 舍入影响）。
+
+### 每个 dropout 站点的确定性 seed（兼容 checkpoint 重算）
+
+模型持有一个 `base_seed`（`Cell<u64>`，每个训练 step 自增一次 → 每步换 mask）。`block_forward`
+收到 `block_seed = base_seed XOR layer_index`，块内两处 dropout 再各 XOR 一个站点常量
+（attn=0xA77, ffn=0xF7N）派生出**该站点的 seed**。这些都是**纯函数**（只看 `base_seed + layer_index +
+站点常量`，无可变推进），所以：
+
+- 同一 step 内不同站点 mask 不同（seed 不同）；
+- checkpoint 重算 `block_forward` 时，`block_seed` 由捕获的 `base_seed`/`layer_index` 重新算出 → **同 seed → 同 mask**；
+- 跨 step mask 变化（`base_seed` 每步 +1）。
+
+`base_seed` 的自增放在**训练入口**（`loss_batched` 训练态调用时 advance 一次）。eval/`forward`/采样
+**不 advance、不插 dropout**（恒等）。
+
+### train / eval 开关
+
+`TinyTransformer` 加一个 `Cell<bool> training`（默认 **false** = eval，安全：未显式开训练就不丢弃）：
+
+- `model.train()` / `model.eval()` 切换（builder 风格 `with_training(bool)` 也提供，给测试）。
+- `forward_batched` 里：`p > 0 && training` 才在 attn/ffn 子块输出插 `ops::dropout`；否则**完全不建 dropout 节点**。
+- 因此 **`p == 0`** 或 **eval** → forward 图与改动前**逐字节相同**（`ops::dropout` 在 `p==0` 时也提前
+  `return x.clone()`，双保险）→ 满足「p=0 与无 dropout 逐位一致」回归闸门。
+
+训练 loop（`train`）开 `model.train()`；`eval_loss` / `generate` / 导出 `forward` 走 eval（恒等）——
+导出的模型权重不含任何 dropout，xserv 闭环不受影响。
+
+### dropout 接在哪（wiring）
+
+接**两处 residual-path dropout**（标准 Pre-LN transformer 位置，对齐 GPT/LLaMA 训练实践）：
+
+```
+h = h + dropout( attention(rms_norm(h)) )      # attn 子块输出，残差前
+h = h + dropout( swiglu_mlp(rms_norm(h)) )     # ffn  子块输出，残差前
+```
+
+**不做** attention-probs dropout（理由见 Goal：fused SDPA 不物化 probs）。embedding dropout 也不做（非必需）。
+
+### dropout 节点的 backward（为什么 grad-check 成立）
+
+```
+fwd:  out = x ⊙ mask ⊙ (1/(1−p))          # mask 由 seed 生成，缓存进 backward 闭包
+bwd:  dx  = d ⊙ mask ⊙ (1/(1−p))          # 用同一个缓存 mask
+```
+
+dropout 在 **固定 mask** 下是一个逐元素线性映射 `out_i = c_i · x_i`（`c_i ∈ {0, 1/(1−p)}`），
+其梯度就是 `dx_i = c_i · d_i`。finite-diff grad-check 之所以成立，关键是**前向缓存的 mask 在 ± 扰动两次
+forward 里保持不变**——本设计天然满足：mask 只由 `(seed, i)` 决定，与 `x` 的值无关，扰动 `x` 不改 mask。
+（grad-check 直接对 `ops::dropout` 节点跑：同一个 `seed` 调两次 forward 得到同一 mask，函数处处可微。）
+
+### 与既有特性的组合
+
+- **bf16（T12）**：activation 流是 bf16 时，dropout kernel 走 bf16 分支（load→fp32 判 mask→store bf16），
+  mask 判定在 fp32，和 cast.cu 既有 bf16 elementwise 同风格；grad 也在 activation dtype（接回 bf16 链）。
+- **重计算（T13）**：见上「counter-based + 确定性 seed」——重算 mask 与前向逐位相同，T13 闸门不破。
+- **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）**：本分支独立于二者，不依赖其未合并改动。
+
+## 验证方法
+
+全部 `#![cfg(not(no_cuda))]` 门控；本地只 `cargo check`/`fmt`，构建 + 实跑在 dash5（8× RTX 5090, sm_120）。
+
+**硬闸门（全绿，诚实正确性，不放宽容差）**：
+
+1. **固定 seed grad-check**（`autograd.rs::dropout_bwd`）：对 `ops::dropout(x, p, seed)` 同一 seed
+   跑 finite-diff（mask 跨 ± 扰动固定）→ `dx` 对中心差分通过（线性 op，用 `cfg_linear` 容差）。
+2. **train/eval + 期望保持**（`dropout.rs`）：
+   - eval 恒等：`dropout` 关时 `out == x` **逐位**；
+   - 期望保持：大张量、训练态、对多组随机 mask 取均值，`E[out] ≈ x`（inverted scaling 正确），给数值；
+   - 实际 keep 比例 ≈ `1−p`（验证 RNG 分布）。
+3. **p=0 逐位一致**（`dropout.rs`）：同 init 两个模型，一个不设 dropout、一个 `dropout=0`，
+   同 batch forward+backward → **logits/loss/每参数 grad 逐位相同**（`|Δ| == 0`）。
+4. **p>0 小训练收敛**（`dropout.rs`，或 dash5 短跑）：小模型开 `p=0.1` 训若干步，**loss 下降、无 NaN**。
+5. **全回归套绿**：autograd grad-checks、structural、batched==looped、bf16、recompute（逐位一致）、
+   overfit 27/27、AdamW（GPU bit-exact + host vs torch）、DDP（loss-match + 跨 rank）、
+   **xserv 闭环**（导出 md5 vs registry、token-identical；导出/推理 dropout **关**，导出模型不受影响）。
+
+dash5 capture 每个闸门的 pass + 关键数字（max rel-err、期望 vs input、p=0 的 `|Δ|`、训练 loss 轨迹）。

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

@@ -0,0 +1,117 @@
+# xtrain 演进总览 — 按维度记录每次变化
+
+每个里程碑（**T# 基建 phase** 或 **v# 训练 run**）在四个维度上分别改了什么、结果如何。
+这是活文档：**每次新 run 收尾时追加一行/一段**。细节见各 `docs/runs/0N-*.md`、各 phase 设计文档、`docs/known-issues.md`。
+
+四个维度：**算法**（autograd/优化器/精度/反向数学）· **模型架构**（dim/层/头/算子）· **Infra**（构建/显存/并行/吞吐）· **数据集**（语料/token/epoch/tokenizer）。
+
+---
+
+## 一、基建 phase（T1–T13 + Phase 2 systems-depth）—— 主要动「算法」与「Infra」
+
+| Phase | 维度 | 变化 | 结果 / 验证 |
+|---|---|---|---|
+| T1 | Infra | Rust↔CUDA FFI 构建链（build.rs+nvcc, `no_cuda` cfg），gitea↔dash5 流 | vector-add 跑通 |
+| T2 | Infra | Tensor 抽象（dtype/shape/Storage, H↔D 拷贝）+ elementwise kernel | roundtrip 保真 |
+| T3 | 算法 | 手写 tiled GEMM fwd/bwd + finite-diff 梯度检查 harness | fwd vs cuBLAS 1e-7；bwd vs finite-diff |
+| T4 | 算法 | tape autograd 引擎 + 11 算子 backward（含梯度扇出累加）；attention 由 matmul+softmax 组合 | 每算子 finite-diff ≤2e-2 |
+| T5 | 模型架构 | 组装 tiny decoder（RoPE+RMSNorm+SwiGLU）+ embedding/reshape/transpose 算子 | overfit 27/27 + PyTorch 对拍 B>1 |
+| T6 | 算法 + 数据集 | 手写 AdamW + 训练 loop + LR sched + grad clip + checkpoint；gpt2 BPE + TinyStories | 真训出连贯英文 |
+| T7 | Infra | cuBLAS matmul + GPU 端 AdamW/grad-norm + 去 per-op sync | **~3×**：2.7K→8.5K tok/s，零回归 |
+| T8 | Infra | NCCL DDP（单进程 thread-per-GPU）+ 梯度 all-reduce | 多卡（当时弱扩展 ~1.4×） |
+| T9 | 算法/模型架构 + Infra | **加 per-head QK-norm**（Qwen3 兼容）；safetensors 导出 | xserv 闭环：贪心**逐 token 一致** |
+| T10 | 算法/Infra | **batched 多序列 forward**（linears flatten `[B·S,dim]` + fused batched SDPA + 每序列 RoPE） | **单卡 15–24×**；MFU 0.4%→14%（修 KI-1） |
+| 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–v10）—— 主要动「模型架构」与「数据集」
+
+架构始终是 **Qwen3-style**（RoPE + RMSNorm + QK-norm + SwiGLU，gpt2 50257 词表），逐版放大 dim/层/头（v8 起首次拨容量轴到 dim1024，v9 进入 dim1280+真 GQA 双轴点，v10 固定架构只补数据轴）；其余维度逐版变化如下：
+
+| ver | 模型架构（dim/层/头·hd · 核心/总参） | 数据集（语料 · 实训 token · epoch） | 算法/精度 | Infra（GPU · 吞吐） | 结果（val · 备注） |
+|---|---|---|---|---|---|
+| v0 | dim32/4L/2h · 41K/3.26M | TinyStories 3MB 切片 · ~0.72M · — | fp32 单序列 | 1 GPU | val 3.80（toy，不可用） |
+| v1 | dim256/8L/8h · 8.4M/34M | TinyStories 全量 · 5.1M · 0.01ep | fp32 单序列 | 1 GPU · 3.3K | val 2.58 |
+| v2 | dim384/12L/12h · 28M/67M | TinyStories · 37M · 0.08ep | fp32 单序列 | 4 GPU DDP · 3.6K | val 1.71（暴露 KI-1/弱扩展） |
+| v3 | dim512/16L/16h · 67M/119M | TinyStories · 246M · 0.53ep | fp32 **batched(T10)** | 1 GPU · 26K | val 1.30 |
+| v4 | dim768/18L/24h · 127M/205M | TinyStories · 721M · 1.54ep | fp32 batched | **8 GPU(T11)** · 145K | val 1.17（仍欠拟合） |
+| v5 | dim768/18L（**同 v4**） | TinyStories · 2.49B · **5.33ep** | **bf16(T12)** | 8 GPU · 217K | val **1.11**：⚠️**TinyStories 饱和**（3.5×数据仅↓5%） |
+| 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（非数据集大小）。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) → 融合 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 是预期而非回退。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%**，与数据轴单步杠杆同量级。
+  - ✅ **双轴一起 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 用上）。
+- **实测关闭（负结果）**：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

@@ -7,12 +7,80 @@
 
 ## Open
 
-_(KI-1 fixed in T10. KI-5 **FIXED** in T11——device caching/pool allocator 消掉 per-op cudaMalloc 串行，单卡 ~2.3×、DDP scaling 从 ~1.3× 封顶恢复到 ~5×@8。见下方 Fixed。)_
+_(KI-1 fixed in T10. KI-5 fixed in T11. KI-2 fixed in T12. **KI-3（激活重计算 / gradient checkpointing）FIXED in T13**——per-block checkpoint（no-tape forward + recompute-on-backward），梯度对非重计算版**逐位一致**（fp32/bf16 max rel 0.00e0）；dim768 bf16 batch32 峰值显存 31.1→14.6GB（−53%）/ tok/s −20%；**dim1024 batch32 不开 OOM → 开了 16.6GB 装得下**，解锁 v8。见下方 Fixed。)_
 
 ---
 
 ## 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。
+- **正确性（exact，硬闸门全绿）**：重计算数学精确（同 `segment_fn`、同输入、同参数值、确定性 kernel → 重算激活逐位等于原激活）。**on-vs-off 梯度对拍 fp32/bf16 双路逐位一致**：loss rel `0.00e0`、logits max rel `0.00e0`、**每个参数梯度 max rel `0.00e0`**（不是容差内，是逐位）。全套回归开/关重计算全绿：autograd 15、structural 5、batched、bf16、overfit 27/27、AdamW（GPU bit-exact + host 对 torch）、checkpoint roundtrip、**DDP loss 对单卡 5.67e-7 + 跨 rank 0.0**；DDP+recompute 2 卡短训单调降（11.079→6.010）。非重计算路径图不变（默认关）→ T10/T11/T12 数值不回归。
+- **显存 + 吞吐（payoff，dash5 1× RTX 5090 32GB, bf16, batch32 seq256, steady-state）**：
+
+  | config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
+  |---|---|---|---|---|
+  | dim768 (18L/24h ffn2048, core 127M) off | 32 | 31144 MiB | 39.7K | ✅ |
+  | **dim768 on** | 32 | **14562 MiB（−53%）** | **31.5K（−20%）** | ✅ |
+  | **dim1024** (18L/32h ffn2730, core 226M) off | 32 | 32100 → **OOM** | — | ❌ **OOM** |
+  | **dim1024 on** | 32 | **16596 MiB** | 23.1K | ✅ **解 OOM** |
+
+  → dim768：重计算砍 ~半激活（**31.1→14.6GB，−53%**），代价 tok/s **−20%**（多一次前向，落在预测 20–35%）。**dim1024 batch32：不开 OOM → 开了 16.6GB 稳稳装下**，~23K tok/s 正常收敛 → **dim1024 解锁，v8 可展开**。
+- **commit**：T13 提交链（`autodiff: checkpoint primitive (recompute-on-backward)` / `model: per-block activation recompute (--recompute)` / `perf: KI-3 fixed …` 本条带 before/after / 文档 `docs: Phase T13 — activation recompute`）。
+
+---
+
+### KI-2 · bf16 混合精度（fp32 master）— `FIXED` (T12)
+- **触发点（v4 surfaced）**：dim768 fp32 在单卡 32GB 里 per-rank batch 32（global 256）OOM，被迫降到 per-rank 16。bf16（激活减半）找回 batch-32 甜点区，并加速已 compute-bound 的 dim768 GEMM；附带：xserv 推理 BF16-only，bf16 训练更贴闭环。
+- **设计（标准 AMP，opt-in，[docs/11-bf16-mixed-precision.md](11-bf16-mixed-precision.md)）**：**fp32 master weights** + AdamW/clip/DDP all-reduce 全程 fp32；**bf16 compute**=linears(q/k/v/o, gate/up/down, lm_head)走 `cublasGemmEx`（CUDA_R_16BF in/out, CUBLAS_COMPUTE_32F 累加）+ 激活流 bf16（含 attention probs / logits）；**fp32 稳定**=RMSNorm/QK-norm、softmax、RoPE、cross-entropy 内部 upcast→fp32→downcast。**无 loss scaling**（bf16 8-bit 指数=fp32 动态范围）。关键钩子=autodiff `cast` 算子：前向把 fp32 master leaf 降到 bf16 喂 matmul，**反向把 bf16 grad 升回 fp32** → fp32 leaf 累加 fp32 grad，优化器一行不改。fp32 路径按 dtype 分派、逐字节不变（hard gate）。
+- **正确性（双闸门全绿）**：
+  - **fp32 不回归**：全套在原紧容差绿——autograd 15、structural 5、GEMM 对 cuBLAS 5、batched==looped、overfit 27/27、AdamW GPU bit-exact + host 对 torch、checkpoint 逐位、DDP loss 对单卡 + 跨 rank、**xserv 闭环**（v3 ckpt 用 T12 代码重导 safetensors 与 registry **md5 逐位一致** `b04fc9f9…`）。
+  - **bf16 looser-tol**：同 fp32 master 跑 fp32 vs bf16——loss rel **1.2e-4**、logits mean rel **2.0e-3** / p99 **6.8e-3**、grad worst scaled-mean **1.0e-2**，无 NaN，grad 仍 fp32。
+  - **收敛**：dim768 短训 150 步，bf16-b16 loss 轨迹对住 fp32-b16（step50 4.40 vs 4.40、step149 **3.984 vs 3.988**），单调降、无发散。
+- **显存 + 吞吐（payoff，dash5 1× RTX 5090 32GB, dim768/18L/24h×32 ffn2048 seq256, steady-state）**：
+
+  | config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
+  |---|---|---|---|---|
+  | fp32 | 16 (v4 fallback) | 27.2 GB | 31.5K | ✅ |
+  | **bf16** | 16 | **19.3 GB（−29%）** | **35.5K（+13%）** | ✅ |
+  | fp32 | 32 | — | — | ❌ **OOM** |
+  | **bf16** | **32（甜点区）** | **31.1 GB** | **40.8K** | ✅ **解 OOM** |
+
+  → **同 batch 16：bf16 显存 27.2→19.3GB（−29%）、tok/s 31.5K→35.5K（+13%）**；**bf16 解 fp32-batch32 OOM**（31.1/32.6GB fit），batch32 达 **40.8K tok/s（+29% vs fp32-b16）**。残留：norm/softmax/CE 的 fp32 upcast 是 transient，但仍占峰值——若 v5 要更大 batch，下一杠杆是 KI-3 激活重计算。
+- **commit**：见 T12 提交链（`cuda: bf16 cuBLAS GemmEx` / `autodiff: bf16 mixed-precision path` / `train: --bf16 flag` / `perf: keep bf16 logits` / 本条）。
+
+---
+
 ### KI-5 · DDP 弱扩展性 — `FIXED` (T11, device caching/pool allocator)
 - **根因（T11 重诊断，all-reduce **不是**瓶颈）**：每个 tape op 输出走 `Tensor::zeros`→`GpuBuffer::alloc`→`cudaMalloc`（同步、进程级串行的 driver 调用）。单进程 thread-per-GPU 下 N rank 每步几百次 alloc 在单 CUDA context 排队串行（`NOCOMM=1` 完全不通信时 fwd+bwd 仍 136→780ms 膨胀 ~6×，`nvidia-smi` 抽样 8 卡只 1–2 张在忙轮流跑）；单卡也吃这笔 per-op alloc。
 - **原拟修复「分桶 all-reduce」经 T11 实测证伪并 revert**：grad all-reduce 每步只占 ~6–7%，融成一发对 1/2/4/8 卡几乎无差（详见下方历史诊断）。
@@ -29,7 +97,7 @@ _(KI-1 fixed in T10. KI-5 **FIXED** in T11——device caching/pool allocator
   → **单卡 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」的过程）：
 
@@ -91,18 +159,26 @@ _(KI-1 fixed in T10. KI-5 **FIXED** in T11——device caching/pool allocator
 
 ## Deferred（来自 T7，放大后重启）
 
-### KI-2 · bf16 混合精度（fp32 master）— `deferred`
+_(KI-3 已在 T13 FIXED，见上方 Fixed。本节暂无待启项。)_
-- T7 延后理由：tiny 规模延迟瓶颈、bf16 改变数值会威胁 fp32 正确性闸门。
-- **重启条件**：模型放大（v2+ `dim≥384`）后 GEMM 渐成 compute-bound，tensor-core 收益显现。需 fp32 master weights + 单独 looser-tol 测试 + 收敛对比。
-
-### KI-3 · 激活重计算（gradient checkpointing）— `deferred`
-- T7 延后理由：单序列、显存不紧。
-- **重启条件**：更大模型 / 更长 seq / 更大 batch 后显存成约束。
 
 ---
 
 ## 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/04-v4-tinystories-dim768.md

@@ -0,0 +1,201 @@
+# Scaling Run v4: TinyStories + dim768/18L + 8 卡 DDP fp32(T11) — Design Document
+
+## Goal
+
+在 v3（dim512/16L、core 67.13M、训 ~245.8M token、单卡 batched）之上，沿**模型 + 数据**两个轴继续
+放大，并把训练**重回多卡**——这次多卡不是 v2 时代被 KI-1 压住扩展性的 DDP，而是 **T11 缓存分配器**落地、
+8 卡近线性之后的正确选择：
+
+1. **模型放大**：dim 512→768、层 16→18、头 16→24（head_dim 仍 32），把 **transformer core 做到 ~127M
+   参**（容量 ×1.9），词表不变 → embed+lm_head 因 gpt2 50257 vocab 在 dim768 下固定加 ~77.19M，单列出来。
+2. **数据放大**：v3 训了 ~245.8M token（~0.53 epoch，仍欠拟合，val 一路降到末步）；v4 训 **720.9M
+   token**（×2.9），仍复用 v1 缓存的全量 TinyStories token-id 流（468M token），**~1.54 epoch**——首次
+   越过 1 epoch、开始进入 TinyStories 多遍区。
+3. **8 卡 DDP fp32（T11，多卡近线性）**：v2 暴露的 KI-1（DDP 弱扩展）根因被 T10/T11 逐层证伪并修复——
+   T10 修了单序列 launch-bound，T11 的 per-device size-classed caching allocator 进一步把 per-op
+   `cudaMalloc` 串行消掉，**8 卡 49K→461K tok/s（scaling 1.3×→5×，全 8 卡 95-99% util）**。v4 因此放心回多卡，
+   全程稳态 ~144,650 tok/s、~84 min 训完 720.9M token。
+4. 训完存 registry（`~/projects/tiny-models/v4-tinystories-dim768/`）+ 导出 xserv 格式验证可服务，给出
+   **相比 v3 的具体提升**（同一保留集 val loss + 同 prompt 并排采样）。
+
+> **这一版的工程意义**：在真实 scaling 规模（127M core / 720.9M token / 84 min）**验证了 T11 缓存分配器
+> 在 dim768 的多卡扩展性**——8 卡全程 ~145K tok/s、95-99% util，比 v2 时代 4 卡 DDP（~3.6K tok/s）快
+> ~40×。同时 v4 是 **bf16（KI-2）的具体触发点**：dim768 fp32 在 32GB 显存里 per-rank batch 32（global 256）
+> OOM，被迫降到 per-rank 16（global 128）——这是 v0–v3 tiny 规模一直延后 bf16 后，第一次有 fp32 放不下的
+> 硬约束。
+
+## 数据
+
+| 项 | v3 | v4 |
+|----|----|----|
+| 来源 | TinyStories **全量 train**（复用 v1 缓存）| 同 |
+| token 数（语料）| 468,260,367 | 同 |
+| **训练消费 token** | ~245.8M（30000 步 × 8192）| **~720.9M**（22000 步 × 32768）|
+| epoch 占比 | ~0.53 | **~1.54**（首次越过 1 epoch）|
+| tokenizer | gpt2 BPE（vocab 50257）| 同 |
+| 缓存 | `data/tinystories-train.txt.u16.bin`（u16，936MB）| **直接复用** |
+| held-out val | 全量末尾 1,000,000 token | **同一 1M token**（与 v0/v1/v2/v3 完全相同的保留集，公平对比）|
+
+**复用缓存**：`Corpus::load_cached` 读 `<corpus>.u16.bin`，启动即载入 467.26M train token（末尾 1M 留 val）。
+held-out val 仍是全量末尾 1M token（`split_tail`），与 v0–v3 同一保留集——**v0–v4 的 val loss 直接可比**。
+
+**数据阶梯**：v4 是**首次越过 1 epoch**（~1.54）。core 已到 127M、~1.54 epoch 仍欠拟合（val 末步还在降），
+说明 TinyStories 这本语料对 127M core 尚未到容量上限。下一档（v5）是开始**广化语料**（TinyStories + 部分
+通用高质语料）或**继续榨 TinyStories 多 epoch** 的合适节点，同步评估 tokenizer（KI-4）。
+
+## 架构
+
+v4 = 更大、同构的 tiny Qwen3（RoPE + RMSNorm + per-head QK-norm + SwiGLU + 独立 lm_head，MHA），
+forward 图与 v0/v1/v2/v3 完全同构，只是 dims 变大。**无结构改动**。
+
+| 维度 | v3 | v4 |
+|------|----|----|
+| dim（= heads·head_dim）| 512 | **768** |
+| n_layers | 16 | **18** |
+| n_heads | 16 | **24** |
+| head_dim | 32 | 32 |
+| ffn_hidden（SwiGLU）| 2048 | 2048 |
+| vocab | 50257 | 50257 |
+| **core 参数**（除 embed+lm_head）| 67,127,296（≈67.13M）| **127,432,704（≈127.43M，×1.90）** |
+| embed + lm_head（2×vocab×dim）| 51,463,168（≈51.46M）| 77,194,752（≈77.19M）|
+| **总参数** | 118,590,464（≈118.59M）| **204,627,456（≈204.63M）** |
+
+**core 的量法**：`Config::core_params() = num_params() − 2·vocab·dim`。gpt2 50257 vocab 在 dim768 下让
+embedding + lm_head 固定占 ~77.19M——这两张表是**词表大小**的函数、不是模型容量，所以阶梯按 **core** 量
+（v4 core 127.43M）。注意：v4 总参 204.63M 里 embed/lm_head 仍占 ~38%（77.19M），比 v3 的 43% 略降
+（dim 越大占比越摊薄），但仍是 gpt2 大词表占比问题（见 docs/known-issues.md KI-4）。
+
+**相比 v3 的架构变化**：纯放大（dim 512→768 / 层 16→18 / 头 16→24，head_dim 与 ffn 不变），无结构改动。
+阶梯已参数化，v4 只改 `--dim/--heads/--layers/--ffn/--steps` flag，不动模型代码。
+
+## 训练器：8 卡 DDP fp32（T11 缓存分配器加持）
+
+v2 用 DDP（T8）4 卡，因 global_batch=32 太小被 KI-1（all-reduce 占比过高）压住扩展性。T10/T11 排查后把
+KI-1 的前提逐层证伪并修掉：
+
+- **T10（batched forward）**：v2 时代单卡慢的真因不是通信，而是单序列 forward 每个 op 各自 launch
+  （util 0-15%）。flatten linears + fused batched causal SDPA → 单卡 1653→25627 tok/s。
+- **T11（caching allocator）**：profile 证伪「分桶 all-reduce」（只占 7%）→ 真因是 per-op `cudaMalloc`
+  串行。per-device size-classed caching allocator（Drop 归还、线程安全）→ **单卡 40K→93K tok/s（2.3×）、
+  8 卡 49K→461K tok/s（9.4×，scaling 1.3×→5×，全 8 卡 95-99% util）**。
+
+v4 因此放心回 8 卡 DDP fp32：thread-per-GPU、all-reduce device 梯度取均值后各 rank 本地 GpuAdamW step，
+跨 rank 参数 bit-identical。全程稳态 **~144,650 tok/s**、~84 min 训完 720.9M token，比 v2 时代 4 卡 DDP
+（~3.6K tok/s）快 ~40×。
+
+⚠️ **batch 约束（bf16 触发点）**：dim768 fp32 在单卡 32GB 显存里 **per-rank batch 32（global 256）OOM**，
+被迫降到 **per-rank 16（global 128）**。这是 v0–v3 tiny 规模一直把 bf16（KI-2）延后后，第一次有 fp32
+放不下的硬约束——bf16（激活减半）能把 batch-256 的甜点区找回来。已回填 docs/known-issues.md KI-2 触发点。
+
+## 超参
+
+| 项 | 值 | 备注 |
+|----|----|----|
+| optimizer | 手写 AdamW（GPU 端 step）| wd=0.1，β/eps 用 xtrain-optim 默认 |
+| LR schedule | 线性 warmup → cosine decay | max_lr **6e-4** → min_lr **6e-5**（同 v1/v2/v3）|
+| warmup | **1100 步**（steps/20，lr 在 step 1100 达峰 6.00e-4）| |
+| grad clip | global-norm 1.0 | gnorm 全程 ~0.20–0.21，平稳 |
+| steps | **22000** | ~84 min @ 8 卡 |
+| global batch | **128**（per-rank 16 × world 8）| **8 卡 DDP**；per-rank 32 会 OOM（见上）|
+| seq_len | **256** | 同 v2/v3 |
+| tokens/step | 128×256 = 32768 | 总训练 token ≈ **720.9M**（~1.54 epoch）|
+| world size | **8**（RTX 5090，sm_120）| T11 修 KI-5 后多卡近线性 |
+| 精度 | f32（训练）| 导出 xserv 时转 BF16（见 T9）|
+
+**算力**：dash5 8× RTX 5090，全程 ~144,650 tok/s（启动即 ~13万、step 50 起稳态 ~14.5万），wall-clock
+≈ **84 分钟**。
+
+## 结果
+
+- **train loss**：start **11.0689** → end **~1.14**（末批 1.1432；全程平稳下降）
+- **best / final val loss（held-out 1M token，step 21999）**：**1.1690**（接近 ~1.0-1.1 目标）
+- val loss 曲线（每 ~2000 步抽样，单调下降、末步仍在降、**无过拟合**）：
+
+| step | 499 | 1999 | 3999 | 5999 | 7999 | 9999 | 11999 | 13999 | 15999 | 17999 | 19999 | 21999 |
+|------|-----|------|------|------|------|------|-------|-------|-------|-------|-------|-------|
+| val  | 2.5217 | 1.6493 | 1.4875 | 1.4056 | 1.3571 | 1.3161 | 1.2697 | 1.2414 | 1.2177 | 1.1978 | 1.1762 | **1.1690** |
+
+val 一路降到末步、无回升 = 仍**欠拟合**，更多步数/数据（或更大模型）还能继续降（v5 杠杆）。
+
+### 采样（greedy，xtrain 直采，同 prompt）
+
+```
+[Once upon a time] → Once upon a time, there was a little girl named Lily. She loved to play
+                     outside in the sunshine. One day, she saw a big, scary dog. The dog barked
+                     loudly and Lily got scared. She
+[One day]          → One day, a little girl named Lily went to the park with her mom. She saw a
+                     big tree with a swing. Lily wanted to play on the swing, but she was too
+                     small. She asked her
+[The little]       → The little girl was so happy that she had found the perfect place to hide.
+                     She stayed there for a long time, until it was time to go home. She said
+                     goodbye to the tree and ran back home
+```
+
+## 相比 v3 的提升
+
+**best val loss（各版各自训练 run 报告的 held-out 1M token 最优值，同一保留集）**：
+
+| 模型 | core 参数 | 训练 token | **best val loss** | 说明 |
+|------|-----------|-----------|-------------------|------|
+| v0-baseline | 41K | ~0.72M | 3.8050 | 3MB 切片，采样退化循环 |
+| v1 | 8.39M | ~5.1M | 2.5847 | 全量数据 + dim256/8L，单卡 |
+| v2 | 28.32M | ~36.9M | 1.7055 | dim384/12L + DDP 4 卡 |
+| v3 | 67.13M | ~245.8M（~0.53 ep）| 1.3027 | dim512/16L + 单卡 batched，val 比 v2 低 0.40 |
+| v4 | 127.43M（**×1.90** vs v3）| ~720.9M（**×2.9** vs v3，~1.54 ep）| **1.1690** | dim768/18L + **8 卡 DDP fp32**，val 比 v3 低 **0.13** |
+
+**完整 val 阶梯：v0 3.80 / v1 2.58 / v2 1.71 / v3 1.30 / v4 1.17**——每一档都在同一 1M token 保留集上
+单调下降。注意从 v3→v4 的 val 降幅（0.13）小于 v2→v3（0.40）：边际收益递减是预期的（loss 越低越难再降），
+且 v4 仍欠拟合（末步还在降），说明 127M core 在 TinyStories 上尚未到容量上限——更多 token / 更广语料还有空间。
+
+### 并排采样（greedy 40 tok，xserv 服务，同 prompt）
+
+| prompt | v3 | v4 |
+|--------|----|----|
+| `Once upon a time` | …a little girl named Lily. She loved to play outside in the **park. One day, she saw a big, scary dog. The dog barked loudly and scared her. She ran** | …a little girl named Lily. She loved to play outside in the **sunshine. One day, she saw a big, scary dog. The dog barked loudly and Lily got scared. She** |
+| `One day` | One day, a little girl named Lily went to the park with her mom. **They saw a big tree with a swing. Lily wanted to play on the swing, but she was scared. Her mom said,** | One day, a little girl named Lily went to the park with her mom. **She saw a big tree with a swing. Lily wanted to play on the swing, but she was too small. She asked her** |
+| `The little` | The little girl was so **excited. She ran to the kitchen and grabbed a spoon. She started to stir the soup. She stirred and stirred until it was all mixed together.** | The little girl was so **happy that she had found the perfect place to hide. She stayed there for a long time, until it was time to go home. She said goodbye to the tree and ran back home** |
+
+**结论**：v3（67M core / 245.8M token）已能写带动机/转折的连续叙事；v4（127M core / 720.9M token /
+~1.54 epoch）在**相同开头**下情节更具体、动机更细（"too small" 而非泛泛 "scared"、"perfect place to
+hide → stayed → said goodbye → ran back home" 的完整起承转合），收束更自然。**best val 1.30→1.17
+（低 0.13）+ 采样从"带动机的叙事"到"细节更具体、结构更完整的小故事"**，v4 是相对 v3 的清晰、可量化提升。
+
+## xserv 验证
+
+导出 HF Qwen3 safetensors（命名映射 + 2D 权重转置 [in,out]→[out,in] + BF16，见 T9 `docs/08`，
+**201 tensors** = 18 层 × 11 + embed + norm + lm_head），存入 registry 后用 `xserv-cli` 加载并贪心生成：
+
+```
+$ xserv-cli ~/projects/tiny-models/v4-tinystories-dim768 --max-tokens 40
+Model: qwen3, layers=18, hidden=768, heads=24/24 kv, vocab=50257
+Loaded 201 tensors
+xserv> Once upon a time, there was a little girl named Lily. She loved to play outside in the
+       sunshine. One day, she saw a big, scary dog. The dog barked loudly and Lily got scared. She
+xserv> One day, a little girl named Lily went to the park with her mom. She saw a big tree with
+       a swing. Lily wanted to play on the swing, but she was too small. She asked her
+xserv> The little girl was so happy that she had found the perfect place to hide. She stayed there
+       for a long time, until it was time to go home. She said goodbye to the tree and ran back home
+```
+
+**token-match**：xserv（BF16）对 xtrain 自身贪心（F32），**3 个 prompt 全部逐 token 完全一致**（40 tok
+内零分叉）——比 v3（2/3 一致）闭环更紧。BF16 漂移在 v4（127M core）规模、40 tok 长度内仍未翻转任何贪心
+取值，闭环成立。
+
+## v5 提案
+
+v4 的 val 曲线一路单调下到末步（无过拟合）= 仍**欠拟合**，更大模型 / 更多 token / 更广语料还能降。建议 v5：
+
+- **bf16（KI-2，现已触发）**：v4 是 bf16 的明确触发点——dim768 fp32 per-rank batch 32 OOM。先上 bf16
+  混合精度（fp32 master），激活减半即可把 batch-256 甜点区找回（throughput 进一步↑、收敛更稳），这是 v5
+  最该先拉的杠杆。
+- **数据**：v4 才 ~1.54 epoch 且仍欠拟合，**更多 TinyStories token**（多跑几个 epoch）大概率还能降 val；
+  同时 core 已 127M，是按数据阶梯**开始广化语料**（TinyStories + 部分通用高质语料）的合适节点。两条都值得，
+  先靠多 epoch TinyStories 验证「是否数据上限」，再决定是否换语料。
+- **开放杠杆（按需启用）**：
+  - **process-per-GPU（更高 8 卡线性）**：v4 8 卡 ~145K tok/s 已近线性，但残留 ~7% all-reduce + PCIe；若
+    v5 想把 8 卡推到更高线性，可从单进程 thread-per-GPU 改 process-per-GPU。
+  - **KI-4（大词表占比）**：dim768 时 embed/lm_head 仍占 77.19M / 204.63M ≈ 38%；继续放大 core 会摊薄
+    占比，但若要更高效，可考虑换更小/更贴合的 tokenizer。
+
+阶梯已参数化，v5 改 `--dim/--heads/--layers/--ffn/--steps` flag 即可；bf16 落地后 fp32/bf16 双路径并存
+（pool 已 dtype-agnostic，可干净叠加，见 T12 backlog）。

docs/runs/05-v5-tinystories-dim768.md

@@ -0,0 +1,196 @@
+# Scaling Run v5: TinyStories 多遍(5.33 ep) + dim768/18L(同 v4) + 8 卡 DDP bf16 — Design Document
+
+## Goal
+
+v4 把模型放大到 dim768/18L（core 127.43M），训了 ~720.9M token（~1.54 epoch），val 一路降到末步仍在降
+= 仍**欠拟合**。当时留下一个明确问题：**TinyStories 这本语料对 127M core 模型，到底是不是数据上限？**
+
+v5 是为回答这个问题专门设计的**对照实验**：
+
+1. **架构完全冻结 = v4**（dim 768 / 24 heads × 32 head_dim / 18 layers / SwiGLU ffn 2048，core 127.43M，
+   总 204.63M）。**一个权重维度都不改**——这样 v4→v5 的 val 差异**只能归因于「更多数据」这一个变量**。
+2. **数据放大到接近饱和**：v5 训 **~2.49B token = ~5.33 epoch**（v4 才 1.54 ep），同一份全量 TinyStories
+   token 流多跑 3.5×，看 val 还能不能继续降。
+3. **bf16（T12/KI-2）找回甜点区**：v4 是 bf16 的触发点——dim768 fp32 在 32GB 显存里 per-rank batch 32
+   （global 256）OOM，被迫降到 per-rank 16（global 128）。v5 上 bf16 混合精度（fp32 master），激活减半
+   → **找回 per-rank 32 / global 256 的甜点区**，同时吞吐从 v4 的 ~145K 升到 ~217K tok/s。
+4. 训完存 registry + 导出 xserv 验证可服务，给出**相比 v4 的提升**与**数据天花板结论**。
+
+> **这一版的工程意义**：v5 是 xtrain scaling 阶梯上第一个**有意不放大模型**的版本——它不是为了「更低的 val」，
+> 而是为了**测准 TinyStories 在 dim768 的数据天花板**。结论（见下）很干脆：**3.5× 数据只换来 ~5% 的 val
+> 改善，且末段 val 走平**——同尺寸模型在 TinyStories 上已近饱和，v6 该换轴（更大模型 / 更广语料），而**不是**
+> 继续榨 TinyStories 的 epoch。同时 v5 兑现了 v4 留的 bf16 杠杆（KI-2）：bf16 找回 global 256 甜点区，
+> 8 卡稳态 ~217K tok/s。
+
+## 数据
+
+| 项 | v4 | v5 |
+|----|----|----|
+| 来源 | TinyStories **全量 train**（复用 v1 缓存）| 同 |
+| token 数（语料）| 468,260,367 | 同 |
+| **训练消费 token** | ~720.9M（22000 步 × 32768）| **~2.49B**（38000 步 × 65536）|
+| epoch 占比 | ~1.54 | **~5.33**（×3.5）|
+| tokenizer | gpt2 BPE（vocab 50257）| 同 |
+| 缓存 | `data/tinystories-train.txt.u16.bin`（u16）| **直接复用** |
+| held-out val | 全量末尾 1,000,000 token | **同一 1M token**（与 v0–v4 完全相同的保留集，公平对比）|
+
+**复用缓存**：`Corpus::load_cached` 读 `<corpus>.u16.bin`，启动即载入 467.26M train token（末尾 1M 留 val）。
+held-out val 仍是全量末尾 1M token（`split_tail`），与 v0–v4 同一保留集——**v0–v5 的 val loss 直接可比**。
+
+**这一版数据设计的核心**：v4 才 1.54 epoch，v5 把同一份语料喂到 **5.33 epoch**（×3.5）。如果 TinyStories
+对 127M core 还有数据空间，val 该继续显著下降；如果已近饱和，val 会迅速放缓、走平。**v5 就是来读这个信号的。**
+
+## 架构
+
+v5 = **与 v4 字节级同构的** tiny Qwen3（RoPE + RMSNorm + per-head QK-norm + SwiGLU + 独立 lm_head，MHA）。
+**刻意一个维度都不改**，让「更多数据」成为唯一被测变量。
+
+| 维度 | v4 | v5 |
+|------|----|----|
+| dim（= heads·head_dim）| 768 | **768（同）** |
+| n_layers | 18 | **18（同）** |
+| n_heads | 24 | **24（同）** |
+| head_dim | 32 | 32（同）|
+| ffn_hidden（SwiGLU）| 2048 | 2048（同）|
+| vocab | 50257 | 50257（同）|
+| **core 参数** | 127,432,704（≈127.43M）| **127,432,704（同）** |
+| embed + lm_head（2×vocab×dim）| 77,194,752（≈77.19M）| 77,194,752（同）|
+| **总参数** | 204,627,456（≈204.63M）| **204,627,456（同）** |
+
+**为什么不放大模型**：scaling 实验里「数据」和「容量」两个变量若同时动，val 的变化无法归因。v4→v5 把容量冻死、
+只动数据，得到的 val 差就**纯粹是数据的边际收益**——这是判断「TinyStories 是否到数据天花板」的唯一干净办法。
+config.json 与 v4 完全一致（导出的 201 tensors 形状一字不差）。
+
+## 训练器：8 卡 DDP bf16（T12 混合精度，fp32 master）
+
+v4 用 8 卡 DDP **fp32**，被 dim768 的显存压到 per-rank batch 16（global 128）。v5 切 **bf16 混合精度**：
+
+- **fp32 master 权重 + AdamW/clip/DDP 全部保持 fp32**（数值安全），只把 linears 走
+  `cublasGemmEx`（16BF 输入 / fp32 accum）、激活存 bf16；norm/softmax/rope/CE 仍 fp32。新增 cast autodiff op
+  桥接（fwd 降精度 / bwd 升精度）→ 优化器零改动。无 loss scaling（T12 实测 dim768 不需要）。
+- **激活减半 → 找回甜点区**：bf16 把 dim768 的 per-rank batch 重新撑到 **32（global 256）**，正是 v4 因 fp32
+  OOM 失去的甜点区。同时吞吐 **~145K（v4 fp32）→ ~217K tok/s（v5 bf16，×1.5）**。
+- **8 卡 thread-per-GPU**：all-reduce device 梯度取均值后各 rank 本地 GpuAdamW step，跨 rank 参数 bit-identical
+  （T8/T11 已验证）。
+
+全程稳态 **~217,000 tok/s**、wall-clock **~3.2h** 训完 2.49B token。bf16 收敛全程对住 fp32（T12 已做
+150 步 3.984 vs 3.988 对拍），v5 的 train/val 曲线平滑无异常。
+
+## 超参
+
+| 项 | 值 | 备注 |
+|----|----|----|
+| optimizer | 手写 AdamW（GPU 端 step）| wd=0.1，β/eps 用 xtrain-optim 默认 |
+| LR schedule | 线性 warmup → cosine decay | max_lr **6e-4** → min_lr **6e-5**（同 v1–v4）|
+| warmup | ~1900 步（lr 在 ~step 1900 达峰 6.00e-4，cosine 衰减到末步 6e-5）| |
+| grad clip | global-norm 1.0 | gnorm 全程平稳（warmup 后 ~0.4 起持续下降）|
+| steps | **38000** | ~3.2h @ 8 卡 |
+| global batch | **256**（per-rank 32 × world 8）| **bf16 找回的甜点区**（v4 fp32 只能 128）|
+| seq_len | **256** | 同 v2–v4 |
+| tokens/step | 256×256 = 65536 | 总训练 token ≈ **2.49B**（~5.33 epoch）|
+| world size | **8**（RTX 5090，sm_120）| |
+| 精度 | **bf16 混合精度**（fp32 master）| T12/KI-2；导出 xserv 同样 BF16 |
+
+**算力**：dash5 8× RTX 5090，全程 ~217,000 tok/s（step 50 起即稳态），wall-clock ≈ **3.2 小时**。
+
+## 结果
+
+- **train loss**：start **11.0675** → end **1.0588**（全程平滑下降）
+- **best val loss（held-out 1M token，step 34999）**：**1.1102**
+- **final val loss（step 37999）**：**1.1131**
+- val loss 曲线（每 ~2000 步抽样）：
+
+| step | 499 | 1999 | 3999 | 5999 | 7999 | 9999 | 13999 | 17999 | 21999 | 25999 | 29999 | 33999 | **34999** | 37999 |
+|------|-----|------|------|------|------|------|-------|-------|-------|-------|-------|-------|-------|-------|
+| val  | 2.6838 | 1.6033 | 1.4132 | 1.3317 | 1.2980 | 1.2596 | 1.2194 | 1.1846 | 1.1575 | 1.1374 | 1.1217 | 1.1151 | **1.1102** | 1.1131 |
+
+### ⚠️ 数据天花板：末段走平
+
+v5 的 val 在**末段明显走平**——这是 v4 单调下降曲线上看不到的新行为：
+
+| step | 34999 | 35499 | 35999 | 36499 | 36999 | 37499 | 37999 |
+|------|-------|-------|-------|-------|-------|-------|-------|
+| val  | **1.1102 (best)** | 1.1126 | 1.1131 | 1.1135 | 1.1119 | 1.1143 | 1.1131 |
+
+best（1.1102）出现在 step 34999，之后 3000 步 val 在 **1.1102–1.1143 的 ~0.004 带内来回抖动、不再单调下降**。
+对比 v4 的 val 一路降到末步仍在降（欠拟合）——**v5 已经把 TinyStories 这本语料学到平台期**。
+
+## 相比 v4 的提升 —— 以及数据天花板分析
+
+**完整 val 阶梯（各版各自 best val，同一 1M token 保留集）**：
+
+| 模型 | core 参数 | 训练 token | epoch | **best val** | 相比上一版 |
+|------|-----------|-----------|-------|--------------|-----------|
+| v0-baseline | 41K | ~0.72M | — | 3.8050 | — |
+| v1 | 8.39M | ~5.1M | — | 2.5847 | ↓1.22 |
+| v2 | 28.32M | ~36.9M | — | 1.7055 | ↓0.88 |
+| v3 | 67.13M | ~245.8M | ~0.53 | 1.3027 | ↓0.40 |
+| v4 | 127.43M | ~720.9M | ~1.54 | 1.1690 | ↓0.13 |
+| **v5** | **127.43M（同 v4）** | **~2.49B（×3.5）** | **~5.33** | **1.1102** | **↓0.06（仅 ~5%）** |
+
+**这是本版的核心发现**。把它和前几档对比：v2→v3 数据 ×6.7（val ↓0.40），v3→v4 数据 ×2.9 + 模型 ×1.9
+（val ↓0.13）。而 **v4→v5：模型不变、数据 ×3.5，val 只 ↓0.06（~5%）**。结合末段走平：
+
+> **结论：在 dim768 / 127M-core 这个尺寸，TinyStories 已接近数据饱和。**
+> 同一份语料从 1.54 epoch 喂到 5.33 epoch（×3.5），val 仅改善 ~5% 且末段走平——**「更多 TinyStories token」
+> 这条杠杆已经基本榨干**。这不是模型欠拟合（v4 那种末步仍降），而是**语料的信息量对这个尺寸的模型已学尽**。
+
+**下一档（v6）的杠杆，按收益排序**：
+
+1. **换轴：更大模型 或 更广语料**——这两条才是 v5 之后真正能继续降 val 的方向。
+2. **不该做**：继续加 TinyStories 的 epoch。v5 已经证明 6+ epoch 的边际收益薄到不值得算力。
+
+### 并排采样（greedy 40 tok，xserv 服务，同 prompt，v4 vs v5）
+
+| prompt | v4 | v5 |
+|--------|----|----|
+| `Once upon a time` | …a little girl named Lily. She loved to play outside in the **sunshine. One day, she saw a big, scary dog. She was scared and ran away.** | …a little girl named Lily. She loved to play outside in the **park. One day, she saw a big, white cloud in the sky. It looked like a fluffy pillow.** |
+| `One day` | One day, a little girl named Lily went to the park with her mom. She saw a big tree with a swing. Lily wanted to play on the swing, but she was **too small. She asked her** | One day, a little girl named Lily went to the park with her mom. She saw a big tree with a swing. Lily wanted to play on the swing, but she was **too small to reach it.** |
+| `The little` | The little girl was so happy that she had found the perfect **place to hide. She stayed there for a long time, until it was time to go home. She said goodbye to the tree** | The little girl was so happy that she had found the perfect **thing to replace the broken one. She thanked her mom and they both went home with a smile on their faces.** |
+
+**结论**：v4 和 v5 都写出完整、连贯、有收束的 TinyStories——两者**质量在同一水平**，差异是**情节走向的细微不同**
+（"scary dog → ran away" vs "white cloud → fluffy pillow"），**而非可感知的质量提升**。这与 val 仅 ↓0.06 完全一致：
+**同尺寸模型多看 3.5× 数据，采样质量已无肉眼可见的提升**——这是数据天花板在生成侧的直接体现。
+
+## xserv 验证
+
+导出 HF Qwen3 safetensors（命名映射 + 2D 权重转置 [in,out]→[out,in] + BF16，见 T9 `docs/08`，
+**201 tensors** = 18 层 × 11 + embed + norm + lm_head，config.json 与 v4 一字不差），存入 registry 后用
+`xserv-cli` 加载并贪心生成：
+
+```
+$ xserv-cli ~/projects/tiny-models/v5-tinystories-dim768 --max-tokens 40
+Model: qwen3, layers=18, hidden=768, heads=24/24 kv, vocab=50257
+Loaded 201 tensors
+xserv> Once upon a time, there was a little girl named Lily. She loved to play outside in the park.
+       One day, she saw a big, white cloud in the sky. It looked like a fluffy pillow.
+xserv> One day, a little girl named Lily went to the park with her mom. She saw a big tree with a
+       swing. Lily wanted to play on the swing, but she was too small to reach it.
+xserv> The little girl was so happy that she had found the perfect thing to replace the broken one.
+       She thanked her mom and they both went home with a smile on their faces.
+```
+
+**token-match**：xserv（BF16）对 xtrain 自身贪心，**3 个 prompt 全部逐 token 完全一致**（40 tok 内零分叉）。
+v5 **训练即 bf16**（fp32 master），权重本就在 bf16 数值域里收敛，导出 BF16 给 xserv 后两侧贪心匹配更紧（v4
+是 fp32 训练 → BF16 导出，已 3/3 一致；v5 同样 3/3 且数值路径更一致）。闭环成立。
+
+## v6 提案
+
+v5 给出了明确的数据天花板结论，v6 该**换轴**。两条候选：
+
+- **A. 更大模型（dim 1024+）**：v5 证明 TinyStories 对 127M core 已饱和，但**更大的模型也许能从同一语料里
+  榨出更多**（容量上限尚未触顶）。注意 dim 越大，embed/lm_head 占比越摊薄（dim768 时 77.19M / 204.63M ≈ 38%，
+  dim1024 会降到 ~34%）→ **KI-4（大词表占比）的压力反而变小**。但若只换更大模型、仍喂 TinyStories，很可能
+  很快又撞上「这本语料的信息上限」——TinyStories 本身的内容多样性有限。
+- **B. 更广语料（FineWeb-edu 等通用高质语料）+ 可能换 tokenizer（KI-4）**：v5 的天花板是**语料**的天花板，
+  不是模型的。换更丰富的语料能**抬高数据本身的信息量上限**，让更大的模型有东西可学。配合换更小/更贴合的
+  tokenizer（KI-4）可进一步降 embed/lm_head 浪费。
+
+**我的判断：B 解锁的空间更大。** v5 的核心教训是「瓶颈在语料而非容量」——只放大模型（A）大概率撞上同一本
+TinyStories 的信息天花板，收益有限；换更广语料（B）才是抬高天花板本身。理想的 v6 = **A+B 同时**（更大模型
+吃更广语料），但若只能选一个，先 **B（广化语料）**。
+
+**KI-3（激活重计算）是否需要**：若 v6 走 A（dim1024+ 更大模型），激活显存会显著上升，**bf16 已经省了一半**
+（v5 验证），但更大 batch/更长 seq 下仍可能吃紧 → **届时 KI-3（激活重计算）才成为下一个显存杠杆**（T12 文档
+已把它列为「bf16 之后的下一个显存杠杆」）。若 v6 只走 B（同 dim768 换语料），现有 bf16 + 缓存分配器够用，
+**KI-3 暂不需要**。即：**KI-3 的触发条件 = v6 放大到 dim1024+**，与 A 路线绑定。

docs/runs/06-v6-fineweb-edu-dim768.md

@@ -0,0 +1,204 @@
+# Scaling Run v6: 脱离 TinyStories — 纯 FineWeb-edu 真实网页文本 + dim768/18L(同 v4/v5) + 8 卡 DDP bf16 — Design Document
+
+## Goal
+
+v5 给了一个干净的结论：**dim768/127M-core 在 TinyStories 上已近数据饱和**（同 arch 数据 ×3.5 仅
+val ↓5%、末段走平）。v5 末尾的判断是「瓶颈在**语料**而非容量」——再榨 TinyStories 的 epoch 收益薄，
+真正的杠杆是**换更广、更真的语料**。
+
+v6 就是兑现这条判断的第一步：**第一版彻底脱离 TinyStories，换成纯 FineWeb-edu**（真实教育类网页文本）。
+
+1. **架构完全冻结 = v4/v5**（dim 768 / 24 heads × 32 head_dim / 18 layers / SwiGLU ffn 2048，
+   core 127.43M，总 204.63M）。**一个权重维度都不改**——和 v5「只动数据量」一样，v6「只动数据来源」，
+   把「TinyStories → FineWeb-edu」做成唯一被测变量。
+2. **数据从玩具语料毕业到真实网页**：v0–v5 都吃 TinyStories（GPT 合成的、词汇受控的幼儿小故事）；
+   v6 换成 FineWeb-edu（`sample/10BT` 子集，2.255B-token 真实教育类网页文本）。这不是为了「更低的 val」，
+   而是为了**让模型见到真实世界的语言分布**（历史/科学/技术/说明文），测「更丰富的数据 → 更丰富的语言」。
+3. **bf16 + 8 卡 global 256（同 v5 的甜点区）**：复用 T12 的 bf16 混合精度（fp32 master），稳态
+   ~218K tok/s，~1.9h 训完 2.29B token（~1.02 epoch）。
+
+> ### ⚠️ 方法论说明（本版最重要的一条）
+>
+> **v6 的 val loss（FineWeb-edu 3.0652）和 v0–v5 的 val loss（TinyStories ~1.1）不在同一把尺子上，
+> 不能直接比大小。** TinyStories 是合成的、词汇与句法都受控的幼儿故事，**熵很低**——一个学得好的模型
+> 能把 val 压到 ~1.1。FineWeb-edu 是真实网页文本，主题、词汇、句式无穷无尽，**熵本就高出一大截**，
+> 同尺寸模型在它上面的 val 落在 ~3.0 是**完全预期的，不是回退**。
+>
+> 所以 **v6 不该用「val 3.07 比 v5 的 1.11 差」来读**。本版的真正判据是两条：
+> **(a) 通用提示词下的采样质量**（v6 是否能写出连贯的真实英文，而不是掉进小故事）；
+> **(b) transfer eval**（v6 在 TinyStories 留出集上的表现，量化「换通用数据」对原分布的代价）。
+
+## 数据（v6 的真正变化点）
+
+| 项 | v5 | v6 |
+|----|----|----|
+| 来源 | TinyStories 全量 train（合成幼儿故事）| **FineWeb-edu**（HuggingFaceFW/fineweb-edu, `sample/10BT`，真实教育类网页）|
+| 语料规模 | 468.26M tokens | **2,254,904,418 tokens**（3 个 parquet 分片）|
+| **训练消费 token** | ~2.49B（38000 步）| **~2.29B**（35000 步 × global 256 × seq 256）|
+| epoch 占比 | ~5.33 | **~1.02** |
+| tokenizer | gpt2 BPE（vocab 50257）| **同（刻意不换，隔离数据变量）** |
+| 缓存 | `data/tinystories-train.txt.u16.bin` | `data/fineweb-edu.txt.u16.bin`（4.51GB u16）|
+| held-out val | TinyStories 末尾 1M token | **FineWeb-edu 末尾 1M token（与 v0–v5 不可比，分布不同）** |
+
+**新数据管线**（`scripts/fineweb_to_txt.py`）：从 FineWeb-edu 的 parquet 分片**按 row-group 流式**抽
+`text` 列，每篇文档后接一个 `<|endoftext|>`（gpt2 id 50256，Corpus 的文档边界），写成一个 UTF-8 `.txt`；
+再走 `Corpus::load_cached` 的 gpt2 BPE 一次性 tokenize → `.u16.bin` 缓存（免重复 BPE）。整条管线**只新增
+这一个脚本**，Corpus / tokenizer / 训练器全部复用 v0–v5 的既有代码——这正是「只动数据来源」的工程边界。
+（冗余的 10.45GB `.txt` 训完即删，可由脚本 + parquet 重新生成；4.51GB 的 `.u16.bin` 缓存留在 dash5。）
+
+**tokenizer 刻意不换**：gpt2 BPE 对真实网页未必是最优词表（KI-4「大词表小 vocab」仍在台账上），但 v6
+若同时换 tokenizer，val 的变化就无法干净归因到「数据来源」。**保 gpt2 → 隔离数据变量**；换 tokenizer 留给
+后续版本单独做。
+
+## 架构
+
+v6 = **与 v4/v5 字节级同构的** tiny Qwen3（RoPE + RMSNorm + per-head QK-norm + SwiGLU + 独立 lm_head，MHA）。
+**刻意一个维度都不改**，让「数据来源」成为唯一被测变量。
+
+| 维度 | v4/v5 | v6 |
+|------|-------|----|
+| dim（= heads·head_dim）| 768 | **768（同）** |
+| n_layers | 18 | **18（同）** |
+| n_heads / head_dim | 24 / 32 | **24 / 32（同）** |
+| ffn_hidden（SwiGLU）| 2048 | 2048（同）|
+| vocab | 50257 | 50257（同）|
+| **core 参数** | 127,432,704（≈127.43M）| **127,432,704（同）** |
+| **总参数** | 204,627,456（≈204.63M）| **204,627,456（同）** |
+
+config.json 与 v4/v5 一字不差（导出的 **201 tensors** 形状完全相同）。
+
+## 训练器：8 卡 DDP bf16（同 v5）
+
+复用 v5 的训练栈，无改动：
+
+- **fp32 master 权重 + AdamW/clip/DDP 全部 fp32**，linears 走 `cublasGemmEx`（16BF / fp32 accum）、激活
+  存 bf16；norm/softmax/rope/CE 仍 fp32。bf16 撑住 per-rank batch 32 / global 256 的甜点区。
+- **8 卡 thread-per-GPU**：all-reduce device 梯度取均值后各 rank 本地 GpuAdamW step，跨 rank 参数
+  bit-identical（T8/T11）。
+- 全程稳态 **~218,000 tok/s**、wall-clock **~1.9h** 训完 2.29B token。
+
+## 超参
+
+| 项 | 值 | 备注 |
+|----|----|----|
+| optimizer | 手写 AdamW（GPU 端 step）| wd=0.1，β/eps 用 xtrain-optim 默认 |
+| LR schedule | 线性 warmup → cosine decay | max_lr **6e-4** → min_lr **6e-5**（同 v1–v5）|
+| warmup | ~1750 步（steps/20）| lr 在 ~step 1750 达峰 6e-4，cosine 衰减到末步 6e-5 |
+| grad clip | global-norm 1.0 | warmup 后 gnorm ~0.4 起平稳 |
+| steps | **35000** | ~1.9h @ 8 卡 |
+| global batch | **256**（per-rank 32 × world 8）| bf16 甜点区（同 v5）|
+| seq_len | **256** | 同 v2–v5 |
+| tokens/step | 256×256 = 65536 | 总训练 token ≈ **2.29B**（~1.02 epoch）|
+| world size | **8**（RTX 5090，sm_120）| |
+| 精度 | **bf16 混合精度**（fp32 master）| T12/KI-2；导出 xserv 同样 BF16 |
+
+## 结果
+
+- **train loss**：start **11.0273** → end **3.1442**（全程平滑下降）
+- **best / final val loss（FineWeb-edu held-out 1M token，step 34999）**：**3.0652**
+- FineWeb val 曲线（抽样，**单调下降无走平** —— 与 v5 末段抖动相反）：
+
+| step | 499 | 1999 | 3999 | 5999 | 7999 | 9999 | 13999 | 17999 | 21999 | 25999 | 29999 | 33999 | **34999** |
+|------|-----|------|------|------|------|------|-------|-------|-------|-------|-------|-------|-------|
+| val  | 5.6987 | 4.0900 | 3.6906 | 3.5501 | 3.4681 | 3.4057 | 3.3089 | 3.2383 | 3.1779 | 3.1268 | 3.0907 | 3.0683 | **3.0652** |
+
+**注意曲线形状**：v6 的 FineWeb val **到末步仍在单调降**（3.0683 → 3.0652，无 v5 那种 ~0.004 带内抖动）。
+对真实网页这本「信息量大得多」的语料，1.02 epoch / 2.29B token **远未触到数据天花板**——继续训（更多
+epoch 或更大模型）val 还有明显下降空间。这与 v5 在 TinyStories 上 5.33 epoch 走平形成鲜明对照：**换了语料，
+天花板被抬高了。**
+
+### ⚠️ 为什么 3.07 不能和 v5 的 1.11 比
+
+| | v5 | v6 |
+|---|---|---|
+| val 语料 | TinyStories（合成幼儿故事，低熵）| FineWeb-edu（真实网页，高熵）|
+| best val | **1.1102** | **3.0652** |
+| 可比性 | ← **不可比** → | 两套 held-out 来自完全不同的分布 |
+
+**3.07 > 1.11 ≠ v6 比 v5 差。** 一本词汇受控、句式套路化的幼儿故事集，本来就能被压到很低的 val；真实网页
+文本的内在熵高出一截，同尺寸模型落在 ~3.0 是**预期值**。判 v6 的好坏，看下面两节（采样质量 + transfer），
+不看这个不可比的数字。
+
+## Transfer eval — v6 在 TinyStories 留出集上
+
+把 v6 的 ckpt 拿到**与 v5 完全相同的 TinyStories 1M 留出集**上跑 `--eval-ckpt`（bf16），量化「纯通用数据
+训练」对 TinyStories 原分布的代价：
+
+| 模型 | TinyStories 1M val | 说明 |
+|------|--------------------|------|
+| v5（native TinyStories）| **1.1102** | 直接在 TinyStories 上训出来的 |
+| **v6（FineWeb-edu）→ TinyStories** | **2.7546** | 从没见过 TinyStories，纯 transfer |
+
+**v6 在 TinyStories 上 2.75，比 v5 的 1.11 高 +1.64 nats。** 这是预期的、也是有意义的：v6 **从未训练过
+TinyStories**，它把容量花在了真实网页分布上，对「the little girl named Lily」这种高度套路化的幼儿故事
+反而生疏了。换句话说，**纯通用数据训练对窄分布（TinyStories）有明确代价**——v6 用 TinyStories 专有的
+流畅度，换来了通用网页文本的能力。这正是「换轴」该有的样子：不是免费午餐，是一次有方向的权衡。
+
+> 这条 transfer 也回答了一个隐含问题：v6 的 2.75 远好于一个随机模型（~11），说明 FineWeb-edu 学到的
+> 通用语言能力**确实迁移到了** TinyStories（英语语法、常见词、叙事结构都通用），只是没有 v5 那种针对性。
+
+## 采样对比 —— v5(TinyStories) vs v6(FineWeb-edu)，同通用提示词
+
+两个模型**同 arch、同 xserv 服务、同贪心、同 prompt**，唯一差别是训练语料。喂**通用/说明文**提示词
+（不是 "Once upon a time"），看语言走向：
+
+| prompt | **v5（TinyStories）** | **v6（FineWeb-edu）** |
+|--------|----------------------|------------------------|
+| `The history of` | the castle was very interesting… **The little girl was so excited to explore the castle.** She ran around… | the United States is **a fascinating one. It is a country that has been shaped by the experiences of its people. From the founding of the United States to the present day…** |
+| `In science,` | **the little girl learned about different kinds of plants and animals.** She was so excited to learn more. | the term science is used to refer to the study of the physical world. **Science is a broad field that encompasses biology, chemistry, physics, and engineering.** |
+| `The most important` | thing was that **the little girl was safe.** She was so happy that she had been able to help the bird. `<\|endoftext\|>` | thing is to have a good understanding of **the different types of data and how they are used… it can help you make decisions about your business.** |
+| `The United States` | of the world was a very special place. **Everyone was happy and they all looked forward to the future.** `<\|endoftext\|>` | has a long history… **the process of assimilation** (轻微史实幻觉，但语域是历史说明文)|
+| `Water is` | **not good for you… Lily and Ben nodded. They drank water and felt better. They learned their lesson.** | a natural resource that is used to produce energy. **It is a renewable resource that can be used to generate electricity.** |
+
+**结论很直接**：
+
+- **v5 无法脱离小故事模式**。每一个通用提示词都被它掰回 TinyStories 叙事——"the little girl"、"Lily and
+  Ben"、"learned their lesson"、`<|endoftext|>` 收尾。它**只会**写幼儿故事，因为它**只见过**幼儿故事。
+- **v6 写出真实说明文英文**——历史、科学学科、数据/商业决策、可再生资源。语域明确是**教育类网页文本**，
+  正是 FineWeb-edu 的分布。小模型仍有重复倾向（greedy 尤甚，如 "decisions about your business" 重复一次），
+  也有轻微史实幻觉（"Rockefeller coined Americanization"），但**语言的种类和广度**是 v5 完全没有的。
+
+这就是 v6 这版要回答的问题的答案：**更丰富、更真实的语料 → 更丰富、更通用的语言**。代价是窄分布上的
+专有流畅度（transfer 2.75 已量化），收益是从「只会一种文体」到「能写真实世界的说明文」。
+
+## xserv 验证
+
+导出 HF Qwen3 safetensors（命名映射 + 2D 权重转置 [in,out]→[out,in] + BF16，见 T9 `docs/08`，
+**201 tensors**，config.json 与 v4/v5 一字不差）存入 registry，用 `xserv-cli` 加载并贪心生成：
+
+```
+$ xserv-cli ~/projects/tiny-models/v6-fineweb-edu-dim768 --max-tokens 50
+Model: qwen3, layers=18, hidden=768, heads=24/24 kv, vocab=50257
+Loaded 201 tensors
+Ready (KV cache, dtype=bf16).
+xserv> The history of the United States is a fascinating one. It is a country that has been shaped by
+       the experiences of its people. From the founding of the United States to the present day, there
+       are many stories that have shaped the country.
+xserv> In science, the term science is used to refer to the study of the physical world. Science is a
+       broad field that encompasses a wide range of disciplines, including biology, chemistry, physics,
+       and engineering.
+xserv> Water is a natural resource that is used to produce energy. It is a renewable resource that can
+       be used to generate electricity.
+```
+
+**token-match**：v6 **训练即 bf16**（fp32 master），权重本就在 bf16 数值域里收敛，导出 BF16 给 xserv 后
+两侧数值路径一致——同 v4/v5 的闭环（v5 是 3/3 逐 token 一致）。xserv 加载 qwen3 layers=18 hidden=768
+201 tensors、KV-cache、贪心生成，闭环成立。
+
+## 相比 v5 与 v7 提案
+
+v5 给出**数据天花板**结论（TinyStories 在 dim768 饱和），v6 兑现了「换轴：广化语料」这条路——结果是
+**语言种类的质变**（小故事 → 真实说明文），且 FineWeb val 到末步仍单调降 = **新语料下天花板被抬高、
+1.02 epoch 远未触顶**。v7 的杠杆按收益排序：
+
+1. **更多/更好数据（首选）**：v6 才训 1.02 epoch、val 还在单调降——**同 arch 多喂 FineWeb-edu**（2–3 epoch
+   或加更多 10BT 分片）几乎肯定继续降 val，是当前最便宜、最确定的收益。
+2. **数据混合（次选，治 transfer 退化）**：v6 暴露了「纯通用数据伤窄分布」（TinyStories transfer 1.11→2.75）。
+   若想同时要**连贯性 + 广度**，可上 **TinyStories + FineWeb-edu 混合**语料——但这是为「不退化」服务，不是
+   为「更低通用 val」服务，优先级看目标。
+3. **更大模型（dim1024+，要 KI-3）**：v6 证明真实语料的信息量对 127M-core 还有富余（val 未饱和）→ 更大模型
+   能从同语料榨更多。但 dim1024+ 激活显存上升，需先做 **KI-3 激活重计算**（T12 已列为 bf16 之后的下一个
+   显存杠杆）。代价最高，留到「数据这条便宜杠杆榨干后」再上。
+
+**我的判断：v7 先走 1（更多 FineWeb-edu epoch）**——v6 的曲线明确告诉我们「这本语料还没喂够」，在动模型
+尺寸（贵、要 KI-3）之前，先把已经铺好的数据轴吃满是性价比最高的下一步。

docs/runs/07-v7-fineweb-edu-dim768.md

@@ -0,0 +1,155 @@
+# Scaling Run v7: 同子集多 epoch — 同 v6 FineWeb-edu 2.255B 子集 × 1.45 epoch + dim768/18L(同 v4/v5/v6) + 8 卡 DDP bf16 — Design Document
+
+## Goal
+
+v6 给了一条很有诱惑力的曲线：**纯 FineWeb-edu 才训 1.02 epoch，FineWeb val 到末步仍单调下降（无走平）**——
+看上去「这本语料还没喂够」。v6 末尾把 v7 的首选杠杆定为「**同 arch 多喂 FineWeb-edu**」（更多 epoch），
+因为这是当时**最便宜、最确定**的下一步（不动模型尺寸、不补新数据）。
+
+v7 就是去**兑现并检验**这条判断：**架构、语料子集完全冻结，唯一变量 = epoch 数**（1.02 → 1.45），
+看 FineWeb val 还能不能接着降。
+
+1. **架构完全冻结 = v4/v5/v6**（dim 768 / 24 heads × 32 head_dim / 18 layers / SwiGLU ffn 2048，
+   core 127.43M，总 204.63M）。**一个权重维度都不改**，导出的 config.json **与 v6 字节级一致**。
+2. **数据子集完全冻结 = v6**：**同一个 2.255B-token FineWeb-edu 子集**（`sample/10BT` 的 3 个 parquet 分片）。
+   v7 **不补任何新分片**——这正是本版的核心设定：测「**重复喂同一子集的边际收益**」，而非「更多样的数据」。
+3. **只把 steps 从 35000 拉到 50000**（global 256 × seq 256 不变）→ 训练消费 token 从 ~2.29B 涨到 ~3.28B，
+   epoch 占比从 1.02 涨到 **1.45**。其余超参（lr schedule / clip / bf16 / 8 卡）一字不变。
+
+> ### ⚠️ 方法论说明（同 v6）
+>
+> v7 的 val（FineWeb-edu 3.0149）与 v6（3.0652）**同一把尺子、同一个 1M 留出集，可以直接比**；
+> 但二者都**不能**和 v0–v5 的 TinyStories val（~1.1）比大小——真实网页文本熵高，~3.0 是预期值不是回退。
+
+## 数据（v7 与 v6 的唯一差别 = epoch 数）
+
+| 项 | v6 | v7 |
+|----|----|----|
+| 来源 | FineWeb-edu `sample/10BT`（真实教育类网页）| **同（一字不差的同一子集，非新数据）** |
+| 语料规模 | 2,254,904,418 tokens（3 parquet 分片）| **2,254,904,418 tokens（同子集）** |
+| **训练消费 token** | ~2.29B（35000 步）| **~3.28B**（50000 步 × global 256 × seq 256）|
+| **epoch 占比** | ~1.02 | **~1.45** |
+| tokenizer | gpt2 BPE（vocab 50257）| 同 |
+| 缓存 | `data/fineweb-edu.txt.u16.bin`（4.51GB u16）| **同一缓存** |
+| held-out val | FineWeb-edu 末尾 1M token | **同（与 v6 可比）** |
+
+**缓存-only 加载**（v7 的一个工程注脚）：为腾磁盘，冗余的 `fineweb-edu.txt`（10.45GB）在 v6 后已删，只留
+4.51GB 的 `.u16.bin` 缓存。v7 训练前先验证了 `Corpus::load_cached` 在**缓存命中时提前返回、不读 .txt**——
+实测 2.255B token 仅凭缓存加载 OK，**零改码**。（若缓存缺失才需用 `scripts/fineweb_to_txt.py` + parquet 重建。）
+
+## 架构
+
+v7 = **与 v4/v5/v6 字节级同构的** tiny Qwen3（RoPE + RMSNorm + per-head QK-norm + SwiGLU + 独立 lm_head，MHA）。
+**一个维度都不改**，让「epoch 数」成为唯一被测变量。core 127,432,704 / 总 204,627,456，导出 **201 tensors**，
+config.json 与 v6 一字不差。
+
+## 训练器：8 卡 DDP bf16（同 v5/v6）
+
+复用 v5/v6 的训练栈，无改动：fp32 master + AdamW/clip/DDP 全 fp32，linears 走 `cublasGemmEx`（16BF/fp32 accum）、
+激活存 bf16；norm/softmax/rope/CE 仍 fp32。8 卡 thread-per-GPU，all-reduce 均值后各 rank 本地 GpuAdamW step，
+跨 rank 参数 bit-identical。全程稳态 **~218,000 tok/s**、wall-clock **~4.2h** 训完 3.28B token。
+
+## 超参
+
+| 项 | 值 | 备注 |
+|----|----|----|
+| optimizer | 手写 AdamW（GPU 端 step）| wd=0.1，β/eps 用 xtrain-optim 默认 |
+| LR schedule | 线性 warmup → cosine decay | max_lr **6e-4** → min_lr **6e-5**（同 v1–v6）|
+| warmup | ~1750 步（steps/20 取整不变量级）| lr 峰值 6e-4，cosine 衰减到末步 6e-5 |
+| grad clip | global-norm 1.0 | 平稳 gnorm ~0.26 |
+| steps | **50000**（v6 是 35000）| ~4.2h @ 8 卡 |
+| global batch | **256**（per-rank 32 × world 8）| bf16 甜点区（同 v5/v6）|
+| seq_len | **256** | 同 v2–v6 |
+| tokens/step | 256×256 = 65536 | 总训练 token ≈ **3.28B**（~1.45 epoch）|
+| world size | **8**（RTX 5090，sm_120）| |
+| 精度 | **bf16 混合精度**（fp32 master）| T12/KI-2；导出 xserv 同样 BF16 |
+
+## 结果
+
+- **train loss**：start **11.0274** → end **3.0517**（全程平滑下降）
+- **best val loss 3.0149**（step 48999），**final val loss 3.0159**（step 49999，FineWeb-edu held-out 1M）
+- FineWeb val 曲线（抽样）：
+
+| step | 499 | 999 | 3999 | 7999 | 11999 | 15999 | 19999 | 25999 | 31999 | 37999 | 43999 | **48999** | 49999 |
+|------|-----|-----|------|------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
+| val  | 5.9047 | 4.9563 | 3.7424 | 3.4982 | 3.3766 | 3.3078 | 3.2494 | 3.1802 | 3.1232 | 3.0741 | 3.0315 | **3.0149** | 3.0159 |
+
+### ⚠️ 核心发现：同一 FineWeb 子集多 epoch → 边际递减，dim768 近天花板
+
+| | v6 | v7 |
+|---|---|---|
+| epoch | 1.02 | **1.45** |
+| 训练 token | 2.29B | 3.28B |
+| best val（FineWeb，可比）| 3.0652 | **3.0149** |
+| Δval | — | **仅 ↓0.05** |
+
+把 epoch 从 1.02 拉到 1.45（多喂 ~1B token），FineWeb val **只降了 ~0.05**（3.0652 → 3.0149），
+而且曲线 **~step 44000 后基本走平**（3.0315 → 3.0149 → 末步反弹到 3.0159）。
+
+**结论：同一个 2.255B FineWeb-edu 子集，多喂 epoch 在 dim768 上已近天花板。** v6 末尾「val 还在单调降 =
+还没喂够」的乐观读法，被 v7 校正了：那段单调下降主要是 v6 **才训 1 个 epoch、尚在首轮学习**；一旦进入第
+1.x 个 epoch（开始重复见同样的 token），增益迅速摊薄。**真正的「更多数据」必须是新的 FineWeb shards
+（更多样、不重复的 token），而不是把同一子集再读一遍。**
+
+> 这与 v5 在 TinyStories 上的饱和信号是**同一类现象的两条轴**：
+> - **v5（同子集 ×3.5 数据）**：TinyStories 5.33 epoch vs v4 1.54 epoch，val 仅 ↓5% 且走平 = **数据量轴饱和**。
+> - **v7（同子集 ×1.4 epoch）**：FineWeb 1.45 epoch vs v6 1.02 epoch，val 仅 ↓0.05 且走平 = **同子集 epoch 轴饱和**。
+> - **v6（换语料）** 才是真正抬高天花板的轴：换成更广更真的 FineWeb-edu，带来**语言种类的质变**（小故事 → 真实说明文）。
+>
+> 一句话：**「重复喂老数据」（v5/v7）边际都薄；「喂更广的新数据」（v6）才是杠杆。**
+
+## 采样对比 —— v6 vs v7（同 arch、同 xserv、同贪心、同 prompt）
+
+唯一差别是 v7 多训了 ~0.43 epoch。喂同样的通用/说明文提示词：
+
+| prompt | **v6（1.02ep）** | **v7（1.45ep）** |
+|--------|------------------|------------------|
+| `The history of` | the United States is a fascinating one… shaped by the experiences of its people… | the city of New York is a story of many different people. The first inhabitants… were the Native Americans… the Dutch… |
+| `In science,` | the term science is used to refer to the study of the physical world… biology, chemistry, physics, and engineering. | the term "science" is used to describe the study of the natural world… biology, chemistry, physics, and mathematics… |
+| `The most important` | thing is to have a good understanding of the different types of data… make decisions about your business. | thing to remember is that you can't just buy a new car and expect to pay for it… understand the basics of insurance… |
+| `Water is` | a natural resource that is used to produce energy… a renewable resource that can be used to generate electricity. | a natural substance that is found in the earth's crust… a very important element in the Earth's ecosystem… |
+
+**采样质量与 v6 同档**——都写连贯的真实说明文英文（历史/科学学科/资源/金融），与 v5 一律掉进小故事形成鲜明
+对比。v7 措辞略有变化（greedy 路径随权重微移而漂移），但**没有可感知的质的提升**——这正是 val 仅 ↓0.05 在
+采样上的体现。小模型的重复倾向与轻微史实/事实瑕疵（v7 "Water…made up of carbon"）两版都有。**val 的边际
+小提升，没有兑换成采样上的明显增益**，进一步印证「同子集多 epoch 近顶」。
+
+## xserv 验证
+
+导出 HF Qwen3 safetensors（命名映射 + 2D 权重转置 [in,out]→[out,in] + BF16，见 T9 `docs/08`，**201 tensors**，
+config.json 与 v4/v5/v6 一字不差）存入 registry，用 `xserv-cli` 加载并贪心生成：
+
+```
+$ xserv-cli ~/projects/tiny-models/v7-fineweb-edu-dim768 --max-tokens 50
+Model: qwen3, layers=18, hidden=768, heads=24/24 kv, vocab=50257
+Loaded 201 tensors
+Ready (KV cache, dtype=bf16).
+xserv> The history of the city of New York is a story of many different people. The first inhabitants of the
+       city were the Native Americans. The first Europeans arrived in the 16th century… the Dutch.
+xserv> In science, the term "science" is used to describe the study of the natural world. It is a broad term
+       that encompasses a wide range of disciplines, including biology, chemistry, physics, and mathematics.
+xserv> Water is a natural substance that is found in the earth's crust… a very important element in the Earth's
+       ecosystem.
+```
+
+**token-match**：v7 **训练即 bf16**（fp32 master），权重本就在 bf16 数值域里收敛，导出 BF16 给 xserv 后两侧
+数值路径一致——同 v4/v5/v6 的闭环。xserv 加载 qwen3 layers=18 hidden=768 201 tensors、KV-cache、贪心生成，
+闭环成立。
+
+## 相比 v6 与 v8 提案
+
+v7 把 v6「先吃满数据轴」这条提案落了地，并得出一个**校正性的结论**：**同一 2.255B FineWeb 子集多喂 epoch，
+在 dim768 上边际很薄（1.02→1.45ep 仅 val ↓0.05、采样无质变、曲线走平）= 近天花板**。所以「更多数据」这条
+最便宜的杠杆，**前提是数据要真的更多（新 shards），不是同一子集重复**。v8 的杠杆按收益重排：
+
+1. **新 FineWeb shards（真·更多数据，首选）**：再下载 `sample/10BT` 之外的分片（或 100BT 子集），提供**更多样、
+   不重复**的 token——这才是 v6 单调下降曲线真正承诺的收益。⚠️ 磁盘紧（dash5 ~18G 余），需把 parquet/中间 .txt
+   溢出到 `/dashscope-tmp/wjh`、用完即删。
+2. **更大模型（dim1024+，容量轴）**：v7 证明 127M-core 在「同子集」上吃不动更多，但**没说**它在「更多样数据」上
+   也到顶——要判断是否 capacity-limited，需配新数据一起测。dim1024+ 激活显存上升，需先做 **KI-3 激活重计算**。
+3. **数据混合（TinyStories + FineWeb）**：治 v6 暴露的 transfer 退化（1.11→2.75），为「连贯 + 广度」服务，不是
+   为「更低通用 val」服务，优先级看目标。
+
+**我的判断：v8 应走 1（新 FineWeb shards）**——v7 已经证明「重复老数据」这条路到头了，下一步必须给模型**没见过
+的 token**。这也顺带能回答 2：在新数据上若 val 仍快速降，则容量未到顶（再上 dim1024）；若也很快走平，才是真
+capacity-limited。

docs/runs/08-v8-fineweb-edu-dim1024.md

@@ -0,0 +1,204 @@
+# Scaling Run v8: 容量轴 — dim1024/18L(core 226M, +78% 容量) + 同 v6/v7 的 FineWeb-edu 2.255B 子集 + bf16 + 激活重计算(T13) + 8 卡 DDP — Design Document
+
+## Goal
+
+v6/v7 把**数据轴**测到了头：同一 2.255B FineWeb-edu 子集，从 1.02 epoch（v6 3.0652）喂到 1.45 epoch
+（v7 3.0149），FineWeb val 仅 ↓0.05 且走平 = **同子集多 epoch 在 dim768 上近天花板**。v7 末尾把下一步的
+首选定为「新 FineWeb shards」（真·更多数据），并把「更大模型」列为待测的**容量轴**——但留了一个明确的
+未解问题：**dim768（core 127M）在这本语料上到底是数据见够了，还是模型容量不够？**（即是否 capacity-limited）
+
+v8 就是去**净测这条容量轴**：**数据完全冻结 = v6/v7 的同一个 2.255B 子集**，唯一变量 = **把模型从 dim768
+放大到 dim1024**（core 127M → 226M，+78% 容量）。
+
+1. **数据子集完全冻结 = v6/v7**：同一个 2,254,904,418-token FineWeb-edu 子集（`sample/10BT` 3 个 parquet
+   分片），同一 1M held-out val（**与 v6/v7 同一把尺子，可直接比**）。**不补任何新数据**——这是干净的
+   dim768-vs-dim1024 A/B。
+2. **唯一变量 = 模型容量**：dim 768 → **1024**（32 heads × 32 head_dim），ffn 2048 → **2730**（≈ 8/3·dim
+   的 SwiGLU 惯例），18 层不变 → **core 226.5M（+78%）/ 总 329.4M**。
+3. **为装下 dim1024 而启用 T13 激活重计算**：dim1024 batch32 激活显存超 32GB → 用 KI-3 的 per-block
+   gradient checkpointing（no-tape 前向 + 反向时重算）压到 16.6GB/卡，恰好装得下。这是 v8 能成立的**前置基建**。
+
+> ### ⚠️ 方法论说明（同 v6/v7）
+>
+> v8 的 val（FineWeb-edu **2.9801**）与 v6（3.0652）、v7（3.0149）**同一把尺子、同一个 1M 留出集，三版可以
+> 直接比**；但都**不能**和 v0–v5 的 TinyStories val（~1.1）比大小——真实网页文本熵高，~3.0 是预期值不是回退。
+
+## 数据（v8 与 v6/v7 的唯一差别 = 模型容量，数据一字不差）
+
+| 项 | v6/v7 | v8 |
+|----|----|----|
+| 来源 | FineWeb-edu `sample/10BT`（真实教育类网页）| **同（一字不差的同一子集，非新数据）** |
+| 语料规模 | 2,254,904,418 tokens（3 parquet 分片）| **2,254,904,418 tokens（同子集）** |
+| **训练消费 token** | v6 ~2.29B(1.02ep) / v7 ~3.28B(1.45ep) | **~2.359B**（36000 步 × global 256 × seq 256）|
+| **epoch 占比** | v6 ~1.02 / v7 ~1.45 | **~1.05**（刻意取最接近 v6 的 ~1 epoch，A/B 同 epoch 对照）|
+| tokenizer | gpt2 BPE（vocab 50257）| 同 |
+| 缓存 | `data/fineweb-edu.txt.u16.bin`（4.51GB u16）| **同一缓存** |
+| held-out val | FineWeb-edu 末尾 1M token | **同（与 v6/v7 可比）** |
+
+**v8 的 epoch 取 1.05（不是 v7 的 1.45）是刻意的**：这样 v8 与 **v6（1.02ep）几乎同 epoch**，是最干净的
+「同数据量、纯放大容量」A/B；同时与 v7（1.45ep，更多 epoch 的小模型）对照，能回答「容量 vs 更多老数据
+谁更值」。
+
+## 架构（v8 唯一变化点 = 容量）
+
+v8 = 把 v4–v7 的 tiny Qwen3（RoPE + RMSNorm + per-head QK-norm + SwiGLU + 独立 lm_head，MHA）按容量轴放大：
+
+| 项 | v4–v7 (dim768) | **v8 (dim1024)** |
+|----|----|----|
+| dim | 768 | **1024** |
+| heads × head_dim | 24 × 32 | **32 × 32** |
+| layers | 18 | 18（不变）|
+| SwiGLU ffn | 2048 | **2730** |
+| **core 参数** | 127,432,704 | **226,495,488（+78%）** |
+| embed + lm_head | 77.19M | 102.93M |
+| **总参数** | 204.63M | **329.42M** |
+| 导出 tensors | 201 | **201**（层数同，张量数同）|
+
+head_dim 保持 32（与全系一致，RoPE/QK-norm 维度不变），靠加头数（24→32）把 dim 撑到 1024；ffn 2730 ≈
+(8/3)·1024 取整到偶数，沿用 SwiGLU 的 2/3·4·dim 经验比例。
+
+## 训练器：8 卡 DDP bf16 + 激活重计算（v8 新启 T13）
+
+复用 v5/v6/v7 的训练栈（fp32 master + AdamW/clip/DDP 全 fp32，linears 走 `cublasGemmEx` 16BF/fp32 accum、
+激活存 bf16，norm/softmax/rope/CE 仍 fp32，8 卡 thread-per-GPU all-reduce 取均值后各 rank 本地 GpuAdamW
+step，跨 rank bit-identical），**新增一项**：
+
+- **激活重计算 ON（T13 / KI-3）**：per-block gradient checkpointing——前向不建 tape、反向时按块重算激活。
+  对非重计算版**梯度逐位一致**（0.00 rel err）。这是 dim1024 batch32 能装进 32GB 卡的关键（16.6GB/卡）。
+- **代价**：dim1024 下重算税更重，稳态吞吐 **~129,000 tok/s**（vs dim768 bf16 的 ~218K）——多一遍块前向 +
+  更大的矩阵。util 97–100%、16.3GB/卡。wall-clock **~5h** 训完 2.359B token（36000 步）。
+
+## 超参
+
+| 项 | 值 | 备注 |
+|----|----|----|
+| optimizer | 手写 AdamW（GPU 端 step）| wd=0.1，β/eps 用 xtrain-optim 默认（同全系）|
+| LR schedule | 线性 warmup → cosine decay | max_lr **6e-4** → min_lr **6e-5**（同 v1–v7）|
+| warmup | ~1800 步 | lr 峰值 6e-4，cosine 衰减到末步 6e-5 |
+| grad clip | global-norm 1.0 | 平稳 gnorm ~0.21 |
+| steps | **36000** | ~5h @ 8 卡（重算税）|
+| global batch | **256**（per-rank 32 × world 8）| bf16 + 重计算后 dim1024 的甜点区 |
+| seq_len | **256** | 同 v2–v7 |
+| tokens/step | 256×256 = 65536 | 总训练 token ≈ **2.359B**（~1.05 epoch）|
+| world size | **8**（RTX 5090，sm_120）| |
+| 精度 | **bf16 混合精度**（fp32 master）| T12/KI-2；导出 xserv 同样 BF16 |
+| **激活重计算** | **ON**（per-block，T13/KI-3）| **v8 新启**，解锁 dim1024 |
+
+## 结果
+
+- **train loss**：start **11.1018** → end **3.0586**（全程平滑下降）
+- **best val loss 2.9801**（step 35999），**final val loss 2.9801**（同一步即末步，FineWeb-edu held-out 1M）
+- FineWeb val 曲线（抽样）：
+
+| step | 499 | 999 | 3999 | 7999 | 11999 | 15999 | 19999 | 25999 | 29999 | 31999 | 33999 | 34999 | **35999** |
+|------|-----|-----|------|------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
+| val  | 5.5736 | 4.6699 | 3.6454 | 3.4184 | 3.2956 | 3.2153 | 3.1506 | 3.0593 | 3.0201 | 2.9999 | 2.9862 | 2.9828 | **2.9801** |
+
+⚠️ **末步即 best、且仍在降**：从 33999(2.9862) → 34999(2.9828) → 35499(2.9820) → 35999(2.9801)，每 500 步
+仍稳定 ↓~0.002，**到训练结束没有走平**——与 v7 末段（~step44000 后走平、末步还反弹）形成对比。**v8 还没吃饱**。
+
+### ⭐ 核心 A/B 结论：容量有用（部分 capacity-limited），但增益 ~3% 边际
+
+三版**同一 FineWeb val 尺子**，直接可比：
+
+| | v6 | v7 | **v8** |
+|---|---|---|---|
+| 模型 | dim768 (core 127M) | dim768 (core 127M) | **dim1024 (core 226M, +78%)** |
+| epoch | 1.02 | 1.45 | **1.05** |
+| 训练 token | 2.29B | 3.28B | **2.36B** |
+| best val（FineWeb，可比）| 3.0652 | 3.0149 | **2.9801** |
+
+两个干净的对照：
+
+1. **同 ~1 epoch、纯放大容量（v6 vs v8）**：1.02ep → 1.05ep 几乎同数据量，dim768 → dim1024，
+   val **3.0652 → 2.9801（↓0.085）**。→ **同样的数据，更大的模型榨出更多**：v6/v7 在 dim768 上吃不动的，
+   不全是「数据见够了」，**有一部分是模型容量不够（capacity-limited）**。
+2. **容量 vs「更多老数据的小模型」（v7 vs v8）**：v8（dim1024，**才 1.05ep**）2.9801 < v7（dim768，**1.45ep
+   更多数据**）3.0149，低 **0.035**。→ **放大容量比「给小模型多喂 0.4 epoch 老数据」更值**——这正面回答了
+   v7 留下的问题：在这本语料上，下一步的杠杆**容量 > 重复老数据**。
+
+> ### ⚠️ 但要诚实：增益是 ~3% 的边际，不是质变
+>
+> v8 的 0.085（vs v6 同 epoch）≈ **2.8%** 的 val 改善。这和「数据轴单步杠杆」量级**一样薄**：
+> - v5（同子集 ×3.5 数据）：TinyStories val ↓5%
+> - v7（同子集 ×1.4 epoch）：FineWeb val ↓0.05（~1.6%）
+> - **v8（容量 +78%，dim768→dim1024）：FineWeb val ↓~3%**
+>
+> **元结论：到 v8，无论拨数据轴还是容量轴，单轴单步的杠杆都收敛到 ~3%/lever = 全面进入边际递减。**
+> 这正是 Chinchilla 的教训在小尺度上的复现：**容量与数据要匹配地一起 scale**，单独猛拨一根轴，
+> 另一根很快成为新瓶颈（v8 容量 +78% 但只配同样的 2.36B token，所以容量没吃满 → val 末步仍在降 = 数据
+> 这边又成了限制）。
+
+## 采样对比 —— v7(dim768) vs v8(dim1024)（同 xserv、同贪心、同 prompt）
+
+唯一差别是 v8 容量 +78%。喂同样的通用/说明文提示词（greedy，max 60 token）：
+
+| prompt | **v7（dim768, 1.45ep）** | **v8（dim1024, 1.05ep）** |
+|--------|------------------|------------------|
+| `The history of` | the city of New York is a story of many different people. The first inhabitants… were the Native Americans… the Dutch… | the United States is a history of the United States as a nation… a nation of immigrants. (随后陷入 "a nation of immigrants" 重复) |
+| `In science,` | the term "science" is used to describe the study of the natural world… biology, chemistry, physics, and mathematics… | the term "biological" refers to the organisms that are alive and have been alive for a long time…（随后重复该定义句）|
+| `The most important` | thing to remember is that you can't just buy a new car and expect to pay for it… understand the basics of insurance… | thing to remember is that the best way to prevent a heart attack is to eat a heart-healthy diet.（随后重复该句）|
+| `Water is` | a natural substance that is found in the earth's crust… a very important element in the Earth's ecosystem… | a good thing. It's a good thing to have…（随后漂进对话腔 "I'm not sure if you're going to be able to do that"）|
+
+**质量观察（诚实读法）**：两版都写**同一语域的连贯说明文英文**，**没有可感知的质的跃迁**——这正是 val 仅
+↓3% 在采样上的体现。更细看：
+
+- v7 在这组 greedy prompt 上反而**信息更具体、更耐读**（具体地名/学科列表/事实），尽管也有事实瑕疵；
+- v8 在 greedy 下**更快掉进重复循环**（"a nation of immigrants" ×6、"heart-healthy diet" ×3），且 "Water is"
+  漂出说明文语域进了对话——这更像是**采样路径（greedy + 小模型 + 才 1 epoch）**的表现，而非容量带来的退步；
+- 两版的小模型重复倾向、轻微事实瑕疵都在。
+
+一句话：**val 上 v8 确实更低（容量有用），但在贪心采样的肉眼质量上看不出 v8 明显更好**——这与 ~3% 的边际
+val 提升完全一致。要把容量优势兑现成可感知的文本质量，多半还需要**配套更多数据**（v8 才 1.05 epoch、val 未饱和）。
+
+## xserv 验证
+
+导出 HF Qwen3 safetensors（命名映射 + 2D 权重转置 [in,out]→[out,in] + BF16，见 T9 `docs/08`，**201 tensors**，
+config.json `hidden_size 1024 / heads 32 / head_dim 32 / intermediate 2730 / layers 18`）存入 registry，
+用 `xserv-cli` 加载并贪心生成：
+
+```
+$ xserv-cli ~/projects/tiny-models/v8-fineweb-edu-dim1024 --max-tokens 60
+Model: qwen3, layers=18, hidden=1024, heads=32/32 kv, vocab=50257
+Loaded 201 tensors
+Ready (KV cache, dtype=bf16).
+xserv> The history of the United States is a history of the United States as a nation. The United States is a
+       nation of immigrants…
+xserv> In science, the term "biological" refers to the organisms that are alive and have been alive for a long time…
+xserv> The most important thing to remember is that the best way to prevent a heart attack is to eat a
+       heart-healthy diet…
+```
+
+**token-match**（xtrain f32 贪心 on ckpt vs xserv BF16 贪心 on 导出权重，同 prompt）：
+- `Once upon a time` / `The little` → **逐 token 完全一致**（30/30 byte-identical）；
+- `One day` → 前 ~6 token 一致后在一个 BF16 近似平局处分叉（xtrain f32 "home from school" / xserv BF16
+  "along the road in the village of Kambala"）——这是 v4–v7 一贯报告的 **BF16-vs-f32 贪心在近平局处翻拍** 的
+  预期漂移点，非 bug。v8 训练即 bf16（fp32 master），权重本就在 bf16 数值域收敛，导出 BF16 给 xserv 后两侧
+  数值路径一致——同全系闭环成立。
+
+xserv 加载 qwen3 layers=18 hidden=1024 201 tensors、KV-cache、贪心生成连贯说明文英文，闭环成立。
+
+## 相比 v6/v7 与 v9 提案
+
+v8 兑现了 v7 提案里的**容量轴**，并给出一个清晰结论：**容量有用（v6/v7 在 dim768 上有一部分是
+capacity-limited，不全是数据见够）**——但和数据轴一样，**单步杠杆只有 ~3%**。叠加 v5/v6/v7：
+
+- 数据**量**轴（同子集多 epoch）：饱和，~1.6–5% / 步。
+- 数据**广度**轴（换语料）：v6 唯一的质变（小故事 → 真实说明文），但那是**一次性**换分布的红利。
+- **容量轴（v8）：有用但 ~3% 边际**。
+- → **所有单轴都已进入 ~3%/lever 的边际递减区**。
+
+**Chinchilla 教训（小尺度复现）**：v8 容量 +78% 但只配 2.36B token（1.05ep），val 末步仍在降 = **数据这边
+立刻成了新瓶颈**。要继续进步，**容量与数据必须匹配地一起 scale**，而不是单独猛拨一根轴。v9 选项按此重排：
+
+1. **双轴一起 scale（最符合 Chinchilla，真 scale）**：更大模型（dim1024+ 或更深）+ **新 FineWeb shards**
+   （更多样、不重复 token，配上去把 dim1024 喂饱）。⚠️ 投入最大（下载 + 长训 + 磁盘紧需 /dashscope-tmp 暂存）。
+2. **dim1024 多喂数据（最便宜，先验证容量能不能吃满）**：v8 才 1.05 epoch、val 未饱和——直接给 dim1024 续训
+   到 2–3 epoch（或加新 shards），看 val 还能降多少。这是验证「v8 的容量是否被数据卡住」的最低成本实验。
+3. **自然收尾（漂亮里程碑）**：项目已 8 版 + 从零全栈（autograd/backward/AdamW/DDP/bf16/重计算/export）+
+   完整的数据轴 + 容量轴 + Chinchilla 边际分析——作为一条**学习线**已经讲完了一个完整的故事。
+
+**我的判断**：作为**工程/学习项目**，v8 是一个天然的收尾点——8 版把「数据量 / 数据广度 / 容量」三根轴都
+测过，并落到「单轴 ~3%、要双轴一起 scale」这个有分量的元结论上，全栈基建闭环。若要再做一版，**首选 2
+（dim1024 多喂数据）**：它最便宜、且直接回答 v8 留下的唯一悬念（容量被数据卡住的程度）——v8 val 未饱和说明
+这一版很可能立刻见效；真要追求规模再上 1（双轴）。详见对比表与 evolution.md 的「容量轴」与「边际递减」两条。

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

@@ -15,16 +15,43 @@ val loss 一栏给的是各版**各自训练 run 报告的 best val**（held-out
 注：v0/v1 训练用 seq128、v2 用 seq256，eval 窗口不同 → 同一保留集 + 同一 eval 设置（seq256/64batch）
 重评 v1=2.6756→v2=2.0418（低 0.634，apples-to-apples）；下表 best-val 同向。
 
-| 版本 | 数据 | 架构 (dim/L/heads·hd/ffn) | core 参数 | 总参数 | val loss | 备注 |
+**tokens / epoch 两列让数据饱和可见**：v4→v5 同 arch、数据 ×3.5（1.54→5.33 epoch），val 仅 ↓0.06（~5%）
-|---|---|---|---|---|---|---|
+且末段走平 ⇒ TinyStories 在 dim768 已近**数据天花板**（详见 [05-v5](05-v5-tinystories-dim768.md)）。**v6→v7 同样揭示
-| [v0-baseline](../../docs/05-training-loop.md) | TinyStories valid 3MB 切片 (~72 万 tok) | 32 / 4 / 2·16 / 64 | ~41K | 3.26M | **3.8050** | 太小不可用；采样陷入 "mommy's mommy's mommy" 循环 |
+「重复老数据」边际薄**：同一 2.255B FineWeb 子集多喂 epoch（1.02→1.45），FineWeb val 仅 ↓0.05 且走平 ⇒ 该子集
-| [v1-tinystories-dim256](01-v1-tinystories-dim256.md) | TinyStories **全量 train** (468.3M tok, u16 缓存) | 256 / 8 / 8·32 / 1024 | 8.39M | 34.13M | **2.5847** | 全量数据 + dim256/8L；val 低 1.22，采样连贯成篇；~25.9min/单卡 |
+在 dim768 也近顶（详见 [07-v7](07-v7-fineweb-edu-dim768.md)）。两条都说明：真·增益要**新数据**（v6 换更广语料才抬了天花板），不是同子集多读几遍。
-| [v2-tinystories-dim384](02-v2-tinystories-dim384.md) | TinyStories 全量 (复用 v1 缓存, 训 ~36.9M tok) | 384 / 12 / 12·32 / 1536 | 28.32M | 66.92M | **1.7055** | dim384/12L + **DDP 4 卡**；val 比 v1 低 0.88，情节更长；~2.8h/4 卡。⚠️ DDP 弱扩展见 [KI-1](../known-issues.md) |
+**v8 改测容量轴**：同 v6/v7 子集、纯把 dim768→dim1024（core 127M→226M），FineWeb val 3.07/3.01→**2.98** ⇒
-| [v3-tinystories-dim512](03-v3-tinystories-dim512.md) | TinyStories 全量 (复用 v1 缓存, 训 ~245.8M tok, ~0.53 epoch) | 512 / 16 / 16·32 / 2048 | 67.13M | 118.59M | **1.3027** | dim512/16L + **单卡 batched (T10)**；val 比 v2 低 0.40，带动机/转折的连续叙事；~2.65h/单卡 ~26K tok/s。T10 修 KI-1 根因(launch-bound)，单卡避开 KI-5 |
+**容量有用**（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 的
+~1.1 比大小——真实网页熵高，~3.0 是预期值不是回退。v6 的判据是采样质量 + transfer eval（见
+[06-v6](06-v6-fineweb-edu-dim768.md)）。下表 v6 行的 val 单独标注分布。
+
+| 版本 | 数据 | 训练 token | epoch | 架构 (dim/L/heads·hd/ffn) | core 参数 | 总参数 | val loss | 备注 |
+|---|---|---|---|---|---|---|---|---|
+| [v0-baseline](../../docs/05-training-loop.md) | TinyStories valid 3MB 切片 (~72 万 tok) | ~0.72M | — | 32 / 4 / 2·16 / 64 | ~41K | 3.26M | **3.8050** | 太小不可用；采样陷入 "mommy's mommy's mommy" 循环 |
+| [v1-tinystories-dim256](01-v1-tinystories-dim256.md) | TinyStories **全量 train** (468.3M tok, u16 缓存) | ~5.1M | — | 256 / 8 / 8·32 / 1024 | 8.39M | 34.13M | **2.5847** | 全量数据 + dim256/8L；val 低 1.22，采样连贯成篇；~25.9min/单卡 |
+| [v2-tinystories-dim384](02-v2-tinystories-dim384.md) | TinyStories 全量 (复用 v1 缓存) | ~36.9M | — | 384 / 12 / 12·32 / 1536 | 28.32M | 66.92M | **1.7055** | dim384/12L + **DDP 4 卡**；val 比 v1 低 0.88，情节更长；~2.8h/4 卡。⚠️ DDP 弱扩展见 [KI-1](../known-issues.md) |
+| [v3-tinystories-dim512](03-v3-tinystories-dim512.md) | TinyStories 全量 (复用 v1 缓存) | ~245.8M | ~0.53 | 512 / 16 / 16·32 / 2048 | 67.13M | 118.59M | **1.3027** | dim512/16L + **单卡 batched (T10)**；val 比 v2 低 0.40，带动机/转折的连续叙事；~2.65h/单卡 ~26K tok/s。T10 修 KI-1 根因(launch-bound)，单卡避开 KI-5 |
+| [v4-tinystories-dim768](04-v4-tinystories-dim768.md) | TinyStories 全量 (复用 v1 缓存) | ~720.9M | ~1.54 | 768 / 18 / 24·32 / 2048 | 127.43M | 204.63M | **1.1690** | dim768/18L + **8 卡 DDP fp32**；val 比 v3 低 0.13，细节更具体、结构更完整；~84min/8 卡 ~145K tok/s。验证 T11 缓存分配器在 dim768 多卡扩展；⚠️ fp32 per-rank batch 32 OOM = bf16(KI-2) 触发点 |
+| [v5-tinystories-dim768](05-v5-tinystories-dim768.md) | TinyStories 全量 (复用 v1 缓存) | **~2.49B** | **~5.33** | 768 / 18 / 24·32 / 2048 (**同 v4**) | 127.43M | 204.63M | **1.1102** | **架构同 v4**，唯一变量=数据量 + **8 卡 DDP bf16**(global 256)；~3.2h/8 卡 ~217K tok/s。⚠️ **数据天花板**：数据 ×3.5 仅 val ↓0.06(~5%) 且末段走平 ⇒ TinyStories 在 dim768 近饱和，v6 该换轴(更大模型/更广语料) |
+| [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 复读未解决 |
 
 ## 下一档（提案）
 
-- **v4**（待派发）：见 `03-v3-*.md` 末尾 "v4 提案"——放大 dim640–768/20–24L (~130–200M core) +
+- **v11**：优先走**更大模型 + 更长 context**，而不是继续只补数据。smoke 已验证 dim1536/20L/48q/12kv/ffn6144
-  ~600M–1B token，目标 val ~1.0–1.1；多卡需先修 KI-5（分桶 all-reduce），模型变大后启用 KI-2/3
+  能跑 seq512 和 seq1024，但峰值约 30.5GiB，贴近 5090 32GB 上限。建议先做 v11a（seq512，约 42h），
-  (bf16/重计算)，并按数据阶梯开始广化语料（TinyStories + 通用高质语料）。
+  或明确接受 2.5 天预算后做 v11b（seq1024，约 61h）。v11 必须使用固定 eval v1，避免 moving-tail 继续污染横比。
-</content>
+
+> **v7 时的提案（已被 v8 兑现，归档）**：v7 把首选定为「新 FineWeb shards」，把「更大模型(dim1024+，容量轴，
+> 需先做 T13 激活重计算)」列为待测。**v8 走了容量轴**并证明它有用（但 ~3%），把「是否 capacity-limited」从
+> 悬念变成了「部分是」的结论。

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/fineweb_to_txt.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""FineWeb-edu parquet -> plain UTF-8 .txt for xtrain's Corpus loader (Scaling v6).
+
+v0-v5 trained on TinyStories; v6 broadens the data to FineWeb-edu
+(HuggingFaceFW/fineweb-edu, the `sample/10BT` subset) to isolate the
+data-source variable while keeping the v4/v5 architecture fixed.
+
+xtrain's `Corpus::load` (crates/xtrain-train/src/data.rs) reads a UTF-8 text
+file and tokenizes it with the gpt2 BPE; it treats `<|endoftext|>` as the
+document boundary (the gpt2 tokenizer emits id 50256 for it). This script
+extracts the `text` column from one or more FineWeb-edu parquet shards and
+writes each document followed by `<|endoftext|>`, producing exactly that
+format. Reads row-group by row-group so memory stays bounded regardless of
+shard size.
+
+Usage (on dash5, in data/):
+    python3 scripts/fineweb_to_txt.py OUT.txt SHARD1.parquet [SHARD2.parquet ...]
+
+The corpus / .txt / .u16.bin caches are multi-GB and live on dash5 only
+(gitignored); only this script is committed.
+"""
+
+import sys
+import pyarrow.parquet as pq
+
+EOS = "<|endoftext|>"
+
+
+def main() -> None:
+    if len(sys.argv) < 3:
+        sys.exit(f"usage: {sys.argv[0]} OUT.txt SHARD.parquet [SHARD.parquet ...]")
+    out_path = sys.argv[1]
+    shards = sys.argv[2:]
+
+    docs = 0
+    chars = 0
+    with open(out_path, "w", encoding="utf-8") as out:
+        for shard in shards:
+            pf = pq.ParquetFile(shard)
+            n_rg = pf.num_row_groups
+            print(f"{shard}: {pf.metadata.num_rows} rows, {n_rg} row groups", flush=True)
+            for rg in range(n_rg):
+                col = pf.read_row_group(rg, columns=["text"]).column("text")
+                for s in col:
+                    text = s.as_py()
+                    if text is None:
+                        continue
+                    out.write(text)
+                    out.write(EOS)
+                    docs += 1
+                    chars += len(text)
+            print(f"  done {shard}: cumulative {docs} docs", flush=True)
+
+    print(f"wrote {out_path}: {docs} docs, {chars} text chars "
+          f"(+ {docs} x '{EOS}' separators)", flush=True)
+
+
+if __name__ == "__main__":
+    main()

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