wip: T10 batched forward (validation)

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));
         }),
     )
 }

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