From ac0ae5b5aea124f98bf1b59be26226154fc3f7e2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 13 Apr 2026 11:45:21 +0200 Subject: [PATCH] =?UTF-8?q?refactor:=20CublasForward=20=E2=86=92=20CublasG?= =?UTF-8?q?emmSet=20with=20shared=20handle?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace per-instance cuBLAS/cublasLt handle ownership in CublasForward with Arc. The struct is renamed to CublasGemmSet and a type alias preserves backward compatibility for external callers. Removed from CublasGemmSet: handle, workspace, lt_handle, lt_workspace, bias kernels (all now in SharedCublasHandle). Kept: branch streams, branch workspaces, GEMM caches, events, network dimensions. External callers (gpu_dqn_trainer, gpu_experience_collector) will be updated in Tasks 4-5 to pass Arc to the constructor. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/batched_forward.rs | 252 ++++-------------- 1 file changed, 45 insertions(+), 207 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 410637ab6..16577d73c 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -59,11 +59,9 @@ use std::collections::HashMap; use std::sync::Arc; use std::mem::ManuallyDrop; -use cudarc::cublas::result as cublas_result; -use cudarc::cublas::sys as cublas_sys; use cudarc::cublaslt::sys as cublaslt_sys; use cudarc::cublaslt::result as cublaslt_result; -use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaEvent, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use crate::MLError; @@ -88,7 +86,7 @@ struct CachedGemmDesc { } /// SAFETY: cublasLt descriptors are opaque pointers owned exclusively by -/// `CublasForward`, which is single-threaded in practice (owned by `GpuDqnTrainer`). +/// `CublasGemmSet`, which is single-threaded in practice (owned by `GpuDqnTrainer`). unsafe impl Send for CachedGemmDesc {} unsafe impl Sync for CachedGemmDesc {} @@ -110,79 +108,20 @@ impl Drop for CachedGemmDesc { /// All forward GEMMs use TRANSA=T, TRANSB=N so transpose flags are implicit. type FwdGemmKey = (usize, usize, usize, usize); -// ── cuBLAS handle wrapper ───────────────────────────────────────────────────── - -/// Wrapper to make `cublasHandle_t` Send + Sync. -/// -/// cuBLAS handles are not thread-safe by default, but `CublasForward` is -/// owned exclusively by `GpuDqnTrainer` which is single-threaded in practice. -struct SendSyncCublasHandle(cublas_sys::cublasHandle_t); - -unsafe impl Send for SendSyncCublasHandle {} -unsafe impl Sync for SendSyncCublasHandle {} - -impl Drop for SendSyncCublasHandle { - fn drop(&mut self) { - unsafe { - let _ = cublas_result::destroy_handle(self.0); - } - } -} - -/// Wrapper to make `cublasLtHandle_t` Send + Sync. -/// -/// Same ownership model as `SendSyncCublasHandle` — exclusively owned by -/// `CublasForward` which is single-threaded in practice. -struct SendSyncCublasLtHandle(cublaslt_sys::cublasLtHandle_t); - -unsafe impl Send for SendSyncCublasLtHandle {} -unsafe impl Sync for SendSyncCublasLtHandle {} - -impl Drop for SendSyncCublasLtHandle { - fn drop(&mut self) { - unsafe { - let _ = cublaslt_result::destroy_handle(self.0); - } - } -} - // ── Batched forward context ─────────────────────────────────────────────────── -/// cuBLAS-based batched forward pass context for the branching dueling DQN. +/// cuBLAS-based batched forward GEMM set for the branching dueling DQN. /// -/// Holds the cuBLAS handle and references to the trainer's pre-allocated -/// activation buffers. GEMM output lands directly in those buffers. +/// Holds an `Arc` (cuBLAS + cublasLt handles, workspaces, +/// bias kernels) and the forward-specific resources: branch streams, per-branch +/// workspaces, cached GEMM descriptors, and network dimensions. /// /// The handle is set to the trainer's stream so all GEMMs execute in order. #[allow(missing_debug_implementations)] // handle is not Debug -pub struct CublasForward { - handle: SendSyncCublasHandle, - /// cuBLAS workspace buffer — must outlive the handle for CUDA Graph replay. - _workspace_buf: CudaSlice, - /// Raw device pointer + size for workspace restoration after cublasSetStream. - /// cublasSetStream resets workspace to the default pool, so we must re-set it. - workspace_ptr: u64, - workspace_size: usize, - - // ── cublasLt handle + workspace (per-call stream, no handle state conflicts) ── - - /// cublasLt handle for forward GEMMs — passes workspace+stream per-call, - /// eliminating the `cublasSetStream` workspace state conflict that hangs - /// CUDA Graph mega-capture on H100. - lt_handle: SendSyncCublasLtHandle, - /// cublasLt workspace buffer — must outlive handle for CUDA Graph replay. - _lt_workspace_buf: CudaSlice, - /// Raw device pointer for cublasLt workspace. - lt_workspace_ptr: u64, - /// Size of cublasLt workspace in bytes. - lt_workspace_size: usize, - - // ── Compiled bias+activation kernels (loaded from precompiled cubin) ── - - /// `add_bias_relu_f32_kernel(output, bias, out_dim, total)` — f32 output + f32 bias + ReLU. - add_bias_relu_f32_kernel: CudaFunction, - /// `add_bias_f32_kernel(output, bias, out_dim, total)` — f32 output + f32 bias (no activation). - add_bias_f32_kernel: CudaFunction, +pub struct CublasGemmSet { + /// Shared cuBLAS + cublasLt handle (owned via Arc, shared with backward pass + /// and experience collector for deterministic single-handle training). + handle: Arc, // ── Network dimensions (baked at construction) ── batch_size: usize, @@ -232,13 +171,20 @@ pub struct CublasForward { gemm_cache_relu_bias: HashMap, } -impl CublasForward { - /// Create a new cuBLAS forward context. +/// Backward-compatible type alias — existing callers (`GpuDqnTrainer`, +/// `GpuExperienceCollector`, etc.) keep using `CublasForward` until Task 4 +/// wires the shared handle through all constructors. +pub type CublasForward = CublasGemmSet; + +impl CublasGemmSet { + /// Create a new cuBLAS forward GEMM set backed by a shared handle. /// - /// Compiles the `add_bias_relu` and `add_bias` kernels from inline NVRTC source. - /// Binds the cuBLAS handle to the provided stream. + /// The shared handle owns the cuBLAS + cublasLt handles, workspaces, and + /// bias kernels. This constructor only creates forward-specific resources: + /// branch streams, per-branch workspaces, cached GEMM descriptors, events, + /// and stores the network dimensions. pub fn new( - stream: &Arc, + shared: Arc, batch_size: usize, state_dim: usize, shared_h1: usize, @@ -252,70 +198,7 @@ impl CublasForward { branch_3_size: usize, s1_input_dim: usize, ) -> Result { - // ── Create cuBLAS handle ──────────────────────────────────────── - let raw_handle = cublas_result::create_handle() - .map_err(|e| MLError::ModelError(format!("cublasCreate: {e:?}")))?; - - // Bind handle to stream — ALL GEMMs will execute on this stream. - // SAFETY: cudarc::driver::sys and cudarc::cublas::sys both typedef - // CUstream to *mut CUstream_st, which is the same underlying type. - // The cast is safe as long as the pointer value is preserved. - unsafe { - let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st; - cublas_result::set_stream(raw_handle, cu_stream) - .map_err(|e| MLError::ModelError(format!("cublasSetStream: {e:?}")))?; - } - - // Deterministic FP32 GEMM — ensures reproducible training across runs. - // TF32 tensor ops (CUBLAS_TF32_TENSOR_OP_MATH) use non-deterministic - // reduction order, causing ~1e-4 per-GEMM variance that compounds over - // thousands of forward+backward passes. DEFAULT_MATH forces IEEE FP32 - // accumulation which is deterministic. Cost: ~1.5× slower GEMM (but GEMM - // is <20% of total step time — most time is in C51 loss/grad kernels). - unsafe { - cublas_sys::cublasSetMathMode( - raw_handle, - cublas_sys::cublasMath_t::CUBLAS_DEFAULT_MATH, - ); - } - - // Allocate explicit cuBLAS workspace for CUDA Graph compatibility. - // CRITICAL: cublasSetStream() resets workspace to the default pool, - // which uses stream-ordered allocation (incompatible with CUDA Graph). - // We must re-set this workspace after every cublasSetStream call. - // 32 MB: TF32 HMMA algorithms on H100 need more workspace than f32. - let workspace_size: usize = 32 * 1024 * 1024; // 32 MB — H100 CUDA 13.0 requires this for TF32 - let workspace_buf = stream.alloc_zeros::(workspace_size) - .map_err(|e| MLError::ModelError(format!("cuBLAS workspace alloc: {e}")))?; - let workspace_ptr = { - let (dev_ptr, _guard) = workspace_buf.device_ptr(stream); - dev_ptr - }; - unsafe { - let ws_ptr = workspace_ptr as *mut std::ffi::c_void; - let status = cublas_sys::cublasSetWorkspace_v2(raw_handle, ws_ptr, workspace_size); - if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { - return Err(MLError::ModelError(format!("cublasSetWorkspace_v2 failed: {status:?}"))); - } - } - - // ── Create cublasLt handle + workspace ──────────────────────────── - // cublasLtMatmul passes workspace+stream per-call (no handle state), - // eliminating the cublasSetStream workspace conflict that hangs CUDA - // Graph mega-capture on H100. - let lt_raw_handle = cublaslt_result::create_handle() - .map_err(|e| MLError::ModelError(format!("cublasLtCreate: {e:?}")))?; - - // 32 MB workspace for cublasLtMatmul. - // cudarc now uses cuMemAlloc_v2 (sync, globally accessible) — see vendor/cudarc comment. - let lt_ws_size: usize = 32 * 1024 * 1024; - let lt_workspace_buf = stream.alloc_zeros::(lt_ws_size) - .map_err(|e| MLError::ModelError(format!("cublasLt workspace alloc: {e}")))?; - let lt_ws_ptr = lt_workspace_buf.raw_ptr(); - - // ── Load f32 bias kernels from precompiled cubin ── - let (add_bias_relu_f32_kernel, add_bias_f32_kernel) = - compile_bias_kernels(stream)?; + let stream = &shared.stream; // ── Fork 4 branch streams for parallel advantage head dispatch ── let branch_streams = [ @@ -364,6 +247,9 @@ impl CublasForward { let s1_ldb = if s1_input_dim < state_dim { s1_input_dim } else { state_dim_padded }; let branch_sizes = [branch_0_size, branch_1_size, branch_2_size, branch_3_size]; + let lt_raw_handle = shared.lt_handle.0; + let lt_ws_size = shared.lt_workspace_size; + let mut unique_shapes: Vec = Vec::new(); // h_s1: M=shared_h1, N=batch, K=s1_input_dim, ldb=s1_ldb unique_shapes.push((shared_h1, batch_size, s1_input_dim, s1_ldb)); @@ -408,16 +294,7 @@ impl CublasForward { } Ok(Self { - handle: SendSyncCublasHandle(raw_handle), - _workspace_buf: workspace_buf, - workspace_ptr, - workspace_size, - lt_handle: SendSyncCublasLtHandle(lt_raw_handle), - _lt_workspace_buf: lt_workspace_buf, - lt_workspace_ptr: lt_ws_ptr, - lt_workspace_size: lt_ws_size, - add_bias_relu_f32_kernel, - add_bias_f32_kernel, + handle: shared, // Dimensions batch_size, state_dim, @@ -458,29 +335,14 @@ impl CublasForward { w_ptr: u64, a_ptr: u64, c_ptr: u64, m: usize, n: usize, k: usize, ) -> Result<(), MLError> { - self.lt_matmul(stream, w_ptr, a_ptr, c_ptr, m, n, k, k, self.lt_workspace_ptr, self.lt_workspace_size, "DIAG_raw") + self.lt_matmul(stream, w_ptr, a_ptr, c_ptr, m, n, k, k, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, "DIAG_raw") } /// CRITICAL: cublasSetStream resets workspace to the default pool. - /// We must re-set our user-owned workspace immediately after. + /// Delegates to `SharedCublasHandle::set_stream` which atomically + /// rebinds the handle and restores the user-owned workspace. pub fn set_stream(&self, stream: &CudaStream) -> Result<(), MLError> { - unsafe { - let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st; - cublas_result::set_stream(self.handle.0, cu_stream) - .map_err(|e| MLError::ModelError(format!("cublasSetStream: {e:?}")))?; - self.restore_workspace(); - } - Ok(()) - } - - /// Restore user-owned workspace after cublasSetStream reset. - /// Must be called after every cublasSetStream to prevent CUDA Graph hang. - unsafe fn restore_workspace(&self) { - cublas_sys::cublasSetWorkspace_v2( - self.handle.0, - self.workspace_ptr as *mut std::ffi::c_void, - self.workspace_size, - ); + self.handle.set_stream(stream) } // ══════════════════════════════════════════════════════════════════════════ @@ -518,8 +380,8 @@ impl CublasForward { // First layer: ldb = state_dim_padded (CUTLASS K-tile alignment). // Try fused GEMM+bias+ReLU epilogue, fall back to separate kernels. - let ws = self.lt_workspace_ptr; - let wss = self.lt_workspace_size; + let ws = self.handle.lt_workspace_ptr; + let wss = self.handle.lt_workspace_size; if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[0], states_ptr, h_s1_ptr, w_ptrs[1], self.shared_h1, b, self.s1_input_dim, self.s1_ldb, ws, wss, "h_s1").is_err() { @@ -587,7 +449,7 @@ impl CublasForward { // Per-branch workspace: eliminates contention between parallel branch streams. // Try fused GEMM+bias+ReLU epilogue, fall back to separate kernels. let bws = self.branch_workspace_ptrs[d]; - let bwss = self.lt_workspace_size; + let bwss = self.handle.lt_workspace_size; if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b, self.shared_h2, self.shared_h2, bws, bwss, "h_bd").is_err() { @@ -945,7 +807,7 @@ impl CublasForward { k: usize, // in_dim _label: &str, ) -> Result<(), MLError> { - self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.lt_workspace_ptr, self.lt_workspace_size, _label) + self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, _label) } /// F32 GEMM on a branch stream with per-branch workspace. @@ -963,7 +825,7 @@ impl CublasForward { branch_idx: usize, _label: &str, ) -> Result<(), MLError> { - self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.branch_workspace_ptrs[branch_idx], self.lt_workspace_size, _label) + self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.branch_workspace_ptrs[branch_idx], self.handle.lt_workspace_size, _label) } /// F32 GEMM with TF32 + explicit `ldb` (for padded states buffer). @@ -980,7 +842,7 @@ impl CublasForward { ldb: usize, _label: &str, ) -> Result<(), MLError> { - self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, self.lt_workspace_ptr, self.lt_workspace_size, _label) + self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, _label) } /// Inner cublasLtMatmul dispatch: uses pre-cached descriptors when available, @@ -1040,7 +902,7 @@ impl CublasForward { let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; tracing::debug!( - handle = ?self.lt_handle.0, + handle = ?self.handle.lt_handle.0, matmul_desc = ?cached.matmul_desc, w_ptr = w_f32_ptr, a_ptr = a_f32_ptr, @@ -1057,7 +919,7 @@ impl CublasForward { ); let matmul_status = cublaslt_sys::cublasLtMatmul( - self.lt_handle.0, + self.handle.lt_handle.0, cached.matmul_desc, &alpha as *const f32 as *const std::ffi::c_void, w_f32_ptr as *const std::ffi::c_void, @@ -1126,7 +988,7 @@ impl CublasForward { unsafe { let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; let status = cublaslt_sys::cublasLtMatmul( - self.lt_handle.0, + self.handle.lt_handle.0, cached.matmul_desc, &alpha as *const f32 as *const std::ffi::c_void, w_ptr as *const std::ffi::c_void, @@ -1220,7 +1082,7 @@ impl CublasForward { ).map_err(|e| MLError::ModelError(format!("set pref ws {_label}: {e:?}")))?; let heuristic = cublaslt_result::get_matmul_algo_heuristic( - self.lt_handle.0, + self.handle.lt_handle.0, matmul_desc, a_layout, b_layout, @@ -1237,7 +1099,7 @@ impl CublasForward { let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; tracing::debug!( - handle = ?self.lt_handle.0, + handle = ?self.handle.lt_handle.0, matmul_desc = ?matmul_desc, w_ptr = w_f32_ptr, a_ptr = a_f32_ptr, @@ -1254,7 +1116,7 @@ impl CublasForward { ); let matmul_status = cublaslt_sys::cublasLtMatmul( - self.lt_handle.0, + self.handle.lt_handle.0, matmul_desc, &alpha as *const f32 as *const std::ffi::c_void, w_f32_ptr as *const std::ffi::c_void, @@ -1317,7 +1179,7 @@ impl CublasForward { let blocks = ((batch * out_dim + 255) / 256) as u32; unsafe { stream - .launch_builder(&self.add_bias_relu_f32_kernel) + .launch_builder(&self.handle.add_bias_relu_f32_kernel) .arg(&out_ptr) .arg(&bias_ptr) .arg(&out_dim_i32) @@ -1347,7 +1209,7 @@ impl CublasForward { unsafe { stream - .launch_builder(&self.add_bias_f32_kernel) + .launch_builder(&self.handle.add_bias_f32_kernel) .arg(&out_ptr) .arg(&bias_ptr) .arg(&out_dim_i32) @@ -1381,30 +1243,6 @@ fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { ptr } -// ── Kernel compilation ──────────────────────────────────────────────────────── - -/// Precompiled bias kernels cubin (build.rs — ZERO runtime nvcc). -static BIAS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bias_kernels.cubin")); - -/// Load f32 bias kernels from precompiled cubin. -fn compile_bias_kernels( - stream: &Arc, -) -> Result<(CudaFunction, CudaFunction), MLError> { - let context = stream.context(); - let module = context - .load_cubin(BIAS_CUBIN.to_vec()) - .map_err(|e| MLError::ModelError(format!("bias cubin load: {e}")))?; - - let add_bias_relu_f32 = module - .load_function("add_bias_relu_f32_kernel") - .map_err(|e| MLError::ModelError(format!("add_bias_relu_f32_kernel load: {e}")))?; - let add_bias_f32 = module - .load_function("add_bias_f32_kernel") - .map_err(|e| MLError::ModelError(format!("add_bias_f32_kernel load: {e}")))?; - - Ok((add_bias_relu_f32, add_bias_f32)) -} - // ── Cached GEMM descriptor factory ────────────────────────────────────────── /// Create a `CachedGemmDesc` for a forward GEMM shape (TRANSA=T, TRANSB=N, f32, FAST_TF32).