docs: Phase T10 — batched forward

docs/09-batched-forward.md: the launch-bound diagnosis recap, the
[B*S,dim]-flatten + fused batched-attention design (RoPE per-seq position +
causal masking inline in softmax), the attention forward/backward via
strided-batched GEMM, autograd implications, the looped-split/merge dead-end
post-mortem (1127 tok/s, host round-trips), verification methods + before→after
throughput, and the v3 recommendation (per-rank batch 16-32, single/small world
until KI-5 bucketed all-reduce lands).

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

crates/xtrain-autodiff/src/ops.rs

@@ -120,15 +120,17 @@ pub fn swiglu(gate: &Var, up: &Var) -> Var {
     mul(&silu(gate), up)
 }
 
-/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]`. Orthogonal map, so the
-/// backward is the inverse rotation of `dy` — no cached forward values needed.
-pub fn rope(x: &Var, theta: f32) -> Var {
-    let out = x.value().rope(theta);
+/// 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
+/// cached forward values needed.
+pub fn rope(x: &Var, theta: f32, period: usize) -> Var {
+    let out = x.value().rope(theta, period);
     Var::from_op(
         out,
         vec![x.clone()],
         Box::new(move |dy, parents| {
-            Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta));
+            Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta, period));
         }),
     )
 }
@@ -190,6 +192,20 @@ pub fn transpose_3d01(x: &Var) -> Var {
     )
 }
 
+/// 4D axis-(1,2) transpose `[a,b,c,d] -> [a,c,b,d]`. Self-inverse structure: the
+/// backward is the same transpose applied to the grad. Lays out the batched
+/// multi-head attention `[B,S,nh,hd] <-> [B,nh,S,hd]`.
+pub fn transpose_4d12(x: &Var) -> Var {
+    let out = x.value().transpose_4d12();
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(|d, parents| {
+            Var::push_grad(&parents[0], d.transpose_4d12());
+        }),
+    )
+}
+
 /// 2D transpose `[r,c] -> [c,r]` as an autograd node (backward transposes the
 /// grad back). Used for `Kᵀ` in attention scores.
 pub fn transpose_2d(x: &Var) -> Var {
@@ -266,6 +282,29 @@ pub fn merge_heads(heads_v: &[Var]) -> Var {
     )
 }
 
+/// Batched causal scaled-dot-product attention. `q`,`k`,`v` are each
+/// `[bh, seq, head_dim]` (bh = batch·n_heads). Returns `[bh, seq, head_dim]`.
+/// One fused op (2 batched GEMMs + 1 causal-softmax kernel forward; 4 batched
+/// GEMMs + 1 softmax-backward kernel in backward) — replaces the per-(batch,head)
+/// matmul/softmax loop, so attention is a handful of launches regardless of bh.
+/// Caches the softmax `probs` for backward.
+pub fn attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
+    let (out, probs) = q.value().attention(&k.value(), &v.value(), scale);
+    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::attention_backward(&q, &k, &v, &probs, dout, scale);
+            Var::push_grad(&parents[0], dq);
+            Var::push_grad(&parents[1], dk);
+            Var::push_grad(&parents[2], dv);
+        }),
+    )
+}
+
 /// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
 /// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
 /// scaled by the upstream scalar grad.

crates/xtrain-autodiff/tests/autograd.rs

@@ -327,12 +327,12 @@ fn rope_bwd() {
     let w = fill(n, 82);
 
     let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim]));
-    let out = ops::rope(&x, theta);
+    let out = ops::rope(&x, theta, tokens);
     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]| weighted_sum(&cuda(v, s).rope(theta), &wf);
+    let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).rope(theta, tokens), &wf);
     report(
         "rope dX",
         &grad_check(
@@ -345,6 +345,38 @@ fn rope_bwd() {
     );
 }
 
+// ---- rope batched (per-sequence position = row % period) ----
+// tokens = B*S laid end to end; period = S. Sequences 2 and 3 re-use positions
+// 0..S, so the kernel's `tok % period` must reset RoPE per sequence.
+#[test]
+fn rope_batched_bwd() {
+    require_gpu();
+    let (b, s, heads, head_dim) = (3, 4, 2, 8);
+    let tokens = b * s;
+    let n = tokens * heads * head_dim;
+    let theta = 10000.0;
+    let x_h = fill(n, 83);
+    let w = fill(n, 84);
+
+    let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim]));
+    let out = ops::rope(&x, theta, s);
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let wf = w.clone();
+    let lx = move |v: &[f32], sh: &[usize]| weighted_sum(&cuda(v, sh).rope(theta, s), &wf);
+    report(
+        "rope batched dX",
+        &grad_check(
+            &x_h,
+            &[tokens, heads, head_dim],
+            &lx,
+            dx.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}
+
 // ---- softmax ----
 #[test]
 fn softmax_bwd() {
@@ -501,6 +533,98 @@ fn attention_composed_bwd() {
     );
 }
 
+// ---- transpose_4d12 ([a,b,c,d] -> [a,c,b,d]) ----
+#[test]
+fn transpose_4d12_bwd() {
+    require_gpu();
+    let (a, b, c, d) = (2, 3, 4, 5);
+    let n = a * b * c * d;
+    let x_h = fill(n, 131);
+    let w = fill(n, 132);
+
+    let x = Var::leaf(cuda(&x_h, &[a, b, c, d]));
+    let out = ops::transpose_4d12(&x);
+    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]| weighted_sum(&cuda(v, s).transpose_4d12(), &wf);
+    report(
+        "transpose_4d12 dX",
+        &grad_check(&x_h, &[a, b, c, d], &lx, dx.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- fused batched causal attention (the T10 op) ----
+// q,k,v: [bh, seq, hd]. Grad-check dq/dk/dv against finite-diff of L = sum(W∘out).
+// bh = 2 (e.g. batch 1 × 2 heads, or 2 sequences × 1 head) exercises the batched
+// GEMM stride; the causal mask is applied inside the op.
+#[test]
+fn 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();
+    let q_h = fill(n, 141);
+    let k_h = fill(n, 142);
+    let v_h = fill(n, 143);
+    let w = fill(n, 144);
+
+    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::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.attention(&kv, &vv, scale);
+        weighted_sum(&o, &w)
+    };
+    let (kf, vf, ff) = (k_h.clone(), v_h.clone(), fwd.clone());
+    let lq = move |x: &[f32], _s: &[usize]| ff(x, &kf, &vf);
+    report(
+        "attn(batched) dQ",
+        &grad_check(
+            &q_h,
+            &[bh, seq, hd],
+            &lq,
+            dq.as_slice::<f32>(),
+            cfg_nonlinear(),
+        ),
+    );
+    let (qf, vf, ff) = (q_h.clone(), v_h.clone(), fwd.clone());
+    let lk = move |x: &[f32], _s: &[usize]| ff(&qf, x, &vf);
+    report(
+        "attn(batched) dK",
+        &grad_check(
+            &k_h,
+            &[bh, seq, hd],
+            &lk,
+            dk.as_slice::<f32>(),
+            cfg_nonlinear(),
+        ),
+    );
+    let (qf, kf, ff) = (q_h.clone(), k_h.clone(), fwd.clone());
+    let lv = move |x: &[f32], _s: &[usize]| ff(&qf, &kf, x);
+    report(
+        "attn(batched) dV",
+        &grad_check(
+            &v_h,
+            &[bh, seq, hd],
+            &lv,
+            dv.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}
+
 // --- test helpers ---
 
 // Scalar loss node L = sum(W ∘ out): wraps a fixed-weight Var and reduces. We

crates/xtrain-cuda/build.rs

@@ -35,6 +35,7 @@ fn main() {
         .file("../../csrc/ops/nn.cu")
         .file("../../csrc/ops/model.cu")
         .file("../../csrc/ops/optim.cu")
+        .file("../../csrc/ops/attention.cu")
         .compile("xtrain_cuda_kernels");
 }
 

crates/xtrain-cuda/src/cublas.rs

@@ -93,3 +93,69 @@ pub fn sgemm(
         assert_eq!(status, 0, "cublasSgemm failed: {status}");
     });
 }
+
+/// Strided-batched row-major SGEMM: for each `i` in `0..batch`,
+/// `C_i[m,n] = alpha·opA(A_i)·opB(B_i) + beta·C_i`, where `A_i`/`B_i`/`C_i` are
+/// consecutive matrices laid `stride_*` elements apart in one contiguous buffer.
+/// Same row-major⟺col-major trick as [`sgemm`] (compute col-major `Cᵀ`), applied
+/// per batch element. Used for the batched attention `QKᵀ` / `PV` GEMMs (and their
+/// backwards), so the whole attention runs as 2 batched-GEMM launches, not a
+/// per-(batch,head) Python loop. `A`/`B`/`C` are device pointers to the first
+/// matrix; strides are in ELEMENTS.
+#[allow(clippy::too_many_arguments)]
+pub fn sgemm_strided_batched(
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: f32,
+    a: *const f32,
+    stride_a: usize,
+    b: *const f32,
+    stride_b: usize,
+    beta: f32,
+    c: *mut f32,
+    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
+    };
+
+    with_handle(|handle| {
+        let status = unsafe {
+            ffi::cublasSgemmStridedBatched(
+                handle,
+                op_b,
+                op_a,
+                n as i32,
+                m as i32,
+                k as i32,
+                &alpha,
+                b,
+                ldb as i32,
+                stride_b as i64,
+                a,
+                lda as i32,
+                stride_a as i64,
+                &beta,
+                c,
+                ldc as i32,
+                stride_c as i64,
+                batch as i32,
+            )
+        };
+        assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
+    });
+}

crates/xtrain-cuda/src/ffi.rs

@@ -125,7 +125,9 @@ unsafe extern "C" {
     pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream);
     pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream);
 
-    // RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index.
+    // RoPE (rotate_half), x:[tokens,heads,head_dim], position = (token index %
+    // period). `period` = sequence length, so a flattened batch of sequences gets
+    // per-sequence positions; period == tokens reproduces the single-sequence case.
     pub fn launch_rope_f32(
         x: *const f32,
         y: *mut f32,
@@ -133,6 +135,7 @@ unsafe extern "C" {
         heads: i32,
         head_dim: i32,
         theta: f32,
+        period: i32,
         s: CudaStream,
     );
     pub fn launch_rope_dx_f32(
@@ -142,6 +145,7 @@ unsafe extern "C" {
         heads: i32,
         head_dim: i32,
         theta: f32,
+        period: i32,
         s: CudaStream,
     );
 
@@ -211,6 +215,31 @@ unsafe extern "C" {
         c: i32,
         s: CudaStream,
     );
+    // 4D axis-(1,2) transpose: in:[a,b,c,d] -> out:[a,c,b,d]. out[i,k,j,l]=in[i,j,k,l].
+    pub fn launch_transpose_4d12_f32(
+        input: *const f32,
+        out: *mut f32,
+        a: i32,
+        b: i32,
+        c: i32,
+        d: i32,
+        s: CudaStream,
+    );
+}
+
+// Batched attention helper (csrc/ops/attention.cu): causal row-wise softmax over
+// score rows [rows, seq] with query position = (row % seq); scales logits by
+// `scale` (= 1/sqrt(head_dim)) and masks future columns to probability 0.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    pub fn launch_softmax_causal_f32(
+        x: *const f32,
+        y: *mut f32,
+        rows: i32,
+        seq: i32,
+        scale: f32,
+        s: CudaStream,
+    );
 }
 
 // GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
@@ -267,6 +296,27 @@ unsafe extern "C" {
         c: *mut f32,
         ldc: i32,
     ) -> i32;
+    #[allow(clippy::too_many_arguments)]
+    pub fn cublasSgemmStridedBatched(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const f32,
+        a: *const f32,
+        lda: i32,
+        stride_a: i64,
+        b: *const f32,
+        ldb: i32,
+        stride_b: i64,
+        beta: *const f32,
+        c: *mut f32,
+        ldc: i32,
+        stride_c: i64,
+        batch_count: i32,
+    ) -> i32;
 }
 
 #[cfg(not(no_cuda))]

crates/xtrain-distributed/src/ddp.rs

@@ -17,7 +17,7 @@ use std::thread;
 use std::time::Instant;
 
 use xtrain_autodiff::tape::Var;
-use xtrain_model::{Config, TinyTransformer, ids_tensor};
+use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
 use xtrain_optim::GpuAdamW;
 use xtrain_tensor::Device;
 use xtrain_train::checkpoint;
@@ -91,10 +91,11 @@ pub fn train_rank(
     let mut losses = Vec::with_capacity(cfg.steps);
     let mut evals = Vec::new();
     let mut best_val: Option<f32> = None;
-    // Each rank reaches the global batch mean as (Σ_global / world) · (1/b_local),
-    // where b_local = batch_size / world (see DdpContext::all_reduce_average_grads).
+    // Each rank runs ONE batched forward over its b_local = batch_size/world
+    // sequences → backward grad = local mean (Σ_local / b_local). all_reduce_average
+    // (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 inv_batch_local = 1.0 / batch_local as f32;
     let start = Instant::now();
     let mut tokens_seen: u64 = 0;
     // Rank 0 owns the held-out eval + best-val checkpoint (params are identical
@@ -105,31 +106,36 @@ pub fn train_rank(
         let lr = cfg.schedule.lr(step);
 
         // Draw the whole global batch from the shared RNG (same on every rank);
-        // run forward+backward only on this rank's shard. The tape SUMs the
-        // shard's grads; the union of shards == the single-GPU batch.
-        let mut local_loss_sum = 0.0f32;
+        // collect only this rank's shard (global index % world == rank) and run it
+        // as ONE batched forward/backward. The union of shards == the single-GPU
+        // batch; each rank's backward yields its local mean (Σ_local / b_local).
+        let mut inputs = Vec::with_capacity(batch_local);
+        let mut targets_v = Vec::with_capacity(batch_local);
         for i in 0..cfg.batch_size {
             let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
-            if i % ctx.world != ctx.rank {
-                continue; // not this rank's sequence
+            if i % ctx.world == ctx.rank {
+                inputs.push(input);
+                targets_v.push(target);
             }
-            let ids = ids_tensor(&input, device);
-            let targets = ids_tensor(&target, device);
-            let loss = model.loss(&ids, &targets);
-            local_loss_sum += read_scalar(&loss);
-            loss.backward();
-            tokens_seen += 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/world.
+        // AllReduce(sum) + /world the grads → every rank holds Σ_global/B_global
+        // (local means summed over ranks, /world = global mean). See note above.
         ctx.all_reduce_average_grads(&params);
-        // The reported loss is the global mean: average local sums across ranks.
-        let step_loss = all_reduce_loss(ctx, local_loss_sum) / cfg.batch_size as f32;
+        // Reported loss = global mean: sum the per-rank local sums (= mean·b_local)
+        // across ranks, /B_global. With equal b_local this is mean over ranks.
+        let step_loss =
+            all_reduce_loss(ctx, local_mean * batch_local as f32) / cfg.batch_size as f32;
         losses.push(step_loss);
 
-        // clip pre_scale = 1/b_local finishes the average to Σ_global/B_global,
-        // identical to the single-GPU clip(pre_scale = 1/B_global).
-        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch_local);
+        // Grads are already the global-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 {
             p.zero_grad();

crates/xtrain-distributed/tests/ddp_correctness.rs

@@ -13,7 +13,7 @@ use std::time::Instant;
 
 use xtrain_cuda::device;
 use xtrain_distributed::{DdpConfig, DdpContext, build_model, get_unique_id, launch, train_rank};
-use xtrain_model::{Config, ids_tensor};
+use xtrain_model::{Config, batched_ids_tensor};
 use xtrain_optim::GpuAdamW;
 use xtrain_tensor::Device;
 use xtrain_train::clip::clip_grad_norm_gpu;
@@ -47,22 +47,25 @@ fn run_single_gpu(cfg: Config, corpus: &Corpus, dcfg: &DdpConfig) -> (Vec<f32>,
     let params = model.params();
     let mut opt = GpuAdamW::new(dcfg.weight_decay);
     let mut rng = dcfg.seed;
-    let inv_batch = 1.0 / dcfg.batch_size as f32;
     let mut losses = Vec::new();
 
     for step in 0..dcfg.steps {
         let lr = dcfg.schedule.lr(step);
-        let mut loss_sum = 0.0f32;
+        // Sample the whole global batch and run it as ONE batched forward/backward
+        // (matches the T10 DDP path: backward yields the global-batch mean grad).
+        let mut inputs = Vec::with_capacity(dcfg.batch_size);
+        let mut targets_v = Vec::with_capacity(dcfg.batch_size);
         for _ in 0..dcfg.batch_size {
             let (input, target) = corpus.sample(dcfg.seq_len, &mut rng);
-            let ids = ids_tensor(&input, device);
-            let targets = ids_tensor(&target, device);
-            let loss = model.loss(&ids, &targets);
-            loss_sum += loss.value().to_device(Device::Cpu).as_slice::<f32>()[0];
-            loss.backward();
+            inputs.push(input);
+            targets_v.push(target);
         }
-        losses.push(loss_sum * inv_batch);
-        clip_grad_norm_gpu(&params, dcfg.max_grad_norm, inv_batch);
+        let ids = batched_ids_tensor(&inputs, device);
+        let targets = batched_ids_tensor(&targets_v, device);
+        let loss = model.loss_batched(&ids, &targets, dcfg.batch_size);
+        losses.push(loss.value().to_device(Device::Cpu).as_slice::<f32>()[0]);
+        loss.backward();
+        clip_grad_norm_gpu(&params, dcfg.max_grad_norm, 1.0);
         opt.step(lr, &params);
         for p in &params {
             p.zero_grad();

crates/xtrain-model/src/lib.rs

@@ -24,4 +24,4 @@ pub use config::Config;
 #[cfg(not(no_cuda))]
 mod model;
 #[cfg(not(no_cuda))]
-pub use model::{TinyTransformer, ids_tensor, param_to_host};
+pub use model::{TinyTransformer, batched_ids_tensor, ids_tensor, param_to_host};

crates/xtrain-model/src/model.rs

@@ -30,7 +30,6 @@ pub struct TinyTransformer {
     blocks: Vec<Block>,
     final_norm: Var, // [dim]
     lm_head: Var,    // [dim, vocab]
-    device: Device,
 }
 
 impl TinyTransformer {
@@ -72,7 +71,6 @@ impl TinyTransformer {
             blocks,
             final_norm,
             lm_head,
-            device,
         }
     }
 
@@ -106,16 +104,34 @@ impl TinyTransformer {
     }
 
     /// Forward over a single sequence of token `ids` (`[seq]` I32 on this
-    /// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`.
+    /// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`. This
+    /// is the batch-1 special case of [`forward_batched`](Self::forward_batched)
+    /// (used by the autoregressive sampler / inference path).
     pub fn forward(&self, ids: &Tensor) -> Var {
-        let seq = ids.shape()[0];
-        let mask = self.causal_mask(seq);
+        self.forward_batched(ids, 1)
+    }
 
-        let mut h = ops::embedding(&self.embed, ids); // [seq, dim]
+    /// Batched forward over `batch` sequences of equal length `seq`, flattened to
+    /// `[batch*seq]` I32 ids in sequence-major order (sequence 0's `seq` tokens,
+    /// then sequence 1's, …). Returns logits `[batch*seq, vocab]` in the SAME flat
+    /// layout. The whole graph runs on the flattened tokens so every linear
+    /// projection is ONE big `[batch*seq, dim] × [dim, out]` GEMM (the
+    /// GPU-filling win); only attention is sequence-aware (per-sequence causal
+    /// mask + RoPE position, NO cross-sequence attention).
+    pub fn forward_batched(&self, ids: &Tensor, batch: usize) -> Var {
+        let total = ids.shape()[0];
+        assert_eq!(
+            total % batch,
+            0,
+            "ids len {total} not divisible by batch {batch}"
+        );
+        let seq = total / batch;
+
+        let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim]
         for b in &self.blocks {
             // --- Attention sub-block (pre-norm + residual) ---
             let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
-            let attn = self.attention(b, &normed, &mask, seq);
+            let attn = self.attention(b, &normed, batch, seq);
             h = ops::add(&h, &attn);
 
             // --- MLP sub-block (pre-norm + residual) ---
@@ -125,7 +141,7 @@ impl TinyTransformer {
         }
 
         let h = ops::rms_norm(&h, &self.final_norm, self.cfg.eps);
-        ops::matmul(&h, &self.lm_head) // [seq, vocab]
+        ops::matmul(&h, &self.lm_head) // [batch*seq, vocab]
     }
 
     /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
@@ -134,76 +150,76 @@ impl TinyTransformer {
         ops::cross_entropy(&logits, targets)
     }
 
-    /// Multi-head causal self-attention. `x`:[seq,dim] (already normed).
-    fn attention(&self, b: &Block, x: &Var, mask: &Var, seq: usize) -> Var {
+    /// Batched cross-entropy mean loss: `forward_batched(ids, batch)` against
+    /// flat `targets` (`[batch*seq]` I32, same sequence-major layout). The CE mean
+    /// is over all `batch*seq` rows — identical to averaging the per-sequence
+    /// losses, so the loss value matches the looped single-sequence path.
+    pub fn loss_batched(&self, ids: &Tensor, targets: &Tensor, batch: usize) -> Var {
+        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]
+    /// (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).
+    fn attention(&self, b: &Block, x: &Var, batch: usize, seq: usize) -> Var {
         let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
+        let total = batch * seq;
+        let bh = batch * nh;
         let scale = 1.0 / (hd as f32).sqrt();
 
-        // Project, then lay out as per-head [seq, head_dim] tensors.
-        //   [seq,dim] @ [dim,dim] = [seq,dim]
-        //   reshape   [seq, nh, hd]
+        // Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
+        //   [B*S,dim] @ [dim,dim] = [B*S,dim]
+        //   reshape   [B*S, nh, hd]
         //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
-        //   rope      (kernel expects exactly [tokens, heads, head_dim])
-        //   transpose [nh, seq, hd]  → split into nh × [seq, hd]
-        let to_heads = |proj: Var, norm: Option<&Var>| -> Vec<Var> {
-            let r = ops::reshape(&proj, &[seq, nh, hd]);
+        //   rope      [B*S, nh, hd] with per-sequence position (period = seq)
+        //   reshape   [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
+        let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
+            let r = ops::reshape(&proj, &[total, nh, hd]);
             let r = match norm {
-                // Per-head RMSNorm: flatten the (seq,nh) head rows, norm over hd,
+                // 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) => {
-                    let flat = ops::reshape(&r, &[seq * nh, hd]);
+                    let flat = ops::reshape(&r, &[total * nh, hd]);
                     let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
-                    let r = ops::reshape(&normed, &[seq, nh, hd]);
-                    ops::rope(&r, self.cfg.rope_theta)
+                    let r = ops::reshape(&normed, &[total, nh, hd]);
+                    ops::rope(&r, self.cfg.rope_theta, seq)
                 }
                 None => r,
             };
-            let t = ops::transpose_3d01(&r); // [nh, seq, hd]
-            ops::split_heads(&t)
+            let r = ops::reshape(&r, &[batch, seq, nh, hd]);
+            let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
+            ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
         };
 
-        let q = to_heads(ops::matmul(x, &b.wq), Some(&b.q_norm));
-        let k = to_heads(ops::matmul(x, &b.wk), Some(&b.k_norm));
-        let v = to_heads(ops::matmul(x, &b.wv), None);
+        let q = to_bh(ops::matmul(x, &b.wq), Some(&b.q_norm));
+        let k = to_bh(ops::matmul(x, &b.wk), Some(&b.k_norm));
+        let v = to_bh(ops::matmul(x, &b.wv), None);
 
-        // Per-head scaled-dot-product attention with causal mask.
-        let heads_out: Vec<Var> = (0..nh)
-            .map(|i| {
-                let kt = ops::transpose_2d(&k[i]); // [hd, seq]
-                let scores = ops::scale(&ops::matmul(&q[i], &kt), scale); // [seq,seq]
-                let scores = ops::add(&scores, mask); // causal
-                let probs = ops::softmax(&scores);
-                ops::matmul(&probs, &v[i]) // [seq, hd]
-            })
-            .collect();
+        // Fused batched causal SDPA over all B*nh (sequence,head) blocks at once
+        // (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
+        let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]
 
-        // Stack heads back: nh × [seq,hd] → [nh,seq,hd] → [seq,nh,hd] → [seq,dim].
-        let merged = ops::merge_heads(&heads_out); // [nh, seq, hd]
-        let t = ops::transpose_3d01(&merged); // [seq, nh, hd]
-        let concat = ops::reshape(&t, &[seq, nh * hd]); // [seq, dim]
+        // Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
+        // [B,S,nh,hd] → [B*S, dim].
+        let out = ops::reshape(&out, &[batch, nh, seq, hd]);
+        let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
+        let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
         ops::matmul(&concat, &b.wo) // out projection
     }
 
-    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[seq,dim].
+    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*seq,dim].
     fn swiglu_mlp(&self, b: &Block, x: &Var) -> Var {
         let gate = ops::matmul(x, &b.w_gate); // [seq, ffn_hidden]
         let up = ops::matmul(x, &b.w_up); // [seq, ffn_hidden]
         let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
         ops::matmul(&act, &b.w_down) // [seq, dim]
     }
-
-    /// Additive causal mask `[seq,seq]`: 0 on/below the diagonal, −1e9 above it
-    /// (so softmax zeros out future positions). A constant leaf (no grad needed,
-    /// but harmless if it accumulates one — it has no consumers downstream of x).
-    fn causal_mask(&self, seq: usize) -> Var {
-        let mut m = vec![0.0f32; seq * seq];
-        for i in 0..seq {
-            for j in (i + 1)..seq {
-                m[i * seq + j] = -1.0e9;
-            }
-        }
-        Var::leaf(Tensor::from_slice(&m, &[seq, seq]).to_device(self.device))
-    }
 }
 
 /// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step
@@ -216,3 +232,17 @@ pub fn param_to_host(v: &Var) -> Vec<f32> {
 pub fn ids_tensor(ids: &[i32], device: Device) -> Tensor {
     Tensor::from_slice(ids, &[ids.len()]).to_device(device)
 }
+
+/// Flatten `batch` equal-length sequences into one `[batch*seq]` I32 tensor in
+/// sequence-major order (the layout `forward_batched` expects). Each row of
+/// `seqs` is one sequence; all must have the same length.
+pub fn batched_ids_tensor(seqs: &[Vec<i32>], device: Device) -> Tensor {
+    assert!(!seqs.is_empty(), "empty batch");
+    let seq = seqs[0].len();
+    let mut flat = Vec::with_capacity(seqs.len() * seq);
+    for s in seqs {
+        assert_eq!(s.len(), seq, "ragged batch: sequences must be equal length");
+        flat.extend_from_slice(s);
+    }
+    Tensor::from_slice(&flat, &[flat.len()]).to_device(device)
+}

crates/xtrain-model/tests/batched.rs

@@ -0,0 +1,142 @@
+// T10 batched-forward equivalence: a batched forward over B sequences must equal
+// the old single-sequence path (run each sequence on its own, concatenate the
+// logits) — both for the forward logits AND every parameter's gradient.
+//
+// This is THE on-GPU correctness gate for batching (no PyTorch needed): if the
+// per-sequence RoPE position, per-sequence causal masking, or any flattened op
+// were wrong, the batched logits/grads would drift from the looped reference.
+//
+// Forward equivalence: batched logits[b*S+i] == single-seq-b logits[i].
+// Gradient equivalence: the batched loss is the mean over all B*S rows, i.e.
+// (1/B)·Σ_b mean_i(loss_b); summing the B single-sequence losses and scaling by
+// 1/B gives the SAME scalar, so their summed grads (tape fan-out) ×1/B match the
+// batched grads. We check that.
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, batched_ids_tensor, ids_tensor};
+use xtrain_tensor::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 batched_matches_looped_single_sequence() {
+    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;
+    let batch = 3usize;
+    let seq = 5usize;
+    // B distinct sequences (sequence-major), within vocab.
+    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 forward: ONE pass over [B*S]. ---
+    let bmodel = build(cfg, device);
+    let bids = batched_ids_tensor(&seqs, device);
+    let blogits = host(&bmodel.forward_batched(&bids, batch).value());
+
+    // --- Looped reference: each sequence on its own, concatenate logits. ---
+    let smodel = build(cfg, device);
+    let mut slogits = Vec::with_capacity(batch * seq * cfg.vocab);
+    for s in &seqs {
+        let ids = ids_tensor(s, device);
+        slogits.extend(host(&smodel.forward(&ids).value()));
+    }
+
+    // Forward equivalence (fp GEMM rounding only differs in summation order).
+    let max_rel = blogits
+        .iter()
+        .zip(&slogits)
+        .map(|(b, s)| (b - s).abs() / s.abs().max(1e-4))
+        .fold(0.0f32, f32::max);
+    println!("batched vs looped: logits max rel err = {max_rel:.3e}");
+    assert!(max_rel < 1e-3, "batched logits diverged: {max_rel:.3e}");
+
+    // --- Gradient equivalence. ---
+    // Batched: loss = mean over B*S rows; one backward.
+    let bparams = bmodel.params();
+    let btgt = batched_ids_tensor(&tgts, device);
+    let bloss = bmodel.loss_batched(&bids, &btgt, batch);
+    let bloss_val = host(&bloss.value())[0];
+    bloss.backward();
+
+    // Looped: Σ_b loss_b (each a per-sequence mean), then grad ×(1/B) == batched.
+    let sparams = smodel.params();
+    let mut sloss_sum = 0.0f32;
+    for (s, t) in seqs.iter().zip(&tgts) {
+        let ids = ids_tensor(s, device);
+        let tg = ids_tensor(t, device);
+        let l = smodel.loss(&ids, &tg);
+        sloss_sum += host(&l.value())[0];
+        l.backward();
+    }
+    println!(
+        "batched loss = {bloss_val:.6}  looped mean = {:.6}",
+        sloss_sum / batch as f32
+    );
+    assert!(
+        (bloss_val - sloss_sum / batch as f32).abs() < 1e-4,
+        "batched loss != looped mean"
+    );
+
+    let mut max_grad_rel = 0.0f32;
+    for (bp, sp) in bparams.iter().zip(&sparams) {
+        let bg = host(&bp.grad().expect("batched grad"));
+        let sg = host(&sp.grad().expect("looped grad"));
+        for (g_b, g_s) in bg.iter().zip(&sg) {
+            // looped grad is the SUM over B sequences; ×(1/B) recovers the mean.
+            let g_s = g_s / batch as f32;
+            let rel = (g_b - g_s).abs() / g_s.abs().max(1e-4);
+            max_grad_rel = max_grad_rel.max(rel);
+        }
+    }
+    println!("batched vs looped: grad max rel err = {max_grad_rel:.3e}");
+    assert!(
+        max_grad_rel < 5e-3,
+        "batched grads diverged: {max_grad_rel:.3e}"
+    );
+}

crates/xtrain-model/tests/parity.py

@@ -55,10 +55,13 @@ NH = int(cfg["n_heads"])
 HD = int(cfg["head_dim"])
 EPS = float(cfg["eps"])
 THETA = float(cfg["rope_theta"])
+# Batched: B sequences of length SEQ, flattened sequence-major to [B*SEQ] ids.
+B = int(cfg.get("batch", "1"))
+SEQ = int(cfg["seq"])
 
 ids = read_ids("ids.txt")
 targets = read_ids("targets.txt")
-SEQ = len(ids)
+assert len(ids) == B * SEQ, f"ids {len(ids)} != B*SEQ {B*SEQ}"
 
 # Load params as leaf tensors requiring grad (float64 for a clean reference).
 P = {}
@@ -76,15 +79,16 @@ def rms_norm(x, gamma):
     return x * torch.rsqrt(ms + EPS) * gamma
 
 
-def rope(x):  # x: [seq, nh, hd], position = token index, matching the kernel
+def rope(x):  # x: [B*SEQ, nh, hd], position = (row % SEQ) — resets per sequence
     half = HD // 2
     out = torch.empty_like(x)
     i = torch.arange(half, dtype=torch.float64)
     freq = THETA ** (-(2.0 * i) / HD)  # [half]
-    pos = torch.arange(SEQ, dtype=torch.float64).reshape(SEQ, 1)  # [seq,1]
-    ang = pos * freq  # [seq, half]
-    c = torch.cos(ang).reshape(SEQ, 1, half)
-    s = torch.sin(ang).reshape(SEQ, 1, half)
+    # Position within each sequence: rows 0..SEQ for seq 0, 0..SEQ for seq 1, ...
+    pos = (torch.arange(B * SEQ, dtype=torch.float64) % SEQ).reshape(B * SEQ, 1)
+    ang = pos * freq  # [B*SEQ, half]
+    c = torch.cos(ang).reshape(B * SEQ, 1, half)
+    s = torch.sin(ang).reshape(B * SEQ, 1, half)
     x0 = x[..., :half]
     x1 = x[..., half:]
     out[..., :half] = x0 * c - x1 * s
@@ -102,26 +106,30 @@ for l in range(NL):
                     "ffn_norm", "w_gate", "w_up", "w_down"]})
 
 idx = torch.tensor(ids, dtype=torch.long)
+# Per-sequence causal mask (broadcast over the batch); NO cross-sequence attention.
 mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float64), diagonal=1)
 
-h = emb[idx]  # [seq, dim]
+h = emb[idx]  # [B*SEQ, dim]  (everything stays flattened, matching the Rust path)
 for L in layers:
     # Attention
     x = rms_norm(h, L["attn_norm"])
-    q = (x @ L["wq"]).reshape(SEQ, NH, HD)
-    k = (x @ L["wk"]).reshape(SEQ, NH, HD)
-    v = (x @ L["wv"]).reshape(SEQ, NH, HD)
+    q = (x @ L["wq"]).reshape(B * SEQ, NH, HD)
+    k = (x @ L["wk"]).reshape(B * SEQ, NH, HD)
+    v = (x @ L["wv"]).reshape(B * SEQ, NH, HD)
     # 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).transpose(0, 1)  # [nh, seq, hd]
-    k = rope(k).transpose(0, 1)
-    v = v.transpose(0, 1)
+    q = rope(q)  # [B*SEQ, nh, hd]
+    k = rope(k)
+    # Reshape to [B, NH, SEQ, HD] so attention runs within each sequence.
+    q = q.reshape(B, SEQ, NH, HD).transpose(1, 2)  # [B, nh, seq, hd]
+    k = k.reshape(B, SEQ, NH, HD).transpose(1, 2)
+    v = v.reshape(B, SEQ, NH, HD).transpose(1, 2)
     scale = 1.0 / math.sqrt(HD)
-    scores = (q @ k.transpose(-1, -2)) * scale + mask  # [nh, seq, seq]
+    scores = (q @ k.transpose(-1, -2)) * scale + mask  # [B, nh, seq, seq]
     probs = torch.softmax(scores, dim=-1)
-    out = probs @ v  # [nh, seq, hd]
-    out = out.transpose(0, 1).reshape(SEQ, DIM)  # [seq, dim]
+    out = probs @ v  # [B, nh, seq, hd]
+    out = out.transpose(1, 2).reshape(B * SEQ, DIM)  # [B*SEQ, dim]
     attn = out @ L["wo"]
     h = h + attn
     # MLP
@@ -133,7 +141,7 @@ for L in layers:
     h = h + mlp
 
 h = rms_norm(h, final_norm)
-logits = h @ lm_head  # [seq, vocab]
+logits = h @ lm_head  # [B*SEQ, vocab]
 
 loss = torch.nn.functional.cross_entropy(
     logits, torch.tensor(targets, dtype=torch.long), reduction="mean")

crates/xtrain-model/tests/parity_dump.rs

@@ -53,12 +53,17 @@ fn dump_for_parity() {
     );
     fs::create_dir_all(&dir).unwrap();
 
-    // Fixed config + ids (independent of any text, for reproducibility).
+    // Fixed config + ids (independent of any text, for reproducibility). B>1 so
+    // the batched forward is exercised: 2 sequences of length 4, flattened
+    // 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.
     let mut cfg = Config::tiny();
     cfg.vocab = 12;
-    let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6];
+    let batch = 2usize;
+    let seq = 4usize;
+    let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major
     let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
-    let seq = ids.len();
 
     // Same deterministic init as the overfit test.
     let mut seed = 1u64;
@@ -83,6 +88,7 @@ fn dump_for_parity() {
         writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
         writeln!(f, "eps {:e}", cfg.eps).unwrap();
         writeln!(f, "rope_theta {:e}", cfg.rope_theta).unwrap();
+        writeln!(f, "batch {batch}").unwrap();
         writeln!(f, "seq {seq}").unwrap();
     }
     {
@@ -105,10 +111,11 @@ fn dump_for_parity() {
         write_vec(&dir, &format!("w_{name}.txt"), &param_to_host(p), &shape);
     }
 
-    // Forward logits + loss, then backward → per-param grads.
+    // Batched forward logits + loss (B sequences as one forward), then backward
+    // → per-param grads.
     let ids_t = ids_tensor(&ids, device);
     let targets_t = ids_tensor(&targets, device);
-    let logits = model.forward(&ids_t);
+    let logits = model.forward_batched(&ids_t, batch);
     write_vec(
         &dir,
         "logits.txt",
@@ -116,7 +123,7 @@ fn dump_for_parity() {
         logits.value().shape(),
     );
 
-    let loss = model.loss(&ids_t, &targets_t);
+    let loss = model.loss_batched(&ids_t, &targets_t, batch);
     let loss_val = param_to_host(&loss)[0];
     {
         let mut f = fs::File::create(dir.join("loss.txt")).unwrap();

crates/xtrain-tensor/src/tensor.rs

@@ -454,13 +454,20 @@ impl Tensor {
         dx
     }
 
-    /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position
-    /// of each token is its row index. Returns the rotated tensor.
+    /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; each token's
+    /// position is `row % period`. `period` = sequence length, so a flattened
+    /// batch `[B*S,heads,head_dim]` gets per-sequence positions (pass `period=S`);
+    /// pass `period=tokens` for a single sequence (position = row). Returns the
+    /// rotated tensor.
     #[cfg(not(no_cuda))]
-    pub fn rope(&self, theta: f32) -> Self {
+    pub fn rope(&self, theta: f32, period: usize) -> Self {
         assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]");
         let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
         assert_eq!(head_dim % 2, 0, "head_dim must be even");
+        assert!(
+            period > 0 && tokens % period == 0,
+            "tokens must be a multiple of period"
+        );
         let out = Tensor::zeros(&self.shape, DType::F32, self.device());
         unsafe {
             xtrain_cuda::ffi::launch_rope_f32(
@@ -470,6 +477,7 @@ impl Tensor {
                 heads as i32,
                 head_dim as i32,
                 theta,
+                period as i32,
                 std::ptr::null_mut(),
             );
         }
@@ -477,9 +485,9 @@ impl Tensor {
     }
 
     /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
-    /// orthogonal map, so it needs no cached forward values, only `theta`.
+    /// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
     #[cfg(not(no_cuda))]
-    pub fn rope_backward(dy: &Tensor, theta: f32) -> Self {
+    pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self {
         let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]);
         let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device());
         unsafe {
@@ -490,6 +498,7 @@ impl Tensor {
                 heads as i32,
                 head_dim as i32,
                 theta,
+                period as i32,
                 std::ptr::null_mut(),
             );
         }
@@ -667,6 +676,202 @@ impl Tensor {
         out
     }
 
+    // --- Batched attention (the T10 fused op) ---
+
+    /// Batched causal scaled-dot-product attention. `self`=Q, `k`, `v` are each
+    /// `[bh, seq, head_dim]` (bh = batch·n_heads), contiguous F32 on one GPU.
+    /// Computes, per batch element, `out = softmax(causal(Q·Kᵀ / √hd)) · V`. The
+    /// two GEMMs run as `cublasSgemmStridedBatched` and the softmax+scale+causal
+    /// mask is one kernel, so the whole attention is 3 launches regardless of bh.
+    /// Returns `(out, probs)` where `probs`:[bh,seq,seq] is cached for backward.
+    #[cfg(not(no_cuda))]
+    pub fn attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> (Tensor, Tensor) {
+        assert_eq!(self.ndim(), 3, "attention Q must be [bh,seq,head_dim]");
+        assert_eq!(self.shape(), k.shape(), "Q/K shape mismatch");
+        assert_eq!(self.shape(), v.shape(), "Q/V shape mismatch");
+        let (bh, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
+        let dev = self.device();
+
+        // scores[bh,seq,seq] = Q[bh,seq,hd] · Kᵀ[bh,hd,seq]
+        let scores = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
+        xtrain_cuda::cublas::sgemm_strided_batched(
+            false,
+            true,
+            seq,
+            seq,
+            hd,
+            1.0,
+            self.data_ptr() as *const f32,
+            seq * hd,
+            k.data_ptr() as *const f32,
+            seq * hd,
+            0.0,
+            scores.data_ptr() as *mut f32,
+            seq * seq,
+            bh,
+        );
+        // probs = softmax(causal(scores · scale)), one block per [bh·seq] row.
+        let probs = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
+        unsafe {
+            xtrain_cuda::ffi::launch_softmax_causal_f32(
+                scores.data_ptr() as *const f32,
+                probs.data_ptr() as *mut f32,
+                (bh * seq) as i32,
+                seq as i32,
+                scale,
+                std::ptr::null_mut(),
+            );
+        }
+        // out[bh,seq,hd] = probs[bh,seq,seq] · V[bh,seq,hd]
+        let out = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        xtrain_cuda::cublas::sgemm_strided_batched(
+            false,
+            false,
+            seq,
+            hd,
+            seq,
+            1.0,
+            probs.data_ptr() as *const f32,
+            seq * seq,
+            v.data_ptr() as *const f32,
+            seq * hd,
+            0.0,
+            out.data_ptr() as *mut f32,
+            seq * hd,
+            bh,
+        );
+        (out, probs)
+    }
+
+    /// Backward of [`attention`](Self::attention). Inputs: forward `q`,`k`,`v`,
+    /// the cached `probs`, the upstream `dout` (all batched `[bh,seq,*]`), and the
+    /// same `scale`. Returns `(dq, dk, dv)`.
+    ///
+    ///   dP = dOut · Vᵀ ; dV = Pᵀ · dOut
+    ///   dScores = softmax_jacobian(P, dP) · scale     (scale folded back in)
+    ///   dQ = dScores · K ; dK = dScoresᵀ · Q
+    ///
+    /// Masked (future) entries of P are 0, so the softmax Jacobian zeros their
+    /// gradient — the causal mask needs no special handling here.
+    #[cfg(not(no_cuda))]
+    pub fn attention_backward(
+        q: &Tensor,
+        k: &Tensor,
+        v: &Tensor,
+        probs: &Tensor,
+        dout: &Tensor,
+        scale: f32,
+    ) -> (Tensor, Tensor, Tensor) {
+        let (bh, seq, hd) = (q.shape[0], q.shape[1], q.shape[2]);
+        let dev = q.device();
+
+        // dP[bh,seq,seq] = dOut[bh,seq,hd] · Vᵀ[bh,hd,seq]
+        let dp = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
+        xtrain_cuda::cublas::sgemm_strided_batched(
+            false,
+            true,
+            seq,
+            seq,
+            hd,
+            1.0,
+            dout.data_ptr() as *const f32,
+            seq * hd,
+            v.data_ptr() as *const f32,
+            seq * hd,
+            0.0,
+            dp.data_ptr() as *mut f32,
+            seq * seq,
+            bh,
+        );
+        // dV[bh,seq,hd] = Pᵀ[bh,seq,seq] · dOut[bh,seq,hd]
+        let dv = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        xtrain_cuda::cublas::sgemm_strided_batched(
+            true,
+            false,
+            seq,
+            hd,
+            seq,
+            1.0,
+            probs.data_ptr() as *const f32,
+            seq * seq,
+            dout.data_ptr() as *const f32,
+            seq * hd,
+            0.0,
+            dv.data_ptr() as *mut f32,
+            seq * hd,
+            bh,
+        );
+        // dScores = softmax Jacobian (per row) applied to dP, then ×scale.
+        // Reuse the row-wise softmax backward over the flattened [bh·seq, seq].
+        let dscores = Tensor::softmax_backward(
+            &probs.reshape(&[bh * seq, seq]),
+            &dp.reshape(&[bh * seq, seq]),
+        )
+        .reshape(&[bh, seq, seq]);
+        let dscores = dscores.scale(scale);
+        // dQ[bh,seq,hd] = dScores[bh,seq,seq] · K[bh,seq,hd]
+        let dq = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        xtrain_cuda::cublas::sgemm_strided_batched(
+            false,
+            false,
+            seq,
+            hd,
+            seq,
+            1.0,
+            dscores.data_ptr() as *const f32,
+            seq * seq,
+            k.data_ptr() as *const f32,
+            seq * hd,
+            0.0,
+            dq.data_ptr() as *mut f32,
+            seq * hd,
+            bh,
+        );
+        // dK[bh,seq,hd] = dScoresᵀ[bh,seq,seq] · Q[bh,seq,hd]
+        let dk = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
+        xtrain_cuda::cublas::sgemm_strided_batched(
+            true,
+            false,
+            seq,
+            hd,
+            seq,
+            1.0,
+            dscores.data_ptr() as *const f32,
+            seq * seq,
+            q.data_ptr() as *const f32,
+            seq * hd,
+            0.0,
+            dk.data_ptr() as *mut f32,
+            seq * hd,
+            bh,
+        );
+        (dq, dk, dv)
+    }
+
+    /// 4D axis-(1,2) transpose: `self`:[a,b,c,d] → [a,c,b,d],
+    /// `out[i,k,j,l]=self[i,j,k,l]`. Lays out batched multi-head attention
+    /// (`[B,S,nh,hd] <-> [B,nh,S,hd]`). Its own backward is the same op (swap b,c).
+    #[cfg(not(no_cuda))]
+    pub fn transpose_4d12(&self) -> Self {
+        assert_eq!(self.dtype, DType::F32, "transpose_4d12 only supports F32");
+        assert_eq!(self.ndim(), 4, "transpose_4d12 requires a 4D tensor");
+        assert!(self.is_contiguous(), "transpose_4d12 requires contiguous");
+        let (a, b, c, d) = (self.shape[0], self.shape[1], self.shape[2], self.shape[3]);
+        let out = Tensor::zeros(&[a, c, b, d], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_transpose_4d12_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                a as i32,
+                b as i32,
+                c as i32,
+                d as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        out
+    }
+
     // Shared validation for same-shape binary elementwise ops.
     #[cfg(not(no_cuda))]
     fn check_binary(&self, other: &Tensor, op: &str) {

crates/xtrain-train/src/train_loop.rs

@@ -1,18 +1,20 @@
-//! The training loop: sample sequences → forward `loss` → backward → grad clip
-//! (with batch averaging) → AdamW step → zero grads; with an LR schedule,
-//! periodic loss logging, and periodic checkpointing.
+//! The training loop: sample a batch of sequences → ONE batched forward `loss` →
+//! backward → grad clip → AdamW step → zero grads; with an LR schedule, periodic
+//! loss logging, and periodic checkpointing.
 //!
-//! The T5 model is single-sequence, so a "batch" of `batch_size` sequences is
-//! handled by running forward+backward on each and letting the tape SUM their
-//! grads (its fan-out rule); the clip pass then multiplies by `1/batch_size` to
-//! recover the batch-mean gradient before clipping + the optimizer step.
+//! Since T10 the model is batched (`loss_batched`): `batch_size` sequences are
+//! flattened to `[batch*seq]` and run as a SINGLE forward/backward, so the linear
+//! projections become big `[batch*seq, dim]` GEMMs that fill the GPU. The
+//! cross-entropy mean is over all `batch*seq` rows — already the batch-mean loss,
+//! so backward yields the batch-mean gradient directly (clip pre-scale = 1.0; no
+//! more "loop B times + SUM + ×1/batch" hack).
 
 #![cfg(not(no_cuda))]
 
 use std::path::PathBuf;
 use std::time::Instant;
 
-use xtrain_model::{TinyTransformer, ids_tensor};
+use xtrain_model::{TinyTransformer, batched_ids_tensor, ids_tensor};
 use xtrain_optim::GpuAdamW;
 use xtrain_tensor::Device;
 
@@ -67,7 +69,6 @@ pub fn train(
     let mut losses = Vec::with_capacity(cfg.steps);
     let mut evals = Vec::new();
     let mut best_val: Option<f32> = None;
-    let inv_batch = 1.0 / cfg.batch_size as f32;
     let start = Instant::now();
     let mut tokens_seen: u64 = 0;
     // Best-val checkpointing only kicks in when we actually evaluate.
@@ -76,22 +77,26 @@ pub fn train(
     for step in 0..cfg.steps {
         let lr = cfg.schedule.lr(step);
 
-        // Accumulate grads over `batch_size` sequences (tape SUMs them).
-        let mut step_loss = 0.0f32;
+        // Sample `batch_size` sequences and run them as ONE batched forward/
+        // backward. The CE mean over all batch*seq rows is the batch-mean loss, so
+        // backward already yields the batch-mean gradient (clip pre-scale = 1.0).
+        let mut inputs = Vec::with_capacity(cfg.batch_size);
+        let mut targets_v = Vec::with_capacity(cfg.batch_size);
         for _ in 0..cfg.batch_size {
             let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
-            let ids = ids_tensor(&input, device);
-            let targets = ids_tensor(&target, device);
-            let loss = model.loss(&ids, &targets);
-            step_loss += read_scalar(&loss);
-            loss.backward();
-            tokens_seen += cfg.seq_len as u64;
+            inputs.push(input);
+            targets_v.push(target);
         }
-        step_loss *= inv_batch;
+        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);
+        let step_loss = read_scalar(&loss);
+        loss.backward();
+        tokens_seen += (cfg.batch_size * cfg.seq_len) as u64;
         losses.push(step_loss);
 
-        // Average the summed grads (×1/batch) and clip to the global norm.
-        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch);
+        // Backward already produced the batch-mean gradient — just clip it.
+        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
         opt.step(lr, &params);
         for p in &params {
             p.zero_grad();

csrc/ops/attention.cu

@@ -0,0 +1,93 @@
+// Batched scaled-dot-product attention helpers (Phase T10).
+//
+// The QKᵀ and PV matmuls run as cublasSgemmStridedBatched in Rust; the only
+// kernel attention needs of its own is a CAUSAL row-wise softmax over the score
+// rows. Scores are [B*nh, S, S] flattened to rows of length S; for a flat row r
+// the query position within its sequence is `r % S`, so columns j > r%S are
+// future positions and get probability 0 (no additive -1e9 mask tensor needed).
+//
+// The forward also folds in the 1/sqrt(head_dim) scale (applied to logits before
+// the max/exp) so we don't need a separate scale pass. Backward is the ordinary
+// softmax Jacobian (csrc/ops/nn.cu launch_softmax_dx_f32): masked entries have
+// y=0, so their contribution vanishes — no causal-specific backward needed.
+//
+// All F32, row-major, contiguous. Reduction helpers mirror nn.cu (inlined so the
+// file is self-contained, matching the csrc/ layout).
+
+#include <math.h>
+
+extern "C" {
+
+__device__ __forceinline__ float att_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 att_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 att_block_sum(float v) {
+    __shared__ float sh[32];
+    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
+    int nwarps = (blockDim.x + 31) >> 5;
+    v = att_warp_sum(v);
+    if (lane == 0) sh[warp] = v;
+    __syncthreads();
+    v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : 0.0f;
+    if (warp == 0) v = att_warp_sum(v);
+    __shared__ float bc;
+    if (threadIdx.x == 0) bc = v;
+    __syncthreads();
+    return bc;
+}
+__device__ __forceinline__ float att_block_max(float v) {
+    __shared__ float sh[32];
+    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
+    int nwarps = (blockDim.x + 31) >> 5;
+    v = att_warp_max(v);
+    if (lane == 0) sh[warp] = v;
+    __syncthreads();
+    v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : -INFINITY;
+    if (warp == 0) v = att_warp_max(v);
+    __shared__ float bc;
+    if (threadIdx.x == 0) bc = v;
+    __syncthreads();
+    return bc;
+}
+
+// One block per score row. rows = B*nh*S total; the query position within its
+// sequence is (blockIdx.x % seq). Logits are scaled by `scale` (= 1/sqrt(hd))
+// before softmax; columns j > qpos are masked to probability 0.
+__global__ void softmax_causal_k(const float* x, float* y, int seq, float scale) {
+    int r = blockIdx.x;
+    int qpos = r % seq;
+    const float* xr = x + (size_t)r * seq;
+    float* yr = y + (size_t)r * seq;
+    int valid = qpos + 1; // attend to columns [0, qpos]
+    float m = -INFINITY;
+    for (int c = threadIdx.x; c < valid; c += blockDim.x)
+        m = fmaxf(m, xr[c] * scale);
+    m = att_block_max(m);
+    float sum = 0.0f;
+    for (int c = threadIdx.x; c < valid; c += blockDim.x) {
+        float e = expf(xr[c] * scale - m);
+        yr[c] = e;
+        sum += e;
+    }
+    sum = att_block_sum(sum);
+    float inv = 1.0f / sum;
+    for (int c = threadIdx.x; c < seq; c += blockDim.x)
+        yr[c] = (c < valid) ? yr[c] * inv : 0.0f;
+}
+void launch_softmax_causal_f32(const float* x, float* y, int rows, int seq,
+                               float scale, void* s) {
+    int blk = seq < 1024 ? seq : 1024;
+    if (blk < 32) blk = 32;
+    softmax_causal_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, y, seq, scale);
+}
+
+} // extern "C"

csrc/ops/model.cu

@@ -63,4 +63,26 @@ void launch_transpose_3d01_f32(const float* in, float* out, int a, int b, int c,
     transpose_3d01_k<<<grid, blk, 0, (cudaStream_t)s>>>(in, out, a, b, c);
 }
 
+// =====================================================================
+// 4D axis-(1,2) transpose: in:[a,b,c,d] -> out:[a,c,b,d]. out[i,k,j,l]=in[i,j,k,l].
+// Lays out batched multi-head attention: [B,S,nh,hd] <-> [B,nh,S,hd], so a
+// flattened [B*nh, S, hd] view feeds the strided-batched-GEMM attention. Its own
+// backward is the same op (swap b,c), so one kernel suffices.
+// =====================================================================
+
+__global__ void transpose_4d12_k(const float* in, float* out, int a, int b, int c, int d) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x; // over a*b*c*d
+    if (idx >= a * b * c * d) return;
+    int l = idx % d;
+    int k = (idx / d) % c;
+    int j = (idx / (d * c)) % b;
+    int i = idx / (d * c * b);
+    // out[i,k,j,l] at ((i*c + k)*b + j)*d + l
+    out[(((i * c + k) * b) + j) * d + l] = in[idx];
+}
+void launch_transpose_4d12_f32(const float* in, float* out, int a, int b, int c, int d, void* s) {
+    int n = a * b * c * d, blk = 256, grid = (n + blk - 1) / blk;
+    transpose_4d12_k<<<grid, blk, 0, (cudaStream_t)s>>>(in, out, a, b, c, d);
+}
+
 } // extern "C"

csrc/ops/nn.cu

@@ -215,14 +215,20 @@ void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void*
 // dx[i+h] = dy[i+h]*cos - dy[i]*sin
 // =====================================================================
 
-__global__ void rope_k(const float* x, float* y, int heads, int head_dim, float theta) {
+// `period` is the sequence length: a flattened batch lays B sequences end to end
+// along the `tokens` axis, so each token's RoPE position is its index WITHIN its
+// own sequence, `tok % period`. With period == tokens (single sequence) this is
+// the original position = row.
+__global__ void rope_k(const float* x, float* y, int heads, int head_dim,
+                       float theta, int period) {
     int tok = blockIdx.x;
     int head = blockIdx.y;
     int half = head_dim / 2;
     int i = threadIdx.x;
     if (i >= half) return;
+    int pos = tok % period;
     float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
-    float angle = (float)tok * freq;
+    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];
@@ -230,20 +236,22 @@ __global__ void rope_k(const float* x, float* y, int heads, int head_dim, float
     y[base + i + half] = x1 * c + x0 * sn;
 }
 void launch_rope_f32(const float* x, float* y, int tokens, int heads,
-                     int head_dim, float theta, void* s) {
+                     int head_dim, float theta, int period, void* s) {
     dim3 grid(tokens, heads);
     int blk = head_dim / 2;
-    rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta);
+    rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, period);
 }
 
-__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, float theta) {
+__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim,
+                         float theta, int period) {
     int tok = blockIdx.x;
     int head = blockIdx.y;
     int half = head_dim / 2;
     int i = threadIdx.x;
     if (i >= half) return;
+    int pos = tok % period;
     float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
-    float angle = (float)tok * freq;
+    float angle = (float)pos * freq;
     float c = cosf(angle), sn = sinf(angle);
     int base = (tok * heads + head) * head_dim;
     float d0 = dy[base + i], d1 = dy[base + i + half];
@@ -251,10 +259,10 @@ __global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, f
     dx[base + i + half] = d1 * c - d0 * sn;
 }
 void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads,
-                        int head_dim, float theta, void* s) {
+                        int head_dim, float theta, int period, void* s) {
     dim3 grid(tokens, heads);
     int blk = head_dim / 2;
-    rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta);
+    rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta, period);
 }
 
 // =====================================================================

docs/09-batched-forward.md

@@ -0,0 +1,187 @@
+# Phase T10: Batched 多序列 Forward — Design Document
+
+## Goal
+
+修 **KI-1 的根因**。v2 暴露、v3 重诊断的结论（见 [docs/known-issues.md](known-issues.md)）：
+吞吐瓶颈**不是** all-reduce，而是 **单序列模型设计 launch-bound**——T5 的模型一次只过**一条**
+序列，每个 op 的 GEMM 都是 `[seq, dim]` 这种小矩阵，喂不饱 GPU；"batch" 是靠**循环 B 次 forward +
+让 tape SUM 梯度**伪造的，于是 GPU util 只有 0–15%、显存占 ~8%，加大 global batch 也只是按比例
+增加串行 kernel-launch（v3 实测 gbatch 32→256 仅 +1.2%）。
+
+T10 的目标：给 model + autograd 加 **batch 维**，把一个 step 的 B 条序列**一次性**过模型，让线性层变成
+**一个大 GEMM** 填满 GPU——这是 launch-bound 的根本解。**硬闸门是正确性**：所有既有 grad-check /
+PyTorch 对拍 / overfit / DDP 跨 rank 一致**必须仍过**（PyTorch 对拍现在用 **B>1**），xserv 推理仍**逐
+token 一致**；在此之上拿吞吐收益。
+
+## 核心设计：linears 摊平为 `[B*S, dim]`，attention 批量化
+
+一个关键观察：**transformer 里只有 attention 需要"序列感知"**，其余全是 per-token 的。
+
+- **线性投影 / 逐元素 / norm / embedding / CE 天然是 2D `[rows, dim]` 上的运算**，不在乎 `rows` 是
+  `seq` 还是 `B*S`。所以把激活**摊平成 `[B*S, dim]`** 喂进去，这些 op **原样复用**：
+  q/k/v/o、gate/up/down、lm_head 全部变成**一个大 `[B*S, dim] × [dim, out]` GEMM**（填 GPU 的收益所在）；
+  embedding 对 `[B*S]` ids gather；rms_norm / qk_norm / swiglu / silu / 逐元素按行（per-token）；
+  cross_entropy 对 `[B*S, vocab]` vs `[B*S]` targets，**mean over B*S 行 == 各序列 loss 的均值**。
+- **attention 是唯一序列感知的 op**：reshape 成 `[B, nh, S, hd]`，**每个序列内**做因果 SDPA（per-seq
+  因果 mask，**绝不跨序列 attend**），写回 `[B*S, dim]`。在 `(B, nh)` 维上批量。
+- **RoPE 位置必须 per-sequence 复位**：位置 = 在**自己序列内**的下标（`row % S`），**不是**摊平后的全局
+  行号。这是正确性陷阱，显式处理（kernel 加 `period` 参数）。
+
+```text
+ids [B*S]
+  └ embedding           → [B*S, dim]
+  每层 block：
+    rms_norm            → [B*S, dim]            (per-token)
+    attention:
+      x@Wq/Wk/Wv        → [B*S, dim]            (大 GEMM)
+      reshape           → [B*S, nh, hd]
+      qk-norm + RoPE(period=S) → [B*S, nh, hd]  (RoPE 位置 = row % S)
+      → [B, nh, S, hd] → [B*nh, S, hd]
+      fused batched SDPA(因果) → [B*nh, S, hd]   (2× 批量GEMM + 1× causal-softmax)
+      → [B*S, dim]; @Wo → [B*S, dim]            (大 GEMM)
+    +residual; rms_norm; SwiGLU MLP(大 GEMM); +residual
+  final rms_norm; @lm_head → [B*S, vocab]
+cross_entropy([B*S, vocab], targets[B*S])       → 标量（B*S 行的 mean）
+```
+
+`forward(ids[seq])` 现在只是 `forward_batched(ids, batch=1)` 的特例 → **采样 / 推理路径 batch=1 不变**。
+
+## Module Layout
+
+```
+csrc/ops/attention.cu                    # 新：causal 行 softmax（含 scale + per-seq 因果 mask）
+csrc/ops/nn.cu                           # rope/rope_dx 加 period（位置 = tok % period）
+csrc/ops/model.cu                        # 加 transpose_4d12（[B,S,nh,hd]<->[B,nh,S,hd]）
+crates/xtrain-cuda/
+├── src/cublas.rs                        # 加 sgemm_strided_batched（批量 GEMM，行主序同 sgemm 套路）
+├── src/ffi.rs                           # +cublasSgemmStridedBatched / launch_softmax_causal / transpose_4d12；rope 加 period
+└── build.rs                             # +attention.cu
+crates/xtrain-tensor/src/tensor.rs       # 加 attention/attention_backward（fused，2/4× 批量GEMM）、transpose_4d12；rope(+period)
+crates/xtrain-autodiff/
+├── src/ops.rs                           # 加 attention 节点、transpose_4d12 节点；rope(+period)
+└── tests/autograd.rs                    # +rope_batched / transpose_4d12 / attention(batched) grad-check
+crates/xtrain-model/
+├── src/model.rs                         # forward_batched(ids,batch)/loss_batched；attention 走 fused 批量 op；删 causal_mask 叶
+├── src/lib.rs                           # +batched_ids_tensor
+├── tests/batched.rs                     # 新：batched == looped 单序列（logits + 梯度）
+├── tests/parity_dump.rs + parity.py     # PyTorch 对拍改 B=2,S=4（per-seq RoPE + per-seq mask）
+crates/xtrain-train/src/train_loop.rs    # 真 batch：一次 batched forward/backward 替代 loop+SUM；clip pre-scale 1.0
+crates/xtrain-distributed/
+├── src/ddp.rs                           # 每 rank 一次 batched forward（b_local 条）；clip pre-scale 1.0
+└── tests/ddp_correctness.rs             # 单卡基线也改 batched（对齐）
+```
+
+## Key Design Decisions
+
+### ① 为什么"摊平 linears"是主要收益
+
+GPU 填不满是因为 GEMM 太小（`[256, 384]` 这种）。摊平成 `[B*S, dim]` 后，B=16/S=256 时左矩阵是
+`[4096, 384]`——同一个 cuBLAS GEMM kernel，但 M 大了 16×，一次 launch 干 16× 的活，启动开销被摊薄，
+SM 占用率上去。**embedding / rms_norm / silu / CE 等本就按行**，摊平后只是行数变多，**零改动复用**。
+
+### ② RoPE 位置 per-sequence 复位（正确性陷阱）
+
+摊平后第 `r` 行是序列 `r/S` 的第 `r%S` 个 token。RoPE 角度 = `pos * freq`，**`pos` 必须是 `r%S`**，
+否则第 2 条及之后的序列会用错位置、与单序列结果不一致、PyTorch 对拍直接挂。kernel 加 `period`
+（= 序列长度 S）参数，`pos = tok % period`；`period == tokens`（单序列）时退化为原"位置=行号"。
+反向同样按 `period` 取角度（RoPE 是正交映射，反向 = 逆旋转，不需缓存 forward）。
+
+### ③ Attention：先摊平到 `[B*nh, S, hd]`，fused 批量 SDPA
+
+q/k/v 投影后是 `[B*S, nh, hd]`（顺序 b,s,head,hd）。要喂批量 GEMM，需排成 `[B*nh, S, hd]`（每个 batch
+元素 = 一个 (b, head) 的整条序列）：
+
+```text
+[B*S, nh, hd] → reshape [B, S, nh, hd] → transpose_4d12 → [B, nh, S, hd] → reshape [B*nh, S, hd]
+```
+
+`transpose_4d12`（轴 1,2 互换，新加的结构 op，自反式 backward）是关键。然后 **fused attention op**：
+
+```text
+scores[B*nh,S,S] = Q · Kᵀ                       (cublasSgemmStridedBatched)
+P[B*nh,S,S]      = softmax(causal(scores / √hd)) (launch_softmax_causal，行内做 mask+scale)
+out[B*nh,S,hd]   = P · V                         (cublasSgemmStridedBatched)
+```
+
+**整个 attention 不论 B*nh 多大，都是 3 次 kernel launch**（2 批量 GEMM + 1 softmax），没有 per-head /
+per-seq 的 Python 循环，也没有 host 往返。
+
+**因果 mask 内联在 softmax kernel**：第 `r` 行的 query 位置 = `r % S`，列 `j > r%S` 是未来 → 概率置 0，
+不需要 `[S,S]` 的 `-1e9` 加性 mask 张量（连 `causal_mask` 叶子都删了）。`1/√hd` 的 scale 也折进
+softmax（取 max/exp 前乘），省一个 scale pass。
+
+> **踩坑记录（重要）**：T10 的第一版 attention 用"per-(batch,head) 循环 + `split_heads`/`merge_heads`"，
+> 而 split/merge_heads 内部走 **host 往返**（`to_device(Cpu)` 再传回）。linears 确实摊平成大 GEMM 了，
+> 但 B=16 时 attention 路径里成千上万次小 kernel + 几 MB 的 host memcpy **反而把吞吐拖到 1127 tok/s
+> （比单序列还慢）**。教训：摊平 linears 是必要非充分——attention 的 host 往返 / launch 风暴必须一起干掉。
+> fused 批量 SDPA 之后才到 **25.6K tok/s**。
+
+### ④ Attention backward（手写，grad-check 兜底）
+
+fused op 的反向是标准 attention 反传，全用批量 GEMM + 复用既有 `softmax_backward`：
+
+```text
+dP      = dOut · Vᵀ ;  dV = Pᵀ · dOut
+dScores = softmax_jacobian(P, dP) · (1/√hd)      (复用行 softmax 反向；mask 处 P=0 → 自动零梯度)
+dQ      = dScores · K ;  dK = dScoresᵀ · Q
+```
+
+新加的 `attention` / `transpose_4d12` 都有 finite-diff grad-check（见 ⑥）。
+
+### ⑤ 训练 loop：真 batch 替代 loop+SUM
+
+旧：循环 B 次 `model.loss(单序列)` + `backward()`，靠 tape 把 B 份梯度 SUM，clip 时 ×`1/B` 还原 batch-mean。
+新：`batched_ids_tensor` 把 B 条序列摊平成 `[B*S]`，**一次** `loss_batched` + **一次** `backward()`。
+CE 已是 B*S 行的 mean = batch-mean loss，**backward 直接给出 batch-mean 梯度**，所以 **clip pre-scale = 1.0**
+（不再有 loop+SUM+×1/B）。
+
+**DDP 同理且保持等价**：每 rank 跑 `b_local = B_global/world` 条的**一次** batched forward → backward
+梯度 = 本地 mean `Σ_local/b_local`；`all_reduce_average`（跨 rank 求和 /world）得
+`Σ_global/(world·b_local) = Σ_global/B_global` = 全局 batch-mean → clip pre-scale 也 = 1.0。
+单卡基线（`ddp_correctness` 里）同步改成对全局 batch 的一次 batched forward，二者对齐。
+
+## 验证方法
+
+**双闸门，都必须过。**
+
+### 正确性（无回归）
+
+- **算子级 finite-diff**（`xtrain-autodiff`，15 个）：新增 `rope_batched`（per-seq 位置）、`transpose_4d12`、
+  `attention(batched)` 的 dQ/dK/dV，连同既有 12 个全过。
+  实测：`attn(batched) dQ 7.5e-3 / dK 1.5e-2 / dV 2.9e-4`，`transpose_4d12 8.2e-5`，`rope_batched 4.5e-4`。
+- **batched == looped 单序列**（`xtrain-model/tests/batched.rs`，新）：同一组权重，batched forward 对
+  "逐序列单独 forward 再拼接" 的 logits + 每参数梯度逐一对比。实测 **logits 完全一致(0.0)**，
+  梯度 max rel **6.4e-4**，loss 完全一致。
+- **PyTorch 对拍 B>1**（`parity.py`，改 B=2/S=4）：等价 PyTorch 模型（per-seq RoPE `pos=row%S`、per-seq
+  因果 mask、QK-norm、SwiGLU）对拍 forward logits + 全部 25 个参数梯度。实测 **loss relerr 5e-8、logits
+  6.9e-6、梯度全在 rtol 2e-2 内**。
+- **overfit**（27/27 token）、**checkpoint 逐位**、**AdamW 对 torch**、**DDP 跨 rank 参数 bit-identical
+  (0.0) + DDP loss 对单卡 5.7e-7**：全过。
+- **xserv 闭环**：短训→导出→xserv serve，对 xtrain 贪心**仍逐 token 一致**；采样路径 batch=1 仍工作。
+
+### 吞吐（收益）
+
+单卡 dim384/12L/12h、batch 16、seq 256，back-to-back：
+
+| 路径 | tok/s | GPU util | 显存 |
+|---|---|---|---|
+| **before**（单序列 launch-bound，KI-1 基线）| ~1653 | 0–15% | ~3 GB |
+| T10 第一版（looped split/merge，host 往返）| 1127 | ~11% | 14.8 GB |
+| **after**（fused 批量 attention，batch 16）| **25627** | **37% 均值 / 54% 峰** | 10.2 GB |
+| **after**（batch 32）| **40263** | — | — |
+
+→ 单卡相对 KI-1 基线 **~15.5×**（batch 16）/ **~24×**（batch 32），GPU util 0–15% → 37–54%。
+
+**DDP 4 卡（dim384, per-rank batch 32, global 128）**：1 卡 40.3K → 4 卡 47.2K tok/s（global），
+仅 **~1.17×**。这是 batched forward **新暴露**的下一层瓶颈：单卡 compute 快了 15–24×后，每步**对全部
+67M 参数的 eager all-reduce + host 侧 optimizer/clip 同步**成了 DDP 的主导开销（不随卡数缩小）。
+注意：**单卡 batch 32 = 40K tok/s 已经把 KI-1 时代的 4 卡 3163 tok/s 干翻 ~12×**——根因（单卡
+launch-bound）已修。DDP 近线性需要 **bucketed / 与 backward overlap 的 all-reduce**（KI-1 修复项 2），
+此前 all-reduce 非瓶颈做了没用，现在才有意义——列为 v3+ 的 follow-up（本 Phase 范围外）。
+
+## 给 v3 的 note（已解锁）
+
+batched forward 就位后，v3（dim512/16L、200–300M tok）建议：**per-rank batch 16–32、seq 256，4 卡
+global batch 64–128**。按单卡 ~25K tok/s（dim384，dim512 略低）、4 卡放大，200–300M tok 的训练时间从
+KI-1 时代估的 17–26h 压到**数小时级**。bucketed/overlapped all-reduce（KI-1 修复项 2）现在才有意义，
+作为 v3 之后的进一步优化项。

docs/known-issues.md

@@ -7,7 +7,35 @@
 
 ## Open
 
-### KI-1 · DDP 弱扩展性（吞吐受单序列 launch-bound 限制）— `P1` · 由 v2 暴露，v3 重新诊断
+_(none — KI-1 fixed in T10; see Fixed below. KI-5 is the DDP-scaling follow-up batching newly exposed.)_
+
+### KI-5 · DDP 弱扩展性（all-reduce 每步全参数，未分桶 / 未与 backward overlap）— `P2` · 由 T10 暴露
+- **现象**：batched forward 修掉单卡 launch-bound 后，dim384/per-rank batch 32：1 卡 40.3K → 4 卡 47.2K tok/s（global），仅 ~1.17×。
+- **根因**：单卡 compute 快了 15–24× 后，每步对全部 ~67M 参数的 **eager all-reduce + host 侧 optimizer/clip 同步**成了 DDP 主导开销，不随卡数缩小。注意**单卡 batch 32 = 40K tok/s 已是 KI-1 时代 4 卡(3163)的 ~12×**——根因已修，这是新的、更上层的瓶颈。
+- **拟修复**：梯度 **bucketed all-reduce + 与 backward 计算 overlap**（即 KI-1 修复项 2，此前 all-reduce 非瓶颈做了无收益，batched 之后才有意义）。可选：optimizer/clip 进一步去 host 同步。
+- **重启条件**：v3 训练若被 DDP 扩展性卡住再做；单卡吞吐已足够，v3 可先单卡 / 小 world 跑。
+
+---
+
+## Fixed
+
+### KI-1 · 单序列 launch-bound（"DDP 弱扩展性"的根因）— `FIXED` (T10, batched forward)
+- **修复**：T10 给 model + autograd 加 batch 维——linears 摊平成 `[B*S, dim]` 一个大 GEMM 填满 GPU；attention 走 fused 批量 SDPA（`cublasSgemmStridedBatched` ×2 + 一个 causal-softmax kernel），RoPE 位置 per-sequence 复位（`row % S`）；训练 loop 用真 batch 一次 forward/backward 替代 "loop B 次 + SUM"。详见 [docs/09-batched-forward.md](09-batched-forward.md)。
+- **before → after**（dim384/12L/12h, batch 16, seq 256, 1 卡, back-to-back A/B）：
+  | | tok/s | GPU util | 显存 |
+  |---|---|---|---|
+  | before（单序列 launch-bound）| ~1653 | 0–15% | ~3 GB |
+  | after（batched）| **25627**（batch16）/ **40263**（batch32）| **37% 均值 / 54% 峰** | ~10 GB |
+
+  → 单卡 **~15.5×（batch16）/ ~24×（batch32）**，util 0–15% → 37–54%。
+- **正确性（全绿，无回归）**：15 算子 grad-check（新增 batched-rope / transpose_4d12 / batched-attention dQ/dK/dV）、batched==looped 单序列（logits 0.0、grad 6.4e-4）、**PyTorch 对拍 B>1**（loss 5e-8 / logits 6.9e-6 / 全参数 grad 在 rtol 2e-2）、overfit 27/27、checkpoint 逐位、AdamW 对 torch、DDP loss 对单卡 5.7e-7 + 跨 rank 参数 bit-identical(0.0)、**xserv 加载导出权重对 xtrain 贪心仍逐 token 一致**（top token 同序、BF16 漂移 ~0.03）。
+- **commit**：见 T10 提交链（`perf: KI-1 fixed — GPU util / tok/s` 那条带 before/after）。
+- **DDP 残留弱扩展性 → KI-5**（这是 batching 后新暴露的 all-reduce 瓶颈，不是 KI-1 的单序列根因）。
+- **历史诊断保留如下**（v2 暴露 → v3 重诊断的过程，证明根因不是 all-reduce）：
+
+---
+
+### KI-1 历史诊断 · DDP 弱扩展性（吞吐受单序列 launch-bound 限制）— v2 暴露，v3 重新诊断
 - **现象**：4 卡 DDP 仅 ~3.2K tok/s，几乎不快于单卡（≈2× over 单卡，远低于近线性；T8 在 tiny micro-bench 为 3.0×@4）。
 - **复现**：`dim384/12L, world=4, seq 256`。
 - **v3 实测（dash5, 4× RTX 5090, dim384, 隔离 back-to-back A/B）**：