diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index a703d2966..0dcc79747 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -237,13 +237,26 @@ impl CublasForward { let lt_raw_handle = cublaslt_result::create_handle() .map_err(|e| MLError::ModelError(format!("cublasLtCreate: {e:?}")))?; - // 32 MB unconditionally — TF32 HMMA algorithms need large workspace - // even on Ampere (SM86). NVIDIA recommends 32 MB for Hopper but - // the TF32 heuristic returns NOT_SUPPORTED with 4 MB at (128,512,64). + // 32 MB workspace for cublasLtMatmul. + // CRITICAL: use cudaMalloc (not cuMemAllocAsync) so the workspace is + // accessible from ALL streams including forked branch streams. Stream-ordered + // allocations are only valid on the allocating stream, causing + // CUBLAS_STATUS_NOT_SUPPORTED when cublasLtMatmul runs on branch streams. let lt_ws_size: usize = 32 * 1024 * 1024; - let lt_ws_buf = stream.alloc_zeros::(lt_ws_size) - .map_err(|e| MLError::ModelError(format!("cublasLt workspace alloc: {e}")))?; - let lt_ws_ptr = lt_ws_buf.raw_ptr(); + let lt_ws_ptr: u64; + unsafe { + let mut dev_ptr: cudarc::driver::sys::CUdeviceptr = 0; + let result = cudarc::driver::sys::cuMemAlloc_v2(&mut dev_ptr, lt_ws_size); + if result != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + return Err(MLError::ModelError(format!("cuMemAlloc workspace: {result:?}"))); + } + // Zero the workspace + cudarc::driver::sys::cuMemsetD8_v2(dev_ptr, 0, lt_ws_size); + lt_ws_ptr = dev_ptr; + } + // Dummy CudaSlice to satisfy struct field (workspace is managed by cuMemAlloc) + let lt_workspace_buf = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("dummy ws alloc: {e}")))?; // ── Load f32 bias kernels from precompiled cubin ── let (add_bias_relu_f32_kernel, add_bias_f32_kernel) = @@ -274,7 +287,7 @@ impl CublasForward { workspace_ptr, workspace_size, lt_handle: SendSyncCublasLtHandle(lt_raw_handle), - _lt_workspace_buf: lt_ws_buf, + _lt_workspace_buf: lt_workspace_buf, lt_workspace_ptr: lt_ws_ptr, lt_workspace_size: lt_ws_size, add_bias_relu_f32_kernel,