model: batched forward [B,S]

forward_batched(ids[B*S], batch)/loss_batched: run B equal-length sequences as
ONE forward over flattened [B*S] ids, so every linear is one big [B*S,dim] GEMM.
Attention reshapes to [B*nh,S,hd], runs the fused batched causal SDPA (per-seq
mask + RoPE period=S, no cross-sequence attention), writes back [B*S,dim]. The
old per-(batch,head) loop + host-round-tripping split/merge_heads + the additive
causal_mask leaf are gone. forward(ids[seq]) is now forward_batched(ids,1), so
the sampler / inference path (batch=1) is unchanged.

+batched_ids_tensor helper. New batched.rs test: batched forward == looped
single-sequence (logits identical 0.0, grads 6.4e-4, loss identical). PyTorch
parity now exercises B>1 (B=2,S=4): loss 5e-8, logits 6.9e-6, all 25 param
grads within rtol — verifying per-seq RoPE position + per-seq causal masking.

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

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();