diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 18644941d..b2d05a37f 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5344,9 +5344,11 @@ extern "C" __global__ void recursive_confidence_forward( /* Kernel: recursive_confidence_backward — MSE grad into trunk */ /* ================================================================== */ /** - * Per-sample trunk gradient + warp-reduced weight gradient. + * Per-sample trunk gradient + per-block weight gradient partials. * d_h_s2[i, k] += d_sigmoid * w_conf[k] (plain write — single writer per (i,k)). - * d_w_conf and d_b_conf use warp+block reduction → one atomicAdd per BLOCK. + * Weight gradients are reduced within each block (warp shuffle + shared mem) + * and written to per-block partial buffers. A separate reduce kernel sums + * the partials deterministically — zero atomicAdd. * * Grid: ceil(B/256), Block: 256. */ @@ -5355,13 +5357,13 @@ extern "C" __global__ void recursive_confidence_backward( const float* __restrict__ predicted_error, /* [B] */ const float* __restrict__ lagged_td_error, /* [1] pinned — target */ const float* __restrict__ w_conf, /* [SH2] */ - float* __restrict__ d_w_conf, /* [SH2] gradient accumulator */ - float* __restrict__ d_b_conf, /* [1] gradient accumulator */ + float* __restrict__ d_conf_partials, /* [num_blocks, SH2 + 1] block partials */ float* __restrict__ d_h_s2, /* [B, SH2] trunk gradient (accumulate) */ int B, int SH2, float loss_weight /* 0.01 */ ) { int i = blockIdx.x * blockDim.x + threadIdx.x; + int stride = SH2 + 1; /* per-block output: SH2 weight grads + 1 bias grad */ float d_sigmoid = 0.0f; if (i < B) { @@ -5376,25 +5378,29 @@ extern "C" __global__ void recursive_confidence_backward( } } - /* Weight gradient reduction: d_b_conf = sum_b(d_sigmoid_b) */ - float my_db = d_sigmoid; - for (int offset = 16; offset > 0; offset >>= 1) - my_db += __shfl_xor_sync(0xFFFFFFFF, my_db, offset); - - __shared__ float ws_db[8]; int warp_id = threadIdx.x / 32; int lane = threadIdx.x % 32; - if (lane == 0) ws_db[warp_id] = my_db; - __syncthreads(); - if (warp_id == 0) { - float val = (lane < blockDim.x / 32) ? ws_db[lane] : 0.0f; - for (int off = 16; off > 0; off >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, off); - if (lane == 0) atomicAdd(d_b_conf, val); + /* Bias gradient: block-reduced partial → d_conf_partials[block, SH2] */ + { + float my_db = d_sigmoid; + for (int offset = 16; offset > 0; offset >>= 1) + my_db += __shfl_xor_sync(0xFFFFFFFF, my_db, offset); + + __shared__ float ws_db[8]; + if (lane == 0) ws_db[warp_id] = my_db; + __syncthreads(); + + if (warp_id == 0) { + float val = (lane < blockDim.x / 32) ? ws_db[lane] : 0.0f; + for (int off = 16; off > 0; off >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, off); + if (lane == 0) + d_conf_partials[blockIdx.x * stride + SH2] = val; + } } - /* Weight gradient reduction: d_w_conf[k] = sum_b(d_sigmoid_b * h_s2[b, k]) */ + /* Weight gradient: block-reduced partial → d_conf_partials[block, k] */ for (int k = 0; k < SH2; k++) { float my_dw = (i < B) ? d_sigmoid * h_s2[(long long)i * SH2 + k] : 0.0f; @@ -5409,12 +5415,44 @@ extern "C" __global__ void recursive_confidence_backward( float val = (lane < blockDim.x / 32) ? ws_dw[lane] : 0.0f; for (int off = 16; off > 0; off >>= 1) val += __shfl_xor_sync(0xFFFFFFFF, val, off); - if (lane == 0) atomicAdd(&d_w_conf[k], val); + if (lane == 0) + d_conf_partials[blockIdx.x * stride + k] = val; } __syncthreads(); } } +/* ================================================================== */ +/* Kernel: recursive_confidence_reduce — deterministic block partial */ +/* reduction for weight gradients. Zero atomicAdd. */ +/* ================================================================== */ +/** + * One thread per output parameter (SH2 + 1 total). + * Loops over block partials in fixed order → fully deterministic. + * Grid: ceil((SH2+1)/256), Block: 256. + */ +extern "C" __global__ void recursive_confidence_reduce( + const float* __restrict__ d_conf_partials, /* [num_blocks, SH2 + 1] */ + float* __restrict__ d_w_conf, /* [SH2] output */ + float* __restrict__ d_b_conf, /* [1] output */ + int num_blocks, int SH2) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = SH2 + 1; + if (idx >= total) return; + + int stride = SH2 + 1; + float sum = 0.0f; + for (int blk = 0; blk < num_blocks; blk++) { + sum += d_conf_partials[blk * stride + idx]; + } + + if (idx < SH2) + d_w_conf[idx] += sum; + else + *d_b_conf += sum; +} + /* ================================================================== */ /* Kernel: isv_feature_gate — ISV embedding modulates h_s2 features */ /* ================================================================== */ diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index c179a8e2a..ec5b95d66 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1375,6 +1375,11 @@ pub struct GpuDqnTrainer { // ── Recursive confidence kernels ── recursive_conf_fwd_kernel: CudaFunction, recursive_conf_bwd_kernel: CudaFunction, + recursive_conf_reduce_kernel: CudaFunction, + /// Per-block partial buffer for recursive_confidence_backward deterministic reduce. + /// Shape: [max_conf_blocks, SH2 + 1]. + recursive_conf_partials: CudaSlice, + max_conf_blocks: usize, // ── Trade plan head ── plan_params_buf: CudaSlice, // [B, 6] @@ -2278,13 +2283,13 @@ impl GpuDqnTrainer { let loss_weight = 0.01_f32; unsafe { + // Stage 1: per-block partial reduction (no atomicAdd) self.stream.launch_builder(&self.recursive_conf_bwd_kernel) .arg(&self.save_h_s2) .arg(&self.predicted_error_buf) .arg(&self.lagged_td_error_dev_ptr) .arg(&w_conf) - .arg(&d_w_conf) - .arg(&d_b_conf) + .arg(&self.recursive_conf_partials) .arg(&self.bw_d_h_s2) .arg(&b_i32) .arg(&sh2) @@ -2295,6 +2300,23 @@ impl GpuDqnTrainer { shared_mem_bytes: 0, }) .map_err(|e| MLError::ModelError(format!("recursive_confidence_backward: {e}")))?; + + // Stage 2: deterministic reduce across block partials + let num_blocks_i32 = blocks as i32; + let total_params = sh2 + 1; + let reduce_blocks = ((total_params as u32 + 255) / 256).max(1); + self.stream.launch_builder(&self.recursive_conf_reduce_kernel) + .arg(&self.recursive_conf_partials) + .arg(&d_w_conf) + .arg(&d_b_conf) + .arg(&num_blocks_i32) + .arg(&sh2) + .launch(LaunchConfig { + grid_dim: (reduce_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("recursive_confidence_reduce: {e}")))?; } Ok(()) } @@ -4425,6 +4447,11 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("recursive_confidence_forward load: {e}")))?; let recursive_conf_bwd_kernel = exp_module_for_mag.load_function("recursive_confidence_backward") .map_err(|e| MLError::ModelError(format!("recursive_confidence_backward load: {e}")))?; + let recursive_conf_reduce_kernel = exp_module_for_mag.load_function("recursive_confidence_reduce") + .map_err(|e| MLError::ModelError(format!("recursive_confidence_reduce load: {e}")))?; + let max_conf_blocks = ((b + 255) / 256).max(1); + let recursive_conf_partials = stream.alloc_zeros::(max_conf_blocks * (config.shared_h2 + 1)) + .map_err(|e| MLError::ModelError(format!("recursive_conf_partials alloc: {e}")))?; let trade_plan_fwd_kernel = exp_module_for_mag.load_function("trade_plan_forward") .map_err(|e| MLError::ModelError(format!("trade_plan_forward load: {e}")))?; let plan_noise_kernel = exp_module_for_mag.load_function("plan_noise_inject") @@ -5970,6 +5997,9 @@ impl GpuDqnTrainer { predicted_error_buf, recursive_conf_fwd_kernel, recursive_conf_bwd_kernel, + recursive_conf_reduce_kernel, + recursive_conf_partials, + max_conf_blocks, plan_params_buf, trade_plan_fwd_kernel, plan_noise_kernel, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 78f6dbe20..cbc03f881 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -53,25 +53,26 @@ pub(crate) struct FusedStepResult { pub grad_norm: f32, } -/// Raw `CUgraph` handle (uninstantiated) for child graph composition. +/// Child sub-graph: owns both the CUgraph topology and an instantiated CUgraphExec. /// -/// Children are NOT individually instantiated -- they are composed into -/// a parent graph via `cuGraphAddChildGraphNode` and instantiated once. -/// The eval forward child also gets a separate instantiated exec for -/// deterministic Q-value replay during evaluation. +/// The exec comes directly from cudarc's `end_capture` instantiation — we do NOT +/// clone + re-instantiate because `cuGraphClone` corrupts cublasLtMatmul HMMA tile +/// scheduling state on Hopper (sm_90), causing graph replay to hang. struct ChildGraph { graph: cuda_sys::CUgraph, exec: cuda_sys::CUgraphExec, } impl ChildGraph { - fn new(graph: cuda_sys::CUgraph) -> Result { + /// Create an additional exec from an existing graph (for eval replays). + /// Does NOT clone — instantiates directly from the original graph. + fn instantiate_extra_exec(graph: cuda_sys::CUgraph) -> Result { let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut(); let result = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0) }; if result != cuda_sys::cudaError_enum::CUDA_SUCCESS { anyhow::bail!("cuGraphInstantiateWithFlags failed: {result:?}"); } - Ok(Self { graph, exec }) + Ok(exec) } fn launch(&self, stream: cuda_sys::CUstream) -> Result<()> { @@ -1389,9 +1390,9 @@ impl FusedTrainingCtx { /// Capture a single child sub-graph using cudarc stream capture. /// - /// Returns a `ChildGraph` with the raw `CUgraph` (uninstantiated). - /// Children are NOT individually instantiated -- they are composed - /// into the parent graph via `cuGraphAddChildGraphNode`. + /// Takes ownership of cudarc's CudaGraph handles directly via `mem::forget`. + /// We do NOT use `cuGraphClone` — cloning corrupts cublasLtMatmul HMMA tile + /// state on Hopper, causing graph replay to hang when temporal ops are present. /// /// Uses `disable_event_tracking` around capture to prevent cudarc corruption. fn capture_child_graph( @@ -1415,8 +1416,6 @@ impl FusedTrainingCtx { let result = ops(self); - // end_capture returns a CudaGraph which instantiates the graph. - // We need the raw CUgraph for child composition, so extract via clone. let graph_result = self.stream.end_capture( cuda_sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ); @@ -1429,17 +1428,14 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("{label} end_capture: {e}"))? .ok_or_else(|| anyhow::anyhow!("{label} capture returned None"))?; - // Clone the raw CUgraph from cudarc's CudaGraph. cuGraphClone gives - // an independent copy that survives CudaGraph drop. - let mut cloned: cuda_sys::CUgraph = std::ptr::null_mut(); - let r = unsafe { cuda_sys::cuGraphClone(&mut cloned, cudarc_graph.cu_graph) }; - if r != cuda_sys::cudaError_enum::CUDA_SUCCESS { - anyhow::bail!("{label} cuGraphClone: {r:?}"); - } - // CudaGraph drops here, destroying the exec + original graph. - drop(cudarc_graph); + // Take ownership of the raw handles from cudarc's CudaGraph. + // cudarc already instantiated the exec with AUTO_FREE_ON_LAUNCH — use + // it directly. mem::forget prevents CudaGraph::drop from destroying them. + let graph = cudarc_graph.cu_graph; + let exec = cudarc_graph.cu_graph_exec; + std::mem::forget(cudarc_graph); - ChildGraph::new(cloned) + Ok(ChildGraph { graph, exec }) } /// Capture all 5 children and compose into parent graph. @@ -1564,46 +1560,23 @@ impl FusedTrainingCtx { /// Create eval-only forward graph exec for deterministic Q-value replay. /// - /// Clones the forward child graph and instantiates it separately. - /// Injects the exec handle into GpuDqnTrainer for replay_forward_for_q_values(). + /// Create eval-only forward graph exec for deterministic Q-value replay. + /// + /// Instantiates directly from the forward child's original CUgraph — no + /// cuGraphClone. Multiple execs from the same graph is safe (NVIDIA docs). fn create_eval_forward_exec(&mut self) -> Result<()> { let forward_child = self.forward_child.as_ref() .ok_or_else(|| anyhow::anyhow!("forward_child not captured"))?; - // Clone the forward child graph for independent instantiation. - let mut eval_graph: cuda_sys::CUgraph = std::ptr::null_mut(); - let r = unsafe { cuda_sys::cuGraphClone(&mut eval_graph, forward_child.graph) }; - if r != cuda_sys::cudaError_enum::CUDA_SUCCESS { - anyhow::bail!("eval forward cuGraphClone: {r:?}"); - } + let eval_exec = ChildGraph::instantiate_extra_exec(forward_child.graph)?; - let mut eval_exec: cuda_sys::CUgraphExec = std::ptr::null_mut(); - let r = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut eval_exec, eval_graph, 0) }; - // Destroy the cloned graph -- only the exec is needed. - unsafe { cuda_sys::cuGraphDestroy(eval_graph); } - if r != cuda_sys::cudaError_enum::CUDA_SUCCESS { - anyhow::bail!("eval forward instantiate: {r:?}"); - } - - // Destroy previous eval exec. if let Some(old) = self.eval_forward_exec.take() { unsafe { cuda_sys::cuGraphExecDestroy(old); } } self.eval_forward_exec = Some(eval_exec); // Inject into GpuDqnTrainer for replay_forward_for_q_values(). - // Create a separate instantiation from the forward child graph. - let mut trainer_eval_graph: cuda_sys::CUgraph = std::ptr::null_mut(); - let r = unsafe { cuda_sys::cuGraphClone(&mut trainer_eval_graph, forward_child.graph) }; - if r != cuda_sys::cudaError_enum::CUDA_SUCCESS { - anyhow::bail!("trainer eval forward cuGraphClone: {r:?}"); - } - let mut trainer_eval_exec: cuda_sys::CUgraphExec = std::ptr::null_mut(); - let r = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut trainer_eval_exec, trainer_eval_graph, 0) }; - unsafe { cuda_sys::cuGraphDestroy(trainer_eval_graph); } - if r != cuda_sys::cudaError_enum::CUDA_SUCCESS { - anyhow::bail!("trainer eval forward instantiate: {r:?}"); - } + let trainer_eval_exec = ChildGraph::instantiate_extra_exec(forward_child.graph)?; self.trainer.set_eval_forward_exec(trainer_eval_exec); Ok(())