diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 135c9ef34..42d859ee3 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -77,6 +77,7 @@ //! here are wired via `GpuDqnTrainer::launch_cublas_backward()` with a //! TODO gate that enables them once the C51 loss kernel is ready. +use std::collections::HashMap; use std::mem::ManuallyDrop; use std::sync::Arc; @@ -89,6 +90,42 @@ use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfi use crate::MLError; use super::gpu_dqn_trainer::GpuDqnTrainConfig; +// ── Cached backward GEMM descriptor ───────────────────────────────────────── + +/// Pre-created cublasLt descriptors + heuristic algo for a fixed backward GEMM shape. +/// +/// Backward GEMMs have variable TRANSA/TRANSB (N/T for dW, N/N for dX), so we +/// key the cache by the full shape tuple including transpose flags and leading dims. +struct CachedBwdGemmDesc { + matmul_desc: cublaslt_sys::cublasLtMatmulDesc_t, + a_layout: cublaslt_sys::cublasLtMatrixLayout_t, + b_layout: cublaslt_sys::cublasLtMatrixLayout_t, + c_layout: cublaslt_sys::cublasLtMatrixLayout_t, + d_layout: cublaslt_sys::cublasLtMatrixLayout_t, + algo: cublaslt_sys::cublasLtMatmulAlgo_t, +} + +/// SAFETY: cublasLt descriptors are opaque pointers owned exclusively by +/// `CublasBackward`, which is single-threaded in practice. +unsafe impl Send for CachedBwdGemmDesc {} +unsafe impl Sync for CachedBwdGemmDesc {} + +impl Drop for CachedBwdGemmDesc { + fn drop(&mut self) { + unsafe { + let _ = cublaslt_result::destroy_matrix_layout(self.d_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.c_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.b_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.a_layout); + let _ = cublaslt_result::destroy_matmul_desc(self.matmul_desc); + } + } +} + +/// Key for the backward GEMM cache: (transa, transb, m, n, k, lda, ldb, ldc). +/// Beta is a runtime parameter passed to cublasLtMatmul, not part of descriptors. +type BwdGemmKey = (i32, i32, i32, i32, i32, i32, i32, i32); + // ── cuBLAS handle wrapper ───────────────────────────────────────────────────── /// Wrapper to make `cublasHandle_t` Send + Sync. @@ -184,6 +221,10 @@ pub struct CublasBackward { branch_1_size: usize, branch_2_size: usize, branch_3_size: usize, + + // ── Cached GEMM descriptors (created once at init, reused per-call) ── + /// Map from (transa, transb, m, n, k, lda, ldb, ldc) → pre-created descriptors + algo. + gemm_cache: HashMap, } impl CublasBackward { @@ -244,6 +285,65 @@ impl CublasBackward { // ── Compile helper kernels ────────────────────────────────── let (relu_mask_kernel, bias_grad_kernel) = compile_backward_kernels(stream)?; + // ── Pre-create backward GEMM descriptors for all unique shapes ── + let batch = config.batch_size; + let sh1 = config.shared_h1; + let sh2 = config.shared_h2; + let vh = config.value_h; + let ah = config.adv_h; + let na = config.num_atoms; + let sd = config.state_dim; + let sd_pad = (config.state_dim + 127) & !127; + let branch_sizes = [config.branch_0_size, config.branch_1_size, config.branch_2_size, config.branch_3_size]; + + let op_n = cublas_sys::cublasOperation_t::CUBLAS_OP_N as i32; + let op_t = cublas_sys::cublasOperation_t::CUBLAS_OP_T as i32; + + // Collect all unique (transa, transb, m, n, k, lda, ldb, ldc) shapes. + let mut unique_shapes: Vec = Vec::new(); + + // For each FC layer with (out_dim, in_dim), backward produces: + // dW: (N, T, in_dim, out_dim, batch, in_dim, out_dim, in_dim) + // dX: (N, N, in_dim, batch, out_dim, in_dim, out_dim, in_dim) + // For lda variants (shared layer 1 with padded input): + // dW: (N, T, in_dim, out_dim, batch, x_lda, out_dim, in_dim) + // dX: (N, N, in_dim, batch, out_dim, in_dim, out_dim, x_lda) + + let add_fc_shapes = |shapes: &mut Vec, out_dim: usize, in_dim: usize| { + // dW: sgemm(N, T, in_dim, out_dim, batch, in_dim, out_dim, in_dim) + shapes.push((op_n, op_t, in_dim as i32, out_dim as i32, batch as i32, in_dim as i32, out_dim as i32, in_dim as i32)); + // dX: sgemm(N, N, in_dim, batch, out_dim, in_dim, out_dim, in_dim) + shapes.push((op_n, op_n, in_dim as i32, batch as i32, out_dim as i32, in_dim as i32, out_dim as i32, in_dim as i32)); + }; + + // Branch output layers: (branch_k*NA, adv_h) + for &bs in &branch_sizes { + add_fc_shapes(&mut unique_shapes, bs * na, ah); + } + // Branch FC layers: (adv_h, shared_h2) + add_fc_shapes(&mut unique_shapes, ah, sh2); + // Value output: (NA, value_h) + add_fc_shapes(&mut unique_shapes, na, vh); + // Value FC: (value_h, shared_h2) + add_fc_shapes(&mut unique_shapes, vh, sh2); + // Shared layer 2: (shared_h2, shared_h1) + add_fc_shapes(&mut unique_shapes, sh2, sh1); + // Shared layer 1 with padded input: (shared_h1, state_dim) — uses x_lda=state_dim_padded + // dW: (N, T, sd, sh1, batch, sd_pad, sh1, sd) + unique_shapes.push((op_n, op_t, sd as i32, sh1 as i32, batch as i32, sd_pad as i32, sh1 as i32, sd as i32)); + // dX: (N, N, sd, batch, sh1, sd, sh1, sd_pad) — only when s1_dx_output != 0 + unique_shapes.push((op_n, op_n, sd as i32, batch as i32, sh1 as i32, sd as i32, sh1 as i32, sd_pad as i32)); + + unique_shapes.sort(); + unique_shapes.dedup(); + + let mut gemm_cache = HashMap::new(); + for &key in &unique_shapes { + let (transa_i32, transb_i32, m, n, k, lda, ldb, ldc) = key; + let desc = create_cached_bwd_gemm_desc(lt_raw_handle, transa_i32, transb_i32, m, n, k, lda, ldb, ldc, lt_ws_size)?; + gemm_cache.insert(key, desc); + } + Ok(Self { handle: SendSyncCublasHandle(raw_handle), _workspace_buf: bw_workspace_buf, @@ -255,7 +355,7 @@ impl CublasBackward { bias_grad_kernel, batch_size: config.batch_size, state_dim: config.state_dim, - state_dim_padded: (config.state_dim + 127) & !127, + state_dim_padded: sd_pad, shared_h1: config.shared_h1, shared_h2: config.shared_h2, value_h: config.value_h, @@ -265,6 +365,7 @@ impl CublasBackward { branch_1_size: config.branch_1_size, branch_2_size: config.branch_2_size, branch_3_size: config.branch_3_size, + gemm_cache, }) } @@ -290,8 +391,8 @@ impl CublasBackward { /// eliminating the `cublasSetStream` workspace conflict that hangs /// CUDA Graph mega-capture on H100. /// - /// Uses CUBLAS_COMPUTE_32F (not FAST_TF32) for portability — TF32 is - /// controlled by the math mode set on the cuBLAS handle at construction. + /// Uses cached descriptors when the shape is pre-cached at init time. + /// Falls back to inline descriptor creation for uncached shapes. fn sgemm_f32( &self, stream: &CudaStream, @@ -304,6 +405,83 @@ impl CublasBackward { beta: f32, c: u64, ldc: i32, label: &str, + ) -> Result<(), MLError> { + let key: BwdGemmKey = (transa as i32, transb as i32, m, n, k, lda, ldb, ldc); + if let Some(cached) = self.gemm_cache.get(&key) { + // ── Fast path: use pre-created descriptors + algo ── + self.sgemm_f32_cached(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, label) + } else { + // ── Slow path: inline descriptor creation ── + self.sgemm_f32_uncached(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, label) + } + } + + /// Fast path: execute backward GEMM with pre-cached descriptors. + #[allow(clippy::too_many_arguments)] + fn sgemm_f32_cached( + &self, + stream: &CudaStream, + cached: &CachedBwdGemmDesc, + alpha: f32, + a: u64, + b_ptr: u64, + beta: f32, + c: u64, + m: i32, n: i32, k: i32, + lda: i32, ldb: i32, ldc: i32, + label: &str, + ) -> Result<(), MLError> { + unsafe { + let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; + + let matmul_status = cublaslt_sys::cublasLtMatmul( + self.lt_handle.0, + cached.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + a as *const std::ffi::c_void, + cached.a_layout, + b_ptr as *const std::ffi::c_void, + cached.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + c as *const std::ffi::c_void, + cached.c_layout, + c as *mut std::ffi::c_void, + cached.d_layout, + &cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, + self.lt_workspace_ptr as *mut std::ffi::c_void, + self.lt_workspace_size, + cu_stream, + ); + + if matmul_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + let e = cublaslt_result::CublasError(matmul_status); + tracing::error!( + m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc, + ws_provided=self.lt_workspace_size, + ws_ptr=self.lt_workspace_ptr, + ?e, + "cublasLtMatmul CACHED FAILED for bw {label}" + ); + return Err(MLError::ModelError(format!("cublasLtMatmul cached bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb}): {e:?}"))); + } + } + Ok(()) + } + + /// Slow path: inline descriptor creation for uncached backward GEMM shapes. + #[allow(clippy::too_many_arguments)] + fn sgemm_f32_uncached( + &self, + stream: &CudaStream, + transa: cublas_sys::cublasOperation_t, + transb: cublas_sys::cublasOperation_t, + m: i32, n: i32, k: i32, + alpha: f32, + a: u64, lda: i32, + b_ptr: u64, ldb: i32, + beta: f32, + c: u64, ldc: i32, + label: &str, ) -> Result<(), MLError> { let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F; // FAST_TF32: explicit TF32 tensor core path. Required on cuBLAS 13.0+ (SM90) @@ -315,7 +493,6 @@ impl CublasBackward { let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type) .map_err(|e| MLError::ModelError(format!("cublasLtMatmulDescCreate bw {label}: {e:?}")))?; - // Set TRANSA from the cublas_sys enum value (N=0, T=1) let transa_i32: i32 = transa as i32; cublaslt_result::set_matmul_desc_attribute( matmul_desc, @@ -324,7 +501,6 @@ impl CublasBackward { std::mem::size_of::(), ).map_err(|e| MLError::ModelError(format!("set TRANSA bw {label}: {e:?}")))?; - // Set TRANSB from the cublas_sys enum value (N=0, T=1) let transb_i32: i32 = transb as i32; cublaslt_result::set_matmul_desc_attribute( matmul_desc, @@ -334,29 +510,21 @@ impl CublasBackward { ).map_err(|e| MLError::ModelError(format!("set TRANSB bw {label}: {e:?}")))?; // ── Matrix layouts ── - // A layout depends on transa: - // N: A is col-major [m, k] with leading dim = lda - // T: A is col-major [k, m] with leading dim = lda (transposed to [m, k]) let a_layout = if transa == cublas_sys::cublasOperation_t::CUBLAS_OP_N { cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64) } else { cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, lda as i64) }.map_err(|e| MLError::ModelError(format!("A layout bw {label}: {e:?}")))?; - // B layout depends on transb: - // N: B is col-major [k, n] with leading dim = ldb - // T: B is col-major [n, k] with leading dim = ldb (transposed to [k, n]) let b_layout = if transb == cublas_sys::cublasOperation_t::CUBLAS_OP_N { cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64) } else { cublaslt_result::create_matrix_layout(f32_type, n as u64, k as u64, ldb as i64) }.map_err(|e| MLError::ModelError(format!("B layout bw {label}: {e:?}")))?; - // C = [m, n] with ldc (input for beta accumulation) let c_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) .map_err(|e| MLError::ModelError(format!("C layout bw {label}: {e:?}")))?; - // D layout must be a SEPARATE descriptor (same dimensions as C) let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) .map_err(|e| MLError::ModelError(format!("D layout bw {label}: {e:?}")))?; @@ -385,8 +553,6 @@ impl CublasBackward { let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo heuristic bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb}): {e:?}")))?; // ── Execute matmul ── - // Cast CUstream (driver::sys) -> cudaStream_t (cublaslt::sys). - // Both are *mut CUstream_st — same underlying type, different module paths. let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; let matmul_result = cublaslt_result::matmul( @@ -414,7 +580,7 @@ impl CublasBackward { ws_needed=heuristic.workspaceSize, ws_ptr=self.lt_workspace_ptr, ?e, - "cublasLtMatmul FAILED for bw {label}" + "cublasLtMatmul UNCACHED FAILED for bw {label}" ); } matmul_result.map_err(|e| MLError::ModelError(format!("cublasLtMatmul bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb},ws={},needed={}): {e:?}", self.lt_workspace_size, heuristic.workspaceSize)))?; @@ -1006,6 +1172,137 @@ impl CublasBackward { } } +// ── Cached backward GEMM descriptor factory ───────────────────────────────── + +/// Create a `CachedBwdGemmDesc` for a backward GEMM shape with given transpose flags. +/// +/// Runs the algorithm heuristic ONCE at init time with `requestedAlgoCount=3` +/// for better algorithm selection. +#[allow(clippy::too_many_arguments)] +fn create_cached_bwd_gemm_desc( + lt_handle: cublaslt_sys::cublasLtHandle_t, + transa_i32: i32, + transb_i32: i32, + m: i32, n: i32, k: i32, + lda: i32, ldb: i32, ldc: i32, + ws_size: usize, +) -> Result { + let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F; + let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32; + + unsafe { + let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type) + .map_err(|e| MLError::ModelError(format!("cached bw MatmulDescCreate (m={m},n={n},k={k}): {e:?}")))?; + + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA, + &transa_i32 as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("cached bw set TRANSA (m={m},n={n},k={k}): {e:?}")))?; + + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB, + &transb_i32 as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("cached bw set TRANSB (m={m},n={n},k={k}): {e:?}")))?; + + // Matrix layouts — match the logic in sgemm_f32_uncached. + let transa_is_n = transa_i32 == 0; // CUBLAS_OP_N + let transb_is_n = transb_i32 == 0; + + let a_layout = if transa_is_n { + cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64) + } else { + cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, lda as i64) + }.map_err(|e| MLError::ModelError(format!("cached bw A layout (m={m},n={n},k={k}): {e:?}")))?; + + let b_layout = if transb_is_n { + cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64) + } else { + cublaslt_result::create_matrix_layout(f32_type, n as u64, k as u64, ldb as i64) + }.map_err(|e| MLError::ModelError(format!("cached bw B layout (m={m},n={n},k={k}): {e:?}")))?; + + let c_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) + .map_err(|e| MLError::ModelError(format!("cached bw C layout (m={m},n={n},k={k}): {e:?}")))?; + + let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) + .map_err(|e| MLError::ModelError(format!("cached bw D layout (m={m},n={n},k={k}): {e:?}")))?; + + // Algorithm heuristic (requestedAlgoCount=3) + let matmul_pref = cublaslt_result::create_matmul_pref() + .map_err(|e| MLError::ModelError(format!("cached bw matmul pref (m={m},n={n},k={k}): {e:?}")))?; + cublaslt_result::set_matmul_pref_attribute( + matmul_pref, + cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_size as *const usize as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("cached bw set pref ws (m={m},n={n},k={k}): {e:?}")))?; + + let mut heuristics: [std::mem::MaybeUninit; 3] = + std::mem::MaybeUninit::uninit().assume_init(); + let mut algo_count: i32 = 0; + + let status = cublaslt_sys::cublasLtMatmulAlgoGetHeuristic( + lt_handle, + matmul_desc, + a_layout, + b_layout, + c_layout, + d_layout, + matmul_pref, + 3, + heuristics[0].as_mut_ptr(), + &mut algo_count, + ); + if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS || algo_count == 0 { + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + let _ = cublaslt_result::destroy_matrix_layout(d_layout); + let _ = cublaslt_result::destroy_matrix_layout(c_layout); + let _ = cublaslt_result::destroy_matrix_layout(b_layout); + let _ = cublaslt_result::destroy_matrix_layout(a_layout); + let _ = cublaslt_result::destroy_matmul_desc(matmul_desc); + return Err(MLError::ModelError(format!( + "cached bw algo heuristic (transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): status={status:?}, count={algo_count}" + ))); + } + + let best = heuristics[0].assume_init(); + if best.state != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + let _ = cublaslt_result::destroy_matrix_layout(d_layout); + let _ = cublaslt_result::destroy_matrix_layout(c_layout); + let _ = cublaslt_result::destroy_matrix_layout(b_layout); + let _ = cublaslt_result::destroy_matrix_layout(a_layout); + let _ = cublaslt_result::destroy_matmul_desc(matmul_desc); + return Err(MLError::ModelError(format!( + "cached bw algo heuristic state invalid (transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): {:?}", + best.state + ))); + } + + tracing::info!( + transa = transa_i32, transb = transb_i32, + m = m, n = n, k = k, lda = lda, ldb = ldb, ldc = ldc, + algo_count = algo_count, + ws_needed = best.workspaceSize, + "cached bwd GEMM desc created" + ); + + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + + Ok(CachedBwdGemmDesc { + matmul_desc, + a_layout, + b_layout, + c_layout, + d_layout, + algo: best.algo, + }) + } +} + // ── Kernel compilation ──────────────────────────────────────────────────────── /// Compile `relu_mask_kernel` and `bias_grad_reduce_kernel` from inline NVRTC. diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 987cd7b03..ed266511f 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -55,6 +55,7 @@ //! state management is needed during multi-stream branch dispatch. Handle //! creation is NOT capturable and must happen before capture begins. +use std::collections::HashMap; use std::sync::Arc; use std::mem::ManuallyDrop; @@ -66,6 +67,49 @@ use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, use crate::MLError; +// ── Cached GEMM descriptor ────────────────────────────────────────────────── + +/// Pre-created cublasLt descriptors + heuristic algo for a fixed GEMM shape. +/// +/// Created once at init time for each unique (M, N, K, ldb) shape in the forward +/// pass. Eliminates per-call `cublasLtMatmulDescCreate`, `cublasLtMatrixLayoutCreate`, +/// `cublasLtMatmulPreferenceCreate`, and `cublasLtMatmulAlgoGetHeuristic` overhead +/// (22+ GEMMs per training step). +/// +/// The algo heuristic runs ONCE at init with `requestedAlgoCount=3` for better +/// algorithm selection, vs the per-call path which only tried 1. +struct CachedGemmDesc { + matmul_desc: cublaslt_sys::cublasLtMatmulDesc_t, + a_layout: cublaslt_sys::cublasLtMatrixLayout_t, + b_layout: cublaslt_sys::cublasLtMatrixLayout_t, + c_layout: cublaslt_sys::cublasLtMatrixLayout_t, + d_layout: cublaslt_sys::cublasLtMatrixLayout_t, + algo: cublaslt_sys::cublasLtMatmulAlgo_t, +} + +/// SAFETY: cublasLt descriptors are opaque pointers owned exclusively by +/// `CublasForward`, which is single-threaded in practice (owned by `GpuDqnTrainer`). +unsafe impl Send for CachedGemmDesc {} +unsafe impl Sync for CachedGemmDesc {} + +impl Drop for CachedGemmDesc { + fn drop(&mut self) { + unsafe { + let _ = cublaslt_result::destroy_matrix_layout(self.d_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.c_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.b_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.a_layout); + let _ = cublaslt_result::destroy_matmul_desc(self.matmul_desc); + // algo is a value type (cublasLtMatmulAlgo_t), no destroy needed. + // matmul_pref is destroyed after heuristic query at init time. + } + } +} + +/// Key for the cached descriptor map: (n=out_dim, b=batch, k=in_dim, ldb). +/// 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. @@ -162,10 +206,21 @@ pub struct CublasForward { /// then the main stream joins all four before consuming the logits. branch_streams: [Arc; 4], + /// Per-branch cublasLt workspace buffers — eliminates workspace contention + /// when 4 branch streams run GEMMs in parallel. Each 32 MB. + /// Main stream GEMMs (trunk, value head) use `lt_workspace_ptr` (shared). + _branch_workspace_bufs: [CudaSlice; 4], + branch_workspace_ptrs: [u64; 4], + // ── Pre-allocated CUDA events for fork-join synchronization ── /// Reused across forward passes — avoids cuEventCreate/Destroy per call. trunk_done_event: CudaEvent, branch_done_events: [CudaEvent; 4], + + // ── Cached GEMM descriptors (created once at init, reused per-call) ── + /// Map from (n, batch, k, ldb) → pre-created descriptors + algo. + /// All forward GEMMs share TRANSA=T, TRANSB=N. + gemm_cache: HashMap, } impl CublasForward { @@ -256,6 +311,25 @@ impl CublasForward { stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 3: {e}")))?, ]; + // ── Allocate per-branch cublasLt workspace buffers (32 MB each) ── + // Eliminates workspace contention when 4 branch streams run GEMMs in parallel. + let branch_ws_size: usize = 32 * 1024 * 1024; + let branch_ws_buf_0 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("branch workspace 0 alloc: {e}")))?; + let branch_ws_buf_1 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("branch workspace 1 alloc: {e}")))?; + let branch_ws_buf_2 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("branch workspace 2 alloc: {e}")))?; + let branch_ws_buf_3 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("branch workspace 3 alloc: {e}")))?; + let branch_workspace_ptrs = [ + branch_ws_buf_0.raw_ptr(), + branch_ws_buf_1.raw_ptr(), + branch_ws_buf_2.raw_ptr(), + branch_ws_buf_3.raw_ptr(), + ]; + let branch_workspace_bufs = [branch_ws_buf_0, branch_ws_buf_1, branch_ws_buf_2, branch_ws_buf_3]; + // ── Pre-allocate CUDA events for fork-join (reused across all forward passes) ── let ctx = stream.context(); let trunk_done_event = ctx.new_event(None) @@ -267,6 +341,36 @@ impl CublasForward { ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 3: {e}")))?, ]; + // ── Pre-create GEMM descriptors for all unique forward shapes ── + // All forward GEMMs: TRANSA=T, TRANSB=N, f32 types, FAST_TF32 compute. + // Shapes keyed by (n=out_dim, b=batch, k=in_dim, ldb). + let state_dim_padded = (state_dim + 127) & !127; + let branch_sizes = [branch_0_size, branch_1_size, branch_2_size, branch_3_size]; + + let mut unique_shapes: Vec = Vec::new(); + // h_s1: M=shared_h1, N=batch, K=state_dim, ldb=state_dim_padded + unique_shapes.push((shared_h1, batch_size, state_dim, state_dim_padded)); + // h_s2: M=shared_h2, N=batch, K=shared_h1, ldb=shared_h1 + unique_shapes.push((shared_h2, batch_size, shared_h1, shared_h1)); + // h_v: M=value_h, N=batch, K=shared_h2, ldb=shared_h2 + unique_shapes.push((value_h, batch_size, shared_h2, shared_h2)); + // v_logits: M=num_atoms, N=batch, K=value_h, ldb=value_h + unique_shapes.push((num_atoms, batch_size, value_h, value_h)); + // h_bd (×4): M=adv_h, N=batch, K=shared_h2, ldb=shared_h2 + unique_shapes.push((adv_h, batch_size, shared_h2, shared_h2)); + // adv_logits (×4): M=branch_k*num_atoms, N=batch, K=adv_h, ldb=adv_h + for &bs in &branch_sizes { + unique_shapes.push((bs * num_atoms, batch_size, adv_h, adv_h)); + } + unique_shapes.sort(); + unique_shapes.dedup(); + + let mut gemm_cache = HashMap::new(); + for &(n, b, k, ldb) in &unique_shapes { + let desc = create_cached_fwd_gemm_desc(lt_raw_handle, n, b, k, ldb, lt_ws_size)?; + gemm_cache.insert((n, b, k, ldb), desc); + } + Ok(Self { handle: SendSyncCublasHandle(raw_handle), _workspace_buf: workspace_buf, @@ -281,7 +385,7 @@ impl CublasForward { // Dimensions batch_size, state_dim, - state_dim_padded: (state_dim + 127) & !127, + state_dim_padded, shared_h1, shared_h2, value_h, @@ -292,8 +396,11 @@ impl CublasForward { branch_2_size, branch_3_size, branch_streams, + _branch_workspace_bufs: branch_workspace_bufs, + branch_workspace_ptrs, trunk_done_event, branch_done_events, + gemm_cache, }) } @@ -312,7 +419,7 @@ 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, "DIAG_raw") + self.lt_matmul(stream, w_ptr, a_ptr, c_ptr, m, n, k, k, self.lt_workspace_ptr, self.lt_workspace_size, "DIAG_raw") } /// CRITICAL: cublasSetStream resets workspace to the default pool. @@ -424,12 +531,12 @@ impl CublasForward { bs.wait(&self.trunk_done_event) .map_err(|e| MLError::ModelError(format!("branch {d} wait trunk: {e}")))?; - // cublasLtMatmul takes stream per-call — dispatch GEMM to branch stream directly. - self.sgemm_f32(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, "h_bd")?; + // Per-branch workspace: eliminates contention between parallel branch streams. + self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, d, "h_bd")?; self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?; let adv_out_ptr = b_logits_ptr + logit_byte_offset; - self.sgemm_f32(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "adv_logits")?; + self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, d, "adv_logits")?; self.launch_add_bias_f32_raw(bs, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?; logit_byte_offset += (b * n_d * na * std::mem::size_of::()) as u64; @@ -517,11 +624,12 @@ impl CublasForward { bs.wait(&self.trunk_done_event) .map_err(|e| MLError::ModelError(format!("f32 branch {d} wait trunk: {e}")))?; - self.sgemm_f32(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, "f32_h_bd")?; + // Per-branch workspace: eliminates contention between parallel branch streams. + self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, d, "f32_h_bd")?; self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?; let adv_out_ptr = b_logits_ptr + logit_byte_offset; - self.sgemm_f32(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "f32_adv_logits")?; + self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, d, "f32_adv_logits")?; self.launch_add_bias_f32_raw(bs, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?; logit_byte_offset += (b * n_d * na * std::mem::size_of::()) as u64; @@ -694,15 +802,15 @@ impl CublasForward { bs.wait(&self.trunk_done_event) .map_err(|e| MLError::ModelError(format!("tg branch {d} wait trunk: {e}")))?; - // Hidden layer - self.sgemm_f32(bs, tg_w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], - self.adv_h, b, self.shared_h2, "tg_h_bd")?; + // Hidden layer — per-branch workspace eliminates contention. + self.sgemm_f32_branch(bs, tg_w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], + self.adv_h, b, self.shared_h2, d, "tg_h_bd")?; self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], tg_w_ptrs[w_fc_idx + 1], self.adv_h, b)?; - // Output layer: f32 cublasLtMatmul + f32 bias + // Output layer: f32 cublasLtMatmul + f32 bias — per-branch workspace. let adv_out_ptr = b_logits_ptr + logit_byte_offset; - self.sgemm_f32(bs, tg_w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, - n_d * na, b, self.adv_h, "tg_adv_logits")?; + self.sgemm_f32_branch(bs, tg_w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, + n_d * na, b, self.adv_h, d, "tg_adv_logits")?; self.launch_add_bias_f32_raw(bs, adv_out_ptr, tg_w_ptrs[w_fc_idx + 3], n_d * na, b)?; logit_byte_offset += (b * n_d * na * std::mem::size_of::()) as u64; @@ -763,6 +871,9 @@ impl CublasForward { /// eliminating the `cublasSetStream` workspace conflict that hangs CUDA /// Graph mega-capture on H100. CUBLAS_COMPUTE_32F_FAST_TF32 gives ~500 /// TFLOPS on H100 (vs ~60 for pure f32 CUDA cores). + /// + /// Uses cached descriptors (pre-created at init) when available — zero + /// per-call descriptor creation overhead. pub(crate) fn sgemm_f32( &self, stream: &CudaStream, @@ -774,7 +885,25 @@ 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, _label) + 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) + } + + /// F32 GEMM on a branch stream with per-branch workspace. + /// Eliminates workspace contention when 4 branches run in parallel. + #[allow(clippy::too_many_arguments)] + pub(crate) fn sgemm_f32_branch( + &self, + stream: &CudaStream, + w_f32_ptr: u64, + a_f32_ptr: u64, + c_f32_ptr: u64, + n: usize, + b: usize, + k: usize, + 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) } /// F32 GEMM with TF32 + explicit `ldb` (for padded states buffer). @@ -791,10 +920,11 @@ 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, _label) + 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) } - /// Inner cublasLtMatmul dispatch: creates descriptors, gets heuristic, calls matmul. + /// Inner cublasLtMatmul dispatch: uses pre-cached descriptors when available, + /// falls back to inline descriptor creation for uncached shapes. /// /// Layout (row-major trick): /// cuBLAS sees A(W) as col-major [K, N] → TRANSA=T @@ -811,6 +941,110 @@ impl CublasForward { b: usize, k: usize, ldb: usize, + ws_ptr: u64, + ws_size: usize, + _label: &str, + ) -> Result<(), MLError> { + let key: FwdGemmKey = (n, b, k, ldb); + if let Some(cached) = self.gemm_cache.get(&key) { + // ── Fast path: use pre-created descriptors + algo ── + self.lt_matmul_cached(stream, cached, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, ws_ptr, ws_size, _label) + } else { + // ── Slow path: inline descriptor creation (backward compat for uncached shapes) ── + self.lt_matmul_uncached(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, ws_ptr, ws_size, _label) + } + } + + /// Fast path: execute cublasLtMatmul with pre-cached descriptors. + /// Zero per-call descriptor creation overhead. + #[allow(clippy::too_many_arguments)] + fn lt_matmul_cached( + &self, + stream: &CudaStream, + cached: &CachedGemmDesc, + w_f32_ptr: u64, + a_f32_ptr: u64, + c_f32_ptr: u64, + n: usize, + b: usize, + k: usize, + ldb: usize, + ws_ptr: u64, + ws_size: usize, + _label: &str, + ) -> Result<(), MLError> { + let alpha = 1.0_f32; + let beta = 0.0_f32; + + unsafe { + let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; + + tracing::debug!( + handle = ?self.lt_handle.0, + matmul_desc = ?cached.matmul_desc, + w_ptr = w_f32_ptr, + a_ptr = a_f32_ptr, + c_ptr = c_f32_ptr, + a_layout = ?cached.a_layout, + b_layout = ?cached.b_layout, + c_layout = ?cached.c_layout, + d_layout = ?cached.d_layout, + ws_ptr = ws_ptr, + ws_size = ws_size, + stream = ?cu_stream, + m = n, n_batch = b, k = k, ldb = ldb, + "lt_matmul CACHED FFI args for {_label}" + ); + + let matmul_status = cublaslt_sys::cublasLtMatmul( + self.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, + cached.a_layout, + a_f32_ptr as *const std::ffi::c_void, + cached.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + c_f32_ptr as *const std::ffi::c_void, + cached.c_layout, + c_f32_ptr as *mut std::ffi::c_void, + cached.d_layout, + &cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, + ws_ptr as *mut std::ffi::c_void, + ws_size, + cu_stream, + ); + + if matmul_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + let e = cublaslt_result::CublasError(matmul_status); + tracing::error!( + m=n, n_batch=b, k=k, ldb=ldb, + ws_provided=ws_size, + ws_ptr=ws_ptr, + ?e, + "cublasLtMatmul CACHED FAILED for {_label}" + ); + return Err(MLError::ModelError(format!("cublasLtMatmul cached {_label} (m={n},n={b},k={k},ldb={ldb}): {e:?}"))); + } + } + Ok(()) + } + + /// Slow path: inline descriptor creation for uncached GEMM shapes. + /// Used for diagnostic calls and any shape not pre-cached at init. + #[allow(clippy::too_many_arguments)] + fn lt_matmul_uncached( + &self, + stream: &CudaStream, + w_f32_ptr: u64, + a_f32_ptr: u64, + c_f32_ptr: u64, + n: usize, + b: usize, + k: usize, + ldb: usize, + ws_ptr: u64, + ws_size: usize, _label: &str, ) -> Result<(), MLError> { let alpha = 1.0_f32; @@ -823,9 +1057,6 @@ impl CublasForward { let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type) .map_err(|e| MLError::ModelError(format!("cublasLtMatmulDescCreate {_label}: {e:?}")))?; - // TRANSA=T (transpose W), TRANSB=N (input as-is) - // Row-major W[N,K] is col-major [K,N] with ld=K. - // TRANSA=T transposes [K,N] → [N,K] for the GEMM. let transa: i32 = 1; // CUBLAS_OP_T cublaslt_result::set_matmul_desc_attribute( matmul_desc, @@ -834,7 +1065,6 @@ impl CublasForward { std::mem::size_of::(), ).map_err(|e| MLError::ModelError(format!("set TRANSA {_label}: {e:?}")))?; - // TRANSB=N (default, but set explicitly for clarity) let transb: i32 = 0; // CUBLAS_OP_N cublaslt_result::set_matmul_desc_attribute( matmul_desc, @@ -843,17 +1073,13 @@ impl CublasForward { std::mem::size_of::(), ).map_err(|e| MLError::ModelError(format!("set TRANSB {_label}: {e:?}")))?; - // ── Matrix layouts (physical storage, before transpose) ── - // A = W: col-major [K, N] with ld=K. TRANSA=T gives logical [N, K]. + // ── Matrix layouts ── let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64) .map_err(|e| MLError::ModelError(format!("A layout {_label}: {e:?}")))?; - // B = input: col-major [K, B] with ld=ldb. TRANSB=N keeps [K, B]. let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64) .map_err(|e| MLError::ModelError(format!("B layout {_label}: {e:?}")))?; - // C = output: col-major [N, B] with ld=N. let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64) .map_err(|e| MLError::ModelError(format!("C layout {_label}: {e:?}")))?; - // D layout (separate descriptor, same dimensions) let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64) .map_err(|e| MLError::ModelError(format!("D layout {_label}: {e:?}")))?; @@ -863,7 +1089,7 @@ impl CublasForward { cublaslt_result::set_matmul_pref_attribute( matmul_pref, cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, - &self.lt_workspace_size as *const usize as *const std::ffi::c_void, + &ws_size as *const usize as *const std::ffi::c_void, std::mem::size_of::(), ).map_err(|e| MLError::ModelError(format!("set pref ws {_label}: {e:?}")))?; @@ -877,16 +1103,13 @@ impl CublasForward { matmul_pref, ); if let Err(ref e) = heuristic { - tracing::error!(m=n, n_batch=b, k=k, ldb=ldb, ws=self.lt_workspace_size, ?e, "cublasLt heuristic FAILED for {_label}"); + tracing::error!(m=n, n_batch=b, k=k, ldb=ldb, ws=ws_size, ?e, "cublasLt heuristic FAILED for {_label}"); } let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo heuristic {_label} (m={n},n={b},k={k},ldb={ldb}): {e:?}")))?; // ── Execute matmul ── - // Cast CUstream (driver::sys) → cudaStream_t (cublaslt::sys). - // Both are *mut CUstream_st — same underlying type, different module paths. let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; - // DEBUG: print all argument addresses for FFI comparison tracing::debug!( handle = ?self.lt_handle.0, matmul_desc = ?matmul_desc, @@ -897,11 +1120,11 @@ impl CublasForward { b_layout = ?b_layout, c_layout = ?c_layout, d_layout = ?d_layout, - ws_ptr = self.lt_workspace_ptr, - ws_size = self.lt_workspace_size, + ws_ptr = ws_ptr, + ws_size = ws_size, stream = ?cu_stream, m = n, n_batch = b, k = k, ldb = ldb, - "lt_matmul FFI args for {_label}" + "lt_matmul UNCACHED FFI args for {_label}" ); let matmul_status = cublaslt_sys::cublasLtMatmul( @@ -918,8 +1141,8 @@ impl CublasForward { c_f32_ptr as *mut std::ffi::c_void, d_layout, &heuristic.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, - self.lt_workspace_ptr as *mut std::ffi::c_void, - self.lt_workspace_size, + ws_ptr as *mut std::ffi::c_void, + ws_size, cu_stream, ); let matmul_result = if matmul_status == cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { @@ -930,16 +1153,16 @@ impl CublasForward { if let Err(ref e) = matmul_result { tracing::error!( m=n, n_batch=b, k=k, ldb=ldb, - ws_provided=self.lt_workspace_size, + ws_provided=ws_size, ws_needed=heuristic.workspaceSize, - ws_ptr=self.lt_workspace_ptr, + ws_ptr=ws_ptr, ?e, - "cublasLtMatmul FAILED for {_label}" + "cublasLtMatmul UNCACHED FAILED for {_label}" ); } - matmul_result.map_err(|e| MLError::ModelError(format!("cublasLtMatmul {_label} (m={n},n={b},k={k},ldb={ldb},ws={},needed={}): {e:?}", self.lt_workspace_size, heuristic.workspaceSize)))?; + matmul_result.map_err(|e| MLError::ModelError(format!("cublasLtMatmul {_label} (m={n},n={b},k={k},ldb={ldb},ws={ws_size},needed={}): {e:?}", heuristic.workspaceSize)))?; - // ── Cleanup ALL descriptors (leaks cause stale handle reuse) ── + // ── Cleanup ALL descriptors ── let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); let _ = cublaslt_result::destroy_matrix_layout(d_layout); let _ = cublaslt_result::destroy_matrix_layout(c_layout); @@ -1056,6 +1279,134 @@ fn compile_bias_kernels( 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). +/// +/// Runs the algorithm heuristic ONCE at init time with `requestedAlgoCount=3` +/// for better algorithm selection than the per-call path (which only tried 1). +fn create_cached_fwd_gemm_desc( + lt_handle: cublaslt_sys::cublasLtHandle_t, + n: usize, // out_dim (M in cuBLAS) + b: usize, // batch (N in cuBLAS) + k: usize, // in_dim + ldb: usize, // leading dimension of B + ws_size: usize, // workspace size for heuristic +) -> Result { + let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F; + let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32; + + unsafe { + // ── Matmul descriptor ── + let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type) + .map_err(|e| MLError::ModelError(format!("cached MatmulDescCreate (n={n},b={b},k={k}): {e:?}")))?; + + let transa: i32 = 1; // CUBLAS_OP_T + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA, + &transa as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("cached set TRANSA (n={n},b={b},k={k}): {e:?}")))?; + + let transb: i32 = 0; // CUBLAS_OP_N + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB, + &transb as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("cached set TRANSB (n={n},b={b},k={k}): {e:?}")))?; + + // ── Matrix layouts (physical storage, before transpose) ── + // A = W: col-major [K, N] with ld=K. TRANSA=T gives logical [N, K]. + let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64) + .map_err(|e| MLError::ModelError(format!("cached A layout (n={n},b={b},k={k}): {e:?}")))?; + // B = input: col-major [K, B] with ld=ldb. TRANSB=N keeps [K, B]. + let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64) + .map_err(|e| MLError::ModelError(format!("cached B layout (n={n},b={b},k={k}): {e:?}")))?; + // C = output: col-major [N, B] with ld=N. + let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64) + .map_err(|e| MLError::ModelError(format!("cached C layout (n={n},b={b},k={k}): {e:?}")))?; + // D layout (separate descriptor, same dimensions as C) + let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64) + .map_err(|e| MLError::ModelError(format!("cached D layout (n={n},b={b},k={k}): {e:?}")))?; + + // ── Algorithm heuristic (requestedAlgoCount=3 for better selection) ── + let matmul_pref = cublaslt_result::create_matmul_pref() + .map_err(|e| MLError::ModelError(format!("cached matmul pref (n={n},b={b},k={k}): {e:?}")))?; + cublaslt_result::set_matmul_pref_attribute( + matmul_pref, + cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_size as *const usize as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("cached set pref ws (n={n},b={b},k={k}): {e:?}")))?; + + // Request top 3 algorithms, pick the best (first valid). + let mut heuristics: [std::mem::MaybeUninit; 3] = + std::mem::MaybeUninit::uninit().assume_init(); + let mut algo_count: i32 = 0; + + let status = cublaslt_sys::cublasLtMatmulAlgoGetHeuristic( + lt_handle, + matmul_desc, + a_layout, + b_layout, + c_layout, + d_layout, + matmul_pref, + 3, // requestedAlgoCount — try 3 for better selection + heuristics[0].as_mut_ptr(), + &mut algo_count, + ); + if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS || algo_count == 0 { + // Cleanup on failure + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + let _ = cublaslt_result::destroy_matrix_layout(d_layout); + let _ = cublaslt_result::destroy_matrix_layout(c_layout); + let _ = cublaslt_result::destroy_matrix_layout(b_layout); + let _ = cublaslt_result::destroy_matrix_layout(a_layout); + let _ = cublaslt_result::destroy_matmul_desc(matmul_desc); + return Err(MLError::ModelError(format!( + "cached algo heuristic (n={n},b={b},k={k},ldb={ldb}): status={status:?}, count={algo_count}" + ))); + } + + // Pick the first valid heuristic result. + let best = heuristics[0].assume_init(); + if best.state != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + let _ = cublaslt_result::destroy_matrix_layout(d_layout); + let _ = cublaslt_result::destroy_matrix_layout(c_layout); + let _ = cublaslt_result::destroy_matrix_layout(b_layout); + let _ = cublaslt_result::destroy_matrix_layout(a_layout); + let _ = cublaslt_result::destroy_matmul_desc(matmul_desc); + return Err(MLError::ModelError(format!( + "cached algo heuristic state invalid (n={n},b={b},k={k},ldb={ldb}): {:?}", + best.state + ))); + } + + tracing::info!( + n = n, b = b, k = k, ldb = ldb, + algo_count = algo_count, + ws_needed = best.workspaceSize, + "cached fwd GEMM desc created" + ); + + // Preference is no longer needed after heuristic query. + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + + Ok(CachedGemmDesc { + matmul_desc, + a_layout, + b_layout, + c_layout, + d_layout, + algo: best.algo, + }) + } +} + // ── Compute BF16 weight pointers from flat params_buf ─────────────────────── /// Compute the 20 raw BF16 device pointers into a flat params_buf at GOFF_* offsets. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 80fb53f72..c95399c03 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -405,8 +405,8 @@ impl DQNTrainer { } if epoch == phase3_start_epoch { if let Some(ref mut fused) = self.fused_ctx { - let alpha = 0.9_f32; - let sigma = 0.01_f32; + let alpha = self.hyperparams.shrink_perturb_alpha as f32; + let sigma = self.hyperparams.shrink_perturb_sigma as f32; if let Err(e) = fused.shrink_and_perturb(alpha, sigma) { tracing::warn!("Shrink-and-Perturb failed at Phase 3 start (non-fatal): {e}"); } else { @@ -422,25 +422,9 @@ impl DQNTrainer { } } - if !in_phase3 && self.hyperparams.epochs > 4 { - let interval = (self.hyperparams.epochs / 4).max(1); - if epoch > 0 && epoch % interval == 0 { - if let Some(ref mut fused) = self.fused_ctx { - let alpha = 0.9_f32; - let sigma = 0.01_f32; - if let Err(e) = fused.shrink_and_perturb(alpha, sigma) { - tracing::warn!("Shrink-and-Perturb failed (non-fatal): {e}"); - } else { - tracing::info!( - epoch, - alpha, - sigma, - "Shrink-and-Perturb applied (plasticity maintenance)" - ); - } - } - } - } + // Plasticity maintenance: use config-defined interval (not epochs/4) + // to avoid firing every epoch on short runs (3-5 epochs). + // The periodic shrink-perturb at line 357 handles this already. // CVaR scales from IQN head if let Some(ref mut fused) = self.fused_ctx {