From 80d871ae1d28328482dd368ead530fb54a803e44 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sat, 18 Apr 2026 16:37:54 +0200
Subject: [PATCH] feat: IQN prepare_buffers extraction for parallel stream
 dispatch

Separates buffer prep (decode_actions, DtoD copies) from training pipeline
so prep runs on main stream and execute_training_pipeline runs on iqn_stream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 57 +++++++++++++--------
 1 file changed, 36 insertions(+), 21 deletions(-)

diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
index 1c03d514e..8a12f5fbb 100644
--- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
@@ -574,30 +574,17 @@ impl GpuIqnHead {
         })
     }
 
-    /// Run one GPU-native IQN training step.
+    /// Prepare IQN buffers: decode flat actions and DtoD copy rewards/dones.
     ///
-    /// All data sourced from DQN trainer's GPU buffers via DtoD copies.
-    /// No CPU arrays, no HtoD uploads for batch data.
-    ///
-    /// Full pipeline:
-    /// 1. Decode flat actions → branch actions on GPU
-    /// 2. DtoD copy rewards/dones from DQN trainer buffers
-    /// 3. Compute target h_s2 (or use cached)
-    /// 4. Forward (cuBLAS GEMMs + element-wise kernels)
-    /// 5. Loss (quantile Huber)
-    /// 6. Backward (cuBLAS GEMMs + element-wise kernels)
-    /// 7. Grad norm + Adam
-    ///
-    /// Returns the batch-mean quantile Huber loss.
-    pub fn train_iqn_step_gpu(
+    /// Runs on the IQN head's own stream (the main training stream). Called
+    /// before `execute_training_pipeline` so that a dedicated override stream
+    /// can safely read `branch_actions`, `rewards_buf`, and `dones_buf`.
+    pub fn prepare_buffers(
         &mut self,
-        online_h_s2: &CudaSlice<f32>,
-        next_states_buf: &CudaSlice<f32>,
-        target_dueling: &DuelingWeightSet,
         dqn_actions_buf: &CudaSlice<i32>,
         dqn_rewards_buf: &CudaSlice<f32>,
         dqn_dones_buf: &CudaSlice<f32>,
-    ) -> Result<f32, MLError> {
+    ) -> Result<(), MLError> {
         let b = self.config.batch_size;
         let shared_h1_i32 = self.config.shared_h1 as i32;
         let hidden_dim_i32 = self.config.hidden_dim as i32;
@@ -607,7 +594,6 @@ impl GpuIqnHead {
         let b2_i32 = self.config.branch_2_size as i32;
         let b3_i32 = self.config.branch_3_size as i32;
 
-        // 1. GPU decode: flat actions [B] → branch actions [B, 4]
         let decode_blocks = (b + 255) / 256;
         let decode_config = LaunchConfig {
             grid_dim: (decode_blocks as u32, 1, 1),
@@ -633,7 +619,6 @@ impl GpuIqnHead {
                 .map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
         }
 
-        // 2. DtoD copy rewards/dones from DQN trainer buffers (f32)
         use super::gpu_dqn_trainer::dtod_copy;
         let f32_bytes = std::mem::size_of::<f32>();
         let r_src = dqn_rewards_buf.raw_ptr();
@@ -644,6 +629,36 @@ impl GpuIqnHead {
         let d_dst = self.dones_buf.raw_ptr();
         dtod_copy(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?;
 
+        Ok(())
+    }
+
+    /// Run one GPU-native IQN training step.
+    ///
+    /// All data sourced from DQN trainer's GPU buffers via DtoD copies.
+    /// No CPU arrays, no HtoD uploads for batch data.
+    ///
+    /// Full pipeline:
+    /// 1. Decode flat actions → branch actions on GPU
+    /// 2. DtoD copy rewards/dones from DQN trainer buffers
+    /// 3. Compute target h_s2 (or use cached)
+    /// 4. Forward (cuBLAS GEMMs + element-wise kernels)
+    /// 5. Loss (quantile Huber)
+    /// 6. Backward (cuBLAS GEMMs + element-wise kernels)
+    /// 7. Grad norm + Adam
+    ///
+    /// Returns the batch-mean quantile Huber loss.
+    pub fn train_iqn_step_gpu(
+        &mut self,
+        online_h_s2: &CudaSlice<f32>,
+        next_states_buf: &CudaSlice<f32>,
+        target_dueling: &DuelingWeightSet,
+        dqn_actions_buf: &CudaSlice<i32>,
+        dqn_rewards_buf: &CudaSlice<f32>,
+        dqn_dones_buf: &CudaSlice<f32>,
+    ) -> Result<f32, MLError> {
+        // 1-2. Decode actions + DtoD copy rewards/dones
+        self.prepare_buffers(dqn_actions_buf, dqn_rewards_buf, dqn_dones_buf)?;
+
         // 3-9. Execute core training pipeline
         self.execute_training_pipeline(online_h_s2, next_states_buf, target_dueling, None, None)
     }