fix(critical): frozen c51_alpha in CUDA Graph + corrupted grad_norm handle

Two critical bugs causing gradient collapse at epoch 3-4 on H100
(batch=16384) while smoke test (batch=58) works:

1. c51_alpha baked as literal scalar in CUDA Graph at capture time.
   set_c51_alpha() updated the CPU field but the graphed SAXPY blend
   kernels kept using the frozen capture-time value (~0.01). The C51/MSE
   gradient blend NEVER ramped — model trained with pure MSE forever.
   Fix: invalidate graph_mega + graph_forward when alpha changes.
   Graph re-captures on next step with the correct blend.

2. grad_norm_standalone loaded but never wired in. clip_grad_buf_inplace
   used the CAPTURED grad_norm_kernel CUfunction outside graph context.
   CUDA forbids launching a captured CUfunction outside its graph —
   silently corrupts kernel state. grad_norm buffer contains garbage,
   Adam reads garbage clip scale, gradients effectively zero.
   Fix: use grad_norm_standalone (separate CUfunction handle) for
   all non-graph gradient norm computations.

Root cause of every H100 gradient collapse this session.

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

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -2053,8 +2053,11 @@ impl GpuDqnTrainer {
         // Disable cudarc event tracking — graph replay leaves stale events
         // that cause INVALID_VALUE on post-graph kernel launches.
         let _evt_guard = EventTrackingGuard::new(self.stream.context());
-        // Two-phase grad_norm — no memset needed
-        self.launch_grad_norm()?;
+        // Two-phase grad_norm using STANDALONE kernel handle (NOT the captured one).
+        // After graph_mega capture, the captured grad_norm_kernel CUfunction cannot be
+        // launched outside the graph context — CUDA corrupts kernel state silently.
+        // grad_norm_standalone is a separately loaded handle of the same kernel.
+        self.launch_grad_norm_standalone()?;
         self.launch_grad_norm_finalize()?;
 
         // Clip in-place
@@ -4697,10 +4700,18 @@ impl GpuDqnTrainer {
     /// `alpha = 1.0`: pure C51 distributional loss (converged).
     /// Intermediate values blend: `grad = alpha*C51 + (1-alpha)*MSE`.
     ///
-    /// No CUDA Graph invalidation — the kernel sequence is always the same
-    /// (both MSE and C51 run every step); only the SAXPY scale factors change.
+    /// INVALIDATES graph_mega — alpha/scale_mse are passed as literal scalar
+    /// arguments to the SAXPY blend kernels inside the graph. CUDA Graph bakes
+    /// scalar args at capture time. Without invalidation, the blend is frozen
+    /// at the capture-time alpha forever. The graph is re-captured on next step.
     pub fn set_c51_alpha(&mut self, alpha: f32) {
-        self.c51_alpha = alpha.clamp(0.0, 1.0);
+        let new_alpha = alpha.clamp(0.0, 1.0);
+        if (new_alpha - self.c51_alpha).abs() > 1e-6 {
+            self.graph_forward = None;
+            self.graph_forward_ddqn = None;
+            // graph_mega invalidated via the caller (FusedTrainingCtx)
+        }
+        self.c51_alpha = new_alpha;
     }
 
     /// Blend MSE and C51 loss values with NaN-safe alpha guard.
@@ -5640,6 +5651,36 @@ impl GpuDqnTrainer {
         Ok(())
     }
 
+    /// Launch grad norm using the STANDALONE kernel handle (not the captured one).
+    /// Must be used outside CUDA Graph context (e.g. clip_grad_buf_inplace).
+    /// After graph_mega capture, the captured CUfunction handle cannot be launched
+    /// outside the graph — CUDA silently corrupts kernel state.
+    fn launch_grad_norm_standalone(&self) -> Result<(), MLError> {
+        let tp = self.total_params as i32;
+        let blocks = self.grad_norm_blocks as u32;
+
+        let grad_ptr = self.ptrs.grad_buf;
+        let partials_ptr = self.grad_norm_partials.raw_ptr();
+
+        unsafe {
+            self.stream
+                .launch_builder(&self.grad_norm_standalone)
+                .arg(&grad_ptr)
+                .arg(&partials_ptr)
+                .arg(&tp)
+                .launch(LaunchConfig {
+                    grid_dim: (blocks, 1, 1),
+                    block_dim: (256, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| {
+                    MLError::ModelError(format!("dqn_grad_norm_standalone launch: {e}"))
+                })?;
+        }
+
+        Ok(())
+    }
+
     /// Launch the Adam update kernel with correct gradient clipping.
     ///
     /// Argument order matches `dqn_adam_update_kernel` in `dqn_training_kernel.cu`:

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

@@ -1626,9 +1626,13 @@ impl FusedTrainingCtx {
     /// Set C51 blend factor for gradual MSE→C51 ramp.
     ///
     /// `alpha = 0.0`: pure MSE (warmup). `alpha = 1.0`: pure C51 (converged).
-    /// No graph invalidation — kernel sequence is fixed (both always run).
+    /// Invalidates graph_mega — alpha is a scalar arg baked at capture time.
     pub(crate) fn set_c51_alpha(&mut self, alpha: f32) {
+        let prev = self.trainer.c51_alpha();
         self.trainer.set_c51_alpha(alpha);
+        if (alpha - prev).abs() > 1e-6 {
+            self.graph_mega = None; // Force recapture with new alpha
+        }
     }
 
     /// Return a reference to the GPU-resident loss scalar (f32) for the training guard.