fix(bf16): standalone grad_norm kernel for non-graph launches

Load separate dqn_grad_norm_kernel instance (grad_norm_standalone)
from fresh cubin for clip_grad_buf_inplace and read_grad_norm_sync.
CUDA graph captures function handles — same CUfunction can't be
used both inside captured graph and outside on same stream.

Smoke test status: experience collection + CUDA graph capture +
forward + loss + backward all pass. Remaining blocker:
non-graph kernel launch after graph replay (needs investigation
of cudarc event tracking interaction with CUDA graphs).

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

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -366,6 +366,10 @@ pub struct GpuDqnTrainer {
     // Dead kernels (forward_loss_kernel, forward_only_kernel, backward_kernel)
     // removed — replaced by cuBLAS forward/backward. See Task 2 for method cleanup.
     grad_norm_kernel: CudaFunction,
+    /// Separate instance of grad_norm for non-graph launches (clip_grad_buf_inplace).
+    /// 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,
     ema_kernel: CudaFunction,
     per_update_kernel: CudaFunction,
@@ -1320,7 +1324,7 @@ impl GpuDqnTrainer {
 
         unsafe {
             self.stream
-                .launch_builder(&self.grad_norm_kernel)
+                .launch_builder(&self.grad_norm_standalone)
                 .arg(&scratch_ptr)
                 .arg(&norm_ptr)
                 .arg(&total)
@@ -1517,10 +1521,12 @@ impl GpuDqnTrainer {
     /// from overwhelming the subsequent auxiliary gradient additions.
     /// All operations are async on the same stream — zero CPU sync.
     pub fn clip_grad_buf_inplace(&mut self, max_norm: f32) -> Result<(), MLError> {
+        // clip_grad_buf_inplace: standalone kernel launch (not captured in CUDA graph)
         // Zero the grad_norm accumulator
         self.stream
             .memset_zeros(&mut self.grad_norm_buf)
             .map_err(|e| MLError::ModelError(format!("zero grad_norm for clip: {e}")))?;
+        // memset_zeros succeeded
 
         // Compute grad_norm (sum of squares)
         self.launch_grad_norm()?;
@@ -1693,6 +1699,16 @@ impl GpuDqnTrainer {
         let (grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) =
             compile_training_kernels(&stream, &config)?;
 
+        // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace).
+        // CUDA graph captures CUfunction handles — the same handle can't be launched
+        // both inside a replayed graph and outside on the same stream.
+        let grad_norm_standalone = {
+            let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec())
+                .map_err(|e| MLError::ModelError(format!("grad_norm_standalone cubin: {e}")))?;
+            module.load_function("dqn_grad_norm_kernel")
+                .map_err(|e| MLError::ModelError(format!("grad_norm_standalone load: {e}")))?
+        };
+
         // ── Compile EMA kernel (standalone — not captured in CUDA Graph) ──
         let ema_kernel = compile_ema_kernel(&stream)?;
         let relu_mask_standalone = compile_relu_mask_standalone(&stream)?;
@@ -2006,6 +2022,7 @@ impl GpuDqnTrainer {
             stream,
             ptrs,
             grad_norm_kernel,
+            grad_norm_standalone,
             adam_update_kernel,
             ema_kernel,
             per_update_kernel,
@@ -3554,6 +3571,7 @@ impl GpuDqnTrainer {
     fn launch_grad_norm(&self) -> Result<(), MLError> {
         let tp = self.total_params as i32;
         let blocks = ((self.total_params + 255) / 256) as u32;
+            tp, blocks, self.ptrs.grad_buf, self.ptrs.grad_norm_buf);
 
         let launch_cfg = LaunchConfig {
             grid_dim: (blocks, 1, 1),
@@ -3568,7 +3586,7 @@ impl GpuDqnTrainer {
         // grad_buf has size = total_params; grad_norm_buf has size = 1.
         unsafe {
             self.stream
-                .launch_builder(&self.grad_norm_kernel)
+                .launch_builder(&self.grad_norm_standalone)
                 .arg(&grad_ptr)
                 .arg(&norm_ptr)
                 .arg(&tp)

crates/ml/src/trainers/dqn/fused_training.rs

@@ -602,6 +602,7 @@ impl FusedTrainingCtx {
             let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac;
             let c51_budget = self.trainer.config().max_grad_norm * c51_frac;
 
+
             self.trainer.clip_grad_buf_inplace(c51_budget)
                 .map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?;
         }