diag: RECAPTURE_DIAG — per-buffer norms after graph recapture

Fires for 3 steps after any CUDA graph recapture. Reports:
- d_val/adv_norm: blended d_logits (C51 × alpha + MSE × (1-alpha))
- mse_val/adv_norm: MSE-only scratch buffers
- grad_norm: backward output grad_buf

Purpose: diagnose C51 grad_norm=0 on H100 at batch=16384 (epoch 2+).

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

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -595,6 +595,7 @@ pub struct GpuDqnTrainer {
     target_params_bf16: CudaSlice<half::bf16>, // [TOTAL_PARAMS + pad] bf16 shadow for GemmEx
     pub(crate) grad_norm_buf: CudaSlice<half::bf16>,    // [1] bf16 L2 norm (written by finalize)
     grad_norm_f32_buf: CudaSlice<f32>,                 // [1] float sum-of-squares output
+    diag_recapture_remaining: u32,                     // diagnostic steps remaining after graph recapture
     grad_norm_partials: CudaSlice<f32>,                // [grad_norm_blocks] per-block partial sums
     grad_norm_blocks: usize,                           // number of blocks for grad_norm kernel
     cql_grad_scratch: CudaSlice<f32>,        // [TOTAL_PARAMS] f32 CQL gradient isolation buffer
@@ -2113,6 +2114,37 @@ impl GpuDqnTrainer {
         );
     }
 
+    /// Extended diagnostics after graph recapture: checks MSE scratch, C51 d_logits,
+    /// blended d_logits, and grad_buf. Fires for first 3 steps after any graph recapture.
+    pub fn run_recapture_diagnostics(&mut self) {
+        if self.diag_recapture_remaining == 0 { return; }
+        self.diag_recapture_remaining -= 1;
+
+        let b = self.config.batch_size;
+        let na = self.config.num_atoms;
+        let tba = (self.config.branch_0_size + self.config.branch_1_size
+            + self.config.branch_2_size + self.config.branch_3_size) * na;
+
+        // After graph replay: d_value/adv_logits contain BLENDED gradients,
+        // MSE scratch contains MSE-only gradients, grad_buf contains backward output.
+        let d_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_buf.raw_ptr(), b * na).unwrap_or(f32::NAN);
+        let d_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN);
+        let mse_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_mse.raw_ptr(), b * na).unwrap_or(f32::NAN);
+        let mse_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_mse.raw_ptr(), b * tba).unwrap_or(f32::NAN);
+        let grad_norm = self.debug_buffer_norm_f32(self.grad_buf.raw_ptr(), self.total_params).unwrap_or(f32::NAN);
+
+        tracing::warn!(
+            remaining = self.diag_recapture_remaining,
+            c51_alpha = self.c51_alpha,
+            d_val_norm,
+            d_adv_norm,
+            mse_val_norm,
+            mse_adv_norm,
+            grad_norm,
+            "RECAPTURE_DIAG: buffer norms after graph recapture replay"
+        );
+    }
+
     /// Compute L2 norm of an arbitrary f32 GPU buffer (diagnostic use only).
     /// Reuses the standalone grad_norm two-phase reduction. Costs 1 stream sync.
     pub fn debug_buffer_norm_f32(&mut self, ptr: u64, n_elems: usize) -> Result<f32, MLError> {
@@ -2997,6 +3029,7 @@ impl GpuDqnTrainer {
             target_params_bf16,
             grad_norm_buf,
             grad_norm_f32_buf,
+            diag_recapture_remaining: 3,
             grad_norm_partials,
             grad_norm_blocks,
             cql_grad_scratch,
@@ -4511,6 +4544,7 @@ impl GpuDqnTrainer {
         self.graph_forward_ddqn = Some(SendSyncGraph(graph_ddqn));
         self.graph_adam = Some(SendSyncGraph(graph_adam));
         self.last_captured_loss_mode = Some(self.loss_mode);
+        self.diag_recapture_remaining = 3; // fire diagnostics for first 3 steps after recapture
         Ok(())
     }
 

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

@@ -834,10 +834,12 @@ impl FusedTrainingCtx {
         }
         self.pending_vaccine_batch = None;
 
-        // ── BUFFER_DIAG: check where gradients die (first 3 steps each epoch) ──
+        // ── BUFFER_DIAG: check where gradients die (first 3 steps each fold) ──
         if self.steps_since_varmap_sync < 3 {
             self.trainer.run_buffer_diagnostics(self.steps_since_varmap_sync);
         }
+        // ── RECAPTURE_DIAG: check C51 gradient path after graph recapture ──
+        self.trainer.run_recapture_diagnostics();
 
         // ── Step 5: Pruning + Adam ───────────────────────────────────────
         self.trainer.apply_pruning_mask()