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.