From 4295cbde795b2f8296d24e5ff7f337bfc4ee07a3 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sat, 18 Apr 2026 14:59:39 +0200
Subject: [PATCH] perf: add per-child graph CUDA event timing breakdown

Records cuEventRecord between each of the 5 child graph launches.
At epoch end, cuEventElapsedTime reports per-child GPU time:
spectral, forward, ddqn, aux, adam.

This tells us exactly which child dominates the 4s/step graph replay
cost, informing Phase 2 multi-stream parallelism priorities.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml/src/trainers/dqn/fused_training.rs | 23 ++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs
index 4e9a7e039..12ece82cc 100644
--- a/crates/ml/src/trainers/dqn/fused_training.rs
+++ b/crates/ml/src/trainers/dqn/fused_training.rs
@@ -960,9 +960,21 @@ impl FusedTrainingCtx {
         // 5 launches instead of 1, ~10µs overhead total, negligible vs step time.
         if self.spectral_child.is_some() {
             self.spectral_child.as_ref().unwrap().launch(cu_stream)?;
+            if let Some(ref pe) = self.phase_events {
+                PhaseEvents::record(pe.spectral_end, cu_stream);
+            }
             self.forward_child.as_ref().unwrap().launch(cu_stream)?;
+            if let Some(ref pe) = self.phase_events {
+                PhaseEvents::record(pe.forward_end, cu_stream);
+            }
             self.ddqn_child.as_ref().unwrap().launch(cu_stream)?;
+            if let Some(ref pe) = self.phase_events {
+                PhaseEvents::record(pe.loss_end, cu_stream);
+            }
             self.aux_child.as_ref().unwrap().launch(cu_stream)?;
+            if let Some(ref pe) = self.phase_events {
+                PhaseEvents::record(pe.backward_end, cu_stream);
+            }
             self.adam_child.as_ref().unwrap().launch(cu_stream)?;
         } else {
             // Ungraphed step 0: run everything, then capture.
@@ -1162,10 +1174,21 @@ impl FusedTrainingCtx {
         let per_update = elapsed(pe.per_update_start, pe.per_update_end);
         let total = upload + fwd_bwd + adam + per_update;
 
+        // Per-child graph breakdown (events recorded between launches)
+        let spectral_ms = elapsed(pe.fwd_bwd_start, pe.spectral_end);
+        let forward_ms = elapsed(pe.spectral_end, pe.forward_end);
+        let ddqn_ms = elapsed(pe.forward_end, pe.loss_end);
+        let aux_ms = elapsed(pe.loss_end, pe.backward_end);
+        let adam_child_ms = elapsed(pe.backward_end, pe.adam_end);
+
         tracing::info!(
             "GPU phase timing ({} batches, last batch): upload={:.1}ms fwd_bwd={:.1}ms adam={:.1}ms per_update={:.1}ms total={:.1}ms",
             self.phase_batch_count, upload, fwd_bwd, adam, per_update, total,
         );
+        tracing::info!(
+            "GPU child graph breakdown: spectral={:.1}ms forward={:.1}ms ddqn={:.1}ms aux={:.1}ms adam={:.1}ms",
+            spectral_ms, forward_ms, ddqn_ms, aux_ms, adam_child_ms,
+        );
 
         self.phase_batch_count = 0;
     }