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

perf(rl): precompute tree_rebuild_levels — eliminate host loop in mega-graph

Browse Source 
The per-step while loop computing (grid, start, nodes_at_level) for
each tree level was host-side work inside the mega-graph replay path.
Now precomputed at init: 16 entries for capacity=65536, 192 bytes.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-28 01:01:53 +02:00
parent 54de55d4bc
commit 346e6670f5
1 changed files with 37 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -1002,6 +1002,10 @@ pub struct IntegratedTrainer {
rl_per_update_priority_fn: CudaFunction,
_rl_per_tree_rebuild_module: Arc<CudaModule>,
rl_per_tree_rebuild_level_fn: CudaFunction,
/// Pre-computed per-level (grid, start, nodes_at_level) for the
/// bottom-up sum-tree rebuild. Avoids recomputing log2(capacity)
/// loop iterations on every step in the non-mega-graph path.
tree_rebuild_levels: Vec<(u32, i32, i32)>,
// ── Bidirectional HER kernel handles ───────────────────────────
pub hindsight: crate::rl::gpu_hindsight::GpuHindsight,
@@ -2136,6 +2140,22 @@ impl IntegratedTrainer {
)
.context("GpuReplayBuffer::new")?;
// Pre-compute per-level launch params for the bottom-up sum-tree
// rebuild. Avoids recomputing log2(capacity) iterations per step.
let tree_rebuild_levels = {
let cap = cfg.per_capacity;
let mut levels = Vec::new();
let mut nodes_at_level = (cap >> 1) as i32;
let mut start = nodes_at_level;
while nodes_at_level >= 1 {
let grid = ((nodes_at_level as u32) + 255) / 256;
levels.push((grid.max(1), start, nodes_at_level));
nodes_at_level >>= 1;
start >>= 1;
}
levels
};
// SP20 P3 FRD head — construct + own per-step buffers. Seed
// distinct from dqn/ppo to keep Xavier draws independent.
let frd_head = crate::rl::frd::FrdHead::new(
@@ -2602,6 +2622,7 @@ impl IntegratedTrainer {
rl_per_update_priority_fn,
_rl_per_tree_rebuild_module: rl_per_tree_rebuild_module,
rl_per_tree_rebuild_level_fn,
tree_rebuild_levels,
hindsight,
_rl_hindsight_track_module: rl_hindsight_track_module,
rl_hindsight_track_fn,
@@ -6950,29 +6971,22 @@ impl IntegratedTrainer {
// barrier that __threadfence() cannot (it is only a memory
// fence, not a synchronization barrier). Graph-safe: each
// level is a separate graph node with implicit dependency.
{
let cap = self.gpu_replay.capacity;
let mut nodes_at_level = (cap >> 1) as i32;
let mut start = nodes_at_level;
while nodes_at_level >= 1 {
let grid = ((nodes_at_level as u32) + 255) / 256;
let mut args = RawArgs::new();
args.push_ptr(self.gpu_replay.priority_tree_d.raw_ptr());
args.push_i32(start);
args.push_i32(nodes_at_level);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_per_tree_rebuild_level_fn.cu_function(),
(grid.max(1), 1, 1), (256, 1, 1), 0,
per_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!(
"rl_per_tree_rebuild_level (start={start}, nodes={nodes_at_level}): {:?}", e
))?;
}
nodes_at_level >>= 1;
start >>= 1;
// Per-level (grid, start, nodes) pre-computed at construction.
for &(grid, start, nodes_at_level) in &self.tree_rebuild_levels {
let mut args = RawArgs::new();
args.push_ptr(self.gpu_replay.priority_tree_d.raw_ptr());
args.push_i32(start);
args.push_i32(nodes_at_level);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_per_tree_rebuild_level_fn.cu_function(),
(grid, 1, 1), (256, 1, 1), 0,
per_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!(
"rl_per_tree_rebuild_level (start={start}, nodes={nodes_at_level}): {:?}", e
))?;
}
}
}