diff --git a/crates/xserv-kernels/build.rs b/crates/xserv-kernels/build.rs index d03e40a..c28d7b8 100644 --- a/crates/xserv-kernels/build.rs +++ b/crates/xserv-kernels/build.rs @@ -16,6 +16,7 @@ fn main() { .include("../../csrc") .file("../../csrc/gemm/naive.cu") .file("../../csrc/gemm/tiled.cu") + .file("../../csrc/gemm/gemv.cu") .file("../../csrc/normalization/rmsnorm.cu") .file("../../csrc/normalization/layernorm.cu") .file("../../csrc/activation/activations.cu") diff --git a/crates/xserv-kernels/src/gemm.rs b/crates/xserv-kernels/src/gemm.rs index 7df34b7..589eb1c 100644 --- a/crates/xserv-kernels/src/gemm.rs +++ b/crates/xserv-kernels/src/gemm.rs @@ -15,6 +15,7 @@ unsafe extern "C" { 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_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); } // --- FFI: cuBLAS --- @@ -124,35 +125,49 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor { } } 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; + // Fast path: custom GEMV for M=1 BF16 (bandwidth-optimal decode) + if m == 1 && dtype == DType::BF16 { + let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap(); + unsafe { + launch_gemv_bf16( + a_ptr, b_ptr, c_ptr, + fp32_buf.as_mut_ptr() as *mut c_void, + k as i32, n as i32, + null_stream, + ); + } + // fp32_buf returned to caching allocator pool on drop + } else { + // 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"), - }; + 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"); + 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"); + } } } } diff --git a/crates/xserv-kernels/tests/gemm_test.rs b/crates/xserv-kernels/tests/gemm_test.rs index 3546d59..f49ccd7 100644 --- a/crates/xserv-kernels/tests/gemm_test.rs +++ b/crates/xserv-kernels/tests/gemm_test.rs @@ -121,6 +121,20 @@ fn test_gemm_cublas_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 4, 4, #[test] fn test_gemm_cublas_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 256, 256, 256); } +// --- Custom GEMV tests (M=1, BF16 fast path) --- + +#[test] +fn test_gemv_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 64, 64); } + +#[test] +fn test_gemv_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 256, 256); } + +#[test] +fn test_gemv_bf16_4096() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 4096, 4096); } + +#[test] +fn test_gemv_bf16_rect() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 512, 4096); } + // --- Larger benchmark-style tests --- #[test] diff --git a/csrc/gemm/gemv.cu b/csrc/gemm/gemv.cu new file mode 100644 index 0000000..060736e --- /dev/null +++ b/csrc/gemm/gemv.cu @@ -0,0 +1,102 @@ +#include +#include + +// Custom GEMV kernel for M=1 decode step (BF16): +// y[n] = sum_k x[k] * W[k * N + n] +// where x: [K] (BF16), W: [K, N] (BF16, row-major), y: [N] (BF16). +// +// Design: K-split for high occupancy on large GPU (170 SMs). +// Grid: (N / TILE_N, K / TILE_K) — each block computes a partial sum +// for TILE_N output columns over a TILE_K slice of K. +// Partial results are atomicAdd'd to an FP32 accumulator, then a +// second kernel converts FP32 -> BF16. +// +// Memory access: adjacent threads read adjacent columns of the same row +// of W, giving perfectly coalesced 128-byte transactions. + +#define GEMV_TILE_N 128 +#define GEMV_TILE_K 256 +#define GEMV_BLOCK 128 // = TILE_N, one thread per output column + +__global__ void gemv_bf16_kernel( + const __nv_bfloat16* __restrict__ x, // [K] + const __nv_bfloat16* __restrict__ W, // [K, N] row-major + float* __restrict__ y_fp32, // [N] accumulator + int K, int N +) { + const int block_n = blockIdx.x; + const int block_k = blockIdx.y; + const int t = threadIdx.x; + const int col = block_n * GEMV_TILE_N + t; + + if (col >= N) return; + + const int k_start = block_k * GEMV_TILE_K; + const int k_end = min(k_start + GEMV_TILE_K, K); + const int k_len = k_end - k_start; + + // Load x[k_start..k_end] into shared memory as FP32 + __shared__ float x_shared[GEMV_TILE_K]; + for (int i = t; i < k_len; i += GEMV_BLOCK) { + x_shared[i] = __bfloat162float(x[k_start + i]); + } + __syncthreads(); + + // Compute partial dot product for this column + float sum = 0.0f; + for (int ki = 0; ki < k_len; ki++) { + sum += x_shared[ki] * __bfloat162float(W[(k_start + ki) * N + col]); + } + + // Atomic accumulate (handles K-split reduction) + atomicAdd(&y_fp32[col], sum); +} + +// Conversion kernel: FP32 accumulator -> BF16 output +__global__ void gemv_fp32_to_bf16_kernel( + const float* __restrict__ src, + __nv_bfloat16* __restrict__ dst, + int n +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + dst[idx] = __float2bfloat16(src[idx]); + } +} + +extern "C" { + +void launch_gemv_bf16( + const void* x, // [K] BF16 + const void* W, // [K, N] BF16 row-major + void* y_bf16, // [N] BF16 output + void* y_fp32_buf, // [N] FP32 temporary (caller-provided) + int K, int N, + void* stream +) { + cudaStream_t s = (cudaStream_t)stream; + + // Zero the FP32 accumulator + cudaMemsetAsync((float*)y_fp32_buf, 0, N * sizeof(float), s); + + // Launch GEMV kernel + dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N, + (K + GEMV_TILE_K - 1) / GEMV_TILE_K); + gemv_bf16_kernel<<>>( + (const __nv_bfloat16*)x, + (const __nv_bfloat16*)W, + (float*)y_fp32_buf, + K, N + ); + + // Convert FP32 -> BF16 + int conv_block = 256; + int conv_grid = (N + conv_block - 1) / conv_block; + gemv_fp32_to_bf16_kernel<<>>( + (const float*)y_fp32_buf, + (__nv_bfloat16*)y_bf16, + N + ); +} + +} // extern "C"