foxhunt/docs/superpowers/plans/2026-04-10-dead-code-gpu-profiling.md

Dead Code Cleanup + GPU Per-Phase Profiling Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Delete dead CPU PER/replay/HER code, simplify agent to use GpuReplayBuffer directly, and add always-on CUDA event profiling to identify the 329ms/batch bottleneck.

Architecture: Part 1 deletes 4 files and simplifies ReplayBufferType enum away — agent.memory() returns &GpuReplayBuffer directly. Part 2 adds 12 CUDA events (6 phase pairs) to FusedTrainingCtx with per-epoch summary logging.

Tech Stack: Rust, cudarc (CUDA events), ml-dqn crate, fused_training.rs

---

File Structure

Action	File	Purpose
Delete	crates/ml-dqn/src/prioritized_replay.rs	Dead CPU segment tree PER
Delete	crates/ml-dqn/src/prioritized_replay_staleness.rs	Dead CPU staleness tracker
Delete	crates/ml-dqn/src/replay_buffer.rs	Dead CPU uniform replay
Delete	crates/ml-dqn/src/hindsight_replay.rs	Dead CPU HER
Rewrite	crates/ml-dqn/src/replay_buffer_type.rs	Delete enum, keep GpuBatch + StagedGpuBuffer
Modify	crates/ml-dqn/src/dqn.rs	Agent memory field → StagedGpuBuffer directly
Modify	crates/ml-dqn/src/lib.rs	Remove dead re-exports
Modify	crates/ml/src/trainers/dqn/trainer/training_loop.rs	Use GpuBatchPtrs directly
Modify	crates/ml/src/trainers/dqn/fused_training.rs	run_full_step takes GpuBatchPtrs, add CUDA events
Modify	crates/ml/src/trainers/dqn/trainer/constructor.rs	Remove CPU HER config
Modify	crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs	upload_batch_gpu takes GpuBatchPtrs
Delete	Tests referencing CPU PER	Dead test code

---

Task 1: Delete Dead Files

Files:

• Delete: crates/ml-dqn/src/prioritized_replay.rs

• Delete: crates/ml-dqn/src/prioritized_replay_staleness.rs

• Delete: crates/ml-dqn/src/replay_buffer.rs

• Delete: crates/ml-dqn/src/hindsight_replay.rs

• Modify: crates/ml-dqn/src/lib.rs

• Step 1: Delete the 4 dead files

rm crates/ml-dqn/src/prioritized_replay.rs
rm crates/ml-dqn/src/prioritized_replay_staleness.rs
rm crates/ml-dqn/src/replay_buffer.rs
rm crates/ml-dqn/src/hindsight_replay.rs

• Step 2: Remove module declarations and re-exports from lib.rs

In crates/ml-dqn/src/lib.rs, remove:

• pub mod prioritized_replay;
• pub mod prioritized_replay_staleness;
• pub mod replay_buffer;
• pub mod hindsight_replay;
• pub use replay_buffer::{ReplayBuffer, ReplayBufferConfig, ReplayBufferStats};
• pub use prioritized_replay::{PrioritizedReplayBuffer, PrioritizedReplayConfig};
• pub use replay_buffer_type::{BatchSample, ReplayBufferType};

Add instead:

• pub use replay_buffer_type::GpuBatch;

• pub use gpu_replay_buffer::GpuBatchPtrs;

• Step 3: Compile check

Run: SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | grep 'error\[' | head -20 Expected: Errors in replay_buffer_type.rs, dqn.rs (they reference deleted types). Fixed in Task 2.

• Step 4: Commit

git add -A
git commit -m "cleanup: delete dead CPU PER, replay buffer, HER, staleness tracker"

---

Task 2: Simplify replay_buffer_type.rs — Delete ReplayBufferType Enum

Files:

• Modify: crates/ml-dqn/src/replay_buffer_type.rs

• Step 1: Strip the file to only GpuBatch + StagedGpuBuffer

Replace the entire file content. Keep GpuBatch struct (it's the batch type used by consumers), StagedGpuBuffer wrapper, and the GPU-only sample/update methods. Delete ReplayBufferType enum, BatchSample, Uniform variant, Prioritized variant, all CPU-path code.

The new file should contain:

//! GPU-resident replay buffer type — the only replay path.

use ml_core::MLError;

/// GPU batch: raw device pointers for sampled transitions.
/// All fields are pre-allocated in GpuReplayBuffer — zero per-batch allocation.
pub struct GpuBatch {
    pub states_ptr: u64,
    pub actions_ptr: u64,
    pub rewards_ptr: u64,
    pub next_states_ptr: u64,
    pub dones_ptr: u64,
    pub weights_ptr: u64,
    pub indices_ptr: u64,
    pub episode_ids_ptr: u64,
    pub batch_size: usize,
}

impl GpuBatch {
    /// Construct from GpuBatchPtrs (the output of sample_proportional)
    pub fn from_ptrs(ptrs: &crate::gpu_replay_buffer::GpuBatchPtrs) -> Self {
        Self {
            states_ptr: ptrs.states_ptr,
            actions_ptr: ptrs.actions_ptr,
            rewards_ptr: ptrs.rewards_ptr,
            next_states_ptr: ptrs.next_states_ptr,
            dones_ptr: ptrs.dones_ptr,
            weights_ptr: ptrs.weights_ptr,
            indices_ptr: ptrs.indices_ptr,
            episode_ids_ptr: ptrs.episode_ids_ptr,
            batch_size: ptrs.batch_size,
        }
    }
}

/// GPU-resident prioritized replay buffer wrapper.
pub struct StagedGpuBuffer {
    pub gpu: crate::gpu_replay_buffer::GpuReplayBuffer,
}

impl StagedGpuBuffer {
    pub fn sample(&mut self, batch_size: usize) -> Result<GpuBatch, MLError> {
        let ptrs = self.gpu.sample_proportional(batch_size)?;
        Ok(GpuBatch::from_ptrs(&ptrs))
    }

    pub fn can_sample(&self, batch_size: usize) -> bool {
        self.gpu.can_sample(batch_size)
    }

    pub fn size(&self) -> usize {
        self.gpu.size()
    }
}

• Step 2: Compile check

Run: SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | grep 'error\[' | head -20 Expected: Errors in dqn.rs (references to deleted types). Fixed in Task 3.

• Step 3: Commit

git add crates/ml-dqn/src/replay_buffer_type.rs
git commit -m "cleanup: delete ReplayBufferType enum — StagedGpuBuffer is the only path"

---

Task 3: Update dqn.rs Agent — Direct StagedGpuBuffer

Files:

• Modify: crates/ml-dqn/src/dqn.rs

• Step 1: Change agent memory field type

The DQNAgent struct has pub memory: super::replay_buffer_type::ReplayBufferType. Change to pub memory: parking_lot::Mutex<super::replay_buffer_type::StagedGpuBuffer>.

Update the constructor (around line 1200) — it currently calls ReplayBufferType::try_gpu_with_halving(...). Replace with direct StagedGpuBuffer construction:

let gpu_buf = crate::gpu_replay_buffer::GpuReplayBuffer::new(
    config.replay_buffer_capacity,
    config.state_dim,
    config.per_alpha,
    config.per_beta_start,
    config.per_beta_max,
    config.per_beta_annealing_steps,
    config.per_max_memory_bytes,
    config.batch_size,
    &stream,
)?;
let memory = parking_lot::Mutex::new(super::replay_buffer_type::StagedGpuBuffer { gpu: gpu_buf });

If the old code had halving retry logic for VRAM, keep that but operating on GpuReplayBuffer::new() directly.

• Step 2: Update memory() accessor

The memory() method should return &Mutex<StagedGpuBuffer>. Callers use agent.memory().lock().sample(batch_size).

• Step 3: Update all methods that call self.memory

Search for self.memory. in dqn.rs. Update insert, update_priorities, can_sample, size, etc. to go through the Mutex.

• Step 4: Remove dead imports

Remove use super::replay_buffer_type::ReplayBufferType; and similar.

• Step 5: Compile check

Run: SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | grep 'error\[' | head -20

• Step 6: Commit

git add crates/ml-dqn/src/dqn.rs
git commit -m "cleanup: agent memory is StagedGpuBuffer directly — no enum dispatch"

---

Task 4: Update Training Loop + Fused Training — Use GpuBatch Directly

Files:

• Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs

• Modify: crates/ml/src/trainers/dqn/fused_training.rs

• Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• Step 1: Update training_loop.rs batch sampling

At line ~1285: let batch = agent.memory().sample(self.current_batch_size)?;

Change to:

let batch = agent.memory().lock().sample(self.current_batch_size)?;

This now returns GpuBatch directly (not BatchSample).

Pass &batch to fused.run_full_step(...) — no .gpu_batch.as_ref().unwrap() needed.

Also update the vaccine batch sampling similarly.

• Step 2: Update fused_training.rs run_full_step signature

Change:

pub(crate) fn run_full_step(
    &mut self,
    batch: &BatchSample,
    ...

To:

pub(crate) fn run_full_step(
    &mut self,
    batch: &crate::dqn::replay_buffer_type::GpuBatch,
    ...

Remove the let gpu_batch = batch.gpu_batch.as_ref().unwrap() line. Use batch directly everywhere that used gpu_batch.

• Step 3: Update gpu_dqn_trainer.rs upload_batch_gpu

Change:

pub(crate) fn upload_batch_gpu(
    &mut self,
    gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,

The signature stays the same type but now callers pass &batch directly instead of unwrapping.

• Step 4: Remove PrioritizedReplayConfig from constructor.rs

Remove the CPU HER construction that used PrioritizedReplayConfig. The GPU HER (gpu_her.rs) is separate and stays.

• Step 5: Remove dead imports across all modified files

Remove use crate::dqn::replay_buffer_type::BatchSample; and similar.

• Step 6: Delete dead test files

Find and delete tests that reference CPU PER:

grep -rn 'PrioritizedReplayBuffer\|ReplayBuffer\b' crates/ml/tests/ --include='*.rs' -l

Delete or update tests as needed.

• Step 7: Full compile check

Run: SQLX_OFFLINE=true cargo check --workspace 2>&1 | grep 'error\[' | head -20 Expected: 0 errors.

• Step 8: Run smoke tests

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget --nocapture Expected: 7/7 pass.

• Step 9: Commit

git add -A
git commit -m "cleanup: training loop uses GpuBatch directly — BatchSample wrapper deleted"

---

Task 5: Add CUDA Event Per-Phase Profiling

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs

• Step 1: Add CUDA event fields to FusedTrainingCtx

Add to the struct:

// Per-phase CUDA event profiling (always-on, zero GPU overhead)
phase_events: Option<PhaseEvents>,
phase_accum: PhaseAccum,

Define the types (in the same file or a small helper):

struct PhaseEvents {
    per_scan_start: cudarc::driver::sys::CUevent,
    per_scan_end: cudarc::driver::sys::CUevent,
    per_sample_start: cudarc::driver::sys::CUevent,
    per_sample_end: cudarc::driver::sys::CUevent,
    upload_start: cudarc::driver::sys::CUevent,
    upload_end: cudarc::driver::sys::CUevent,
    fwd_bwd_start: cudarc::driver::sys::CUevent,
    fwd_bwd_end: cudarc::driver::sys::CUevent,
    adam_start: cudarc::driver::sys::CUevent,
    adam_end: cudarc::driver::sys::CUevent,
    per_update_start: cudarc::driver::sys::CUevent,
    per_update_end: cudarc::driver::sys::CUevent,
}

#[derive(Default)]
struct PhaseAccum {
    per_scan_ms: f64,
    per_sample_ms: f64,
    upload_ms: f64,
    fwd_bwd_ms: f64,
    adam_ms: f64,
    per_update_ms: f64,
    count: usize,
}

• Step 2: Create events at FusedTrainingCtx construction

In new(), create 12 CUDA events using cudarc::driver::result::event_create:

let create = || -> Result<cudarc::driver::sys::CUevent, _> {
    let mut event = std::mem::MaybeUninit::uninit();
    unsafe {
        cudarc::driver::sys::cuEventCreate(
            event.as_mut_ptr(),
            cudarc::driver::sys::CUevent_flags_enum::CU_EVENT_DEFAULT as u32,
        ).result()?;
        Ok(event.assume_init())
    }
};
let phase_events = PhaseEvents {
    per_scan_start: create()?,
    per_scan_end: create()?,
    // ... all 12
};

• Step 3: Record events around each phase in run_full_step

The PER sample happens OUTSIDE run_full_step (in training_loop.rs). But the GPU PER kernels run on the replay buffer's stream. We need events on THAT stream for per_scan and per_sample.

For phases inside run_full_step (upload, fwd_bwd, adam, per_update), record on self.stream:

unsafe { cudarc::driver::sys::cuEventRecord(self.phase_events.upload_start, self.stream.cu_stream()); }
self.trainer.upload_batch_gpu(batch)?;
unsafe { cudarc::driver::sys::cuEventRecord(self.phase_events.upload_end, self.stream.cu_stream()); }

For the graph_mega replay:

unsafe { cudarc::driver::sys::cuEventRecord(self.phase_events.fwd_bwd_start, self.stream.cu_stream()); }
self.graph_mega.as_ref().unwrap().launch(self.stream.cu_stream())?;
unsafe { cudarc::driver::sys::cuEventRecord(self.phase_events.fwd_bwd_end, self.stream.cu_stream()); }

Same pattern for graph_adam and per_update.

• Step 4: Compute elapsed times at epoch end

Add a method log_phase_timing(&mut self) called at epoch end (after cuStreamSynchronize):

fn log_phase_timing(&mut self) {
    if self.phase_accum.count == 0 { return; }
    let n = self.phase_accum.count as f64;
    // cuEventElapsedTime requires both events completed — we're after stream sync
    let elapsed = |start, end| -> f32 {
        let mut ms = 0.0_f32;
        unsafe {
            cudarc::driver::sys::cuEventElapsedTime(&mut ms, start, end);
        }
        ms
    };
    // Log last-batch timing as representative (events get overwritten each batch)
    let pe = self.phase_events.as_ref().unwrap();
    tracing::info!(
        "GPU phase timing ({} batches, last batch): upload={:.1}ms fwd_bwd={:.1}ms adam={:.1}ms per_update={:.1}ms",
        self.phase_accum.count,
        elapsed(pe.upload_start, pe.upload_end),
        elapsed(pe.fwd_bwd_start, pe.fwd_bwd_end),
        elapsed(pe.adam_start, pe.adam_end),
        elapsed(pe.per_update_start, pe.per_update_end),
    );
    self.phase_accum = PhaseAccum::default();
}

Note: CUDA events get overwritten each batch, so we can only report the LAST batch timing (not averages). For averages, we'd need 12 events per batch × 178 batches = too many. Last-batch timing is sufficient to identify the bottleneck.

• Step 5: Call log_phase_timing from the epoch end

In training_loop.rs, after the training step loop completes and before validation, call:

fused.log_phase_timing();

• Step 6: Compile check + smoke test

Run: SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | grep 'error\[' | head -10 Run: SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget --nocapture Expected: compiles clean, 7/7 pass, phase timing appears in log output.

• Step 7: Commit

git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "perf: add always-on CUDA event per-phase profiling — identifies training bottleneck"

---

Task 6: Run Full Smoke Suite + Deploy Baseline

Files:

• No modifications — verification only

• Step 1: Run all 19 smoke tests

SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -5

Expected: 19/19 pass.

• Step 2: Push and deploy H100 baseline

git push origin main
./scripts/argo-train.sh dqn --baseline --epochs 5 --watch

Watch for the new GPU phase timing line in the logs — this reveals the actual per-kernel bottleneck.

• Step 3: Analyze timing output

The timing will show one of:

• fwd_bwd=25ms total=30ms → 329ms is pipeline stall (investigate cudarc sync, cross-stream waits)
• fwd_bwd=280ms → actual GPU kernel bottleneck (investigate occupancy, memory bandwidth)