wip: T10 batched forward (validation)

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