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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user