jgrusewski 2c39d100a0 perf(ml-dqn): eliminate CPU roundtrips in PER sample_proportional ...

Eliminate the 3 CPU roundtrips per sample_proportional() call (called
300x/epoch), which was the #1 performance bottleneck:

1. Replace memcpy_dtoh(cs_total) with DtoD copy to pre-allocated
   total_sum_buf — total_sum never leaves GPU
2. Replace CPU rand() + memcpy_htod(thresholds) with GPU-resident
   Philox PRNG kernel — thresholds generated directly on device
3. Replace memcpy_htod([0.0], max_weight) with memset_zeros —
   async GPU memset, zero host staging
4. is_weights_f32 kernel now reads total_sum from GPU pointer
   instead of scalar argument

Pre-allocate 13 PER sampling buffers on the GpuReplayBuffer struct
(thresholds, indices, gathered data, weights, max_weight, total_sum_buf,
rng_step counter). All intermediate computation uses these pre-allocated
buffers. Output GpuBatchSlices are DtoD-cloned for ownership transfer.

Add max_batch_size field to GpuReplayBufferConfig (defaults to 1024).
Delete the cs_total() method entirely.

Net result: zero memcpy_dtoh, zero memcpy_htod, zero CPU synchronization
points in the PER sampling hot path.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

2026-03-21 10:38:33 +01:00

13 KiB

Raw Blame History

View Raw