From c0b43c9fad01722baa804dba4ef16e0ebc9dfe4a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 22 Mar 2026 00:19:44 +0100 Subject: [PATCH] perf: fix PER prefix sum single-block fallback at 500K buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of H100 training time growth (12ms→23ms/step across epochs): pfx_sum() fell back to single-block O(n) scan when num_blocks > 1024. With block_dim=256 and buffer_size=500K: num_blocks = 1953 > 1024 → fallback. Fix: increase block_dim to 1024 (capped at device max_threads_per_block). Now: num_blocks = 500000/1024 = 489 < 1024 → multi-block 3-phase scan works. Expected impact: PER sampling cost stays constant as buffer fills instead of growing linearly. Training step time should be stable across epochs. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/gpu_replay_buffer.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 06a1ed34e..5a71e04df 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -644,7 +644,7 @@ impl GpuReplayBuffer { if bs == 0 { return Ok(()); } let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32); let mut bm = a32f(&self.stream, 1, "bm")?; - self.stream.memcpy_htod(&[0.0_f32], &mut bm).map_err(|e| MLError::ModelError(format!("{e}")))?; + self.stream.memset_zeros(&mut bm).map_err(|e| MLError::ModelError(format!("{e}")))?; // SAFETY: td_errors, indices, priorities, bm are valid device allocations. bs <= buffer size. unsafe { self.stream.launch_builder(&self.kernels.priority_update_f32) @@ -653,7 +653,6 @@ impl GpuReplayBuffer { } self.pending_max_priority = Some(match self.pending_max_priority.take() { Some(prev) => { - // GPU-only max of two scalars (zero CPU readback) let mut r = a32f(&self.stream, 1, "mm")?; // SAFETY: r, prev, bm are valid 1-element device allocations on same context. unsafe { @@ -774,8 +773,11 @@ impl GpuReplayBuffer { } }); - // Multi-block scan uses 256 threads/block. - let block_dim = 256_u32; + // Multi-block scan: use 1024 threads/block to maximize elements per block. + // This ensures num_blocks stays below max_threads_per_block (1024) for phase-2 + // block_sums scan. At buffer_size=500K: num_blocks = 500000/1024 = 489 < 1024. + // With 256 threads: num_blocks = 1953 > 1024 → falls back to single-block O(n). + let block_dim = 1024_u32.min(mt); let num_blocks = (n as u32).div_ceil(block_dim); if num_blocks <= 1 {