gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: cuBLAS matmul fwd/bwd

Browse Source 
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>
This commit is contained in:
Gahow Wang
2026-06-15 16:48:35 +08:00
parent 5df1d4d57b
commit 0e5c7d22e2
4 changed files with 150 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
crates/xtrain-tensor/src/tensor.rs
View File

@@ -161,8 +161,9 @@ impl Tensor {
/// Matrix multiply: `C = self @ other`. `self`:[M,K], `other`:[K,N] → [M,N].
///
/// Runs the tiled `gemm_tiled_f32` CUDA kernel. Requires contiguous F32
/// tensors on the same GPU. Available only when CUDA is compiled in.
/// Routes through cuBLAS `Sgemm` (Phase T7). fp32, so it is the same GEMM as
/// the T3 tiled kernel up to rounding order. Requires contiguous F32 tensors
/// on the same GPU. Available only when CUDA is compiled in.
#[cfg(not(no_cuda))]
pub fn matmul(&self, other: &Tensor) -> Self {
assert_eq!(self.dtype, DType::F32, "matmul only supports F32");
@@ -188,17 +189,18 @@ impl Tensor {
let k = self.shape[1];
let n = other.shape[1];
let out = Tensor::zeros(&[m, n], DType::F32, self.device());
unsafe {
xtrain_cuda::ffi::launch_gemm_tiled_f32(
self.data_ptr() as *const f32,
other.data_ptr() as *const f32,
out.data_ptr() as *mut f32,
m as i32,
n as i32,
k as i32,
std::ptr::null_mut(),
);
}
xtrain_cuda::cublas::sgemm(
false,
false,
m,
n,
k,
1.0,
self.data_ptr() as *const f32,
other.data_ptr() as *const f32,
0.0,
out.data_ptr() as *mut f32,
);
xtrain_cuda::device::synchronize().expect("matmul kernel sync failed");
out
}
@@ -234,6 +236,9 @@ impl Tensor {
/// Backward of `C = A @ B` given the upstream gradient `dC` (shape [M,N]).
/// Returns `(dA, dB)` where `dA = dC @ Bᵀ` ([M,K]) and `dB = Aᵀ @ dC`
/// ([K,N]). All tensors contiguous F32 on the same GPU.
///
/// Phase T7: cuBLAS applies the transposes internally via its op flags, so we
/// avoid the two transpose kernels (and their allocations) the T3 version ran.
#[cfg(not(no_cuda))]
pub fn matmul_backward(a: &Tensor, b: &Tensor, dc: &Tensor) -> (Tensor, Tensor) {
assert_eq!(a.ndim(), 2, "matmul_backward requires 2D A");
@@ -243,8 +248,36 @@ impl Tensor {
assert_eq!(dc.shape[0], a.shape[0], "dC rows != A rows (M)");
assert_eq!(dc.shape[1], b.shape[1], "dC cols != B cols (N)");
let da = dc.matmul(&b.transpose_2d()); // [M,N] @ [N,K] = [M,K]
let db = a.transpose_2d().matmul(dc); // [K,M] @ [M,N] = [K,N]
let (m, k, n) = (a.shape[0], a.shape[1], b.shape[1]);
// dA[M,K] = dC[M,N] · Bᵀ (B stored [K,N], transposed by cuBLAS)
let da = Tensor::zeros(&[m, k], DType::F32, a.device());
xtrain_cuda::cublas::sgemm(
false,
true,
m,
k,
n,
1.0,
dc.data_ptr() as *const f32,
b.data_ptr() as *const f32,
0.0,
da.data_ptr() as *mut f32,
);
// dB[K,N] = Aᵀ · dC[M,N] (A stored [M,K], transposed by cuBLAS)
let db = Tensor::zeros(&[k, n], DType::F32, a.device());
xtrain_cuda::cublas::sgemm(
true,
false,
k,
n,
m,
1.0,
a.data_ptr() as *const f32,
dc.data_ptr() as *const f32,
0.0,
db.data_ptr() as *mut f32,
);
xtrain_cuda::device::synchronize().expect("matmul_backward sync failed");
(da, db)
}