diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 5c0d2d06f..747fc84cc 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -84,7 +84,7 @@ use std::sync::Arc; use cudarc::cublas::sys as cublas_sys; use cudarc::cublaslt::sys as cublaslt_sys; use cudarc::cublaslt::result as cublaslt_result; -use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use crate::MLError; use super::gpu_dqn_trainer::{GpuDqnTrainConfig, compute_param_sizes, padded_byte_offset}; @@ -197,6 +197,35 @@ pub struct CublasBackwardSet { /// Used to compute `padded_byte_offset` for gradient buffer offsets. param_sizes: [usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], + // ── Multi-stream branch dispatch ── + /// 4 forked CUDA streams for parallel advantage branch backward execution. + /// Each branch submits dW/db GEMMs and KAN backward to its own stream, + /// then the main stream joins all four before the dX (Step 5) accumulation. + branch_streams: [Arc; 4], + + /// Per-branch cublasLt workspace buffers — eliminates workspace contention + /// when 4 branch streams run GEMMs in parallel. Each 32 MB. + _branch_workspace_bufs: [CudaSlice; 4], + branch_workspace_ptrs: [u64; 4], + + /// Reused CUDA events for fork-join synchronization. + trunk_done_event: CudaEvent, + branch_done_events: [CudaEvent; 4], + + // ── Per-branch scratch buffers (safe for parallel execution) ── + /// Per-branch GLU value gradient scratch [B, AH] each — avoids cross-branch contention. + branch_d_glu_value: [CudaSlice; 4], + /// Per-branch GLU gate_pre gradient scratch [B, AH] each. + branch_d_glu_gate: [CudaSlice; 4], + + // ── Per-branch KAN backward intermediate buffers ── + /// Per-element B-spline basis contributions [B*AH, 4] per branch — f32. + branch_kan_d_coeff_per_elem: [CudaSlice; 4], + /// Per-element span index [B*AH] per branch — i32. + branch_kan_d_coeff_span: [CudaSlice; 4], + /// Per-element residual_w gradient contribution [B*AH] per branch — f32. + branch_kan_d_residual_per_elem: [CudaSlice; 4], + // ── 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, @@ -308,6 +337,86 @@ impl CublasBackwardSet { let kan_d_residual_per_elem = stream.alloc_zeros::(kan_n) .map_err(|e| MLError::ModelError(format!("kan_d_residual_per_elem alloc [{kan_n}]: {e}")))?; + // ── Fork 4 branch streams for parallel backward dispatch ──── + let branch_streams = [ + stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 0: {e}")))?, + stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 1: {e}")))?, + stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 2: {e}")))?, + stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 3: {e}")))?, + ]; + + // ── Per-branch cublasLt workspace buffers (32 MB each) ────── + let branch_ws_size: usize = 32 * 1024 * 1024; + let bw_ws_buf_0 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("bw branch workspace 0 alloc: {e}")))?; + let bw_ws_buf_1 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("bw branch workspace 1 alloc: {e}")))?; + let bw_ws_buf_2 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("bw branch workspace 2 alloc: {e}")))?; + let bw_ws_buf_3 = stream.alloc_zeros::(branch_ws_size) + .map_err(|e| MLError::ModelError(format!("bw branch workspace 3 alloc: {e}")))?; + let branch_workspace_ptrs = [ + bw_ws_buf_0.raw_ptr(), + bw_ws_buf_1.raw_ptr(), + bw_ws_buf_2.raw_ptr(), + bw_ws_buf_3.raw_ptr(), + ]; + let branch_workspace_bufs = [bw_ws_buf_0, bw_ws_buf_1, bw_ws_buf_2, bw_ws_buf_3]; + + // ── Pre-allocate CUDA events for fork-join ────────────────── + let ctx = stream.context(); + let trunk_done_event = ctx.new_event(None) + .map_err(|e| MLError::DeviceError(format!("bw alloc trunk_done_event: {e}")))?; + let branch_done_events = [ + ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 0: {e}")))?, + ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 1: {e}")))?, + ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 2: {e}")))?, + ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 3: {e}")))?, + ]; + + // ── Per-branch GLU scratch buffers ────────────────────────── + let b_ah = config.batch_size * config.adv_h; + let alloc_f32 = |n: usize, label: &str| -> Result, MLError> { + stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("{label} alloc [{n}]: {e}"))) + }; + let branch_d_glu_value = [ + alloc_f32(b_ah, "branch_d_glu_value_0")?, + alloc_f32(b_ah, "branch_d_glu_value_1")?, + alloc_f32(b_ah, "branch_d_glu_value_2")?, + alloc_f32(b_ah, "branch_d_glu_value_3")?, + ]; + let branch_d_glu_gate = [ + alloc_f32(b_ah, "branch_d_glu_gate_0")?, + alloc_f32(b_ah, "branch_d_glu_gate_1")?, + alloc_f32(b_ah, "branch_d_glu_gate_2")?, + alloc_f32(b_ah, "branch_d_glu_gate_3")?, + ]; + + // ── Per-branch KAN backward intermediate buffers ──────────── + let alloc_i32 = |n: usize, label: &str| -> Result, MLError> { + stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("{label} alloc [{n}]: {e}"))) + }; + let branch_kan_d_coeff_per_elem = [ + alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_0")?, + alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_1")?, + alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_2")?, + alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_3")?, + ]; + let branch_kan_d_coeff_span = [ + alloc_i32(kan_n, "branch_kan_d_coeff_span_0")?, + alloc_i32(kan_n, "branch_kan_d_coeff_span_1")?, + alloc_i32(kan_n, "branch_kan_d_coeff_span_2")?, + alloc_i32(kan_n, "branch_kan_d_coeff_span_3")?, + ]; + let branch_kan_d_residual_per_elem = [ + alloc_f32(kan_n, "branch_kan_d_resid_pe_0")?, + alloc_f32(kan_n, "branch_kan_d_resid_pe_1")?, + alloc_f32(kan_n, "branch_kan_d_resid_pe_2")?, + alloc_f32(kan_n, "branch_kan_d_resid_pe_3")?, + ]; + Ok(Self { handle: shared, relu_mask_kernel, @@ -330,6 +439,16 @@ impl CublasBackwardSet { branch_1_size: config.branch_1_size, branch_2_size: config.branch_2_size, branch_3_size: config.branch_3_size, + branch_streams, + _branch_workspace_bufs: branch_workspace_bufs, + branch_workspace_ptrs, + trunk_done_event, + branch_done_events, + branch_d_glu_value, + branch_d_glu_gate, + branch_kan_d_coeff_per_elem, + branch_kan_d_coeff_span, + branch_kan_d_residual_per_elem, param_sizes, gemm_cache, }) @@ -355,6 +474,8 @@ impl CublasBackwardSet { /// /// Uses cached descriptors when the shape is pre-cached at init time. /// Falls back to inline descriptor creation for uncached shapes. + /// + /// Uses the shared (main) workspace — NOT safe for branch streams. fn sgemm_f32( &self, stream: &CudaStream, @@ -367,18 +488,40 @@ impl CublasBackwardSet { beta: f32, c: u64, ldc: i32, label: &str, + ) -> Result<(), MLError> { + self.sgemm_f32_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, self.handle.lt_workspace_ptr, label) + } + + /// F32 GEMM with explicit workspace pointer. + /// + /// Branch streams MUST use per-branch workspace pointers to avoid + /// contention with the shared workspace used by the main stream. + #[allow(clippy::too_many_arguments)] + fn sgemm_f32_ws( + &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, + ws_ptr: u64, + 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) + self.sgemm_f32_cached_ws(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, ws_ptr, 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) + self.sgemm_f32_uncached_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, ws_ptr, label) } } - /// Fast path: execute backward GEMM with pre-cached descriptors. + /// Fast path: execute backward GEMM with pre-cached descriptors (shared workspace). #[allow(clippy::too_many_arguments)] fn sgemm_f32_cached( &self, @@ -392,6 +535,25 @@ impl CublasBackwardSet { m: i32, n: i32, k: i32, lda: i32, ldb: i32, ldc: i32, label: &str, + ) -> Result<(), MLError> { + self.sgemm_f32_cached_ws(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, self.handle.lt_workspace_ptr, label) + } + + /// Fast path: execute backward GEMM with pre-cached descriptors and explicit workspace. + #[allow(clippy::too_many_arguments)] + fn sgemm_f32_cached_ws( + &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, + ws_ptr: u64, + label: &str, ) -> Result<(), MLError> { unsafe { let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t; @@ -410,7 +572,7 @@ impl CublasBackwardSet { c as *mut std::ffi::c_void, cached.d_layout, &cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, - self.handle.lt_workspace_ptr as *mut std::ffi::c_void, + ws_ptr as *mut std::ffi::c_void, self.handle.lt_workspace_size, cu_stream, ); @@ -420,7 +582,7 @@ impl CublasBackwardSet { tracing::error!( m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc, ws_provided=self.handle.lt_workspace_size, - ws_ptr=self.handle.lt_workspace_ptr, + ws_ptr=ws_ptr, ?e, "cublasLtMatmul CACHED FAILED for bw {label}" ); @@ -430,7 +592,7 @@ impl CublasBackwardSet { Ok(()) } - /// Slow path: inline descriptor creation for uncached backward GEMM shapes. + /// Slow path: inline descriptor creation for uncached backward GEMM shapes (shared workspace). #[allow(clippy::too_many_arguments)] fn sgemm_f32_uncached( &self, @@ -444,6 +606,25 @@ impl CublasBackwardSet { beta: f32, c: u64, ldc: i32, label: &str, + ) -> Result<(), MLError> { + self.sgemm_f32_uncached_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, self.handle.lt_workspace_ptr, label) + } + + /// Slow path: inline descriptor creation with explicit workspace pointer. + #[allow(clippy::too_many_arguments)] + fn sgemm_f32_uncached_ws( + &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, + ws_ptr: u64, + 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) @@ -531,7 +712,7 @@ impl CublasBackwardSet { c as *mut std::ffi::c_void, d_layout, &heuristic.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, - self.handle.lt_workspace_ptr as *mut std::ffi::c_void, + ws_ptr as *mut std::ffi::c_void, self.handle.lt_workspace_size, cu_stream, ); @@ -540,7 +721,7 @@ impl CublasBackwardSet { m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc, ws_provided=self.handle.lt_workspace_size, ws_needed=heuristic.workspaceSize, - ws_ptr=self.handle.lt_workspace_ptr, + ws_ptr=ws_ptr, ?e, "cublasLtMatmul UNCACHED FAILED for bw {label}" ); @@ -796,6 +977,388 @@ impl CublasBackwardSet { Ok(()) } + // ══════════════════════════════════════════════════════════════════════════ + // Per-branch backward (Steps 1-4: dW/db only, no dX) + // ══════════════════════════════════════════════════════════════════════════ + + /// KAN spline gate backward using per-branch intermediate buffers. + /// + /// Identical to `launch_kan_gate_backward` but uses branch-specific + /// `d_coeff_per_elem`, `d_coeff_span`, and `d_residual_per_elem` buffers, + /// making it safe for concurrent execution on branch streams. + #[allow(clippy::too_many_arguments)] + fn launch_kan_gate_backward_branch( + &self, + stream: &Arc, + d: usize, // branch index 0..3 + d_output: u64, // [B, AH] + gate_pre: u64, // [B, AH] + value: u64, // [B, AH] + spline_coeff: u64, // [AH, 8] + residual_w: u64, // [AH] + d_gate_pre: u64, // [B, AH] output + d_value: u64, // [B, AH] output + d_spline_coeff: u64, // [AH, 8] output (in grad_buf) + d_residual_w: u64, // [AH] output (in grad_buf) + n: usize, + adv_h: usize, + ) -> Result<(), MLError> { + let n_i32 = n as i32; + let adv_h_i32 = adv_h as i32; + let blocks = ((n + 255) / 256) as u32; + + // Use per-branch intermediate buffers (no cross-branch contention). + let d_coeff_pe_ptr = raw_f32_ptr(&self.branch_kan_d_coeff_per_elem[d], stream); + let d_span_ptr = { + let (ptr, guard) = self.branch_kan_d_coeff_span[d].device_ptr(stream); + let _no_drop = ManuallyDrop::new(guard); + ptr + }; + let d_resid_pe_ptr = raw_f32_ptr(&self.branch_kan_d_residual_per_elem[d], stream); + + // ── Stage 1: per-element scatter (no atomicAdd) ── + unsafe { + stream + .launch_builder(&self.kan_gate_backward_kernel) + .arg(&d_output) + .arg(&gate_pre) + .arg(&value) + .arg(&spline_coeff) + .arg(&residual_w) + .arg(&d_gate_pre) + .arg(&d_value) + .arg(&d_coeff_pe_ptr) + .arg(&d_span_ptr) + .arg(&d_resid_pe_ptr) + .arg(&n_i32) + .arg(&adv_h_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("kan_gate_backward_kernel branch {d}: {e}")))?; + } + + // ── Stage 2: deterministic reduction ── + let total_params = adv_h * 9; + let reduce_blocks = ((total_params + 255) / 256) as u32; + + unsafe { + stream + .launch_builder(&self.kan_grad_reduce_kernel) + .arg(&d_coeff_pe_ptr) + .arg(&d_span_ptr) + .arg(&d_resid_pe_ptr) + .arg(&d_spline_coeff) + .arg(&d_residual_w) + .arg(&n_i32) + .arg(&adv_h_i32) + .launch(LaunchConfig { + grid_dim: (reduce_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("kan_grad_reduce_kernel branch {d}: {e}")))?; + } + + Ok(()) + } + + /// Backward Steps 1-4 for a single branch on a branch stream. + /// + /// Computes dW/db for the branch output layer, KAN spline backward, and + /// value/gate path weight gradients. All writes go to non-overlapping + /// regions of `grad_buf`, making this safe for parallel execution. + /// + /// Does NOT compute Step 5 (dX to d_h_s2 or d_concat). That runs on the + /// main stream after all 4 branches join. + #[allow(clippy::too_many_arguments)] + fn backward_branch_dw( + &self, + stream: &Arc, + d: usize, + ws_ptr: u64, + // Per-branch data + d_adv_logits_d: u64, // [B, n_d] + save_h_b_d: u64, // [B, AH] (KAN gate output) + scratch_d_h_b_d: u64, // [B, AH] scratch for d_h_bd + glu_gate_pre_d: u64, // [B, AH] saved gate pre-activation + glu_value_d: u64, // [B, AH] saved value path output + // Weight pointers (read-only) + w_out: u64, // W_bdk_out [n_d, AH] + w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], + // Grad buf offsets + grad_buf_base: u64, + goff_w_bout_d: u64, + goff_b_bout_d: u64, + goff_w_bfc_d: u64, + goff_b_bfc_d: u64, + goff_w_gate_d: u64, + goff_b_gate_d: u64, + goff_kan_coeff_d: u64, + goff_kan_resid_d: u64, + // FC input for dW (branches 1,2,3 may use concat buffers) + fc_input: u64, + fc_in_dim: usize, + // Branch output dim + n_d: usize, + b: usize, + ) -> Result<(), MLError> { + // ── Step 1: Branch output layer (no activation — logit layer) ── + // dW += dY^T @ X, db += sum(dY), dX = dY @ W^T → scratch_d_h_b[d] + self.backward_fc_layer_ws( + stream, + d_adv_logits_d, + save_h_b_d, + w_out, + grad_buf_base + goff_w_bout_d, + grad_buf_base + goff_b_bout_d, + scratch_d_h_b_d, + n_d, + self.adv_h, + b, + ws_ptr, + )?; + + // Per-branch GLU scratch pointers + let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream); + let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream); + + // ── Step 2: KAN spline backward ── + let kan_coeff_idx = 42 + d * 2; + let kan_resid_idx = 42 + d * 2 + 1; + self.launch_kan_gate_backward_branch( + stream, + d, + scratch_d_h_b_d, + glu_gate_pre_d, + glu_value_d, + w_ptrs[kan_coeff_idx], + w_ptrs[kan_resid_idx], + d_glu_gate_ptr, + d_glu_value_ptr, + grad_buf_base + goff_kan_coeff_d, + grad_buf_base + goff_kan_resid_d, + b * self.adv_h, + self.adv_h, + )?; + + // ── Step 3: Value path weight gradients ── + self.launch_dw_only_ws( + stream, + d_glu_value_ptr, + fc_input, + grad_buf_base + goff_w_bfc_d, + grad_buf_base + goff_b_bfc_d, + self.adv_h, + fc_in_dim, + b, + ws_ptr, + )?; + + // ── Step 4: Gate path weight gradients ── + self.launch_dw_only_ws( + stream, + d_glu_gate_ptr, + fc_input, + grad_buf_base + goff_w_gate_d, + grad_buf_base + goff_b_gate_d, + self.adv_h, + fc_in_dim, + b, + ws_ptr, + )?; + + Ok(()) + } + + /// Backward Step 5 for a single branch on the main stream. + /// + /// Computes dX (upstream gradient to VSN input / d_h_s2 / concat buffers). + /// Must run on the main stream AFTER all branches finish Steps 1-4 so that + /// d_h_s2 accumulation with beta is ordered correctly. + #[allow(clippy::too_many_arguments)] + fn backward_branch_dx( + &self, + stream: &Arc, + d: usize, + w_fc: u64, + w_gate: u64, + scratch_d_h_s2: u64, + mag_concat_ptr: u64, + d_mag_concat_ptr: u64, + ord_concat_ptr: u64, + d_ord_concat_ptr: u64, + urg_concat_ptr: u64, + d_urg_concat_ptr: u64, + b: usize, + ) -> Result<(), MLError> { + // Read per-branch GLU scratch (produced by Steps 1-4 on branch stream, + // now readable by main stream after join). + let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream); + let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream); + + if d == 1 && mag_concat_ptr != 0 { + // Magnitude: dX writes to d_mag_concat [B, SH2+3]. + self.launch_dx_only( + stream, + d_glu_value_ptr, + w_fc, + d_mag_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 0.0_f32, + )?; + self.launch_dx_only( + stream, + d_glu_gate_ptr, + w_gate, + d_mag_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 1.0_f32, + )?; + } else if d == 2 && ord_concat_ptr != 0 { + // Order: dX writes to d_ord_concat [B, SH2+3]. + self.launch_dx_only( + stream, + d_glu_value_ptr, + w_fc, + d_ord_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 0.0_f32, + )?; + self.launch_dx_only( + stream, + d_glu_gate_ptr, + w_gate, + d_ord_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 1.0_f32, + )?; + } else if d == 3 && urg_concat_ptr != 0 { + // Urgency: dX writes to d_urg_concat [B, SH2+3]. + self.launch_dx_only( + stream, + d_glu_value_ptr, + w_fc, + d_urg_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 0.0_f32, + )?; + self.launch_dx_only( + stream, + d_glu_gate_ptr, + w_gate, + d_urg_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 1.0_f32, + )?; + } else { + // Direction (d==0) or any branch with no concat: accumulate into d_h_s2 directly. + let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 }; + self.launch_dx_only( + stream, + d_glu_value_ptr, + w_fc, + scratch_d_h_s2, + self.adv_h, + self.shared_h2, + b, + beta_s2, + )?; + self.launch_dx_only( + stream, + d_glu_gate_ptr, + w_gate, + scratch_d_h_s2, + self.adv_h, + self.shared_h2, + b, + 1.0_f32, + )?; + } + + Ok(()) + } + + /// Workspace-aware `backward_fc_layer`: dW + db + dX using explicit workspace. + #[allow(clippy::too_many_arguments)] + fn backward_fc_layer_ws( + &self, + stream: &Arc, + dy: u64, x: u64, w: u64, dw: u64, db: u64, dx: u64, + out_dim: usize, in_dim: usize, batch: usize, + ws_ptr: u64, + ) -> Result<(), MLError> { + // dW += dY^T @ X + self.sgemm_f32_ws( + stream, + cublas_sys::cublasOperation_t::CUBLAS_OP_N, + cublas_sys::cublasOperation_t::CUBLAS_OP_T, + in_dim as i32, out_dim as i32, batch as i32, + 1.0, x, in_dim as i32, + dy, out_dim as i32, + 1.0, dw, in_dim as i32, + ws_ptr, + "fc_dW_ws", + )?; + + self.launch_bias_grad(stream, dy, db, out_dim, batch)?; + + // dX = dY @ W^T + if dx != 0 { + self.sgemm_f32_ws( + stream, + cublas_sys::cublasOperation_t::CUBLAS_OP_N, + cublas_sys::cublasOperation_t::CUBLAS_OP_N, + in_dim as i32, batch as i32, out_dim as i32, + 1.0, w, in_dim as i32, + dy, out_dim as i32, + 0.0, dx, in_dim as i32, + ws_ptr, + "fc_dX_ws", + )?; + } + + Ok(()) + } + + /// Workspace-aware `launch_dw_only`: dW + db using explicit workspace. + #[allow(clippy::too_many_arguments)] + fn launch_dw_only_ws( + &self, + stream: &Arc, + dy: u64, x: u64, dw: u64, db: u64, + out_dim: usize, in_dim: usize, batch: usize, + ws_ptr: u64, + ) -> Result<(), MLError> { + self.sgemm_f32_ws( + stream, + cublas_sys::cublasOperation_t::CUBLAS_OP_N, + cublas_sys::cublasOperation_t::CUBLAS_OP_T, + in_dim as i32, out_dim as i32, batch as i32, + 1.0, x, in_dim as i32, + dy, out_dim as i32, + 1.0, dw, in_dim as i32, + ws_ptr, + "dW_only_ws", + )?; + + self.launch_bias_grad(stream, dy, db, out_dim, batch) + } + // ══════════════════════════════════════════════════════════════════════════ // Full backward pass // ══════════════════════════════════════════════════════════════════════════ @@ -876,9 +1439,10 @@ impl CublasBackwardSet { // GLU saved forward activations for backward pass glu_gate_pre: &[u64; 4], // [B, AH] each — saved pre-sigmoid gate activations glu_value: &[u64; 4], // [B, AH] each — saved value path outputs - // GLU backward scratch buffers (reused across branches) - scratch_d_glu_value: u64, // [B, AH] — d_value output from GLU backward - scratch_d_glu_gate: u64, // [B, AH] — d_gate_pre output from GLU backward + // GLU backward scratch buffers (kept for backward compatibility; per-branch + // versions in branch_d_glu_value/branch_d_glu_gate are used for parallel dispatch) + _scratch_d_glu_value: u64, // [B, AH] — superseded by per-branch scratch + _scratch_d_glu_gate: u64, // [B, AH] — superseded by per-branch scratch ) -> Result<(), MLError> { let b = self.batch_size; let na = self.num_atoms; @@ -975,57 +1539,28 @@ impl CublasBackwardSet { ]; // ══════════════════════════════════════════════════════════════════ - // BRANCHES (all 4, with KAN spline backward replacing ReLU mask) + // BRANCHES: parallel dW/db (Steps 1-4), sequential dX (Step 5) + // + // Steps 1-4 write to non-overlapping regions of grad_buf and + // per-branch scratch buffers, so they run safely in parallel on + // branch_streams[0..3]. Step 5 (dX) may accumulate into the + // shared scratch_d_h_s2 with beta ordering, so it runs + // sequentially on the main stream after all branches join. // ══════════════════════════════════════════════════════════════════ + // ── Fork: record trunk completion, dispatch Steps 1-4 on branch streams ── + self.trunk_done_event.record(stream) + .map_err(|e| MLError::DeviceError(format!("bw trunk_done_event record: {e}")))?; + for d in 0..4 { + let bs = &self.branch_streams[d]; + bs.wait(&self.trunk_done_event) + .map_err(|e| MLError::DeviceError(format!("bw branch stream {d} wait trunk: {e}")))?; + let n_d = branch_n[d]; - let w_out = w_ptrs[w_bout_idx[d]]; // W_bdk_out[n_d, AH] - let w_fc = w_ptrs[w_bfc_idx[d]]; // W_bdk_fc[AH, SH2] (value path weight) - let w_gate = w_ptrs[w_gate_idx[d]]; // W_gate_d[AH, SH2] (gate path weight) + let w_out = w_ptrs[w_bout_idx[d]]; - // KAN spline weight pointers (from params_buf, read-only) - let kan_coeff_idx = 42 + d * 2; // kan_coeff_d - let kan_resid_idx = 42 + d * 2 + 1; // kan_resid_d - - // ── Step 1: Branch output layer (no activation — logit layer) ── - // dY = d_adv_logits[d] [B, n_d] - // X = save_h_b[d] [B, AH] (KAN gate output) - // W = W_bdk_out [n_d, AH] - // dW += dY^T @ X, db += sum(dY), dX_b = dY @ W^T → scratch_d_h_b[d] - self.backward_fc_layer( - stream, - d_adv_logits[d], // dY [B, n_d] - save_h_b[d], // X [B, AH] - w_out, // W [n_d, AH] - grad_buf_base + goff_w_bout[d], // dW - grad_buf_base + goff_b_bout[d], // db - scratch_d_h_b[d], // dX [B, AH] → d_h_bd (gradient into KAN gate output) - n_d, // out_dim - self.adv_h, // in_dim - b, - )?; - - // ── Step 2: KAN spline backward (replaces GLU backward) ────── - // d_h_bd → d_value + d_gate_pre + d_coeff + d_resid_w - self.launch_kan_gate_backward( - stream, - scratch_d_h_b[d], // d_output = d_h_bd [B, AH] - glu_gate_pre[d], // saved gate_pre [B, AH] - glu_value[d], // saved value [B, AH] - w_ptrs[kan_coeff_idx], // spline_coeff [AH, 8] (current weights) - w_ptrs[kan_resid_idx], // residual_w [AH] (current weights) - scratch_d_glu_gate, // d_gate_pre output [B, AH] - scratch_d_glu_value, // d_value output [B, AH] - grad_buf_base + goff_kan_coeff[d], // d_spline_coeff [AH, 8] (accumulated in grad_buf) - grad_buf_base + goff_kan_resid[d], // d_residual_w [AH] (accumulated in grad_buf) - b * self.adv_h, - self.adv_h, - )?; - - // ── Step 3: Value path weight gradients ──────────────────── - // dW_bdf += d_value^T @ vsn_input, db_bdf += sum(d_value) - // Branches 1,2,3 use wider concat input [B, SH2+3] for dW. + // FC input for dW (branches 1,2,3 may use wider concat buffers) let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 { (mag_concat_ptr, self.shared_h2 + 3) } else if d == 2 && ord_concat_ptr != 0 { @@ -1035,135 +1570,66 @@ impl CublasBackwardSet { } else { (save_h_s2, self.shared_h2) }; - self.launch_dw_only( - stream, - scratch_d_glu_value, // dY = d_value [B, AH] — f32 - fc_input, // X [B, fc_in_dim] (vsn_input) - grad_buf_base + goff_w_bfc[d], // dW [AH, fc_in_dim] - grad_buf_base + goff_b_bfc[d], // db [AH] - self.adv_h, // out_dim - fc_in_dim, // in_dim + + self.backward_branch_dw( + bs, + d, + self.branch_workspace_ptrs[d], + d_adv_logits[d], + save_h_b[d], + scratch_d_h_b[d], + glu_gate_pre[d], + glu_value[d], + w_out, + w_ptrs, + grad_buf_base, + goff_w_bout[d], + goff_b_bout[d], + goff_w_bfc[d], + goff_b_bfc[d], + goff_w_gate[d], + goff_b_gate[d], + goff_kan_coeff[d], + goff_kan_resid[d], + fc_input, + fc_in_dim, + n_d, b, )?; - // ── Step 4: Gate path weight gradients ───────────────────── - // dW_gate += d_gate_pre^T @ vsn_input, db_gate += sum(d_gate_pre) - self.launch_dw_only( + self.branch_done_events[d].record(bs) + .map_err(|e| MLError::DeviceError(format!("bw branch_done_event {d} record: {e}")))?; + } + + // ── Join: main stream waits for all 4 branches to finish Steps 1-4 ── + for d in 0..4 { + stream.wait(&self.branch_done_events[d]) + .map_err(|e| MLError::DeviceError(format!("bw main stream wait branch {d}: {e}")))?; + } + + // ── Step 5: dX accumulation runs sequentially on main stream ── + // d_h_s2 requires ordered beta accumulation (d==0 overwrites, d>0 accumulates). + // Concat branches (mag/ord/urg) write to separate buffers, so ordering + // among them does not matter — but they must complete before the trunk + // backward reads d_h_s2. + for d in 0..4 { + let w_fc = w_ptrs[w_bfc_idx[d]]; + let w_gate = w_ptrs[w_gate_idx[d]]; + + self.backward_branch_dx( stream, - scratch_d_glu_gate, // dY = d_gate_pre [B, AH] — f32 - fc_input, // X [B, fc_in_dim] (same vsn_input) - grad_buf_base + goff_w_gate[d], // dW_gate [AH, fc_in_dim] - grad_buf_base + goff_b_gate[d], // db_gate [AH] - self.adv_h, // out_dim - fc_in_dim, // in_dim + d, + w_fc, + w_gate, + scratch_d_h_s2, + mag_concat_ptr, + d_mag_concat_ptr, + ord_concat_ptr, + d_ord_concat_ptr, + urg_concat_ptr, + d_urg_concat_ptr, b, )?; - - // ── Step 5: Upstream gradient to VSN input ───────────────── - // d_vsn_input = W_bdf^T @ d_value + W_gate^T @ d_gate_pre - // Sum of both paths (value + gate) using beta accumulation. - if d == 1 && mag_concat_ptr != 0 { - // Magnitude: dX writes to d_mag_concat [B, SH2+3]. - // Value path: beta=0 (overwrite) - self.launch_dx_only( - stream, - scratch_d_glu_value, // dY = d_value [B, AH] - w_fc, // W_bdf [AH, SH2+3] - d_mag_concat_ptr, // dX [B, SH2+3] - self.adv_h, - self.shared_h2 + 3, - b, - 0.0_f32, // overwrite - )?; - // Gate path: beta=1 (accumulate on top of value path) - self.launch_dx_only( - stream, - scratch_d_glu_gate, // dY = d_gate_pre [B, AH] - w_gate, // W_gate [AH, SH2+3] - d_mag_concat_ptr, // dX [B, SH2+3] (accumulate) - self.adv_h, - self.shared_h2 + 3, - b, - 1.0_f32, // accumulate - )?; - // Caller will call accumulate_d_h_s2_from_concat(d_mag_concat, d_h_s2, B, 1.0) - // after backward_full returns. - } else if d == 2 && ord_concat_ptr != 0 { - // Order: dX writes to d_ord_concat [B, SH2+3]. - // Value path: beta=0 (overwrite) - self.launch_dx_only( - stream, - scratch_d_glu_value, - w_fc, - d_ord_concat_ptr, - self.adv_h, - self.shared_h2 + 3, - b, - 0.0_f32, - )?; - // Gate path: beta=1 (accumulate) - self.launch_dx_only( - stream, - scratch_d_glu_gate, - w_gate, - d_ord_concat_ptr, - self.adv_h, - self.shared_h2 + 3, - b, - 1.0_f32, - )?; - // Caller will call accumulate_d_h_s2_from_concat(d_ord_concat, d_h_s2, B, 1.0) - } else if d == 3 && urg_concat_ptr != 0 { - // Urgency: dX writes to d_urg_concat [B, SH2+3]. - // Value path: beta=0 (overwrite) - self.launch_dx_only( - stream, - scratch_d_glu_value, - w_fc, - d_urg_concat_ptr, - self.adv_h, - self.shared_h2 + 3, - b, - 0.0_f32, - )?; - // Gate path: beta=1 (accumulate) - self.launch_dx_only( - stream, - scratch_d_glu_gate, - w_gate, - d_urg_concat_ptr, - self.adv_h, - self.shared_h2 + 3, - b, - 1.0_f32, - )?; - // Caller will call accumulate_d_h_s2_from_concat(d_urg_concat, d_h_s2, B, 1.0) - } else { - // Other branches: accumulate into d_h_s2 directly. - let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 }; - // Value path - self.launch_dx_only( - stream, - scratch_d_glu_value, // dY = d_value [B, AH] - w_fc, // W_bdf [AH, SH2] - scratch_d_h_s2, // dX [B, SH2] - self.adv_h, - self.shared_h2, - b, - beta_s2, // d==0: overwrite, d>0: accumulate - )?; - // Gate path: always beta=1 (accumulate on top of value path) - self.launch_dx_only( - stream, - scratch_d_glu_gate, // dY = d_gate_pre [B, AH] - w_gate, // W_gate [AH, SH2] - scratch_d_h_s2, // dX [B, SH2] (accumulate) - self.adv_h, - self.shared_h2, - b, - 1.0_f32, // always accumulate - )?; - } } // ══════════════════════════════════════════════════════════════════