From 55acb45f3baf0eb2fc361323242b774189cf9515 Mon Sep 17 00:00:00 2001 From: pasha Date: Thu, 16 Jul 2026 17:32:32 +0300 Subject: [PATCH 1/8] Add asynchronous cuBLAS BF16 GEMM dispatch --- Cargo.toml | 1 + crates/cubecl-cuda/examples/cublas_gemm.rs | 290 ++++++++++++++ .../cubecl-cuda/examples/cublas_gemm_check.rs | 375 ++++++++++++++++++ crates/cubecl-cuda/src/compute/cublas.rs | 233 +++++++++++ crates/cubecl-cuda/src/compute/mod.rs | 1 + crates/cubecl-cuda/src/compute/server.rs | 50 ++- crates/cubecl-cuda/src/runtime.rs | 5 + crates/cubecl-ir/src/features.rs | 4 +- crates/cubecl-runtime/src/client.rs | 15 +- crates/cubecl-runtime/src/server/base.rs | 54 +++ 10 files changed, 1022 insertions(+), 6 deletions(-) create mode 100644 crates/cubecl-cuda/examples/cublas_gemm.rs create mode 100644 crates/cubecl-cuda/examples/cublas_gemm_check.rs create mode 100644 crates/cubecl-cuda/src/compute/cublas.rs diff --git a/Cargo.toml b/Cargo.toml index 27f0b25a41..3d4dbd6e38 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -153,6 +153,7 @@ cudarc = { version = "0.19.6", features = [ "std", "driver", "nvrtc", + "cublas", "nccl", "fallback-dynamic-loading", "cuda-version-from-build-system", diff --git a/crates/cubecl-cuda/examples/cublas_gemm.rs b/crates/cubecl-cuda/examples/cublas_gemm.rs new file mode 100644 index 0000000000..f0ba17d549 --- /dev/null +++ b/crates/cubecl-cuda/examples/cublas_gemm.rs @@ -0,0 +1,290 @@ +use cubecl_common::future; +use cubecl_core::{ + Runtime, + ir::{ElemType, FloatKind}, + server::{GemmDescriptor, GemmMatrix}, +}; +use cubecl_cuda::{CudaDevice, CudaRuntime}; +use half::bf16; +use std::time::Instant; + +#[derive(Clone, Copy)] +struct Shape { + m: usize, + n: usize, + k: usize, + batch: usize, + lhs_t: bool, + rhs_t: bool, + lhs_broadcast: bool, + rhs_broadcast: bool, + repeats: usize, +} + +fn main() { + let args = std::env::args().skip(1).collect::>(); + if args.first().is_some_and(|arg| arg == "check") { + correctness_matrix(); + return; + } + if args.first().is_some_and(|arg| arg == "enqueue") { + enqueue_overhead(args.get(1).map_or(10_000, |value| parse(value))); + return; + } + if args.len() != 9 { + eprintln!( + "usage: cublas_gemm M N K BATCH LHS_T RHS_T LHS_BROADCAST RHS_BROADCAST REPEATS\n cublas_gemm check\n cublas_gemm enqueue [REPEATS]" + ); + std::process::exit(2); + } + let shape = Shape { + m: parse(&args[0]), + n: parse(&args[1]), + k: parse(&args[2]), + batch: parse(&args[3]), + lhs_t: parse_flag(&args[4]), + rhs_t: parse_flag(&args[5]), + lhs_broadcast: parse_flag(&args[6]), + rhs_broadcast: parse_flag(&args[7]), + repeats: parse(&args[8]), + }; + let ms = benchmark(shape); + let flops = 2.0 * shape.m as f64 * shape.n as f64 * shape.k as f64 * shape.batch as f64; + println!( + "m={} n={} k={} batch={} lhs_t={} rhs_t={} lhs_broadcast={} rhs_broadcast={} ms={ms:.4} tflops={:.2}", + shape.m, + shape.n, + shape.k, + shape.batch, + shape.lhs_t, + shape.rhs_t, + shape.lhs_broadcast, + shape.rhs_broadcast, + flops / (ms * 1.0e9), + ); +} + +fn enqueue_overhead(repeats: usize) { + let client = CudaRuntime::client(&CudaDevice::default()); + let one = bf16::ONE.to_bits(); + let lhs = client.create_from_slice(bytemuck::cast_slice(&[one])); + let rhs = client.create_from_slice(bytemuck::cast_slice(&[one])); + let out = client.empty(2); + let descriptor = descriptor( + Shape { + m: 1, + n: 1, + k: 1, + batch: 1, + lhs_t: false, + rhs_t: false, + lhs_broadcast: false, + rhs_broadcast: false, + repeats, + }, + lhs.binding(), + rhs.binding(), + out.binding(), + ); + client.gemm(descriptor.clone()); + future::block_on(client.sync()).unwrap(); + + let start = Instant::now(); + for _ in 0..repeats { + client.gemm(descriptor.clone()); + } + let enqueue = start.elapsed(); + let drain_start = Instant::now(); + future::block_on(client.sync()).unwrap(); + let drain = drain_start.elapsed(); + println!( + "repeats={repeats} enqueue_us_per_call={:.3} final_gpu_drain_ms={:.3}", + enqueue.as_secs_f64() * 1.0e6 / repeats as f64, + drain.as_secs_f64() * 1.0e3, + ); +} + +fn benchmark(shape: Shape) -> f64 { + let client = CudaRuntime::client(&CudaDevice::default()); + let lhs_batches = if shape.lhs_broadcast { 1 } else { shape.batch }; + let rhs_batches = if shape.rhs_broadcast { 1 } else { shape.batch }; + let lhs_elems = lhs_batches * shape.m * shape.k; + let rhs_elems = rhs_batches * shape.k * shape.n; + let out_elems = shape.batch * shape.m * shape.n; + let lhs = client.create_from_slice(bytemuck::cast_slice(&values(lhs_elems, 3))); + let rhs = client.create_from_slice(bytemuck::cast_slice(&values(rhs_elems, 7))); + let out = client.empty(out_elems * 2); + let descriptor = descriptor(shape, lhs.binding(), rhs.binding(), out.binding()); + + for _ in 0..5 { + client.gemm(descriptor.clone()); + } + future::block_on(client.sync()).unwrap(); + + let start = Instant::now(); + for _ in 0..shape.repeats { + client.gemm(descriptor.clone()); + } + future::block_on(client.sync()).unwrap(); + start.elapsed().as_secs_f64() * 1_000.0 / shape.repeats as f64 +} + +fn correctness_matrix() { + for lhs_t in [false, true] { + for rhs_t in [false, true] { + for batch in [1, 3] { + for lhs_broadcast in [false, true] { + for rhs_broadcast in [false, true] { + if batch == 1 && (lhs_broadcast || rhs_broadcast) { + continue; + } + check(Shape { + m: 5, + n: 7, + k: 3, + batch, + lhs_t, + rhs_t, + lhs_broadcast, + rhs_broadcast, + repeats: 1, + }); + } + } + } + } + } + println!("cuBLAS BF16 layout correctness matrix passed"); +} + +fn check(shape: Shape) { + let client = CudaRuntime::client(&CudaDevice::default()); + let lhs_batches = if shape.lhs_broadcast { 1 } else { shape.batch }; + let rhs_batches = if shape.rhs_broadcast { 1 } else { shape.batch }; + let lhs_bits = values(lhs_batches * shape.m * shape.k, 3); + let rhs_bits = values(rhs_batches * shape.k * shape.n, 7); + let lhs = client.create_from_slice(bytemuck::cast_slice(&lhs_bits)); + let rhs = client.create_from_slice(bytemuck::cast_slice(&rhs_bits)); + let out = client.empty(shape.batch * shape.m * shape.n * 2); + client.gemm(descriptor( + shape, + lhs.binding(), + rhs.binding(), + out.clone().binding(), + )); + let bytes = client.read_one_unchecked(out); + let actual = bytemuck::cast_slice::(&bytes); + + for batch in 0..shape.batch { + for row in 0..shape.m { + for col in 0..shape.n { + let expected = (0..shape.k) + .map(|inner| { + get( + &lhs_bits, + shape.m, + shape.k, + shape.lhs_t, + if shape.lhs_broadcast { 0 } else { batch }, + row, + inner, + ) * get( + &rhs_bits, + shape.k, + shape.n, + shape.rhs_t, + if shape.rhs_broadcast { 0 } else { batch }, + inner, + col, + ) + }) + .sum::(); + let index = (batch * shape.m + row) * shape.n + col; + let actual = bf16::from_bits(actual[index]).to_f32(); + assert!( + (actual - expected).abs() <= 0.06, + "layout failed: lhs_t={} rhs_t={} lhs_broadcast={} rhs_broadcast={} batch={batch} row={row} col={col}: {actual} != {expected}", + shape.lhs_t, + shape.rhs_t, + shape.lhs_broadcast, + shape.rhs_broadcast, + ); + } + } + } +} + +fn descriptor( + shape: Shape, + lhs: cubecl_core::server::Binding, + rhs: cubecl_core::server::Binding, + out: cubecl_core::server::Binding, +) -> GemmDescriptor { + GemmDescriptor::new( + GemmMatrix::new( + lhs, + if shape.lhs_t { shape.m } else { shape.k } as u32, + if shape.lhs_broadcast { + 0 + } else { + (shape.m * shape.k) as u64 + }, + shape.lhs_t, + ), + GemmMatrix::new( + rhs, + if shape.rhs_t { shape.k } else { shape.n } as u32, + if shape.rhs_broadcast { + 0 + } else { + (shape.k * shape.n) as u64 + }, + shape.rhs_t, + ), + GemmMatrix::new(out, shape.n as u32, (shape.m * shape.n) as u64, false), + shape.m as u32, + shape.n as u32, + shape.k as u32, + shape.batch as u32, + ElemType::Float(FloatKind::BF16), + ) +} + +fn values(len: usize, offset: usize) -> Vec { + (0..len) + .map(|index| { + let value = ((index + offset) % 11) as f32 / 32.0 - 0.15; + bf16::from_f32(value).to_bits() + }) + .collect() +} + +fn get( + bits: &[u16], + rows: usize, + cols: usize, + transposed: bool, + batch: usize, + row: usize, + col: usize, +) -> f32 { + let base = batch * rows * cols; + let index = if transposed { + base + col * rows + row + } else { + base + row * cols + col + }; + bf16::from_bits(bits[index]).to_f32() +} + +fn parse(value: &str) -> usize { + value.parse().unwrap() +} + +fn parse_flag(value: &str) -> bool { + match value { + "0" => false, + "1" => true, + _ => panic!("flag must be 0 or 1"), + } +} diff --git a/crates/cubecl-cuda/examples/cublas_gemm_check.rs b/crates/cubecl-cuda/examples/cublas_gemm_check.rs new file mode 100644 index 0000000000..f7e71fd9f9 --- /dev/null +++ b/crates/cubecl-cuda/examples/cublas_gemm_check.rs @@ -0,0 +1,375 @@ +use cubecl_common::{future, stream_id::StreamId}; +use cubecl_core as cubecl; +use cubecl_core::{ + Runtime, + ir::{ElemType, FloatKind}, + prelude::*, + server::{Binding, GemmDescriptor, GemmMatrix, Handle}, +}; +use cubecl_cuda::{CudaDevice, CudaRuntime}; +use half::bf16; + +#[cube(launch)] +fn copy_bf16(input: &[bf16], output: &mut [bf16]) { + if ABSOLUTE_POS < output.len() { + output[ABSOLUTE_POS] = input[ABSOLUTE_POS]; + } +} + +#[derive(Clone, Copy, Debug)] +struct Problem { + m: usize, + n: usize, + k: usize, + batch: usize, + lhs_t: bool, + rhs_t: bool, + lhs_broadcast: bool, + rhs_broadcast: bool, + padding: usize, + batch_gap: usize, + offset: usize, +} + +struct Matrix { + bits: Vec, + rows: usize, + cols: usize, + batches: usize, + transposed: bool, + ld: usize, + batch_stride: usize, + offset: usize, +} + +fn main() { + correctness_matrix(); + cross_stream_ordering(); + overlapping_output_is_rejected(); + prior_error_is_not_bypassed(); + println!("cuBLAS BF16 padded/offset/batched/multistream checks passed"); +} + +fn correctness_matrix() { + for lhs_t in [false, true] { + for rhs_t in [false, true] { + for lhs_broadcast in [false, true] { + for rhs_broadcast in [false, true] { + check(Problem { + m: 5, + n: 7, + k: 3, + batch: 3, + lhs_t, + rhs_t, + lhs_broadcast, + rhs_broadcast, + padding: 5, + batch_gap: 11, + offset: 13, + }); + } + } + } + } +} + +fn check(problem: Problem) { + let client = CudaRuntime::client(&CudaDevice::default()); + let lhs = Matrix::new( + problem.m, + problem.k, + if problem.lhs_broadcast { + 1 + } else { + problem.batch + }, + problem.lhs_t, + problem.padding, + problem.batch_gap, + problem.offset, + 3, + ); + let rhs = Matrix::new( + problem.k, + problem.n, + if problem.rhs_broadcast { + 1 + } else { + problem.batch + }, + problem.rhs_t, + problem.padding + 2, + problem.batch_gap + 3, + problem.offset + 5, + 7, + ); + let out = Matrix::zeros( + problem.m, + problem.n, + problem.batch, + false, + problem.padding + 4, + problem.batch_gap + 7, + problem.offset + 9, + ); + + let lhs_base = client.create_from_slice(bytemuck::cast_slice(&lhs.bits)); + let rhs_base = client.create_from_slice(bytemuck::cast_slice(&rhs.bits)); + let out_base = client.empty(out.bits.len() * 2); + let lhs_view = view(&lhs_base, lhs.offset, lhs.bits.len()); + let rhs_view = view(&rhs_base, rhs.offset, rhs.bits.len()); + let out_view = view(&out_base, out.offset, out.bits.len()); + let descriptor = GemmDescriptor::new( + matrix_arg(&lhs, lhs_view, problem.lhs_broadcast), + matrix_arg(&rhs, rhs_view, problem.rhs_broadcast), + matrix_arg(&out, out_view, false), + problem.m as u32, + problem.n as u32, + problem.k as u32, + problem.batch as u32, + ElemType::Float(FloatKind::BF16), + ); + client.gemm(descriptor); + let bytes = client.read_one_unchecked(out_base); + let actual = bytemuck::cast_slice::(&bytes); + + for batch in 0..problem.batch { + for row in 0..problem.m { + for col in 0..problem.n { + let expected = (0..problem.k) + .map(|inner| { + lhs.get(if problem.lhs_broadcast { 0 } else { batch }, row, inner) + * rhs.get(if problem.rhs_broadcast { 0 } else { batch }, inner, col) + }) + .sum::(); + let index = out.index(batch, row, col); + let actual = bf16::from_bits(actual[index]).to_f32(); + assert!( + (actual - expected).abs() <= 0.06, + "{problem:?}, b={batch}, row={row}, col={col}: {actual} != {expected}" + ); + } + } + } +} + +fn cross_stream_ordering() { + let mut producer = CudaRuntime::client(&CudaDevice::default()); + let mut gemm = producer.clone(); + let mut consumer = producer.clone(); + unsafe { + producer.set_stream(StreamId { value: 100 }); + gemm.set_stream(StreamId { value: 101 }); + consumer.set_stream(StreamId { value: 102 }); + } + + let problem = Problem { + m: 8, + n: 6, + k: 4, + batch: 2, + lhs_t: false, + rhs_t: true, + lhs_broadcast: false, + rhs_broadcast: true, + padding: 0, + batch_gap: 0, + offset: 0, + }; + let lhs = Matrix::new(8, 4, 2, false, 0, 0, 0, 3); + let rhs = Matrix::new(4, 6, 1, true, 0, 0, 0, 7); + let lhs_handle = producer.create_from_slice(bytemuck::cast_slice(&lhs.bits)); + let rhs_handle = producer.create_from_slice(bytemuck::cast_slice(&rhs.bits)); + let gemm_out = gemm.empty(problem.batch * problem.m * problem.n * 2); + let descriptor = GemmDescriptor::new( + matrix_arg(&lhs, lhs_handle.binding(), false), + matrix_arg(&rhs, rhs_handle.binding(), true), + GemmMatrix::new( + gemm_out.clone().binding(), + problem.n as u32, + (problem.m * problem.n) as u64, + false, + ), + problem.m as u32, + problem.n as u32, + problem.k as u32, + problem.batch as u32, + ElemType::Float(FloatKind::BF16), + ); + gemm.gemm(descriptor); + + let consumed = consumer.empty(problem.batch * problem.m * problem.n * 2); + copy_bf16::launch::( + &consumer, + CubeCount::Static(1, 1, 1), + CubeDim::new(&consumer, 128), + unsafe { BufferArg::from_raw_parts(gemm_out, problem.batch * problem.m * problem.n) }, + unsafe { + BufferArg::from_raw_parts(consumed.clone(), problem.batch * problem.m * problem.n) + }, + ); + let bytes = consumer.read_one_unchecked(consumed); + let actual = bytemuck::cast_slice::(&bytes); + for batch in 0..problem.batch { + for row in 0..problem.m { + for col in 0..problem.n { + let expected = (0..problem.k) + .map(|inner| lhs.get(batch, row, inner) * rhs.get(0, inner, col)) + .sum::(); + let index = (batch * problem.m + row) * problem.n + col; + let actual = bf16::from_bits(actual[index]).to_f32(); + assert!((actual - expected).abs() <= 0.06); + } + } + } +} + +fn overlapping_output_is_rejected() { + let mut client = CudaRuntime::client(&CudaDevice::default()); + unsafe { client.set_stream(StreamId { value: 103 }) }; + + let m = 2; + let n = 3; + let k = 4; + let batches = 2; + let lhs = client.create_from_slice(bytemuck::cast_slice(&vec![ + bf16::ONE.to_bits(); + batches * m * k + ])); + let rhs = client.create_from_slice(bytemuck::cast_slice(&vec![ + bf16::ONE.to_bits(); + batches * k * n + ])); + let out = client.empty(batches * m * n * 2); + let descriptor = GemmDescriptor::new( + GemmMatrix::new(lhs.binding(), k as u32, (m * k) as u64, false), + GemmMatrix::new(rhs.binding(), n as u32, (k * n) as u64, false), + GemmMatrix::new(out.binding(), n as u32, (m * n - 1) as u64, false), + m as u32, + n as u32, + k as u32, + batches as u32, + ElemType::Float(FloatKind::BF16), + ); + client.gemm(descriptor); + assert!(future::block_on(client.sync()).is_err()); +} + +fn prior_error_is_not_bypassed() { + let client = CudaRuntime::client(&CudaDevice::default()); + let input = client.create_from_slice(bytemuck::cast_slice(&[bf16::ONE.to_bits()])); + let output = client.empty(2); + + // The CUDA launch limit is at most 1024 units per cube. The fire-and-forget + // launch records a server error; the following GEMM must observe that + // unhealthy stream instead of enqueueing through it. + copy_bf16::launch::( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(2048), + unsafe { BufferArg::from_raw_parts(input.clone(), 1) }, + unsafe { BufferArg::from_raw_parts(output.clone(), 1) }, + ); + let descriptor = GemmDescriptor::new( + GemmMatrix::new(input.clone().binding(), 1, 0, false), + GemmMatrix::new(input.binding(), 1, 0, false), + GemmMatrix::new(output.binding(), 1, 0, false), + 1, + 1, + 1, + 1, + ElemType::Float(FloatKind::BF16), + ); + client.gemm(descriptor); + // Surfaces the original launch error and restores stream health. + assert!(future::block_on(client.sync()).is_err()); +} + +impl Matrix { + #[allow(clippy::too_many_arguments)] + fn new( + rows: usize, + cols: usize, + batches: usize, + transposed: bool, + padding: usize, + batch_gap: usize, + offset: usize, + seed: usize, + ) -> Self { + let mut matrix = Self::zeros(rows, cols, batches, transposed, padding, batch_gap, offset); + for batch in 0..batches { + for row in 0..rows { + for col in 0..cols { + let logical = (batch * rows + row) * cols + col; + let value = ((logical + seed) % 11) as f32 / 32.0 - 0.15; + let index = matrix.index(batch, row, col); + matrix.bits[index] = bf16::from_f32(value).to_bits(); + } + } + } + matrix + } + + #[allow(clippy::too_many_arguments)] + fn zeros( + rows: usize, + cols: usize, + batches: usize, + transposed: bool, + padding: usize, + batch_gap: usize, + offset: usize, + ) -> Self { + let ld = (if transposed { rows } else { cols }) + padding; + let span = (if transposed { cols } else { rows }) * ld; + let batch_stride = span + batch_gap; + let suffix = 17; + Self { + bits: vec![0; offset + batches * batch_stride + suffix], + rows, + cols, + batches, + transposed, + ld, + batch_stride, + offset, + } + } + + fn index(&self, batch: usize, row: usize, col: usize) -> usize { + assert!(batch < self.batches && row < self.rows && col < self.cols); + let matrix = self.offset + batch * self.batch_stride; + if self.transposed { + matrix + col * self.ld + row + } else { + matrix + row * self.ld + col + } + } + + fn get(&self, batch: usize, row: usize, col: usize) -> f32 { + bf16::from_bits(self.bits[self.index(batch, row, col)]).to_f32() + } +} + +fn matrix_arg(matrix: &Matrix, binding: Binding, broadcast: bool) -> GemmMatrix { + GemmMatrix::new( + binding, + matrix.ld as u32, + if broadcast { + 0 + } else { + matrix.batch_stride as u64 + }, + matrix.transposed, + ) +} + +fn view(base: &Handle, offset: usize, total: usize) -> Binding { + let suffix = total - offset; + base.clone() + .offset_start((offset * 2) as u64) + .offset_end((suffix.min(17) * 2) as u64) + .binding() +} diff --git a/crates/cubecl-cuda/src/compute/cublas.rs b/crates/cubecl-cuda/src/compute/cublas.rs new file mode 100644 index 0000000000..00a1a760ef --- /dev/null +++ b/crates/cubecl-cuda/src/compute/cublas.rs @@ -0,0 +1,233 @@ +use cubecl_common::backtrace::BackTrace; +use cubecl_core::{ + ir::{ElemType, FloatKind}, + server::{GemmDescriptor, GemmMatrix, ServerError}, +}; +use cudarc::cublas::{result, sys}; +use std::ffi::c_void; + +use super::storage::gpu::GpuResource; + +#[derive(Default)] +pub(crate) struct CublasState { + handle: Option, + stream: Option, +} + +impl core::fmt::Debug for CublasState { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("CublasState") + .field("initialized", &self.handle.is_some()) + .finish() + } +} + +impl CublasState { + pub(crate) fn launch( + &mut self, + descriptor: &GemmDescriptor, + lhs: &GpuResource, + rhs: &GpuResource, + out: &GpuResource, + stream: cudarc::driver::sys::CUstream, + ) -> Result<(), ServerError> { + if descriptor.elem != ElemType::Float(FloatKind::BF16) { + return Err(validation_error("cuBLAS GEMM currently supports only BF16")); + } + validate(descriptor, lhs, rhs, out)?; + + if descriptor.m == 0 + || descriptor.n == 0 + || descriptor.k == 0 + || descriptor.batch_count == 0 + { + return Ok(()); + } + + let handle = match self.handle { + Some(handle) => handle, + None => { + let handle = result::create_handle().map_err(cublas_error)?; + self.handle = Some(handle); + handle + } + }; + if self.stream != Some(stream) { + // SAFETY: both handles are owned by this CUDA server. Changing the + // cuBLAS stream only changes where subsequent work is enqueued. + unsafe { result::set_stream(handle, stream.cast()) }.map_err(cublas_error)?; + self.stream = Some(stream); + } + + let alpha = 1.0f32; + let beta = 0.0f32; + let op_rhs = operation(&descriptor.rhs); + let op_lhs = operation(&descriptor.lhs); + + // cuBLAS is column-major. Swapping the row-major operands computes + // C^T = rhs^T @ lhs^T without copies. The call is asynchronous on the + // exact stream resolved by CubeCL for these bindings. + unsafe { + result::gemm_strided_batched_ex( + handle, + op_rhs, + op_lhs, + descriptor.n as i32, + descriptor.m as i32, + descriptor.k as i32, + (&alpha as *const f32).cast::(), + rhs.ptr as *const c_void, + sys::cudaDataType::CUDA_R_16BF, + descriptor.rhs.leading_dimension as i32, + descriptor.rhs.batch_stride as i64, + lhs.ptr as *const c_void, + sys::cudaDataType::CUDA_R_16BF, + descriptor.lhs.leading_dimension as i32, + descriptor.lhs.batch_stride as i64, + (&beta as *const f32).cast::(), + out.ptr as *mut c_void, + sys::cudaDataType::CUDA_R_16BF, + descriptor.out.leading_dimension as i32, + descriptor.out.batch_stride as i64, + descriptor.batch_count as i32, + sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, + ) + } + .map_err(cublas_error)?; + + Ok(()) + } + + pub(crate) fn destroy(&mut self) { + if let Some(handle) = self.handle.take() { + // SAFETY: this state uniquely owns the handle and destroys it once. + if let Err(err) = unsafe { result::destroy_handle(handle) } { + log::warn!("Unable to destroy cuBLAS handle: {err}"); + } + } + } +} + +fn operation(matrix: &GemmMatrix) -> sys::cublasOperation_t { + if matrix.transposed { + sys::cublasOperation_t::CUBLAS_OP_T + } else { + sys::cublasOperation_t::CUBLAS_OP_N + } +} + +fn validate( + descriptor: &GemmDescriptor, + lhs: &GpuResource, + rhs: &GpuResource, + out: &GpuResource, +) -> Result<(), ServerError> { + let max_i32 = i32::MAX as u32; + if descriptor.m > max_i32 + || descriptor.n > max_i32 + || descriptor.k > max_i32 + || descriptor.batch_count > max_i32 + { + return Err(validation_error( + "GEMM dimensions exceed the cuBLAS i32 API", + )); + } + if descriptor.out.transposed { + return Err(validation_error("GEMM output must be row-major")); + } + + validate_matrix( + "lhs", + &descriptor.lhs, + descriptor.m, + descriptor.k, + descriptor.batch_count, + lhs, + true, + )?; + validate_matrix( + "rhs", + &descriptor.rhs, + descriptor.k, + descriptor.n, + descriptor.batch_count, + rhs, + true, + )?; + validate_matrix( + "out", + &descriptor.out, + descriptor.m, + descriptor.n, + descriptor.batch_count, + out, + false, + ) +} + +#[allow(clippy::too_many_arguments)] +fn validate_matrix( + name: &str, + matrix: &GemmMatrix, + rows: u32, + cols: u32, + batches: u32, + resource: &GpuResource, + allow_broadcast: bool, +) -> Result<(), ServerError> { + let required_ld = if matrix.transposed { rows } else { cols }; + if matrix.leading_dimension < required_ld || matrix.leading_dimension > i32::MAX as u32 { + return Err(validation_error(&format!( + "{name} leading dimension {} is invalid for logical shape [{rows}, {cols}]", + matrix.leading_dimension + ))); + } + if matrix.batch_stride > i64::MAX as u64 { + return Err(validation_error(&format!( + "{name} batch stride exceeds the cuBLAS i64 API" + ))); + } + let matrix_elems = if rows == 0 || cols == 0 { + 0 + } else if matrix.transposed { + (cols as u64 - 1) * matrix.leading_dimension as u64 + rows as u64 + } else { + (rows as u64 - 1) * matrix.leading_dimension as u64 + cols as u64 + }; + if !allow_broadcast && batches > 1 && matrix.batch_stride < matrix_elems { + return Err(validation_error("GEMM output batches may not overlap")); + } + let batch_offset = if batches <= 1 || matrix.batch_stride == 0 { + 0 + } else { + (batches as u64 - 1) + .checked_mul(matrix.batch_stride) + .ok_or_else(|| validation_error("GEMM batch stride overflow"))? + }; + let required_bytes = batch_offset + .checked_add(matrix_elems) + .and_then(|elements| elements.checked_mul(2)) + .ok_or_else(|| validation_error("GEMM buffer size overflow"))?; + if required_bytes > resource.size { + return Err(validation_error(&format!( + "{name} requires {required_bytes} bytes but its binding contains {}", + resource.size + ))); + } + Ok(()) +} + +fn validation_error(message: &str) -> ServerError { + ServerError::Validation { + message: message.into(), + backtrace: BackTrace::capture(), + } +} + +fn cublas_error(error: result::CublasError) -> ServerError { + ServerError::Generic { + reason: format!("cuBLAS GEMM failed: {error}"), + backtrace: BackTrace::capture(), + } +} diff --git a/crates/cubecl-cuda/src/compute/mod.rs b/crates/cubecl-cuda/src/compute/mod.rs index c595738e3a..a3ed1a0944 100644 --- a/crates/cubecl-cuda/src/compute/mod.rs +++ b/crates/cubecl-cuda/src/compute/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod command; pub(crate) mod communication; pub(crate) mod context; +pub(crate) mod cublas; pub(crate) mod io; pub(crate) mod storage; pub(crate) mod stream; diff --git a/crates/cubecl-cuda/src/compute/server.rs b/crates/cubecl-cuda/src/compute/server.rs index d8f3107cc7..ea5a6db6b4 100644 --- a/crates/cubecl-cuda/src/compute/server.rs +++ b/crates/cubecl-cuda/src/compute/server.rs @@ -5,6 +5,7 @@ use crate::{ command::Command, communication::{get_nccl_comm_id, get_nccl_dtype_count, to_nccl_op}, context::CudaContext, + cublas::CublasState, stream::CudaStreamBackend, sync::Fence, }, @@ -16,9 +17,9 @@ use cubecl_core::{ ir::{ElemType, FloatKind, IntKind, MemoryDeviceProperties, StorageType, UIntKind}, prelude::*, server::{ - Binding, CommunicationId, CopyDescriptor, Handle, KernelArguments, LaunchError, - ProfileError, ProfilingToken, ReduceOperation, ServerCommunication, ServerError, - ServerUtilities, StreamErrorMode, TensorMapBinding, TensorMapMeta, + Binding, CommunicationId, CopyDescriptor, GemmDescriptor, Handle, KernelArguments, + LaunchError, ProfileError, ProfilingToken, ReduceOperation, ServerCommunication, + ServerError, ServerUtilities, StreamErrorMode, TensorMapBinding, TensorMapMeta, }, }; use cubecl_environment::backtrace::BackTrace; @@ -56,6 +57,7 @@ pub struct CudaServer { utilities: Arc>, comm_stream: *mut CUstream_st, communicators: HashMap, + cublas: CublasState, } // SAFETY: `CudaServer` is only accessed from one thread at a time via the `DeviceHandle`, @@ -166,6 +168,16 @@ impl ComputeServer for CudaServer { } } + fn gemm(&mut self, descriptor: GemmDescriptor, stream_id: StreamId) { + if let Err(err) = self.gemm_checked(descriptor, stream_id) { + let mut stream = match self.streams.resolve(stream_id, [].into_iter(), false) { + Ok(stream) => stream, + Err(err) => unreachable!("{err}"), + }; + stream.current().errors.push(err); + } + } + fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError> { let mut command = self.command_no_inputs( stream_id, @@ -597,6 +609,7 @@ impl CudaServer { utilities: Arc::new(utilities), comm_stream, communicators: HashMap::default(), + cublas: CublasState::default(), } } @@ -657,6 +670,30 @@ impl CudaServer { errors } + fn gemm_checked( + &mut self, + descriptor: GemmDescriptor, + stream_id: StreamId, + ) -> Result<(), ServerError> { + let bindings = [ + &descriptor.lhs.binding, + &descriptor.rhs.binding, + &descriptor.out.binding, + ]; + self.unsafe_set_current(); + let streams = self + .streams + .resolve(stream_id, bindings.into_iter(), true)?; + let cublas = &mut self.cublas; + let mut command = Command::new(&mut self.ctx, streams); + let lhs = command.resource(descriptor.lhs.binding.clone())?; + let rhs = command.resource(descriptor.rhs.binding.clone())?; + let out = command.resource(descriptor.out.binding.clone())?; + let stream = command.streams.current().sys; + + cublas.launch(&descriptor, &lhs, &rhs, &out, stream) + } + fn launch_checked( &mut self, kernel: Box>, @@ -919,6 +956,13 @@ impl CudaServer { } } +impl Drop for CudaServer { + fn drop(&mut self) { + self.unsafe_set_current(); + self.cublas.destroy(); + } +} + fn elem_to_tensor_map_type(ty: StorageType) -> CUtensorMapDataType { use cudarc::driver::sys::CUtensorMapDataType::*; match ty { diff --git a/crates/cubecl-cuda/src/runtime.rs b/crates/cubecl-cuda/src/runtime.rs index 2a32feb782..5efb5bfaf0 100644 --- a/crates/cubecl-cuda/src/runtime.rs +++ b/crates/cubecl-cuda/src/runtime.rs @@ -212,6 +212,11 @@ impl DeviceService for CudaServer { if arch_version >= 80 { device_props.features.copy_async = true; + device_props + .features + .matmul + .accelerated_gemm + .insert(ElemType::Float(FloatKind::BF16)); } // NOTE: I commented that since I observed synchronisation issues with atomic add for bf16. diff --git a/crates/cubecl-ir/src/features.rs b/crates/cubecl-ir/src/features.rs index 54fb027d8b..9b487ce38f 100644 --- a/crates/cubecl-ir/src/features.rs +++ b/crates/cubecl-ir/src/features.rs @@ -1,4 +1,4 @@ -use crate::{AddressType, OpaqueType, SemanticType, StorageType, Type}; +use crate::{AddressType, ElemType, OpaqueType, SemanticType, StorageType, Type}; use alloc::collections::{BTreeMap, BTreeSet}; use enumset::EnumSetType; @@ -53,6 +53,8 @@ pub struct Types { /// Matrix multiplication-related features #[derive(Debug, Clone, PartialEq, Eq, Default, Hash)] pub struct MatmulFeatures { + /// Element types supported by a backend-native GEMM implementation. + pub accelerated_gemm: BTreeSet, /// The cmma feature enables cooperative matrix-multiply and accumulate operations. pub cmma: BTreeSet, /// Cube MMA is like cmma but at the cube level, rather than the plane level. diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index b213f1f10a..bd5ff939da 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -7,8 +7,8 @@ use crate::{ memory_management::{MemoryAllocationMode, MemoryConfiguration, MemoryUsage}, runtime::Runtime, server::{ - CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, Handle, IoError, - KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, + CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, GemmDescriptor, + Handle, IoError, KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, MemoryLayoutStrategy, ProfileError, ReduceOperation, ServerCommunication, ServerError, ServerUtilities, }, @@ -366,6 +366,17 @@ impl ComputeClient { .unwrap_or_resume() } + /// Enqueue a backend-accelerated GEMM. + /// + /// The operation is ordered on this client's CubeCL stream and returns + /// immediately, like a kernel launch. Callers must first check the device's + /// `features.matmul.accelerated_gemm` capability. + pub fn gemm(&self, descriptor: GemmDescriptor) { + let stream_id = self.stream_id(); + self.device + .submit(move |server| server.gemm(descriptor, stream_id)); + } + fn do_create_from_slices( &self, descriptors: Vec, diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index cc940dbec6..dccb959432 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -352,6 +352,49 @@ pub struct StreamErrorMode { pub flush: bool, } +/// Describes one matrix operand of an accelerated GEMM. +/// +/// Strides are expressed in elements. A matrix is either row-major +/// (`transposed == false`) or a transposed view of a row-major matrix. The +/// latter has a unit row stride and `leading_dimension` as its column stride. +#[derive(new, Clone, Debug)] +pub struct GemmMatrix { + /// Buffer containing the matrix data. + pub binding: Binding, + /// Leading dimension of the underlying row-major allocation. + pub leading_dimension: u32, + /// Distance between consecutive batch matrices, in elements. Zero + /// broadcasts this operand across batches. + pub batch_stride: u64, + /// Whether the logical matrix is a transposed row-major view. + pub transposed: bool, +} + +/// Backend-optional accelerated GEMM descriptor. +/// +/// Computes `out = lhs @ rhs` for `batch_count` matrices. All matrices use +/// `elem`; accumulation precision is selected by the backend. The output must +/// be row-major. +#[derive(new, Clone, Debug)] +pub struct GemmDescriptor { + /// Left-hand matrix, logically shaped `[m, k]`. + pub lhs: GemmMatrix, + /// Right-hand matrix, logically shaped `[k, n]`. + pub rhs: GemmMatrix, + /// Row-major output matrix, logically shaped `[m, n]`. + pub out: GemmMatrix, + /// Number of rows in the output. + pub m: u32, + /// Number of columns in the output. + pub n: u32, + /// Contracting dimension. + pub k: u32, + /// Number of independent matrix products. + pub batch_count: u32, + /// Scalar element type shared by all operands. + pub elem: ElemType, +} + /// The compute server is responsible for handling resources and computations over resources. /// /// Everything in the server is mutable, therefore it should be solely accessed through the @@ -435,6 +478,17 @@ where launch_mode: LaunchMode, ); + /// Enqueue an accelerated GEMM when the backend provides one. + /// + /// Callers must first check + /// [`Features::matmul.accelerated_gemm`](cubecl_ir::features::MatmulFeatures). + /// The enqueue returns immediately and records failures on the stream so + /// they surface at its next synchronization point. + fn gemm(&mut self, descriptor: GemmDescriptor, stream_id: StreamId) { + let _ = (descriptor, stream_id); + panic!("Compute server advertised an accelerated GEMM without implementing it") + } + /// Flush all outstanding tasks in the server. fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>; From 1e62e7eef27b442bd64b3bebba38f228f3dbd758 Mon Sep 17 00:00:00 2001 From: pasha Date: Thu, 16 Jul 2026 18:44:19 +0300 Subject: [PATCH 2/8] Harden accelerated GEMM contracts --- .../cubecl-cuda/examples/cublas_gemm_check.rs | 75 ++++++++++++++++++- crates/cubecl-cuda/src/compute/cublas.rs | 26 +++++-- crates/cubecl-cuda/src/compute/server.rs | 10 +++ crates/cubecl-runtime/src/server/base.rs | 4 +- 4 files changed, 105 insertions(+), 10 deletions(-) diff --git a/crates/cubecl-cuda/examples/cublas_gemm_check.rs b/crates/cubecl-cuda/examples/cublas_gemm_check.rs index f7e71fd9f9..98ab3ea0c7 100644 --- a/crates/cubecl-cuda/examples/cublas_gemm_check.rs +++ b/crates/cubecl-cuda/examples/cublas_gemm_check.rs @@ -46,6 +46,9 @@ fn main() { correctness_matrix(); cross_stream_ordering(); overlapping_output_is_rejected(); + aliased_output_is_rejected(); + foreign_output_is_rejected(); + zero_k_is_rejected(); prior_error_is_not_bypassed(); println!("cuBLAS BF16 padded/offset/batched/multistream checks passed"); } @@ -256,6 +259,70 @@ fn overlapping_output_is_rejected() { assert!(future::block_on(client.sync()).is_err()); } +fn aliased_output_is_rejected() { + let mut client = CudaRuntime::client(&CudaDevice::default()); + unsafe { client.set_stream(StreamId { value: 104 }) }; + + let values = vec![bf16::ONE.to_bits(); 4]; + let lhs_and_out = client.create_from_slice(bytemuck::cast_slice(&values)); + let rhs = client.create_from_slice(bytemuck::cast_slice(&values)); + client.gemm(GemmDescriptor::new( + GemmMatrix::new(lhs_and_out.clone().binding(), 2, 0, false), + GemmMatrix::new(rhs.binding(), 2, 0, false), + GemmMatrix::new(lhs_and_out.binding(), 2, 0, false), + 2, + 2, + 2, + 1, + ElemType::Float(FloatKind::BF16), + )); + assert!(future::block_on(client.sync()).is_err()); +} + +fn foreign_output_is_rejected() { + let mut origin = CudaRuntime::client(&CudaDevice::default()); + let mut execution = origin.clone(); + unsafe { + origin.set_stream(StreamId { value: 105 }); + execution.set_stream(StreamId { value: 106 }); + } + + let values = vec![bf16::ONE.to_bits(); 4]; + let lhs = execution.create_from_slice(bytemuck::cast_slice(&values)); + let rhs = execution.create_from_slice(bytemuck::cast_slice(&values)); + let foreign_out = origin.empty(8); + execution.gemm(GemmDescriptor::new( + GemmMatrix::new(lhs.binding(), 2, 0, false), + GemmMatrix::new(rhs.binding(), 2, 0, false), + GemmMatrix::new(foreign_out.binding(), 2, 0, false), + 2, + 2, + 2, + 1, + ElemType::Float(FloatKind::BF16), + )); + assert!(future::block_on(execution.sync()).is_err()); +} + +fn zero_k_is_rejected() { + let mut client = CudaRuntime::client(&CudaDevice::default()); + unsafe { client.set_stream(StreamId { value: 107 }) }; + + let placeholder = client.empty(2); + let out = client.empty(12); + client.gemm(GemmDescriptor::new( + GemmMatrix::new(placeholder.clone().binding(), 1, 0, false), + GemmMatrix::new(placeholder.binding(), 3, 0, false), + GemmMatrix::new(out.binding(), 3, 0, false), + 2, + 3, + 0, + 1, + ElemType::Float(FloatKind::BF16), + )); + assert!(future::block_on(client.sync()).is_err()); +} + fn prior_error_is_not_bypassed() { let client = CudaRuntime::client(&CudaDevice::default()); let input = client.create_from_slice(bytemuck::cast_slice(&[bf16::ONE.to_bits()])); @@ -274,16 +341,20 @@ fn prior_error_is_not_bypassed() { let descriptor = GemmDescriptor::new( GemmMatrix::new(input.clone().binding(), 1, 0, false), GemmMatrix::new(input.binding(), 1, 0, false), - GemmMatrix::new(output.binding(), 1, 0, false), + GemmMatrix::new(output.clone().binding(), 1, 0, false), 1, 1, 1, 1, ElemType::Float(FloatKind::BF16), ); - client.gemm(descriptor); + client.gemm(descriptor.clone()); // Surfaces the original launch error and restores stream health. assert!(future::block_on(client.sync()).is_err()); + + client.gemm(descriptor); + let bytes = client.read_one_unchecked(output); + assert_eq!(bytemuck::cast_slice::(&bytes), &[bf16::ONE.to_bits()]); } impl Matrix { diff --git a/crates/cubecl-cuda/src/compute/cublas.rs b/crates/cubecl-cuda/src/compute/cublas.rs index 00a1a760ef..57a6c68d7b 100644 --- a/crates/cubecl-cuda/src/compute/cublas.rs +++ b/crates/cubecl-cuda/src/compute/cublas.rs @@ -34,15 +34,15 @@ impl CublasState { if descriptor.elem != ElemType::Float(FloatKind::BF16) { return Err(validation_error("cuBLAS GEMM currently supports only BF16")); } - validate(descriptor, lhs, rhs, out)?; - - if descriptor.m == 0 - || descriptor.n == 0 - || descriptor.k == 0 - || descriptor.batch_count == 0 - { + if descriptor.m == 0 || descriptor.n == 0 || descriptor.batch_count == 0 { return Ok(()); } + if descriptor.k == 0 { + return Err(validation_error( + "cuBLAS GEMM does not initialize a nonempty zero-K output", + )); + } + validate(descriptor, lhs, rhs, out)?; let handle = match self.handle { Some(handle) => handle, @@ -136,6 +136,9 @@ fn validate( if descriptor.out.transposed { return Err(validation_error("GEMM output must be row-major")); } + if resources_overlap(out, lhs) || resources_overlap(out, rhs) { + return Err(validation_error("GEMM output may not overlap either input")); + } validate_matrix( "lhs", @@ -166,6 +169,15 @@ fn validate( ) } +fn resources_overlap(lhs: &GpuResource, rhs: &GpuResource) -> bool { + if lhs.size == 0 || rhs.size == 0 { + return false; + } + let lhs_end = lhs.ptr.saturating_add(lhs.size); + let rhs_end = rhs.ptr.saturating_add(rhs.size); + lhs.ptr < rhs_end && rhs.ptr < lhs_end +} + #[allow(clippy::too_many_arguments)] fn validate_matrix( name: &str, diff --git a/crates/cubecl-cuda/src/compute/server.rs b/crates/cubecl-cuda/src/compute/server.rs index ea5a6db6b4..92a489b6c8 100644 --- a/crates/cubecl-cuda/src/compute/server.rs +++ b/crates/cubecl-cuda/src/compute/server.rs @@ -675,6 +675,16 @@ impl CudaServer { descriptor: GemmDescriptor, stream_id: StreamId, ) -> Result<(), ServerError> { + // A binding records its allocation stream, which is also the stream + // later consumers use for dependency tracking. Executing a write on a + // different stream would leave that metadata stale and let a consumer + // on the allocation stream race the GEMM. + if descriptor.out.binding.stream != stream_id { + return Err(ServerError::Validation { + message: "GEMM output must be allocated on the execution stream".into(), + backtrace: BackTrace::capture(), + }); + } let bindings = [ &descriptor.lhs.binding, &descriptor.rhs.binding, diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index dccb959432..e0965d54bc 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -374,7 +374,9 @@ pub struct GemmMatrix { /// /// Computes `out = lhs @ rhs` for `batch_count` matrices. All matrices use /// `elem`; accumulation precision is selected by the backend. The output must -/// be row-major. +/// be row-major, allocated on the execution stream, and disjoint from both +/// inputs. Backends may reject a nonempty zero-K problem rather than provide a +/// separate output-initialization path. #[derive(new, Clone, Debug)] pub struct GemmDescriptor { /// Left-hand matrix, logically shaped `[m, k]`. From bb7961e4a59d6cd3ea80ef7c64cf6dbd026afd99 Mon Sep 17 00:00:00 2001 From: pasha Date: Fri, 17 Jul 2026 13:50:19 +0300 Subject: [PATCH 3/8] Route the accelerated GEMM dispatch through cublasLt cublasGemmEx with CUBLAS_GEMM_DEFAULT_TENSOR_OP leaves algorithm selection to the legacy heuristic. cublasLt exposes the real heuristic API, so each GEMM shape now gets a cached execution plan (descriptor, layouts, and the heuristic-selected algorithm) plus a 32 MiB workspace so split-K algorithms are eligible - the same machinery PyTorch's linear uses. Workspaces are per CUDA stream since concurrent matmuls must not share scratch. Shapes the heuristic rejects surface as validation errors, which the autotune arbitration already treats as candidate unavailable. --- Cargo.toml | 1 + .../cubecl-cuda/examples/cublas_gemm_check.rs | 5 +- crates/cubecl-cuda/src/compute/cublas.rs | 357 ++++++++++++++++-- 3 files changed, 322 insertions(+), 41 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3d4dbd6e38..88508bc41f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -154,6 +154,7 @@ cudarc = { version = "0.19.6", features = [ "driver", "nvrtc", "cublas", + "cublaslt", "nccl", "fallback-dynamic-loading", "cuda-version-from-build-system", diff --git a/crates/cubecl-cuda/examples/cublas_gemm_check.rs b/crates/cubecl-cuda/examples/cublas_gemm_check.rs index 98ab3ea0c7..844c2329bc 100644 --- a/crates/cubecl-cuda/examples/cublas_gemm_check.rs +++ b/crates/cubecl-cuda/examples/cublas_gemm_check.rs @@ -354,7 +354,10 @@ fn prior_error_is_not_bypassed() { client.gemm(descriptor); let bytes = client.read_one_unchecked(output); - assert_eq!(bytemuck::cast_slice::(&bytes), &[bf16::ONE.to_bits()]); + assert_eq!( + bytemuck::cast_slice::(&bytes), + &[bf16::ONE.to_bits()] + ); } impl Matrix { diff --git a/crates/cubecl-cuda/src/compute/cublas.rs b/crates/cubecl-cuda/src/compute/cublas.rs index 57a6c68d7b..05b48c38c4 100644 --- a/crates/cubecl-cuda/src/compute/cublas.rs +++ b/crates/cubecl-cuda/src/compute/cublas.rs @@ -1,23 +1,82 @@ +use std::collections::HashMap; +use std::ffi::c_void; + use cubecl_common::backtrace::BackTrace; use cubecl_core::{ ir::{ElemType, FloatKind}, server::{GemmDescriptor, GemmMatrix, ServerError}, }; -use cudarc::cublas::{result, sys}; -use std::ffi::c_void; +use cudarc::cublas::sys::cublasOperation_t; +use cudarc::cublaslt::{result as lt, sys}; use super::storage::gpu::GpuResource; +/// Workspace given to every cublasLt matmul. 32 MiB matches the size the +/// cuBLAS documentation recommends for Ampere+ so the heuristic can pick +/// split-K and other workspace-hungry algorithms. +const WORKSPACE_BYTES: usize = 32 * 1024 * 1024; + +/// A cached execution plan: the descriptor/layout objects plus the algorithm +/// the cublasLt heuristic selected for one GEMM shape. +struct MatmulPlan { + desc: sys::cublasLtMatmulDesc_t, + a_layout: sys::cublasLtMatrixLayout_t, + b_layout: sys::cublasLtMatrixLayout_t, + d_layout: sys::cublasLtMatrixLayout_t, + algo: sys::cublasLtMatmulAlgo_t, +} + +#[derive(PartialEq, Eq, Hash)] +struct MatrixKey { + leading_dimension: u32, + batch_stride: u64, + transposed: bool, +} + +#[derive(PartialEq, Eq, Hash)] +struct PlanKey { + m: u32, + n: u32, + k: u32, + batch_count: u32, + lhs: MatrixKey, + rhs: MatrixKey, + out: MatrixKey, +} + +impl PlanKey { + fn new(descriptor: &GemmDescriptor) -> Self { + let matrix = |m: &GemmMatrix| MatrixKey { + leading_dimension: m.leading_dimension, + batch_stride: m.batch_stride, + transposed: m.transposed, + }; + Self { + m: descriptor.m, + n: descriptor.n, + k: descriptor.k, + batch_count: descriptor.batch_count, + lhs: matrix(&descriptor.lhs), + rhs: matrix(&descriptor.rhs), + out: matrix(&descriptor.out), + } + } +} + #[derive(Default)] pub(crate) struct CublasState { - handle: Option, - stream: Option, + handle: Option, + /// One workspace per CUDA stream: concurrent matmuls on different + /// streams must not share scratch memory. + workspaces: HashMap, + plans: HashMap, } impl core::fmt::Debug for CublasState { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("CublasState") .field("initialized", &self.handle.is_some()) + .field("plans", &self.plans.len()) .finish() } } @@ -47,51 +106,60 @@ impl CublasState { let handle = match self.handle { Some(handle) => handle, None => { - let handle = result::create_handle().map_err(cublas_error)?; + let handle = lt::create_handle().map_err(cublas_error)?; self.handle = Some(handle); handle } }; - if self.stream != Some(stream) { - // SAFETY: both handles are owned by this CUDA server. Changing the - // cuBLAS stream only changes where subsequent work is enqueued. - unsafe { result::set_stream(handle, stream.cast()) }.map_err(cublas_error)?; - self.stream = Some(stream); + let workspace = match self.workspaces.get(&(stream as usize)) { + Some(workspace) => *workspace, + None => { + // SAFETY: the server made its CUDA context current before + // resolving the streams for this launch. + let workspace = unsafe { cudarc::driver::result::malloc_sync(WORKSPACE_BYTES) } + .map_err(|err| ServerError::Generic { + reason: format!("cublasLt workspace allocation failed: {err}"), + backtrace: BackTrace::capture(), + })?; + self.workspaces.insert(stream as usize, workspace); + workspace + } + }; + + let key = PlanKey::new(descriptor); + if !self.plans.contains_key(&key) { + let plan = build_plan(handle, descriptor)?; + self.plans.insert(PlanKey::new(descriptor), plan); } + let plan = self + .plans + .get(&key) + .expect("cublasLt plan was just inserted"); let alpha = 1.0f32; let beta = 0.0f32; - let op_rhs = operation(&descriptor.rhs); - let op_lhs = operation(&descriptor.lhs); // cuBLAS is column-major. Swapping the row-major operands computes - // C^T = rhs^T @ lhs^T without copies. The call is asynchronous on the - // exact stream resolved by CubeCL for these bindings. + // D^T = rhs^T @ lhs^T without copies, so A carries `rhs` and B + // carries `lhs`. The call is asynchronous on the given stream. unsafe { - result::gemm_strided_batched_ex( + lt::matmul( handle, - op_rhs, - op_lhs, - descriptor.n as i32, - descriptor.m as i32, - descriptor.k as i32, + plan.desc, (&alpha as *const f32).cast::(), + (&beta as *const f32).cast::(), rhs.ptr as *const c_void, - sys::cudaDataType::CUDA_R_16BF, - descriptor.rhs.leading_dimension as i32, - descriptor.rhs.batch_stride as i64, + plan.a_layout, lhs.ptr as *const c_void, - sys::cudaDataType::CUDA_R_16BF, - descriptor.lhs.leading_dimension as i32, - descriptor.lhs.batch_stride as i64, - (&beta as *const f32).cast::(), + plan.b_layout, + out.ptr as *const c_void, + plan.d_layout, out.ptr as *mut c_void, - sys::cudaDataType::CUDA_R_16BF, - descriptor.out.leading_dimension as i32, - descriptor.out.batch_stride as i64, - descriptor.batch_count as i32, - sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, - sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, + plan.d_layout, + &plan.algo, + workspace as *mut c_void, + WORKSPACE_BYTES, + stream.cast(), ) } .map_err(cublas_error)?; @@ -100,20 +168,229 @@ impl CublasState { } pub(crate) fn destroy(&mut self) { + for (_, plan) in self.plans.drain() { + // SAFETY: each plan uniquely owns its descriptor objects and is + // destroyed exactly once. + unsafe { + let _ = lt::destroy_matmul_desc(plan.desc); + let _ = lt::destroy_matrix_layout(plan.a_layout); + let _ = lt::destroy_matrix_layout(plan.b_layout); + let _ = lt::destroy_matrix_layout(plan.d_layout); + } + } + for (_, workspace) in self.workspaces.drain() { + // SAFETY: the workspace was allocated by this state on the + // server's context and freed exactly once. + if let Err(err) = unsafe { cudarc::driver::result::free_sync(workspace) } { + log::warn!("Unable to free cublasLt workspace: {err}"); + } + } if let Some(handle) = self.handle.take() { // SAFETY: this state uniquely owns the handle and destroys it once. - if let Err(err) = unsafe { result::destroy_handle(handle) } { - log::warn!("Unable to destroy cuBLAS handle: {err}"); + if let Err(err) = unsafe { lt::destroy_handle(handle) } { + log::warn!("Unable to destroy cublasLt handle: {err}"); } } } } -fn operation(matrix: &GemmMatrix) -> sys::cublasOperation_t { +/// Build the descriptor, layouts, and heuristic-selected algorithm for one +/// GEMM shape. All dimensions below are in cuBLAS column-major terms, i.e. +/// the row-major operands are swapped: `m_lt = n`, `n_lt = m`, `A = rhs`, +/// `B = lhs`. +fn build_plan( + handle: sys::cublasLtHandle_t, + descriptor: &GemmDescriptor, +) -> Result { + let m_lt = descriptor.n as u64; + let n_lt = descriptor.m as u64; + let k_lt = descriptor.k as u64; + let op_a = operation(&descriptor.rhs); + let op_b = operation(&descriptor.lhs); + + let desc = lt::create_matmul_desc( + sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + sys::cudaDataType::CUDA_R_32F, + ) + .map_err(cublas_error)?; + let destroy_desc = || { + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matmul_desc(desc); + } + }; + for (attr, op) in [ + ( + sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA, + op_a, + ), + ( + sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB, + op_b, + ), + ] { + let value = op as i32; + // SAFETY: `desc` is live and the attribute is an i32 by contract. + if let Err(err) = unsafe { + lt::set_matmul_desc_attribute( + desc, + attr, + (&value as *const i32).cast::(), + core::mem::size_of::(), + ) + } { + destroy_desc(); + return Err(cublas_error(err)); + } + } + + // Stored (pre-transpose) dimensions of each column-major operand. + let a_dims = if matches!(op_a, cublasOperation_t::CUBLAS_OP_N) { + (m_lt, k_lt) + } else { + (k_lt, m_lt) + }; + let b_dims = if matches!(op_b, cublasOperation_t::CUBLAS_OP_N) { + (k_lt, n_lt) + } else { + (n_lt, k_lt) + }; + let layout = |rows: u64, cols: u64, matrix: &GemmMatrix| -> Result<_, ServerError> { + let layout = lt::create_matrix_layout( + sys::cudaDataType::CUDA_R_16BF, + rows, + cols, + matrix.leading_dimension as i64, + ) + .map_err(cublas_error)?; + let batch_count = descriptor.batch_count as i32; + let batch_stride = matrix.batch_stride as i64; + // SAFETY: `layout` is live; both attributes take the given widths. + let result = unsafe { + lt::set_matrix_layout_attribute( + layout, + sys::cublasLtMatrixLayoutAttribute_t::CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + (&batch_count as *const i32).cast::(), + core::mem::size_of::(), + ) + .and_then(|_| { + lt::set_matrix_layout_attribute( + layout, + sys::cublasLtMatrixLayoutAttribute_t:: + CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, + (&batch_stride as *const i64).cast::(), + core::mem::size_of::(), + ) + }) + }; + if let Err(err) = result { + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(layout); + } + return Err(cublas_error(err)); + } + Ok(layout) + }; + + let a_layout = match layout(a_dims.0, a_dims.1, &descriptor.rhs) { + Ok(layout) => layout, + Err(err) => { + destroy_desc(); + return Err(err); + } + }; + let b_layout = match layout(b_dims.0, b_dims.1, &descriptor.lhs) { + Ok(layout) => layout, + Err(err) => { + destroy_desc(); + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(a_layout); + } + return Err(err); + } + }; + let d_layout = match layout(m_lt, n_lt, &descriptor.out) { + Ok(layout) => layout, + Err(err) => { + destroy_desc(); + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(a_layout); + let _ = lt::destroy_matrix_layout(b_layout); + } + return Err(err); + } + }; + let cleanup = || { + destroy_desc(); + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(a_layout); + let _ = lt::destroy_matrix_layout(b_layout); + let _ = lt::destroy_matrix_layout(d_layout); + } + }; + + let pref = match lt::create_matmul_pref() { + Ok(pref) => pref, + Err(err) => { + cleanup(); + return Err(cublas_error(err)); + } + }; + let workspace_bytes = WORKSPACE_BYTES as u64; + // SAFETY: `pref` is live and the attribute is a u64 by contract. + if let Err(err) = unsafe { + lt::set_matmul_pref_attribute( + pref, + sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + (&workspace_bytes as *const u64).cast::(), + core::mem::size_of::(), + ) + } { + // SAFETY: created above. + unsafe { + let _ = lt::destroy_matmul_pref(pref); + } + cleanup(); + return Err(cublas_error(err)); + } + + // SAFETY: every descriptor is live; the C layout equals the D layout + // because beta is always zero. + let heuristic = unsafe { + lt::get_matmul_algo_heuristic(handle, desc, a_layout, b_layout, d_layout, d_layout, pref) + }; + // SAFETY: created above. + unsafe { + let _ = lt::destroy_matmul_pref(pref); + } + let heuristic = match heuristic { + Ok(heuristic) => heuristic, + Err(_) => { + cleanup(); + return Err(validation_error( + "cublasLt has no algorithm for this GEMM shape", + )); + } + }; + + Ok(MatmulPlan { + desc, + a_layout, + b_layout, + d_layout, + algo: heuristic.algo, + }) +} + +fn operation(matrix: &GemmMatrix) -> cublasOperation_t { if matrix.transposed { - sys::cublasOperation_t::CUBLAS_OP_T + cublasOperation_t::CUBLAS_OP_T } else { - sys::cublasOperation_t::CUBLAS_OP_N + cublasOperation_t::CUBLAS_OP_N } } @@ -237,9 +514,9 @@ fn validation_error(message: &str) -> ServerError { } } -fn cublas_error(error: result::CublasError) -> ServerError { +fn cublas_error(error: lt::CublasError) -> ServerError { ServerError::Generic { - reason: format!("cuBLAS GEMM failed: {error}"), + reason: format!("cuBLAS error: {error:?}"), backtrace: BackTrace::capture(), } } From d840eb752cbd583fa34dca4d1d17a432e964a226 Mon Sep 17 00:00:00 2001 From: pasha Date: Thu, 16 Jul 2026 07:39:09 +0300 Subject: [PATCH 4/8] Fix allocation retry after cached-page OOM --- .../src/memory_management/memory_manage.rs | 121 ++++++++++++++++-- 1 file changed, 110 insertions(+), 11 deletions(-) diff --git a/crates/cubecl-runtime/src/memory_management/memory_manage.rs b/crates/cubecl-runtime/src/memory_management/memory_manage.rs index 4f6b396acb..ef2207550e 100644 --- a/crates/cubecl-runtime/src/memory_management/memory_manage.rs +++ b/crates/cubecl-runtime/src/memory_management/memory_manage.rs @@ -938,20 +938,35 @@ impl MemoryManagement { ); // Find first pool that fits this allocation - let pool = self - .pools - .iter_mut() - .find(|p| p.accept(size)) - .ok_or(IoError::BufferTooBig { - size, - backtrace: BackTrace::capture(), - })?; + let pool_index = + self.pools + .iter() + .position(|p| p.accept(size)) + .ok_or(IoError::BufferTooBig { + size, + backtrace: BackTrace::capture(), + })?; - if let Some(slice) = pool.try_reserve(size) { + if let Some(slice) = self.pools[pool_index].try_reserve(size) { return Ok(slice); } - let allocated = pool.alloc(&mut self.storage, size); + let allocated = self.pools[pool_index].alloc(&mut self.storage, size); + + // Autotuning can leave one-off pages cached in several size buckets. + // If the device is full, release every unused page and retry once + // before surfacing the allocation error. + let allocated = match allocated { + Err(IoError::BufferTooBig { .. }) => { + self.cleanup(true); + if let Some(slice) = self.pools[pool_index].try_reserve(size) { + Ok(slice) + } else { + self.pools[pool_index].alloc(&mut self.storage, size) + } + } + result => result, + }; self.logger.log_memory( |level| matches!(level, MemoryLogLevel::Full), @@ -1070,8 +1085,12 @@ impl core::fmt::Debug for MemoryManagement { #[cfg(test)] mod tests { use super::*; - use crate::{memory_management::MemoryManagement, storage::BytesStorage}; + use crate::{ + memory_management::MemoryManagement, + storage::{BytesResource, BytesStorage, ComputeStorage, StorageHandle, StorageId}, + }; use alloc::vec; + use hashbrown::HashMap; const DUMMY_MEM_PROPS: MemoryDeviceProperties = MemoryDeviceProperties { max_page_size: 128 * 1024 * 1024, @@ -1085,6 +1104,86 @@ mod tests { } } + #[derive(Debug)] + struct LimitedStorage { + inner: BytesStorage, + capacity: u64, + allocated: u64, + sizes: HashMap, + } + + impl LimitedStorage { + fn new(capacity: u64) -> Self { + Self { + inner: BytesStorage::default(), + capacity, + allocated: 0, + sizes: HashMap::new(), + } + } + } + + impl ComputeStorage for LimitedStorage { + type Resource = BytesResource; + + fn alignment(&self) -> usize { + self.inner.alignment() + } + + fn get(&mut self, handle: &StorageHandle) -> Self::Resource { + self.inner.get(handle) + } + + fn alloc(&mut self, size: u64) -> Result { + if self.allocated + size > self.capacity { + return Err(IoError::BufferTooBig { + size, + backtrace: BackTrace::capture(), + }); + } + let handle = self.inner.alloc(size)?; + self.allocated += size; + self.sizes.insert(handle.id, size); + Ok(handle) + } + + fn dealloc(&mut self, id: StorageId) { + if let Some(size) = self.sizes.remove(&id) { + self.allocated -= size; + } + self.inner.dealloc(id); + } + + fn flush(&mut self) { + self.inner.flush(); + } + } + + #[test_log::test] + fn allocation_oom_reclaims_cached_pages_and_retries() { + let mut memory_management = MemoryManagement::from_configuration( + LimitedStorage::new(1024), + &DUMMY_MEM_PROPS, + MemoryConfiguration::Custom { + pool_options: vec![MemoryPoolOptions { + pool_type: PoolType::ExclusivePages { + max_alloc_size: 1024, + }, + dealloc_period: None, + }], + }, + Arc::new(ServerLogger::default()), + options(), + ); + + let cached = memory_management.reserve(768).unwrap(); + drop(cached); + + let _live = memory_management.reserve(1024).unwrap(); + assert_eq!(memory_management.memory_usage().bytes_in_use, 1024); + assert_eq!(memory_management.memory_usage().bytes_reserved, 1024); + } + // Test pools with slices. #[test_log::test] #[cfg(not(exclusive_memory_only))] From 43eef09280b33517764dabfb0b076d0e4cafcc4b Mon Sep 17 00:00:00 2001 From: pasha Date: Tue, 21 Jul 2026 14:08:49 +0300 Subject: [PATCH 5/8] feat(cuda): add native grouped GEMM dispatch --- .../cubecl-cuda/examples/cublas_gemm_check.rs | 65 +++- crates/cubecl-cuda/src/compute/cublas.rs | 352 +++++++++++++++++- crates/cubecl-cuda/src/compute/server.rs | 53 +++ crates/cubecl-cuda/src/runtime.rs | 8 + crates/cubecl-ir/src/features.rs | 2 + crates/cubecl-runtime/src/client.rs | 19 +- crates/cubecl-runtime/src/server/base.rs | 23 ++ 7 files changed, 512 insertions(+), 10 deletions(-) diff --git a/crates/cubecl-cuda/examples/cublas_gemm_check.rs b/crates/cubecl-cuda/examples/cublas_gemm_check.rs index 844c2329bc..6624a80fdd 100644 --- a/crates/cubecl-cuda/examples/cublas_gemm_check.rs +++ b/crates/cubecl-cuda/examples/cublas_gemm_check.rs @@ -4,7 +4,7 @@ use cubecl_core::{ Runtime, ir::{ElemType, FloatKind}, prelude::*, - server::{Binding, GemmDescriptor, GemmMatrix, Handle}, + server::{Binding, GemmDescriptor, GemmMatrix, GroupedGemmDescriptor, Handle}, }; use cubecl_cuda::{CudaDevice, CudaRuntime}; use half::bf16; @@ -44,13 +44,14 @@ struct Matrix { fn main() { correctness_matrix(); + grouped_correctness(); cross_stream_ordering(); overlapping_output_is_rejected(); aliased_output_is_rejected(); foreign_output_is_rejected(); zero_k_is_rejected(); prior_error_is_not_bypassed(); - println!("cuBLAS BF16 padded/offset/batched/multistream checks passed"); + println!("cuBLAS BF16 padded/offset/batched/grouped/multistream checks passed"); } fn correctness_matrix() { @@ -157,6 +158,66 @@ fn check(problem: Problem) { } } +fn grouped_correctness() { + let client = CudaRuntime::client(&CudaDevice::default()); + let elem = ElemType::Float(FloatKind::BF16); + if !client + .features() + .matmul + .accelerated_grouped_gemm + .contains(&elem) + { + return; + } + let problems = [(3, 5, 4), (7, 2, 3), (4, 6, 5)]; + let mut matrices = Vec::with_capacity(problems.len()); + let mut groups = Vec::with_capacity(problems.len()); + + for (index, (m, n, k)) in problems.into_iter().enumerate() { + let lhs = Matrix::new(m, k, 1, index % 2 == 0, 3 + index, 0, 7, index + 2); + let rhs = Matrix::new(k, n, 1, index % 2 != 0, 5 + index, 0, 11, index + 5); + let out = Matrix::zeros(m, n, 1, false, 2 + index, 0, 13); + let lhs_base = client.create_from_slice(bytemuck::cast_slice(&lhs.bits)); + let rhs_base = client.create_from_slice(bytemuck::cast_slice(&rhs.bits)); + let out_base = client.empty(out.bits.len() * 2); + groups.push(GemmDescriptor::new( + matrix_arg(&lhs, view(&lhs_base, lhs.offset, lhs.bits.len()), false), + matrix_arg(&rhs, view(&rhs_base, rhs.offset, rhs.bits.len()), false), + matrix_arg(&out, view(&out_base, out.offset, out.bits.len()), false), + m as u32, + n as u32, + k as u32, + 1, + elem, + )); + matrices.push((lhs, rhs, out, lhs_base, rhs_base, out_base)); + } + + let descriptor = GroupedGemmDescriptor::new(groups); + for _ in 0..16 { + client.grouped_gemm(descriptor.clone()); + } + for (lhs, rhs, out, _lhs_base, _rhs_base, out_base) in matrices { + let bytes = client.read_one_unchecked(out_base); + let actual = bytemuck::cast_slice::(&bytes); + for row in 0..out.rows { + for col in 0..out.cols { + let expected = (0..lhs.cols) + .map(|inner| lhs.get(0, row, inner) * rhs.get(0, inner, col)) + .sum::(); + let actual = bf16::from_bits(actual[out.index(0, row, col)]).to_f32(); + assert!( + (actual - expected).abs() <= 0.06, + "grouped m={}, n={}, k={}, row={row}, col={col}: {actual} != {expected}", + out.rows, + out.cols, + lhs.cols + ); + } + } + } +} + fn cross_stream_ordering() { let mut producer = CudaRuntime::client(&CudaDevice::default()); let mut gemm = producer.clone(); diff --git a/crates/cubecl-cuda/src/compute/cublas.rs b/crates/cubecl-cuda/src/compute/cublas.rs index 05b48c38c4..7581ff4d53 100644 --- a/crates/cubecl-cuda/src/compute/cublas.rs +++ b/crates/cubecl-cuda/src/compute/cublas.rs @@ -2,11 +2,15 @@ use std::collections::HashMap; use std::ffi::c_void; use cubecl_common::backtrace::BackTrace; +#[cfg(cuda_12050)] +use cubecl_core::server::GroupedGemmDescriptor; use cubecl_core::{ ir::{ElemType, FloatKind}, server::{GemmDescriptor, GemmMatrix, ServerError}, }; use cudarc::cublas::sys::cublasOperation_t; +#[cfg(cuda_12050)] +use cudarc::cublas::{result as blas, sys as blas_sys}; use cudarc::cublaslt::{result as lt, sys}; use super::storage::gpu::GpuResource; @@ -26,6 +30,96 @@ struct MatmulPlan { algo: sys::cublasLtMatmulAlgo_t, } +/// Pinned-host and device storage for one in-flight grouped-GEMM pointer list. +/// +/// cuBLAS consumes the matrix-pointer arrays on the device. The completion +/// event prevents either side of this staging pair from being overwritten +/// while an earlier grouped launch is still reading it. +#[cfg(cuda_12050)] +struct GroupedPointerStaging { + host: *mut std::ffi::c_void, + device: cudarc::driver::sys::CUdeviceptr, + event: cudarc::driver::sys::CUevent, + capacity: usize, + in_flight: bool, + /// Captured graph nodes retain both staging addresses for replay, so a + /// captured slot must remain immutable until the server is destroyed. + captured: bool, +} + +#[cfg(cuda_12050)] +impl GroupedPointerStaging { + fn new(capacity: usize) -> Result { + let bytes = capacity + .checked_mul(core::mem::size_of::()) + .ok_or_else(|| validation_error("grouped GEMM pointer staging size overflow"))?; + // SAFETY: both allocations are owned by the returned staging slot and + // released exactly once in `destroy`. + let host = unsafe { + cudarc::driver::result::malloc_host( + bytes, + cudarc::driver::sys::CU_MEMHOSTALLOC_WRITECOMBINED, + ) + } + .map_err(cuda_error)?; + let device = match unsafe { cudarc::driver::result::malloc_sync(bytes) } { + Ok(device) => device, + Err(err) => { + // SAFETY: `host` was allocated immediately above and has not + // been exposed or freed. + unsafe { + let _ = cudarc::driver::result::free_host(host); + } + return Err(cuda_error(err)); + } + }; + let event = match cudarc::driver::result::event::create( + cudarc::driver::sys::CUevent_flags_enum::CU_EVENT_DISABLE_TIMING, + ) { + Ok(event) => event, + Err(err) => { + // SAFETY: both allocations are uniquely owned here. + unsafe { + let _ = cudarc::driver::result::free_sync(device); + let _ = cudarc::driver::result::free_host(host); + } + return Err(cuda_error(err)); + } + }; + Ok(Self { + host, + device, + event, + capacity, + in_flight: false, + captured: false, + }) + } + + fn available(&self) -> bool { + !self.captured + && (!self.in_flight + // SAFETY: `event` remains live for this slot's entire lifetime. + || unsafe { cudarc::driver::result::event::query(self.event) }.is_ok()) + } + + fn destroy(self) { + // The server is shutting down. Waiting here protects the pinned host + // buffer if a final pointer upload is still in flight. + if self.in_flight { + // SAFETY: `event` is live and was recorded after the last grouped + // launch using this slot. + let _ = unsafe { cudarc::driver::result::event::synchronize(self.event) }; + } + // SAFETY: all three resources are uniquely owned by this slot. + unsafe { + let _ = cudarc::driver::result::event::destroy(self.event); + let _ = cudarc::driver::result::free_sync(self.device); + let _ = cudarc::driver::result::free_host(self.host); + } + } +} + #[derive(PartialEq, Eq, Hash)] struct MatrixKey { leading_dimension: u32, @@ -66,18 +160,26 @@ impl PlanKey { #[derive(Default)] pub(crate) struct CublasState { handle: Option, + #[cfg(cuda_12050)] + grouped_handle: Option, /// One workspace per CUDA stream: concurrent matmuls on different /// streams must not share scratch memory. workspaces: HashMap, plans: HashMap, + #[cfg(cuda_12050)] + grouped_staging: Vec, } impl core::fmt::Debug for CublasState { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - f.debug_struct("CublasState") - .field("initialized", &self.handle.is_some()) - .field("plans", &self.plans.len()) - .finish() + let mut debug = f.debug_struct("CublasState"); + debug.field("initialized", &self.handle.is_some()); + #[cfg(cuda_12050)] + debug.field("grouped_initialized", &self.grouped_handle.is_some()); + debug.field("plans", &self.plans.len()); + #[cfg(cuda_12050)] + debug.field("grouped_staging", &self.grouped_staging.len()); + debug.finish() } } @@ -167,6 +269,207 @@ impl CublasState { Ok(()) } + #[cfg(cuda_12050)] + pub(crate) fn launch_grouped( + &mut self, + descriptor: &GroupedGemmDescriptor, + lhs: &[GpuResource], + rhs: &[GpuResource], + out: &[GpuResource], + stream: cudarc::driver::sys::CUstream, + ) -> Result<(), ServerError> { + let Some(elem) = descriptor.groups.first().map(|group| group.elem) else { + return Ok(()); + }; + if elem != ElemType::Float(FloatKind::BF16) { + return Err(validation_error( + "cuBLAS grouped GEMM currently supports only BF16", + )); + } + if descriptor.groups.len() != lhs.len() + || descriptor.groups.len() != rhs.len() + || descriptor.groups.len() != out.len() + { + return Err(validation_error( + "grouped GEMM descriptor/resource count mismatch", + )); + } + + let mut active = Vec::with_capacity(descriptor.groups.len()); + let mut pointer_count = 0usize; + for (index, group) in descriptor.groups.iter().enumerate() { + if group.elem != elem { + return Err(validation_error( + "every grouped GEMM entry must use the descriptor element type", + )); + } + if group.m == 0 || group.n == 0 || group.batch_count == 0 { + continue; + } + if group.k == 0 { + return Err(validation_error( + "cuBLAS grouped GEMM does not initialize a nonempty zero-K output", + )); + } + validate(group, &lhs[index], &rhs[index], &out[index])?; + pointer_count = pointer_count + .checked_add(group.batch_count as usize) + .ok_or_else(|| validation_error("grouped GEMM batch count overflow"))?; + active.push(index); + } + if active.is_empty() { + return Ok(()); + } + for (position, &index) in active.iter().enumerate() { + for &other in &active[position + 1..] { + if resources_overlap(&out[index], &out[other]) + || resources_overlap(&out[index], &lhs[other]) + || resources_overlap(&out[index], &rhs[other]) + || resources_overlap(&out[other], &lhs[index]) + || resources_overlap(&out[other], &rhs[index]) + { + return Err(validation_error( + "grouped GEMM outputs may not overlap another group", + )); + } + } + } + let group_count: i32 = active + .len() + .try_into() + .map_err(|_| validation_error("grouped GEMM has more than i32::MAX groups"))?; + let staging_values = pointer_count + .checked_mul(3) + .ok_or_else(|| validation_error("grouped GEMM pointer count overflow"))?; + + let slot_index = match self + .grouped_staging + .iter() + .position(|slot| slot.capacity >= staging_values && slot.available()) + { + Some(index) => index, + None => { + let capacity = staging_values + .checked_next_power_of_two() + .ok_or_else(|| validation_error("grouped GEMM staging capacity overflow"))?; + self.grouped_staging + .push(GroupedPointerStaging::new(capacity)?); + self.grouped_staging.len() - 1 + } + }; + let slot = &mut self.grouped_staging[slot_index]; + // CUDA graphs replay the captured host-to-device copy, so both the + // pinned source and device destination addresses must remain stable. + // Retaining a tiny slot per captured grouped launch provides that + // lifetime without imposing synchronization on ordinary launches. + // SAFETY: `stream` is the live execution stream resolved by the server. + slot.captured = !matches!( + unsafe { cudarc::driver::result::stream::is_capturing(stream) }.map_err(cuda_error)?, + cudarc::driver::sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_NONE + ); + // SAFETY: the pinned allocation contains `capacity` u64 values and the + // selected slot is no longer in use by a previous launch. + let host = + unsafe { core::slice::from_raw_parts_mut(slot.host.cast::(), slot.capacity) }; + + let mut trans_a = Vec::with_capacity(active.len()); + let mut trans_b = Vec::with_capacity(active.len()); + let mut m = Vec::with_capacity(active.len()); + let mut n = Vec::with_capacity(active.len()); + let mut k = Vec::with_capacity(active.len()); + let mut lda = Vec::with_capacity(active.len()); + let mut ldb = Vec::with_capacity(active.len()); + let mut ldc = Vec::with_capacity(active.len()); + let mut group_sizes = Vec::with_capacity(active.len()); + let alpha = vec![1.0_f32; active.len()]; + let beta = vec![0.0_f32; active.len()]; + let elem_size = core::mem::size_of::() as u64; + let (a_values, rest) = host[..staging_values].split_at_mut(pointer_count); + let (b_values, c_values) = rest.split_at_mut(pointer_count); + let mut pointer = 0; + + for &index in &active { + let group = &descriptor.groups[index]; + // cuBLAS is column-major. Swapping the row-major operands computes + // D^T = rhs^T @ lhs^T without materialization. + trans_a.push(operation(&group.rhs)); + trans_b.push(operation(&group.lhs)); + m.push(group.n as i32); + n.push(group.m as i32); + k.push(group.k as i32); + lda.push(group.rhs.leading_dimension as i32); + ldb.push(group.lhs.leading_dimension as i32); + ldc.push(group.out.leading_dimension as i32); + group_sizes.push(group.batch_count as i32); + for batch in 0..group.batch_count as u64 { + a_values[pointer] = batch_pointer(&rhs[index], &group.rhs, batch, elem_size)?; + b_values[pointer] = batch_pointer(&lhs[index], &group.lhs, batch, elem_size)?; + c_values[pointer] = batch_pointer(&out[index], &group.out, batch, elem_size)?; + pointer += 1; + } + } + debug_assert_eq!(pointer, pointer_count); + + // SAFETY: the host slice is pinned and remains untouched until the + // completion event recorded below. The device allocation is large + // enough for exactly `staging_values` pointers. + unsafe { + cudarc::driver::result::memcpy_htod_async(slot.device, &host[..staging_values], stream) + } + .map_err(cuda_error)?; + + let handle = match self.grouped_handle { + Some(handle) => handle, + None => { + let handle = blas::create_handle().map_err(grouped_cublas_error)?; + self.grouped_handle = Some(handle); + handle + } + }; + // SAFETY: the handle and CubeCL stream are live for the server. + unsafe { blas::set_stream(handle, stream.cast()) }.map_err(grouped_cublas_error)?; + let a_device = slot.device as *const *const c_void; + let b_device = (slot.device + (pointer_count * core::mem::size_of::()) as u64) + as *const *const c_void; + let c_device = (slot.device + (2 * pointer_count * core::mem::size_of::()) as u64) + as *const *mut c_void; + // SAFETY: dimensions and leading dimensions were validated; all + // pointer arrays reside in the staging device allocation and refer to + // live CubeCL resources ordered on `stream`. + let launch = unsafe { + blas_sys::cublasGemmGroupedBatchedEx( + handle, + trans_a.as_ptr(), + trans_b.as_ptr(), + m.as_ptr(), + n.as_ptr(), + k.as_ptr(), + alpha.as_ptr().cast::(), + a_device, + blas_sys::cudaDataType_t::CUDA_R_16BF, + lda.as_ptr(), + b_device, + blas_sys::cudaDataType_t::CUDA_R_16BF, + ldb.as_ptr(), + beta.as_ptr().cast::(), + c_device, + blas_sys::cudaDataType_t::CUDA_R_16BF, + ldc.as_ptr(), + group_count, + group_sizes.as_ptr(), + blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + ) + .result() + }; + // Record even when cuBLAS rejects the launch: the pointer upload is + // already queued and must complete before the slot is reused. + // SAFETY: `event` and `stream` are live. + unsafe { cudarc::driver::result::event::record(slot.event, stream) }.map_err(cuda_error)?; + slot.in_flight = true; + launch.map_err(grouped_cublas_error)?; + Ok(()) + } + pub(crate) fn destroy(&mut self) { for (_, plan) in self.plans.drain() { // SAFETY: each plan uniquely owns its descriptor objects and is @@ -185,12 +488,23 @@ impl CublasState { log::warn!("Unable to free cublasLt workspace: {err}"); } } + #[cfg(cuda_12050)] + for staging in self.grouped_staging.drain(..) { + staging.destroy(); + } if let Some(handle) = self.handle.take() { // SAFETY: this state uniquely owns the handle and destroys it once. if let Err(err) = unsafe { lt::destroy_handle(handle) } { log::warn!("Unable to destroy cublasLt handle: {err}"); } } + #[cfg(cuda_12050)] + if let Some(handle) = self.grouped_handle.take() { + // SAFETY: this state uniquely owns the handle and destroys it once. + if let Err(err) = unsafe { blas::destroy_handle(handle) } { + log::warn!("Unable to destroy cuBLAS grouped-GEMM handle: {err}"); + } + } } } @@ -455,6 +769,20 @@ fn resources_overlap(lhs: &GpuResource, rhs: &GpuResource) -> bool { lhs.ptr < rhs_end && rhs.ptr < lhs_end } +#[cfg(cuda_12050)] +fn batch_pointer( + resource: &GpuResource, + matrix: &GemmMatrix, + batch: u64, + elem_size: u64, +) -> Result { + batch + .checked_mul(matrix.batch_stride) + .and_then(|offset| offset.checked_mul(elem_size)) + .and_then(|offset| resource.ptr.checked_add(offset)) + .ok_or_else(|| validation_error("grouped GEMM batch pointer overflow")) +} + #[allow(clippy::too_many_arguments)] fn validate_matrix( name: &str, @@ -520,3 +848,19 @@ fn cublas_error(error: lt::CublasError) -> ServerError { backtrace: BackTrace::capture(), } } + +#[cfg(cuda_12050)] +fn grouped_cublas_error(error: blas::CublasError) -> ServerError { + ServerError::Generic { + reason: format!("cuBLAS grouped GEMM error: {error:?}"), + backtrace: BackTrace::capture(), + } +} + +#[cfg(cuda_12050)] +fn cuda_error(error: cudarc::driver::DriverError) -> ServerError { + ServerError::Generic { + reason: format!("CUDA grouped GEMM staging error: {error:?}"), + backtrace: BackTrace::capture(), + } +} diff --git a/crates/cubecl-cuda/src/compute/server.rs b/crates/cubecl-cuda/src/compute/server.rs index 92a489b6c8..c1717e6101 100644 --- a/crates/cubecl-cuda/src/compute/server.rs +++ b/crates/cubecl-cuda/src/compute/server.rs @@ -11,6 +11,8 @@ use crate::{ }, }; use cubecl_common::{bytes::Bytes, profile::ProfileDuration}; +#[cfg(cuda_12050)] +use cubecl_core::server::GroupedGemmDescriptor; use cubecl_core::{ MemoryConfiguration, device::DeviceId, @@ -178,6 +180,17 @@ impl ComputeServer for CudaServer { } } + #[cfg(cuda_12050)] + fn grouped_gemm(&mut self, descriptor: GroupedGemmDescriptor, stream_id: StreamId) { + if let Err(err) = self.grouped_gemm_checked(descriptor, stream_id) { + let mut stream = match self.streams.resolve(stream_id, [].into_iter(), false) { + Ok(stream) => stream, + Err(err) => unreachable!("{err}"), + }; + stream.current().errors.push(err); + } + } + fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError> { let mut command = self.command_no_inputs( stream_id, @@ -704,6 +717,46 @@ impl CudaServer { cublas.launch(&descriptor, &lhs, &rhs, &out, stream) } + #[cfg(cuda_12050)] + fn grouped_gemm_checked( + &mut self, + descriptor: GroupedGemmDescriptor, + stream_id: StreamId, + ) -> Result<(), ServerError> { + if descriptor + .groups + .iter() + .any(|group| group.out.binding.stream != stream_id) + { + return Err(ServerError::Validation { + message: "grouped GEMM outputs must be allocated on the execution stream".into(), + backtrace: BackTrace::capture(), + }); + } + let bindings = descriptor + .groups + .iter() + .flat_map(|group| [&group.lhs.binding, &group.rhs.binding, &group.out.binding]) + .collect::>(); + self.unsafe_set_current(); + let streams = self + .streams + .resolve(stream_id, bindings.into_iter(), true)?; + let cublas = &mut self.cublas; + let mut command = Command::new(&mut self.ctx, streams); + let mut lhs = Vec::with_capacity(descriptor.groups.len()); + let mut rhs = Vec::with_capacity(descriptor.groups.len()); + let mut out = Vec::with_capacity(descriptor.groups.len()); + for group in &descriptor.groups { + lhs.push(command.resource(group.lhs.binding.clone())?); + rhs.push(command.resource(group.rhs.binding.clone())?); + out.push(command.resource(group.out.binding.clone())?); + } + let stream = command.streams.current().sys; + + cublas.launch_grouped(&descriptor, &lhs, &rhs, &out, stream) + } + fn launch_checked( &mut self, kernel: Box>, diff --git a/crates/cubecl-cuda/src/runtime.rs b/crates/cubecl-cuda/src/runtime.rs index 5efb5bfaf0..b6136e1d5a 100644 --- a/crates/cubecl-cuda/src/runtime.rs +++ b/crates/cubecl-cuda/src/runtime.rs @@ -217,6 +217,14 @@ impl DeviceService for CudaServer { .matmul .accelerated_gemm .insert(ElemType::Float(FloatKind::BF16)); + #[cfg(cuda_12050)] + { + device_props + .features + .matmul + .accelerated_grouped_gemm + .insert(ElemType::Float(FloatKind::BF16)); + } } // NOTE: I commented that since I observed synchronisation issues with atomic add for bf16. diff --git a/crates/cubecl-ir/src/features.rs b/crates/cubecl-ir/src/features.rs index 9b487ce38f..d42de0bd77 100644 --- a/crates/cubecl-ir/src/features.rs +++ b/crates/cubecl-ir/src/features.rs @@ -55,6 +55,8 @@ pub struct Types { pub struct MatmulFeatures { /// Element types supported by a backend-native GEMM implementation. pub accelerated_gemm: BTreeSet, + /// Element types supported by a backend-native grouped GEMM implementation. + pub accelerated_grouped_gemm: BTreeSet, /// The cmma feature enables cooperative matrix-multiply and accumulate operations. pub cmma: BTreeSet, /// Cube MMA is like cmma but at the cube level, rather than the plane level. diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index bd5ff939da..fe68d170fc 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -8,9 +8,9 @@ use crate::{ runtime::Runtime, server::{ CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, GemmDescriptor, - Handle, IoError, KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, - MemoryLayoutStrategy, ProfileError, ReduceOperation, ServerCommunication, ServerError, - ServerUtilities, + GroupedGemmDescriptor, Handle, IoError, KernelArguments, MemoryLayout, + MemoryLayoutDescriptor, MemoryLayoutPolicy, MemoryLayoutStrategy, ProfileError, + ReduceOperation, ServerCommunication, ServerError, ServerUtilities, }, storage::{ComputeStorage, ManagedResource}, throughput::{ @@ -368,7 +368,7 @@ impl ComputeClient { /// Enqueue a backend-accelerated GEMM. /// - /// The operation is ordered on this client's CubeCL stream and returns + /// The operation is ordered on this client's `CubeCL` stream and returns /// immediately, like a kernel launch. Callers must first check the device's /// `features.matmul.accelerated_gemm` capability. pub fn gemm(&self, descriptor: GemmDescriptor) { @@ -377,6 +377,17 @@ impl ComputeClient { .submit(move |server| server.gemm(descriptor, stream_id)); } + /// Enqueue a backend-accelerated grouped GEMM. + /// + /// The operation is ordered on this client's `CubeCL` stream and returns + /// immediately, like a kernel launch. Callers must first check the device's + /// `features.matmul.accelerated_grouped_gemm` capability. + pub fn grouped_gemm(&self, descriptor: GroupedGemmDescriptor) { + let stream_id = self.stream_id(); + self.device + .submit(move |server| server.grouped_gemm(descriptor, stream_id)); + } + fn do_create_from_slices( &self, descriptors: Vec, diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index e0965d54bc..04195d9e2e 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -377,6 +377,7 @@ pub struct GemmMatrix { /// be row-major, allocated on the execution stream, and disjoint from both /// inputs. Backends may reject a nonempty zero-K problem rather than provide a /// separate output-initialization path. +#[allow(clippy::too_many_arguments)] #[derive(new, Clone, Debug)] pub struct GemmDescriptor { /// Left-hand matrix, logically shaped `[m, k]`. @@ -397,6 +398,17 @@ pub struct GemmDescriptor { pub elem: ElemType, } +/// Backend-optional accelerated grouped GEMM descriptor. +/// +/// Every entry computes one strided batch of `out = lhs @ rhs`. Entries may +/// have different matrix dimensions and batch counts, but must share an +/// element type. A backend may execute the entries concurrently in one launch. +#[derive(new, Clone, Debug)] +pub struct GroupedGemmDescriptor { + /// Independent GEMM groups. An empty collection is a no-op. + pub groups: Vec, +} + /// The compute server is responsible for handling resources and computations over resources. /// /// Everything in the server is mutable, therefore it should be solely accessed through the @@ -491,6 +503,17 @@ where panic!("Compute server advertised an accelerated GEMM without implementing it") } + /// Enqueue an accelerated grouped GEMM when the backend provides one. + /// + /// Callers must first check + /// [`Features::matmul.accelerated_grouped_gemm`](cubecl_ir::features::MatmulFeatures). + /// The enqueue returns immediately and records failures on the stream so + /// they surface at its next synchronization point. + fn grouped_gemm(&mut self, descriptor: GroupedGemmDescriptor, stream_id: StreamId) { + let _ = (descriptor, stream_id); + panic!("Compute server advertised an accelerated grouped GEMM without implementing it") + } + /// Flush all outstanding tasks in the server. fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>; From dcdb2b7d531278a5a83126763594d467f82f8ba1 Mon Sep 17 00:00:00 2001 From: pasha Date: Sat, 1 Aug 2026 10:45:22 +0300 Subject: [PATCH 6/8] feat(runtime): expose nonblocking dispatch queue flush --- crates/cubecl-common/src/device/handle/tests.rs | 15 +++++++++++++++ crates/cubecl-runtime/src/client.rs | 10 ++++++++++ 2 files changed, 25 insertions(+) diff --git a/crates/cubecl-common/src/device/handle/tests.rs b/crates/cubecl-common/src/device/handle/tests.rs index 480bf44fa9..db8da7a072 100644 --- a/crates/cubecl-common/src/device/handle/tests.rs +++ b/crates/cubecl-common/src/device/handle/tests.rs @@ -28,6 +28,21 @@ fn test_concurrent_increment() { let count = context.submit_blocking(move |state| state.counter).unwrap(); assert_eq!(count, thread_count); } + +#[test] +fn test_flush_queue_dispatches_partial_batch() { + let device = TestDevice::<3>::new(0); + let context = DeviceHandle::>::new(device.to_id()); + let (sender, receiver) = std::sync::mpsc::channel(); + + context.submit(move |_state| sender.send(()).unwrap()); + context.flush_queue(); + + receiver + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("partial batch was not dispatched"); +} + #[test] fn test_recursive_execution_different_state() { let device_id = DeviceId { diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index fe68d170fc..af3200f31d 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -388,6 +388,16 @@ impl ComputeClient { .submit(move |server| server.grouped_gemm(descriptor, stream_id)); } + /// Dispatches the current partial client-side command batch without waiting for it. + /// + /// This does not flush or synchronize the device stream. It only makes commands queued in + /// the runtime's batching channel available to the device service thread immediately. It is + /// useful when a host-side dependency has just drained the channel and the next command would + /// otherwise wait for a full command batch before being dispatched. + pub fn flush_dispatch_queue(&self) { + self.device.flush_queue(); + } + fn do_create_from_slices( &self, descriptors: Vec, From 45ed466abf28e404e1d3d1f3c25263941fb440e9 Mon Sep 17 00:00:00 2001 From: pasha Date: Sat, 1 Aug 2026 10:48:52 +0300 Subject: [PATCH 7/8] fix(cuda): follow environment module migration --- crates/cubecl-cuda/examples/cublas_gemm.rs | 2 +- crates/cubecl-cuda/examples/cublas_gemm_check.rs | 2 +- crates/cubecl-cuda/src/compute/cublas.rs | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/cubecl-cuda/examples/cublas_gemm.rs b/crates/cubecl-cuda/examples/cublas_gemm.rs index f0ba17d549..54ba64c2ed 100644 --- a/crates/cubecl-cuda/examples/cublas_gemm.rs +++ b/crates/cubecl-cuda/examples/cublas_gemm.rs @@ -1,10 +1,10 @@ -use cubecl_common::future; use cubecl_core::{ Runtime, ir::{ElemType, FloatKind}, server::{GemmDescriptor, GemmMatrix}, }; use cubecl_cuda::{CudaDevice, CudaRuntime}; +use cubecl_environment::future; use half::bf16; use std::time::Instant; diff --git a/crates/cubecl-cuda/examples/cublas_gemm_check.rs b/crates/cubecl-cuda/examples/cublas_gemm_check.rs index 6624a80fdd..7ae6e5287e 100644 --- a/crates/cubecl-cuda/examples/cublas_gemm_check.rs +++ b/crates/cubecl-cuda/examples/cublas_gemm_check.rs @@ -1,4 +1,3 @@ -use cubecl_common::{future, stream_id::StreamId}; use cubecl_core as cubecl; use cubecl_core::{ Runtime, @@ -7,6 +6,7 @@ use cubecl_core::{ server::{Binding, GemmDescriptor, GemmMatrix, GroupedGemmDescriptor, Handle}, }; use cubecl_cuda::{CudaDevice, CudaRuntime}; +use cubecl_environment::{future, stream::StreamId}; use half::bf16; #[cube(launch)] diff --git a/crates/cubecl-cuda/src/compute/cublas.rs b/crates/cubecl-cuda/src/compute/cublas.rs index 7581ff4d53..8ef2df633a 100644 --- a/crates/cubecl-cuda/src/compute/cublas.rs +++ b/crates/cubecl-cuda/src/compute/cublas.rs @@ -1,13 +1,13 @@ use std::collections::HashMap; use std::ffi::c_void; -use cubecl_common::backtrace::BackTrace; #[cfg(cuda_12050)] use cubecl_core::server::GroupedGemmDescriptor; use cubecl_core::{ ir::{ElemType, FloatKind}, server::{GemmDescriptor, GemmMatrix, ServerError}, }; +use cubecl_environment::backtrace::BackTrace; use cudarc::cublas::sys::cublasOperation_t; #[cfg(cuda_12050)] use cudarc::cublas::{result as blas, sys as blas_sys}; From c0efe74d3d4b820fbf992dbbe443ed7125c5205c Mon Sep 17 00:00:00 2001 From: pasha Date: Sat, 1 Aug 2026 11:57:56 +0300 Subject: [PATCH 8/8] Revert "feat(runtime): expose nonblocking dispatch queue flush" This reverts commit dcdb2b7d531278a5a83126763594d467f82f8ba1. --- crates/cubecl-common/src/device/handle/tests.rs | 15 --------------- crates/cubecl-runtime/src/client.rs | 10 ---------- 2 files changed, 25 deletions(-) diff --git a/crates/cubecl-common/src/device/handle/tests.rs b/crates/cubecl-common/src/device/handle/tests.rs index db8da7a072..480bf44fa9 100644 --- a/crates/cubecl-common/src/device/handle/tests.rs +++ b/crates/cubecl-common/src/device/handle/tests.rs @@ -28,21 +28,6 @@ fn test_concurrent_increment() { let count = context.submit_blocking(move |state| state.counter).unwrap(); assert_eq!(count, thread_count); } - -#[test] -fn test_flush_queue_dispatches_partial_batch() { - let device = TestDevice::<3>::new(0); - let context = DeviceHandle::>::new(device.to_id()); - let (sender, receiver) = std::sync::mpsc::channel(); - - context.submit(move |_state| sender.send(()).unwrap()); - context.flush_queue(); - - receiver - .recv_timeout(std::time::Duration::from_secs(1)) - .expect("partial batch was not dispatched"); -} - #[test] fn test_recursive_execution_different_state() { let device_id = DeviceId { diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index af3200f31d..fe68d170fc 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -388,16 +388,6 @@ impl ComputeClient { .submit(move |server| server.grouped_gemm(descriptor, stream_id)); } - /// Dispatches the current partial client-side command batch without waiting for it. - /// - /// This does not flush or synchronize the device stream. It only makes commands queued in - /// the runtime's batching channel available to the device service thread immediately. It is - /// useful when a host-side dependency has just drained the channel and the next command would - /// otherwise wait for a full command batch before being dispatched. - pub fn flush_dispatch_queue(&self) { - self.device.flush_queue(); - } - fn do_create_from_slices( &self, descriptors: Vec,