# Dead Code Cleanup + GPU Per-Phase Profiling ## Part 1: Dead Code Cleanup ### Problem Three replay buffer implementations exist: CPU uniform, CPU PER (segment tree), GPU PER (prefix scan). Only GPU PER is used. The CPU paths are dead code that confuses the architecture and wastes the `ReplayBufferType` enum indirection. ### Delete - `crates/ml-dqn/src/prioritized_replay.rs` — CPU segment tree PER - `crates/ml-dqn/src/prioritized_replay_staleness.rs` — CPU staleness tracker - `crates/ml-dqn/src/replay_buffer.rs` — CPU uniform replay buffer - `crates/ml-dqn/src/hindsight_replay.rs` — CPU HER (GPU HER is `gpu_her.rs`) ### Simplify - `crates/ml-dqn/src/replay_buffer_type.rs` — delete `ReplayBufferType` enum. Agent holds `GpuReplayBuffer` directly. - `crates/ml-dqn/src/dqn.rs` — `memory` field becomes `GpuReplayBuffer` (was `ReplayBufferType`). `memory()` returns `&GpuReplayBuffer`. `sample()` returns `GpuBatchPtrs` directly. - `crates/ml-dqn/src/lib.rs` — remove dead re-exports (`PrioritizedReplayBuffer`, `ReplayBuffer`, `ReplayBufferConfig`, etc.) - `crates/ml/src/trainers/dqn/trainer/training_loop.rs` — `agent.memory().sample()` returns `GpuBatchPtrs` directly. Remove `BatchSample` wrapper with empty `experiences` vec. Remove `.gpu_batch.as_ref().unwrap()`. - `crates/ml/src/trainers/dqn/fused_training.rs` — `run_full_step` takes `&GpuBatchPtrs` directly (not `&BatchSample`). - `crates/ml/src/trainers/dqn/trainer/constructor.rs` — remove `PrioritizedReplayConfig` import used for CPU HER. - Tests referencing CPU PER — update or delete. ### Non-goals - Do not change `gpu_replay_buffer.rs` internals. - Do not change `gpu_her.rs` (GPU HER stays). - Do not change PER kernel code (`per_kernels.cu`). --- ## Part 2: GPU Per-Phase CUDA Event Profiling ### Problem H100 shows 329ms/batch for a 336K param DQN with batch=16384. Expected ~30-40ms. CPU-side wall-clock timers conflate GPU pipeline stalls with actual kernel work. We need GPU-side per-phase timing to identify the real bottleneck. ### Design Pre-allocate CUDA event pairs in `FusedTrainingCtx`. Record `cuEventRecord` before and after each training phase. At epoch end, compute `cuEventElapsedTime` for each phase and log averages. ### Phases to measure | Phase | What it covers | |-------|---------------| | `per_scan` | `per_prefix_scan` kernel (prefix sum over 24.7M priorities) | | `per_sample` | `per_sample` + `i64_to_u32` + `gather_*` + `is_weights` kernels | | `upload` | `upload_batch_gpu` (DtoD copy from replay to trainer buffers) | | `fwd_bwd` | `graph_mega` or `graph_forward` + backward replay | | `adam` | `graph_adam` replay (grad norm + optimizer + unflatten) | | `per_update` | `per_update_pa` kernel (priority scatter) | ### Implementation - 12 CUDA events (2 per phase) allocated at `FusedTrainingCtx` construction - `cuEventRecord(start, stream)` before each phase - `cuEventRecord(end, stream)` after each phase - Per-batch: accumulate elapsed times into f64 accumulators - Per-epoch: divide by batch count, log one summary line - Zero GPU overhead — events are hardware timestamp queries, not sync points ### Output format ``` GPU phase timing (178 batches avg): per_scan=1.8ms per_sample=3.2ms upload=0.9ms fwd_bwd=22.1ms adam=1.8ms per_update=0.3ms total=30.1ms ``` ### Constraints - Always-on (no CLI flag). CUDA events have zero overhead. - Events must be on the SAME stream as the kernels they bracket. - `cuEventElapsedTime` requires both events to have completed — call at epoch end after `cuStreamSynchronize`. - Do not add `cuStreamSynchronize` between phases — that would serialize the pipeline and change the behavior we're measuring. ### Success criteria The timing output reveals where the 329ms goes: - If `fwd_bwd` is ~25ms and `total` is ~30ms → the 329ms wall-clock is a pipeline stall (cross-stream sync, cudarc overhead, or CPU scheduling) - If `fwd_bwd` is ~280ms → actual GPU kernel bottleneck (occupancy, register pressure, or memory bandwidth) - Either answer tells us exactly what to fix next