wip: T10 batched forward (validation)

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