fix: CUfunction isolation for post_aux child + raw memset in IQN/IQL/Attention

ROOT CAUSE of 3100ms adam_child: adam_update_kernel and grad_norm_finalize_kernel
were shared between post_aux_child and adam_update/forward children. On Hopper,
CUfunction sharing across child graphs corrupts kernel state.

Fix: load utility cubin from separate CUmodule for post_aux child, giving it
isolated adam_update_post_aux, grad_norm_finalize_post_aux, and scale_f32_post_aux
handles. Zero cross-child CUfunction sharing for these kernels.

Also: convert memset_zeros to raw cuMemsetD8Async in IQN (3 calls), IQL (1),
Attention (1) — eliminates cudarc device_ptr_mut() overhead during graph capture.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/gpu_attention.rs

@@ -606,9 +606,13 @@ impl GpuAttention {
         let f32_sz = std::mem::size_of::<f32>();
         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;

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(&param_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(&params_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(&params_ptr)
                 .arg(&grad_ptr)
                 .arg(&m_ptr)
@@ -10250,7 +10260,7 @@ impl GpuDqnTrainer {
 fn compile_training_kernels(
     stream: &Arc<CudaStream>,
     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.

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;

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;