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**

```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<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**

```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<super::replay_buffer_type::StagedGpuBuffer>`.

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<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**

```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<PhaseEvents>,
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<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`:

```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)