gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(cuda): GPU-resident action routing in select_actions_batch_gpu

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-10 14:44:21 +01:00
parent ce534c26a5
commit 3dde6d285c
1 changed files with 42 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -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::<u32>()
.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::<u32>()
.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));
}