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