phase 3: GEMM kernels (naive, tiled, cuBLAS)

- Naive GEMM kernel: one thread per output element (F32 + BF16)
- Tiled GEMM kernel: 32x32 shared memory tiles (F32 + BF16)
- cuBLAS wrapper: cublasGemmEx with row-major trick
- GemmBackend enum for runtime backend selection
- CublasContext RAII handle
- Made error::check public for cross-crate use
- 17 GEMM tests: small/medium/rect sizes, all backends, F32+BF16
- Cross-backend consistency verified (naive vs tiled vs cuBLAS)
- All 44 tests pass across all crates

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/xserv-kernels/src/gemm.rs

@@ -0,0 +1,151 @@
+use std::ffi::c_void;
+use xserv_cuda::error::{self, Result};
+use xserv_tensor::{DType, Device, Tensor};
+
+#[derive(Debug, Clone, Copy)]
+pub enum GemmBackend {
+    Naive,
+    Tiled,
+    CuBlas,
+}
+
+// --- FFI: custom CUDA kernels ---
+unsafe extern "C" {
+    fn launch_gemm_naive_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+}
+
+// --- FFI: cuBLAS ---
+type CublasHandle = *mut c_void;
+
+#[allow(non_upper_case_globals)]
+const CUBLAS_OP_N: i32 = 0;
+
+// cudaDataType
+const CUDA_R_32F: i32 = 0;
+const CUDA_R_16BF: i32 = 14;
+
+// cublasComputeType
+const CUBLAS_COMPUTE_32F: i32 = 68;
+
+unsafe extern "C" {
+    fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
+    fn cublasDestroy_v2(handle: CublasHandle) -> i32;
+    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
+    fn cublasGemmEx(
+        handle: CublasHandle,
+        transa: i32, transb: i32,
+        m: i32, n: i32, k: i32,
+        alpha: *const c_void,
+        a: *const c_void, a_type: i32, lda: i32,
+        b: *const c_void, b_type: i32, ldb: i32,
+        beta: *const c_void,
+        c: *mut c_void, c_type: i32, ldc: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+}
+
+pub struct CublasContext {
+    handle: CublasHandle,
+}
+
+impl CublasContext {
+    pub fn new() -> Result<Self> {
+        let mut handle = std::ptr::null_mut();
+        error::check(unsafe { cublasCreate_v2(&mut handle) })?;
+        Ok(Self { handle })
+    }
+}
+
+impl Drop for CublasContext {
+    fn drop(&mut self) {
+        if !self.handle.is_null() {
+            unsafe { cublasDestroy_v2(self.handle) };
+        }
+    }
+}
+
+/// Matrix multiplication: C = A @ B
+/// A: [M, K], B: [K, N], C: [M, N]
+/// All tensors must be contiguous and on the same GPU.
+pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
+    assert_eq!(a.ndim(), 2);
+    assert_eq!(b.ndim(), 2);
+    assert_eq!(a.shape()[1], b.shape()[0], "inner dimension mismatch");
+    assert_eq!(a.dtype(), b.dtype(), "dtype mismatch");
+    assert!(a.is_contiguous() && b.is_contiguous(), "matmul requires contiguous tensors");
+    assert!(matches!(a.device(), Device::Cuda(_)), "matmul requires GPU tensors");
+
+    let m = a.shape()[0];
+    let k = a.shape()[1];
+    let n = b.shape()[1];
+    let dtype = a.dtype();
+
+    let c = Tensor::zeros(&[m, n], dtype, a.device());
+
+    let a_ptr = a.data_ptr() as *const c_void;
+    let b_ptr = b.data_ptr() as *const c_void;
+    let c_ptr = c.data_ptr() as *mut c_void;
+    let null_stream = std::ptr::null_mut();
+
+    match backend {
+        GemmBackend::Naive => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_naive_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_naive_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for naive GEMM"),
+                }
+            }
+            xserv_cuda::device::synchronize().unwrap();
+        }
+        GemmBackend::Tiled => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_tiled_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_tiled_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for tiled GEMM"),
+                }
+            }
+            xserv_cuda::device::synchronize().unwrap();
+        }
+        GemmBackend::CuBlas => {
+            // cuBLAS uses column-major, but we have row-major tensors.
+            // Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
+            // cuBLAS sees our row-major data as column-major transposed.
+            let ctx = CublasContext::new().unwrap();
+            let alpha = 1.0f32;
+            let beta = 0.0f32;
+
+            let (a_type, b_type, c_type) = match dtype {
+                DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
+                DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
+                _ => panic!("unsupported dtype for cuBLAS GEMM"),
+            };
+
+            unsafe {
+                cublasSetStream_v2(ctx.handle, null_stream);
+                // Row-major trick: swap A/B and transpose flags
+                // C(row-major) = A @ B  <=>  C^T(col-major) = B^T @ A^T
+                error::check(cublasGemmEx(
+                    ctx.handle,
+                    CUBLAS_OP_N, CUBLAS_OP_N,
+                    n as i32, m as i32, k as i32,
+                    &alpha as *const f32 as *const c_void,
+                    b_ptr, b_type, n as i32,  // B as col-major = B^T
+                    a_ptr, a_type, k as i32,  // A as col-major = A^T
+                    &beta as *const f32 as *const c_void,
+                    c_ptr, c_type, n as i32,  // C as col-major = C^T
+                    CUBLAS_COMPUTE_32F,
+                    -1, // default algo
+                )).expect("cuBLAS GEMM failed");
+            }
+            xserv_cuda::device::synchronize().unwrap();
+        }
+    }
+
+    c
+}

crates/xserv-kernels/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod gemm;
+
+pub use gemm::{GemmBackend, matmul};