perf: cuBLAS matmul fwd/bwd

Route Tensor::matmul and matmul_backward through cuBLAS Sgemm instead of
the hand-written tiled kernel. fp32 → same GEMM up to rounding order, so
the T3 cuBLAS tolerance and downstream grad-checks are preserved.

- cublas.rs: thread-local persistent handle + row-major sgemm helper with
  transpose flags (col-major⟺row-major as the T3 oracle does).
- matmul_backward: dA/dB via cuBLAS OP_T, dropping the two transpose
  kernels + their allocations the T3 version ran.

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

crates/xtrain-cuda/src/cublas.rs

@@ -0,0 +1,95 @@
+//! cuBLAS GEMM backend (Phase T7).
+//!
+//! The hand-written tiled kernel (csrc/ops/gemm.cu) is kept as the T3 learning
+//! artifact + correctness oracle's counterpart, but the forward + both backward
+//! matmuls now route through cuBLAS `Sgemm` — fp32, so the result is numerically
+//! the same GEMM (only the rounding order changes), which is why the T3 tolerance
+//! against cuBLAS holds unchanged.
+//!
+//! **Layout.** cuBLAS is column-major; our tensors are row-major. A row-major
+//! `[r,c]` matrix handed to cuBLAS with leading dim `c` is read as its transpose
+//! (col-major `[c,r]`). To get a row-major result `C[m,n] = opA(A)·opB(B)` we
+//! compute the col-major transpose `Cᵀ[n,m] = opB(B)ᵀ·opA(A)ᵀ`; the bytes of
+//! col-major `Cᵀ` are exactly row-major `C`. See [`sgemm`] for the index algebra.
+//!
+//! **Handle.** cuBLAS handle creation is expensive (T3's oracle made one per
+//! call). We cache one handle per thread for the lifetime of the process.
+
+#![cfg(not(no_cuda))]
+
+use crate::ffi::{self, CublasHandle};
+use std::cell::RefCell;
+
+thread_local! {
+    static HANDLE: RefCell<Option<CublasHandle>> = const { RefCell::new(None) };
+}
+
+/// Run `f` with the thread's cached cuBLAS handle, creating it on first use.
+fn with_handle<R>(f: impl FnOnce(CublasHandle) -> R) -> R {
+    HANDLE.with(|h| {
+        let mut slot = h.borrow_mut();
+        if slot.is_none() {
+            let mut handle: CublasHandle = std::ptr::null_mut();
+            let status = unsafe { ffi::cublasCreate_v2(&mut handle) };
+            assert_eq!(status, 0, "cublasCreate failed: {status}");
+            *slot = Some(handle);
+        }
+        f(slot.unwrap())
+    })
+}
+
+/// Row-major single-precision GEMM: `C[m,n] = opA(A) · opB(B)` with
+/// `C = alpha·(…) + beta·C`. `A`/`B`/`C` are device pointers to row-major fp32
+/// matrices; `trans_a`/`trans_b` request the transpose of the *logical* operand.
+///
+/// `m,n,k` are the dims of the math (`opA(A)` is `[m,k]`, `opB(B)` is `[k,n]`).
+/// The stored, untransposed shapes are: `A` is `[m,k]` (or `[k,m]` if `trans_a`),
+/// `B` is `[k,n]` (or `[n,k]` if `trans_b`). Their row-major leading dims are the
+/// stored column counts, derived below.
+///
+/// We ask cuBLAS for col-major `Cᵀ[n,m] = opB(B)ᵀ · opA(A)ᵀ`. Since a row-major
+/// `[r,c]` buffer is col-major `[c,r]`, a row-major operand already *is* its own
+/// transpose to cuBLAS — so `opB(B)ᵀ` over the row-major bytes of `B` is obtained
+/// by passing `B` with the OPPOSITE op flag of what `opB` would suggest. Working
+/// it through: first cuBLAS arg = `B` with op `trans_b ? N : T`, second = `A` with
+/// op `trans_a ? N : T`, sizes (m=n, n=m, k=k).
+#[allow(clippy::too_many_arguments)]
+pub fn sgemm(
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: f32,
+    a: *const f32,
+    b: *const f32,
+    beta: f32,
+    c: *mut f32,
+) {
+    // Leading dims = stored (row-major) column count of each untransposed matrix.
+    let lda = if trans_a { m } else { k }; // A stored [m,k] or [k,m]
+    let ldb = if trans_b { k } else { n }; // B stored [k,n] or [n,k]
+    let ldc = n; // Cᵀ is [n,m] col-major with ld n  (== row-major C[m,n])
+
+    let op_b = if trans_b {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let op_a = if trans_a {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+
+    with_handle(|handle| {
+        let status = unsafe {
+            ffi::cublasSgemm_v2(
+                handle, op_b, op_a, n as i32, // rows of Cᵀ
+                m as i32, // cols of Cᵀ
+                k as i32, &alpha, b, ldb as i32, a, lda as i32, &beta, c, ldc as i32,
+            )
+        };
+        assert_eq!(status, 0, "cublasSgemm failed: {status}");
+    });
+}

crates/xtrain-cuda/src/ffi.rs

@@ -212,8 +212,9 @@ unsafe extern "C" {
     );
 }
 
-// cuBLAS — used ONLY as a correctness reference for the hand-written GEMM in
-// tests. Declared (and linked, see build.rs) only when CUDA is compiled in.
+// cuBLAS — the production GEMM backend (Phase T7) and the correctness oracle the
+// T3 GEMM tests still compare against. Declared (and linked, see build.rs) only
+// when CUDA is compiled in.
 #[cfg(not(no_cuda))]
 pub type CublasHandle = *mut c_void;
 
@@ -241,3 +242,5 @@ unsafe extern "C" {
 
 #[cfg(not(no_cuda))]
 pub const CUBLAS_OP_N: i32 = 0;
+#[cfg(not(no_cuda))]
+pub const CUBLAS_OP_T: i32 = 1;

crates/xtrain-cuda/src/lib.rs

@@ -1,3 +1,5 @@
+#[cfg(not(no_cuda))]
+pub mod cublas;
 pub mod device;
 pub mod error;
 pub mod ffi;