diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index 3b6953507..5f50a34f3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -606,9 +606,13 @@ impl GpuAttention { let f32_sz = std::mem::size_of::(); let inv_batch = 1.0_f32 / b as f32; - // Zero grad_buf before accumulation - stream.memset_zeros(&mut self.d_params) - .map_err(|e| MLError::ModelError(format!("attention zero d_params: {e}")))?; + // Zero grad_buf — raw cuMemsetD8Async, no cudarc overhead during graph capture + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.d_params.raw_ptr(), 0, self.d_params.num_bytes(), + stream.cu_stream(), + ); + } // Weight layout offsets let w_q_off = 0; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 172472aad..1dfe6e962 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -615,11 +615,18 @@ pub struct GpuDqnTrainer { /// Separate finalize instance for adam_child — cannot share CUfunction across /// child graphs on Hopper (causes 3100ms replay from kernel state corruption). grad_norm_finalize_adam: CudaFunction, + /// Separate finalize instance for post_aux_child — same Hopper CUfunction isolation. + grad_norm_finalize_post_aux: CudaFunction, /// Separate instance of grad_norm for non-graph launches (outside CUDA graph). /// CUDA doesn't allow the same CUfunction to be launched both inside a captured /// graph and outside on the same stream. grad_norm_standalone: CudaFunction, adam_update_kernel: CudaFunction, + /// Separate instances for post_aux_child — each from its own CUmodule. + /// CUfunction sharing across child graphs corrupts kernel state on Hopper, + /// causing 3100ms replay. Every child gets isolated CUfunction handles. + adam_update_post_aux: CudaFunction, + scale_f32_post_aux: CudaFunction, ema_kernel: CudaFunction, saxpy_kernel: CudaFunction, saxpy_f32_kernel: CudaFunction, @@ -2532,7 +2539,7 @@ impl GpuDqnTrainer { let risk_ptr = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 64); unsafe { - self.stream.launch_builder(&self.scale_f32_kernel) + self.stream.launch_builder(&self.scale_f32_post_aux) .arg(&risk_ptr) .arg(&decay) .arg(&(risk_aligned_count as i32)) @@ -4081,7 +4088,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile utility kernels (grad_norm, adam_update, etc.) ─ - let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, adam_update_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel) = + let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (outside CUDA graph). @@ -5680,8 +5687,11 @@ impl GpuDqnTrainer { grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, + grad_norm_finalize_post_aux, grad_norm_standalone, adam_update_kernel, + adam_update_post_aux, + scale_f32_post_aux, ema_kernel, saxpy_kernel, saxpy_f32_kernel, @@ -7834,7 +7844,7 @@ impl GpuDqnTrainer { let nb = 1_i32; unsafe { self.stream - .launch_builder(&self.grad_norm_finalize_kernel) + .launch_builder(&self.grad_norm_finalize_post_aux) .arg(&partials_ptr) .arg(&norm_out_ptr) .arg(&nb) @@ -7865,7 +7875,7 @@ impl GpuDqnTrainer { let blocks = ((DENOISE_TOTAL_PARAMS as u32 + 255) / 256).max(1); unsafe { self.stream - .launch_builder(&self.adam_update_kernel) + .launch_builder(&self.adam_update_post_aux) .arg(¶ms_ptr) .arg(&grad_ptr) .arg(&m_ptr) @@ -10179,7 +10189,7 @@ impl GpuDqnTrainer { let nb = 1_i32; unsafe { self.stream - .launch_builder(&self.grad_norm_finalize_kernel) + .launch_builder(&self.grad_norm_finalize_post_aux) .arg(&partials_ptr) .arg(&norm_out_ptr) .arg(&nb) @@ -10210,7 +10220,7 @@ impl GpuDqnTrainer { let blocks = ((sel_dim as u32 + 255) / 256).max(1); unsafe { self.stream - .launch_builder(&self.adam_update_kernel) + .launch_builder(&self.adam_update_post_aux) .arg(¶ms_ptr) .arg(&grad_ptr) .arg(&m_ptr) @@ -10250,7 +10260,7 @@ impl GpuDqnTrainer { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -10272,6 +10282,16 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("adam finalize cubin load: {e}")))?; let grad_norm_finalize_adam = adam_finalize_module.load_function("dqn_grad_norm_finalize") .map_err(|e| MLError::ModelError(format!("adam grad_norm_finalize load: {e}")))?; + // Separate CUmodule for post_aux_child — adam_update + grad_norm_finalize + // must NOT share CUfunction handles with adam_child or forward_child on Hopper. + let post_aux_module = context.load_cubin(DQN_UTILITY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("post_aux cubin load: {e}")))?; + let grad_norm_finalize_post_aux = post_aux_module.load_function("dqn_grad_norm_finalize") + .map_err(|e| MLError::ModelError(format!("post_aux grad_norm_finalize load: {e}")))?; + let adam_update_post_aux = post_aux_module.load_function("dqn_adam_update_kernel") + .map_err(|e| MLError::ModelError(format!("post_aux adam_update load: {e}")))?; + let scale_f32_post_aux = post_aux_module.load_function("dqn_scale_f32_kernel") + .map_err(|e| MLError::ModelError(format!("post_aux scale_f32 load: {e}")))?; let adam_update = module.load_function("dqn_adam_update_kernel") .map_err(|e| MLError::ModelError(format!("dqn_adam_update_kernel load: {e}")))?; // f32_to_bf16_kernel / bf16_to_f32_kernel DELETED — pure f32 pipeline. @@ -10333,7 +10353,7 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("popart_normalize_kernel load: {e}")))?; info!("GpuDqnTrainer: 30 utility kernels loaded from precompiled cubin"); - Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, adam_update, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize)) + Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 9d446869a..d664dc250 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -670,9 +670,13 @@ impl GpuIqlTrainer { // ── Backward pass (5 GEMMs + 2 SiLU backward + bias grad reduce) ── - // Zero grad_buf before accumulating weight gradients - self.stream.memset_zeros(&mut self.grad_buf) - .map_err(|e| MLError::ModelError(format!("IQL zero grad_buf: {e}")))?; + // Zero grad_buf — raw cuMemsetD8Async, no cudarc overhead during graph capture + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.grad_buf.raw_ptr(), 0, self.grad_buf.num_bytes(), + self.stream.cu_stream(), + ); + } // inv_batch for 1/N mean reduction let inv_batch: f32 = 1.0 / b as f32; diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 8a12f5fbb..a8c4b41f1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -990,9 +990,14 @@ impl GpuIqnHead { } // ── Step 6: Quantile Huber loss + dq gradients ────────────────── - // Zero d_branch_logits_buf before loss kernel writes to it - effective_stream.memset_zeros(&mut self.d_branch_logits_buf) - .map_err(|e| MLError::ModelError(format!("IQN zero d_branch_logits: {e}")))?; + // Zero d_branch_logits_buf — raw cuMemsetD8Async, no cudarc overhead during graph capture + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.d_branch_logits_buf.raw_ptr(), 0, + self.d_branch_logits_buf.num_bytes(), + effective_stream.cu_stream(), + ); + } let kappa = self.config.kappa; let gamma = self.config.gamma; @@ -1094,12 +1099,17 @@ impl GpuIqnHead { // Now backward through branches: let inv_bq = 1.0_f32 / bq as f32; - // Zero grad_buf before accumulating weight gradients - effective_stream.memset_zeros(&mut self.grad_buf) - .map_err(|e| MLError::ModelError(format!("IQN zero grad_buf: {e}")))?; - // Zero d_combined_buf for accumulation - effective_stream.memset_zeros(&mut self.d_combined_buf) - .map_err(|e| MLError::ModelError(format!("IQN zero d_combined: {e}")))?; + // Zero grad_buf + d_combined_buf — raw cuMemsetD8Async, no cudarc overhead + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.grad_buf.raw_ptr(), 0, self.grad_buf.num_bytes(), + effective_stream.cu_stream(), + ); + cudarc::driver::sys::cuMemsetD8Async( + self.d_combined_buf.raw_ptr(), 0, self.d_combined_buf.num_bytes(), + effective_stream.cu_stream(), + ); + } { let mut row_off = 0usize;