Files
foxhunt/docs/superpowers/specs/2026-04-10-dead-code-cleanup-gpu-profiling.md

4.0 KiB

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.rsmemory 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.rsagent.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.rsrun_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