perf: capture spectral norm in graph_forward + EMA in graph_ema

Spectral norm (10 weight matrices):
- Moved into graph_forward capture (submit_spectral_norm_ops)
- Runs before cuBLAS forward — graph ensures normalized weights
- All device_ptr() → raw_ptr() (no cudarc events)
- ~10 kernel launches eliminated from per-step ungraphed ops

EMA target update:
- Captured as graph_ema (EMA kernel + f32→bf16 cast + 20 unflatten DtoD)
- tau read from GPU-resident tau_buf (async HtoD before replay)
- ~22 kernel launches (EMA + cast + 20 DtoD) → 1 graph replay

Per-step kernel launches now: 7 graph replays + ~4 ungraphed
(was: 5 graph replays + ~14 ungraphed)

Remaining ungraphed:
- upload_batch_gpu (8 ops — batch ptrs change per step)
- HER relabel (2 ops — donor indices change)
- IQN trunk gradient (2 ops — tau changes)
- IQN→PER DtoD (1 op)
- causal/vaccine (amortized, rare)
- regime_scale + PER update (2 ops)

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

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -929,6 +929,8 @@ impl GpuDqnTrainer {
     ///
     /// Shape: `[B]` f32 — rewards from the batch. Valid after `train_step()` or
     /// `train_step_gpu()`. Used by IQN for Bellman target computation on GPU.
+    pub fn tau_buf_ptr(&self) -> u64 { self.tau_buf.raw_ptr() }
+
     pub fn rewards_buf(&self) -> &CudaSlice<f32> {
         &self.rewards_buf
     }
@@ -2833,8 +2835,8 @@ impl GpuDqnTrainer {
     pub fn train_step_gpu(
         &mut self,
         gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
-        online_dueling: &DuelingWeightSet,
-        online_branching: &BranchingWeightSet,
+        online_dueling: &mut DuelingWeightSet,
+        online_branching: &mut BranchingWeightSet,
         _target_dueling: &DuelingWeightSet,
         _target_branching: &BranchingWeightSet,
     ) -> Result<FusedTrainScalars, MLError> {
@@ -3026,9 +3028,10 @@ impl GpuDqnTrainer {
     /// - (injection point: external code adds IQN/attention/ensemble gradients to grad_buf)
     /// - `graph_adam`: grad_norm → Adam → unflatten
     fn execute_train_and_readback(
+
         &mut self,
-        online_dueling: &DuelingWeightSet,
-        online_branching: &BranchingWeightSet,
+        online_dueling: &mut DuelingWeightSet,
+        online_branching: &mut BranchingWeightSet,
     ) -> Result<FusedTrainResult, MLError> {
         let b = self.config.batch_size;
 
@@ -3426,8 +3429,8 @@ impl GpuDqnTrainer {
     /// by `update_priorities_cuda()` which reads them directly via device pointer.
     fn execute_train_scalars_only(
         &mut self,
-        online_dueling: &DuelingWeightSet,
-        online_branching: &BranchingWeightSet,
+        online_dueling: &mut DuelingWeightSet,
+        online_branching: &mut BranchingWeightSet,
     ) -> Result<FusedTrainScalars, MLError> {
         // ── Update Adam step counter on device (OUTSIDE graph) ───────
         self.adam_step += 1;
@@ -3877,8 +3880,8 @@ impl GpuDqnTrainer {
     /// they are pure element-wise ops with fixed control flow.
     fn capture_training_graphs(
         &mut self,
-        online_d: &DuelingWeightSet,
-        online_b: &BranchingWeightSet,
+        online_d: &mut DuelingWeightSet,
+        online_b: &mut BranchingWeightSet,
     ) -> Result<(), MLError> {
         self.stream.synchronize()
             .map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?;
@@ -3895,6 +3898,7 @@ impl GpuDqnTrainer {
             return Err(MLError::ModelError(format!("CUDA graph_forward begin_capture: {e}")));
         }
 
+        self.submit_spectral_norm_ops(online_d, online_b)?;
         let submit_fwd_result = self.submit_forward_ops();
 
         let graph_fwd_result = self.stream.end_capture(
@@ -4005,6 +4009,51 @@ impl GpuDqnTrainer {
         Ok(())
     }
 
+    /// Submit spectral norm ops for graph capture.
+    /// Normalizes all weight matrices and syncs back to params_buf.
+    /// Must run BEFORE submit_forward_ops so cuBLAS reads normalized weights.
+    fn submit_spectral_norm_ops(
+        &mut self,
+        online_dueling: &mut DuelingWeightSet,
+        online_branching: &mut BranchingWeightSet,
+    ) -> Result<(), MLError> {
+        self.apply_spectral_norm(online_dueling, online_branching)
+    }
+
+    /// Submit EMA + bf16 cast + unflatten ops for graph capture.
+    /// tau comes from GPU-resident tau_buf (async HtoD before replay).
+    pub fn submit_ema_ops(
+        &mut self,
+        target_d: &mut DuelingWeightSet,
+        target_b: &mut BranchingWeightSet,
+    ) -> Result<(), MLError> {
+        // EMA kernel: target_params_buf = (1-tau)*target + tau*online
+        let n = self.total_params as i32;
+        let blocks = ((self.total_params + 255) / 256) as u32;
+        let t_ptr = self.target_params_buf.raw_ptr();
+        let o_ptr = self.params_buf.raw_ptr();
+        let tau_ptr = self.tau_buf.raw_ptr();
+        unsafe {
+            self.stream
+                .launch_builder(&self.ema_kernel)
+                .arg(&t_ptr)
+                .arg(&o_ptr)
+                .arg(&tau_ptr)
+                .arg(&n)
+                .launch(LaunchConfig {
+                    grid_dim: (blocks, 1, 1),
+                    block_dim: (256, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("ema capture: {e}")))?;
+        }
+        // f32 → bf16 cast
+        self.launch_f32_to_bf16_target_cast()?;
+        // Unflatten to individual weight tensors
+        self.unflatten_target_weights(target_d, target_b)?;
+        Ok(())
+    }
+
     fn submit_forward_ops(&mut self) -> Result<(), MLError> {
         // ── Zero accumulators (capturable: memset_zeros uses cuMemsetD32Async) ─
         self.stream

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

@@ -109,6 +109,8 @@ pub(crate) struct FusedTrainingCtx {
     /// CUDA Graph for IQN training (decode + fwd/loss + bwd + adam + trunk grad).
     /// Captured lazily on first step.
     graph_iqn: Option<crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph>,
+    /// CUDA Graph for EMA target update + bf16 cast + unflatten.
+    graph_ema: Option<crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph>,
     /// Ensemble extra heads (heads 1..K-1). Head 0 lives inside the CUDA Graph.
     /// Each element is an independent (DuelingWeightSet, BranchingWeightSet) pair
     /// that shares the same trunk but has perturbed value/advantage weights.
@@ -510,6 +512,7 @@ impl FusedTrainingCtx {
             graph_attention: None,
             graph_iql: None,
             graph_iqn: None,
+            graph_ema: None,
             ensemble_extra_heads,
             ensemble_diversity_weight: hyperparams.ensemble_diversity_weight as f32,
             ensemble_aggregate_kernel,
@@ -576,17 +579,14 @@ impl FusedTrainingCtx {
         let gpu_batch = batch.gpu_batch.as_ref()
             .ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?;
 
-        // ── Step 1: Spectral normalization BEFORE forward pass ─────────
-        self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching)
-            .map_err(|e| anyhow::anyhow!("Spectral norm (pre-forward): {e}"))?;
 
         // ── Step 2: Upload batch + replay graph_forward ONLY ─────────────
         // Upload the RAW PER batch into trainer staging buffers (DtoD).
         // HER relabeling happens AFTER upload, modifying the staging buffers in-place.
         // This eliminates the intermediate GpuBatch (was 17+ GPU allocs per step).
-        let _fused_placeholder = self.trainer.train_step_gpu(
+        let _step_scalars = self.trainer.train_step_gpu(
             gpu_batch,
-            &self.online_dueling, &self.online_branching,
+            &mut self.online_dueling, &mut self.online_branching,
             &self.target_dueling, &self.target_branching,
         ).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?;
 
@@ -635,21 +635,54 @@ impl FusedTrainingCtx {
         // C51 gradient clip now captured in graph_adam — no per-step call needed.
 
         // ── Step 3: GPU-native Polyak EMA target update ──────────────────
+        // tau uploaded async to GPU buffer before graph replay.
         {
             let dqn = agent.primary_dqn_mut();
-
             let tau = compute_cosine_annealed_tau(
                 dqn.get_training_steps(),
-                dqn.config.tau,
-                dqn.config.tau_final,
-                dqn.config.tau_anneal_steps,
+                dqn.config.tau, dqn.config.tau_final, dqn.config.tau_anneal_steps,
             );
+            // Async HtoD for tau (graph reads from tau_buf device address)
+            unsafe {
+                cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
+                    self.trainer.tau_buf_ptr(),
+                    (&(tau as f32) as *const f32).cast(),
+                    std::mem::size_of::<f32>(),
+                    self.stream.cu_stream(),
+                );
+            }
 
-            self.trainer.target_ema_update(
-                &self.online_dueling, &self.online_branching,
-                &mut self.target_dueling, &mut self.target_branching,
-                tau as f32,
-            ).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?;
+            if self.graph_ema.is_none() {
+                // First step: run EMA ungraphed to initialize
+                self.trainer.target_ema_update(
+                    &self.online_dueling, &self.online_branching,
+                    &mut self.target_dueling, &mut self.target_branching,
+                    tau as f32,
+                ).map_err(|e| anyhow::anyhow!("EMA first step: {e}"))?;
+
+                // Capture EMA graph
+                self.stream.synchronize()
+                    .map_err(|e| anyhow::anyhow!("sync before ema capture: {e}"))?;
+                unsafe { self.stream.context().disable_event_tracking(); }
+                if let Ok(()) = self.stream.begin_capture(
+                    cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
+                ) {
+                    let _ = self.trainer.submit_ema_ops(
+                        &mut self.target_dueling, &mut self.target_branching,
+                    );
+                    if let Ok(Some(graph)) = self.stream.end_capture(
+                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
+                    ) {
+                        use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph;
+                        self.graph_ema = Some(SendSyncGraph(graph));
+                        tracing::info!("Captured EMA CUDA graph (ema + bf16 cast + unflatten)");
+                    }
+                }
+                unsafe { self.stream.context().enable_event_tracking(); }
+                let _ = self.stream.context().check_err();
+            } else if let Some(ref graph) = self.graph_ema {
+                graph.0.launch().map_err(|e| anyhow::anyhow!("EMA graph replay: {e}"))?;
+            }
         }
 
         // ── Step 3b-5b: Auxiliary heads (attention + IQL + IQN) ─────────