wip: T10 batched forward (validation)

crates/xtrain-cuda/build.rs

@@ -35,6 +35,7 @@ fn main() {
         .file("../../csrc/ops/nn.cu")
         .file("../../csrc/ops/model.cu")
         .file("../../csrc/ops/optim.cu")
+        .file("../../csrc/ops/attention.cu")
         .compile("xtrain_cuda_kernels");
 }
 

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}");
+    });
+}

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))]