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/ffi.rs

@@ -125,7 +125,9 @@ unsafe extern "C" {
     pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream);
     pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream);
 
-    // RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index.
+    // RoPE (rotate_half), x:[tokens,heads,head_dim], position = (token index %
+    // period). `period` = sequence length, so a flattened batch of sequences gets
+    // per-sequence positions; period == tokens reproduces the single-sequence case.
     pub fn launch_rope_f32(
         x: *const f32,
         y: *mut f32,
@@ -133,6 +135,7 @@ unsafe extern "C" {
         heads: i32,
         head_dim: i32,
         theta: f32,
+        period: i32,
         s: CudaStream,
     );
     pub fn launch_rope_dx_f32(
@@ -142,6 +145,7 @@ unsafe extern "C" {
         heads: i32,
         head_dim: i32,
         theta: f32,
+        period: i32,
         s: CudaStream,
     );
 
@@ -211,6 +215,31 @@ unsafe extern "C" {
         c: i32,
         s: CudaStream,
     );
+    // 4D axis-(1,2) transpose: in:[a,b,c,d] -> out:[a,c,b,d]. out[i,k,j,l]=in[i,j,k,l].
+    pub fn launch_transpose_4d12_f32(
+        input: *const f32,
+        out: *mut f32,
+        a: i32,
+        b: i32,
+        c: i32,
+        d: i32,
+        s: CudaStream,
+    );
+}
+
+// Batched attention helper (csrc/ops/attention.cu): causal row-wise softmax over
+// score rows [rows, seq] with query position = (row % seq); scales logits by
+// `scale` (= 1/sqrt(head_dim)) and masks future columns to probability 0.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    pub fn launch_softmax_causal_f32(
+        x: *const f32,
+        y: *mut f32,
+        rows: i32,
+        seq: i32,
+        scale: f32,
+        s: CudaStream,
+    );
 }
 
 // GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
@@ -267,6 +296,27 @@ unsafe extern "C" {
         c: *mut f32,
         ldc: i32,
     ) -> i32;
+    #[allow(clippy::too_many_arguments)]
+    pub fn cublasSgemmStridedBatched(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const f32,
+        a: *const f32,
+        lda: i32,
+        stride_a: i64,
+        b: *const f32,
+        ldb: i32,
+        stride_b: i64,
+        beta: *const f32,
+        c: *mut f32,
+        ldc: i32,
+        stride_c: i64,
+        batch_count: i32,
+    ) -> i32;
 }
 
 #[cfg(not(no_cuda))]