perf: fix PER prefix sum single-block fallback at 500K buffer

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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-22 00:19:44 +01:00
parent 49444c19d8
commit c0b43c9fad

View File

@@ -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 {