From 3dde6d285c5ccb9143ba8d89504b383c9837cfd7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 10 Mar 2026 14:44:21 +0100 Subject: [PATCH] feat(cuda): GPU-resident action routing in select_actions_batch_gpu MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire route_exposure_to_factored() into both select_actions_batch_gpu and select_actions_batch GPU paths, eliminating per-item CPU routing. Previously, exposure indices (0-4) were downloaded from GPU and routed to factored indices (0-44) one-by-one on CPU via route_action(). Now the exposure→factored mapping runs entirely on GPU via the routing kernel, with a single batch readback of the final factored indices. Branching DQN path unchanged (already produces factored indices 0-44). Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/trainers/dqn/trainer.rs | 71 ++++++++++++++++----------- 1 file changed, 42 insertions(+), 29 deletions(-) diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 9f72d0555..095e46c6b 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -4116,17 +4116,26 @@ impl DQNTrainer { let action_indices_tensor = selector .select_actions(&batch_q_values, epsilon, batch_size, 5) .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {}", e))?; - // Single GPU->CPU readback for route_action() dispatch - let action_indices = action_indices_tensor + + // GPU-resident action routing: exposure (0-4) → factored (0-44) on-device. + // Single batch readback of factored indices — no per-item CPU routing. + let factored_tensor = selector.route_exposure_to_factored( + &action_indices_tensor, + batch_size, + self.hyperparams.avg_spread as f32, + self.hyperparams.avg_spread as f32, + self.vol_ema as f32, + self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))?; + + let factored_indices = factored_tensor .to_vec1::() - .map_err(|e| anyhow::anyhow!("Failed to readback action indices: {}", e))?; + .map_err(|e| anyhow::anyhow!("Factored indices readback: {e}"))?; let mut actions = Vec::with_capacity(batch_size); - for idx in &action_indices { - let action_idx = *idx as usize; - let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) - .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; - let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); + for idx in factored_indices { + let action = FactoredAction::from_index(idx as usize) + .map_err(|e| anyhow::anyhow!("Invalid factored index {}: {}", idx, e))?; actions.push(action); } return Ok(actions); @@ -4307,30 +4316,34 @@ impl DQNTrainer { .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {}", e))? }; - let action_indices = action_indices_tensor + // GPU-resident action routing: exposure (0-4) → factored (0-44) on-device. + // Branching already produces factored indices, so pass through. + // Single batch readback of factored indices — no per-item CPU routing. + let factored_tensor = if self.hyperparams.use_branching { + // Branching: already factored (0-44), pass through + action_indices_tensor + } else { + // Route exposure (0-4) → factored (0-44) on GPU + selector.route_exposure_to_factored( + &action_indices_tensor, + batch_size, + self.hyperparams.avg_spread as f32, + self.hyperparams.avg_spread as f32, + self.vol_ema as f32, + self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? + }; + + // Single batch readback (experience collection path, not training hot path) + let factored_indices = factored_tensor .to_vec1::() - .map_err(|e| anyhow::anyhow!("Failed to readback action indices: {}", e))?; + .map_err(|e| anyhow::anyhow!("Factored indices readback: {e}"))?; let mut actions = Vec::with_capacity(batch_size); - if self.hyperparams.use_branching { - // Branching: factored indices 0-44 - for &idx in &action_indices { - let action = FactoredAction::from_index(idx as usize) - .map_err(|e| anyhow::anyhow!("Invalid factored index {}: {}", idx, e))?; - actions.push(action); - } - } else { - // Standard: exposure indices 0-4 - for &idx in &action_indices { - let action_idx = idx as usize; - let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) - .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; - // When routed kernel was used, exposure is already post-fill - // (unfilled → Flat override done on GPU). Route to get the - // FactoredAction struct but the exposure level is final. - let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); - actions.push(action); - } + for idx in factored_indices { + let action = FactoredAction::from_index(idx as usize) + .map_err(|e| anyhow::anyhow!("Invalid factored index {}: {}", idx, e))?; + actions.push(action); } return Ok((actions, use_routed)); }