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));
         }),
     )
 }
@@ -190,6 +192,20 @@ pub fn transpose_3d01(x: &Var) -> Var {
     )
 }
 
+/// 4D axis-(1,2) transpose `[a,b,c,d] -> [a,c,b,d]`. Self-inverse structure: the
+/// backward is the same transpose applied to the grad. Lays out the batched
+/// multi-head attention `[B,S,nh,hd] <-> [B,nh,S,hd]`.
+pub fn transpose_4d12(x: &Var) -> Var {
+    let out = x.value().transpose_4d12();
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(|d, parents| {
+            Var::push_grad(&parents[0], d.transpose_4d12());
+        }),
+    )
+}
+
 /// 2D transpose `[r,c] -> [c,r]` as an autograd node (backward transposes the
 /// grad back). Used for `Kᵀ` in attention scores.
 pub fn transpose_2d(x: &Var) -> Var {
@@ -266,6 +282,29 @@ pub fn merge_heads(heads_v: &[Var]) -> Var {
     )
 }
 
+/// Batched causal scaled-dot-product attention. `q`,`k`,`v` are each
+/// `[bh, seq, head_dim]` (bh = batch·n_heads). Returns `[bh, seq, head_dim]`.
+/// One fused op (2 batched GEMMs + 1 causal-softmax kernel forward; 4 batched
+/// GEMMs + 1 softmax-backward kernel in backward) — replaces the per-(batch,head)
+/// matmul/softmax loop, so attention is a handful of launches regardless of bh.
+/// Caches the softmax `probs` for backward.
+pub fn attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
+    let (out, probs) = q.value().attention(&k.value(), &v.value(), scale);
+    Var::from_op(
+        out,
+        vec![q.clone(), k.clone(), v.clone()],
+        Box::new(move |dout, parents| {
+            let q = parents[0].value();
+            let k = parents[1].value();
+            let v = parents[2].value();
+            let (dq, dk, dv) = Tensor::attention_backward(&q, &k, &v, &probs, dout, scale);
+            Var::push_grad(&parents[0], dq);
+            Var::push_grad(&parents[1], dk);
+            Var::push_grad(&parents[2], dv);
+        }),
+    )
+}
+
 /// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
 /// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
 /// scaled by the upstream scalar grad.

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