diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index e5886933f..e7d2e9b38 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -183,6 +183,10 @@ pub struct GpuReplayBuffer { sample_max_weight: CudaSlice, sample_episode_ids: CudaSlice, total_sum_buf: CudaSlice, + // Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step) + update_td_f32: CudaSlice, // [max_batch_size] bf16→f32 conversion + update_batch_max: CudaSlice, // [1] per-batch atomicMax accumulator + update_max_merge: CudaSlice, // [1] merge pending + batch max rng_step: u32, /// Size of the most recent `sample_proportional` call. Used by `sample_weights_ref` / /// `sample_indices_ref` to return correctly-sized views (pre-allocated buffers are @@ -244,6 +248,11 @@ impl GpuReplayBuffer { let ep_ids = a32i(stream, cap, "episode_ids")?; let s_ep = a32i(stream, mbs, "s_episode_ids")?; + // Pre-allocate PER update scratch buffers (zero cuMemAlloc per step) + let update_td_f32 = a32f(stream, mbs, "update_td_f32")?; + let update_batch_max = a32f(stream, 1, "update_batch_max")?; + let update_max_merge = a32f(stream, 1, "update_max_merge")?; + Ok(Self { config, stream: Arc::clone(stream), kernels: k, states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p, @@ -258,6 +267,7 @@ impl GpuReplayBuffer { sample_priorities: sp, sample_weights: sw, sample_episode_ids: s_ep, sample_max_weight: smw, total_sum_buf: tsb, + update_td_f32, update_batch_max, update_max_merge, rng_step: 0, last_batch_size: 0, }) @@ -684,62 +694,73 @@ impl GpuReplayBuffer { /// writes to priorities buffer AND tree leaves, propagates sums to root. /// O(log n) per update instead of O(n) prefix sum rebuild. /// Raw-pointer variant for tests and callers without CudaSlice access. - pub fn update_priorities_gpu_raw(&mut self, _indices_ptr: u64, td_errors: &CudaSlice, bs: usize) -> Result<(), MLError> { + pub fn update_priorities_gpu_raw(&mut self, _indices_ptr: u64, td_errors: &CudaSlice, bs: usize, ext_stream: Option<&Arc>) -> Result<(), MLError> { let idx_slice = unsafe { &*(&self.sample_indices_u32 as *const CudaSlice) }; - self.update_priorities_gpu(idx_slice, td_errors, bs) + self.update_priorities_gpu(idx_slice, td_errors, bs, ext_stream) } + /// Update PER priorities from TD errors using the segment tree kernel. + /// + /// If `ext_stream` is provided, all GPU work runs on that stream instead of + /// the replay buffer's own stream. This avoids cross-stream synchronization + /// when the TD errors were computed on a different stream (e.g., the trainer's + /// main stream). All pre-allocated scratch buffers are used — zero per-step + /// GPU allocations. pub fn update_priorities_gpu( &mut self, indices: &CudaSlice, td_errors: &CudaSlice, bs: usize, + ext_stream: Option<&Arc>, ) -> Result<(), MLError> { let _nvtx = NvtxRange::new("per_update_priorities"); if bs == 0 { return Ok(()); } - // Convert bf16 td_errors to f32 for the CUDA kernel (which reads float*) - let kernels = get_cast_kernels(&self.stream)?; - let td_errors_f32 = { - let n = bs; - let ni = n as i32; - let out = self.stream.alloc_zeros::(n) - .map_err(|e| MLError::ModelError(format!("td bf16->f32 alloc: {e}")))?; - // SAFETY: out, td_errors are valid device allocations of at least n elements. - unsafe { - self.stream.launch_builder(&kernels.bf16_to_f32) - .arg(&out).arg(td_errors).arg(&ni) - .launch(lcfg(n)) - .map_err(|e| MLError::ModelError(format!("td bf16_to_f32: {e}")))?; - } - out - }; + let stream = ext_stream.unwrap_or(&self.stream); + // Convert bf16 td_errors to f32 using pre-allocated scratch buffer + let kernels = get_cast_kernels(stream)?; + let ni = bs as i32; + // SAFETY: update_td_f32, td_errors are valid device allocations of at least bs elements. + unsafe { + stream.launch_builder(&kernels.bf16_to_f32) + .arg(&self.update_td_f32).arg(td_errors).arg(&ni) + .launch(lcfg(bs)) + .map_err(|e| MLError::ModelError(format!("td bf16_to_f32: {e}")))?; + } let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32); let cap_pow2_i = self.capacity_pow2 as i32; - let mut bm = a32f(&self.stream, 1, "bm")?; - self.stream.memset_zeros(&mut bm).map_err(|e| MLError::ModelError(format!("{e}")))?; - // seg_tree_update: writes priorities[], tree leaves, propagates, atomicMax. - // SAFETY: seg_tree, priorities, bm, indices, td_errors_f32 are valid device allocations. + stream.memset_zeros(&mut self.update_batch_max) + .map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?; + // seg_tree_update: writes priorities[], tree leaves, propagates via atomicAdd, atomicMax. + // SAFETY: seg_tree, priorities, update_batch_max, indices, update_td_f32 are valid device allocations. unsafe { - self.stream.launch_builder(&self.kernels.seg_tree_update) - .arg(&self.seg_tree).arg(&self.priorities).arg(&mut bm) - .arg(indices).arg(&td_errors_f32) + stream.launch_builder(&self.kernels.seg_tree_update) + .arg(&self.seg_tree).arg(&self.priorities).arg(&self.update_batch_max) + .arg(indices).arg(&self.update_td_f32) .arg(&al).arg(&ep).arg(&cap_pow2_i).arg(&bsi) .launch(lcfg(bs)).map_err(|e| MLError::ModelError(format!("st update: {e}")))?; } - self.pending_max_priority = Some(match self.pending_max_priority.take() { + // Merge batch max into pending max (GPU-only). + // Uses pre-allocated update_max_merge as output scratch. + match self.pending_max_priority.take() { Some(prev) => { - let mut r = a32f(&self.stream, 1, "mm")?; - // SAFETY: r, prev, bm are valid 1-element device allocations on same context. + // SAFETY: update_max_merge, prev, update_batch_max are valid 1-element device allocations. unsafe { - self.stream.launch_builder(&self.kernels.max_of_two_f32) - .arg(&mut r).arg(&prev).arg(&bm) + stream.launch_builder(&self.kernels.max_of_two_f32) + .arg(&self.update_max_merge).arg(&prev).arg(&self.update_batch_max) .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 }) .map_err(|e| MLError::ModelError(format!("max_of_two: {e}")))?; } - r + // Rotate: update_max_merge becomes pending, prev becomes the new merge scratch. + self.pending_max_priority = Some(std::mem::replace(&mut self.update_max_merge, prev)); } - None => bm, - }); + None => { + // First update this epoch: batch_max IS the pending. + // Rotate: update_batch_max becomes pending, allocate new batch_max scratch. + // This allocation happens once per epoch (after flush_max_priority), not per step. + let fresh = a32f(stream, 1, "bm_epoch")?; + self.pending_max_priority = Some(std::mem::replace(&mut self.update_batch_max, fresh)); + } + } Ok(()) } diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 305838ea0..b9ba4c033 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -366,7 +366,7 @@ impl ReplayBufferType { let mut buf = buffer.lock(); if bs == 0 { return Ok(()); } // indices already CudaSlice, td_errors already CudaSlice - buf.gpu.update_priorities_gpu(indices, td_errors, bs) + buf.gpu.update_priorities_gpu(indices, td_errors, bs, None) } Self::Uniform(_) | Self::Prioritized(_) => Ok(()), // No-op for non-GPU buffers } @@ -390,16 +390,22 @@ impl ReplayBufferType { /// Uses the segment tree update kernel which writes priorities, updates /// tree leaves, and propagates sums to root. The sample indices are stored /// internally from the last `sample()` call. + /// Update PER priorities using td_errors and the indices from the last sample. + /// + /// `ext_stream`: the stream on which `td_errors` was computed. All GPU work + /// runs on this stream to avoid cross-stream synchronization. Pass the + /// trainer's main stream here. pub fn update_priorities_from_td( &self, td_errors: &cudarc::driver::CudaSlice, batch_size: usize, + ext_stream: &std::sync::Arc, ) -> Result<(), MLError> { match self { Self::GpuPrioritized(buffer) => { let mut buf = buffer.lock(); if batch_size == 0 { return Ok(()); } - buf.gpu.update_priorities_gpu_raw(0, td_errors, batch_size) + buf.gpu.update_priorities_gpu_raw(0, td_errors, batch_size, Some(ext_stream)) } Self::Uniform(_) | Self::Prioritized(_) => Ok(()), } diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 04926b2b5..10e1f269f 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -643,8 +643,9 @@ impl DQNAgentType { &self, td_errors: &cudarc::driver::CudaSlice, batch_size: usize, + stream: &std::sync::Arc, ) -> Result<(), crate::MLError> { - self.memory().update_priorities_from_td(td_errors, batch_size) + self.memory().update_priorities_from_td(td_errors, batch_size, stream) } /// Set max_priority from a CPU scalar (epoch-boundary flush from CUDA kernel). diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 806fd6d2f..0d32e2b35 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -783,6 +783,7 @@ impl FusedTrainingCtx { agent.update_priorities_from_td( self.trainer.td_errors_buf(), self.batch_size, + &self.stream, ).map_err(|e| anyhow::anyhow!("PER seg_tree_update: {e}"))?; if step < 3 { eprintln!("H100_HANG4: step {step} — PER update done"); } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index 7a00484d2..2432d0870 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -190,7 +190,7 @@ fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> { // Sample and get indices let batch = buf.sample_proportional(16)?; let td_errors = random_cuda_bf16(16, &stream)?; - let idx = batch.indices_ptr; buf.update_priorities_gpu_raw(idx, &td_errors, 16)?; + let idx = batch.indices_ptr; buf.update_priorities_gpu_raw(idx, &td_errors, 16, None)?; // Verify sampling still works with updated priorities (proves update succeeded) buf.sample_proportional(16)?; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 0edaf6eec..1c5daa541 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1558,7 +1558,7 @@ impl DQNTrainer { let mut gpu_buf = agent.memory().as_gpu_buffer(); if let Some(ref mut buf) = gpu_buf { if let Err(e) = buf.gpu.update_priorities_gpu( - &idx_slice, &td_slice, batch_size_refresh, + &idx_slice, &td_slice, batch_size_refresh, None, ) { debug!("M2: GPU priority refresh failed (non-fatal): {}", e); } else { diff --git a/crates/ml/tests/gpu_per_integration_test.rs b/crates/ml/tests/gpu_per_integration_test.rs index 3866713fb..7f90359e6 100644 --- a/crates/ml/tests/gpu_per_integration_test.rs +++ b/crates/ml/tests/gpu_per_integration_test.rs @@ -157,7 +157,7 @@ fn test_gpu_per_training_loop_cycle() { let td_errors: CudaSlice = stream.clone_htod(&td_host).unwrap(); let idx_ptr = batch.indices_ptr; - buf.update_priorities_gpu_raw(idx_ptr, &td_errors, batch_size).unwrap(); + buf.update_priorities_gpu_raw(idx_ptr, &td_errors, batch_size, None).unwrap(); buf.step(); } @@ -181,7 +181,7 @@ fn test_gpu_per_priorities_affect_sampling() { // Set experience 0 to very high priority let high_idx: CudaSlice = stream.clone_htod(&[0_u32]).unwrap(); let high_td: CudaSlice = stream.clone_htod(&[half::bf16::from_f32(100.0)]).unwrap(); - buf.update_priorities_gpu(&high_idx, &high_td, 1).unwrap(); + buf.update_priorities_gpu(&high_idx, &high_td, 1, None).unwrap(); // Sample 1000 times in batches of 10, count how often index 0 appears // Download indices to host, count matches @@ -225,7 +225,7 @@ fn test_gpu_per_sampling_distribution_differs_from_uniform() { for i in 0..50 { let idx: CudaSlice = stream.clone_htod(&[i as u32]).unwrap(); let td: CudaSlice = stream.clone_htod(&[half::bf16::from_f32(i as f32 + 1.0)]).unwrap(); - buf.update_priorities_gpu(&idx, &td, 1).unwrap(); + buf.update_priorities_gpu(&idx, &td, 1, None).unwrap(); } // Sample 5000 indices via proportional sampling @@ -267,7 +267,7 @@ fn test_gpu_per_is_weights_correct_range() { for i in 0..200 { let idx: CudaSlice = stream.clone_htod(&[i as u32]).unwrap(); let td: CudaSlice = stream.clone_htod(&[half::bf16::from_f32((i % 10) as f32 + 1.0)]).unwrap(); - buf.update_priorities_gpu(&idx, &td, 1).unwrap(); + buf.update_priorities_gpu(&idx, &td, 1, None).unwrap(); } let batch = buf.sample_proportional(64).unwrap();