autograd: batch dim for ops (flatten linears, batched attention)

Add the batched-forward primitives. Linears/norms/elementwise/embedding/CE
already act on flat [rows,dim], so they work unchanged on [B*S,dim]; only
attention + RoPE need sequence awareness:

- RoPE: kernel takes a `period` (= seq len) so position = row % period, i.e.
  per-sequence position on a flattened batch (period == tokens = single seq).
- Fused batched causal attention: new `Tensor::attention`/`attention_backward`
  + ops node, running QKᵀ and PV as cublasSgemmStridedBatched over the B*nh
  (sequence,head) blocks (new sgemm_strided_batched binding) and a causal
  softmax kernel (scale + per-row causal mask inline) — the whole attention is
  3 launches regardless of B*nh, no per-head/per-seq loop, no host round-trip.
- transpose_4d12 ([B,S,nh,hd] <-> [B,nh,S,hd]) to lay out the batched heads.

grad-checks: new batched-rope, transpose_4d12, batched-attention dQ/dK/dV all
pass finite-diff (attn dK 1.5e-2, dQ 7.5e-3, dV 2.9e-4; rest tighter) alongside
the existing 12.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-cuda/src/cublas.rs

@@ -93,3 +93,69 @@ pub fn sgemm(
         assert_eq!(status, 0, "cublasSgemm failed: {status}");
     });
 }
+
+/// Strided-batched row-major SGEMM: for each `i` in `0..batch`,
+/// `C_i[m,n] = alpha·opA(A_i)·opB(B_i) + beta·C_i`, where `A_i`/`B_i`/`C_i` are
+/// consecutive matrices laid `stride_*` elements apart in one contiguous buffer.
+/// Same row-major⟺col-major trick as [`sgemm`] (compute col-major `Cᵀ`), applied
+/// per batch element. Used for the batched attention `QKᵀ` / `PV` GEMMs (and their
+/// backwards), so the whole attention runs as 2 batched-GEMM launches, not a
+/// per-(batch,head) Python loop. `A`/`B`/`C` are device pointers to the first
+/// matrix; strides are in ELEMENTS.
+#[allow(clippy::too_many_arguments)]
+pub fn sgemm_strided_batched(
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: f32,
+    a: *const f32,
+    stride_a: usize,
+    b: *const f32,
+    stride_b: usize,
+    beta: f32,
+    c: *mut f32,
+    stride_c: usize,
+    batch: usize,
+) {
+    let lda = if trans_a { m } else { k };
+    let ldb = if trans_b { k } else { n };
+    let ldc = n;
+    let op_a = if trans_a {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let op_b = if trans_b {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+
+    with_handle(|handle| {
+        let status = unsafe {
+            ffi::cublasSgemmStridedBatched(
+                handle,
+                op_b,
+                op_a,
+                n as i32,
+                m as i32,
+                k as i32,
+                &alpha,
+                b,
+                ldb as i32,
+                stride_b as i64,
+                a,
+                lda as i32,
+                stride_a as i64,
+                &beta,
+                c,
+                ldc as i32,
+                stride_c as i64,
+                batch as i32,
+            )
+        };
+        assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
+    });
+}