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

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

232
crates/xserv-kernels/src/gemm.rs
View File

@@ -1,14 +1,22 @@
use std::cell::RefCell;
use std::ffi::c_void;
use xserv_cuda::error::{self, Result};
use xserv_cuda::GpuBuffer;
use xserv_cuda::error::{self, Result};
use xserv_tensor::{DType, Device, Tensor};
const CUBLAS_WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
// GEMV: single-kernel, no FP32 temp buffer needed
unsafe extern "C" {
fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
fn launch_gemv_bf16(
x: *const c_void,
w: *const c_void,
y_bf16: *mut c_void,
y_fp32_buf: *mut c_void,
k: i32,
n: i32,
stream: *mut c_void,
);
}
#[derive(Debug, Clone, Copy)]
@@ -20,10 +28,42 @@ pub enum GemmBackend {
// --- 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);
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 ---
@@ -46,25 +86,46 @@ unsafe extern "C" {
fn cublasSetWorkspace_v2(handle: CublasHandle, workspace: *mut c_void, size: usize) -> i32;
fn cublasGemmEx(
handle: CublasHandle,
transa: i32, transb: i32,
m: i32, n: i32, k: i32,
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,
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,
c: *mut c_void,
c_type: i32,
ldc: i32,
compute_type: i32,
algo: i32,
) -> i32;
fn cublasGemmStridedBatchedEx(
handle: CublasHandle,
transa: i32, transb: i32,
m: i32, n: i32, k: i32,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const c_void,
a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
a: *const c_void,
a_type: i32,
lda: i32,
stride_a: i64,
b: *const c_void,
b_type: i32,
ldb: i32,
stride_b: i64,
beta: *const c_void,
c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
c: *mut c_void,
c_type: i32,
ldc: i32,
stride_c: i64,
batch_count: i32,
compute_type: i32,
algo: i32,
@@ -89,9 +150,16 @@ impl CublasContext {
// set, so we keep the GpuBuffer in this struct.
let mut workspace = GpuBuffer::alloc(CUBLAS_WORKSPACE_BYTES)?;
error::check(unsafe {
cublasSetWorkspace_v2(handle, workspace.as_mut_ptr() as *mut c_void, CUBLAS_WORKSPACE_BYTES)
cublasSetWorkspace_v2(
handle,
workspace.as_mut_ptr() as *mut c_void,
CUBLAS_WORKSPACE_BYTES,
)
})?;
Ok(Self { handle, _workspace: Some(workspace) })
Ok(Self {
handle,
_workspace: Some(workspace),
})
}
}
@@ -123,9 +191,7 @@ where
/// Get the thread-local cuBLAS handle for use with dispatch module.
pub fn cublas_handle() -> CublasHandle {
CUBLAS_CTX.with(|cell| {
cell.borrow().handle
})
CUBLAS_CTX.with(|cell| cell.borrow().handle)
}
/// Matrix multiplication: C = A @ B
@@ -136,8 +202,14 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
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");
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];
@@ -154,32 +226,63 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
let null_stream = xserv_cuda::current_stream_raw();
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"),
}
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"),
}
}
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"),
}
},
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"),
}
}
},
GemmBackend::CuBlas => {
if m == 1 && dtype == DType::BF16 && n >= 256 {
let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
unsafe {
launch_gemv_bf16(
a_ptr, b_ptr, c_ptr,
a_ptr,
b_ptr,
c_ptr,
fp32_buf.as_mut_ptr() as *mut c_void,
k as i32, n as i32,
k as i32,
n as i32,
null_stream,
);
}
@@ -197,16 +300,26 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
cublasSetStream_v2(handle, null_stream);
error::check(cublasGemmEx(
handle,
CUBLAS_OP_N, CUBLAS_OP_N,
n as i32, m as i32, k as i32,
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,
a_ptr, a_type, k as i32,
b_ptr,
b_type,
n as i32,
a_ptr,
a_type,
k as i32,
&beta as *const f32 as *const c_void,
c_ptr, c_type, n as i32,
c_ptr,
c_type,
n as i32,
CUBLAS_COMPUTE_32F,
-1,
)).expect("cuBLAS GEMM failed");
))
.expect("cuBLAS GEMM failed");
});
}
}
@@ -264,17 +377,30 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
// Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
error::check(cublasGemmStridedBatchedEx(
handle,
CUBLAS_OP_N, CUBLAS_OP_N,
n as i32, m as i32, k as i32,
CUBLAS_OP_N,
CUBLAS_OP_N,
n as i32,
m as i32,
k as i32,
&alpha as *const f32 as *const c_void,
b.data_ptr() as _, b_type, n as i32, stride_b,
a.data_ptr() as _, a_type, k as i32, stride_a,
b.data_ptr() as _,
b_type,
n as i32,
stride_b,
a.data_ptr() as _,
a_type,
k as i32,
stride_a,
&beta as *const f32 as *const c_void,
c.data_ptr() as *mut c_void, c_type, n as i32, stride_c,
c.data_ptr() as *mut c_void,
c_type,
n as i32,
stride_c,
batch as i32,
CUBLAS_COMPUTE_32F,
-1,
)).expect("cuBLAS batched GEMM failed");
))
.expect("cuBLAS batched GEMM failed");
});
c
}