# 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** ```bash 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** ```bash 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: ```rust //! 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 { 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** ```bash 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`. Update the constructor (around line 1200) — it currently calls `ReplayBufferType::try_gpu_with_halving(...)`. Replace with direct `StagedGpuBuffer` construction: ```rust 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`. 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** ```bash 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: ```rust 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: ```rust pub(crate) fn run_full_step( &mut self, batch: &BatchSample, ... ``` To: ```rust 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: ```rust 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: ```bash 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** ```bash 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: ```rust // Per-phase CUDA event profiling (always-on, zero GPU overhead) phase_events: Option, phase_accum: PhaseAccum, ``` Define the types (in the same file or a small helper): ```rust 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`: ```rust let create = || -> Result { 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`: ```rust 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: ```rust 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`): ```rust 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: ```rust 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** ```bash 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** ```bash 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** ```bash 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)