diff --git a/crates/ml-dqn/src/noisy_layers.rs b/crates/ml-dqn/src/noisy_layers.rs index 6d239c617..b5526979a 100644 --- a/crates/ml-dqn/src/noisy_layers.rs +++ b/crates/ml-dqn/src/noisy_layers.rs @@ -359,17 +359,17 @@ impl NoisyLinear { /// Copies `weight_mu`, `bias_mu`, `weight_sigma`, `bias_sigma`, and epsilon buffers. /// Zero CPU involvement — all copies are async device-to-device. pub fn copy_params_from(&mut self, src: &NoisyLinear) -> Result<(), MLError> { - Self::dtod_copy(&src.weight_mu, &mut self.weight_mu, "weight_mu", &self.stream)?; - Self::dtod_copy(&src.bias_mu, &mut self.bias_mu, "bias_mu", &self.stream)?; - Self::dtod_copy(&src.weight_sigma, &mut self.weight_sigma, "weight_sigma", &self.stream)?; - Self::dtod_copy(&src.bias_sigma, &mut self.bias_sigma, "bias_sigma", &self.stream)?; - Self::dtod_copy(&src.weight_epsilon, &mut self.weight_epsilon, "weight_epsilon", &self.stream)?; - Self::dtod_copy(&src.bias_epsilon, &mut self.bias_epsilon, "bias_epsilon", &self.stream)?; + Self::dtod_copy_async(&src.weight_mu, &mut self.weight_mu, "weight_mu", &self.stream)?; + Self::dtod_copy_async(&src.bias_mu, &mut self.bias_mu, "bias_mu", &self.stream)?; + Self::dtod_copy_async(&src.weight_sigma, &mut self.weight_sigma, "weight_sigma", &self.stream)?; + Self::dtod_copy_async(&src.bias_sigma, &mut self.bias_sigma, "bias_sigma", &self.stream)?; + Self::dtod_copy_async(&src.weight_epsilon, &mut self.weight_epsilon, "weight_epsilon", &self.stream)?; + Self::dtod_copy_async(&src.bias_epsilon, &mut self.bias_epsilon, "bias_epsilon", &self.stream)?; Ok(()) } /// Async device-to-device memcpy for a single `CudaSlice` buffer. - fn dtod_copy( + fn dtod_copy_async( src: &CudaSlice, dst: &mut CudaSlice, label: &str, diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 6f3069d20..03ec09827 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -8397,7 +8397,7 @@ impl GpuDqnTrainer { // td_errors → readback[3..3+B] (only td_errors needs DtoD → DtoH) let td_src = self.td_errors_buf.raw_ptr(); let td_bytes = b * f32_bytes; - dtod_copy(readback_base + (3 * f32_bytes) as u64, td_src, td_bytes, &self.stream, 3, "readback_gather")?; + dtod_copy_async(readback_base + (3 * f32_bytes) as u64, td_src, td_bytes, &self.stream, 3, "readback_gather")?; // ── Single DtoH transfer (only td_errors portion matters) ───── super::dtoh_f32(&self.stream, &self.readback_buf, &mut self.readback_host)?; @@ -10851,26 +10851,26 @@ impl GpuDqnTrainer { )?; // Rewards, dones: f32 DtoD - dtod_copy( + dtod_copy_async( self.rewards_buf.raw_ptr(), gpu_batch.rewards_ptr, b * f32_size, &self.stream, 2, "rewards", )?; - dtod_copy( + dtod_copy_async( self.dones_buf.raw_ptr(), gpu_batch.dones_ptr, b * f32_size, &self.stream, 3, "dones", )?; // IS-weights: f32 DtoD - dtod_copy( + dtod_copy_async( self.is_weights_buf.raw_ptr(), gpu_batch.weights_ptr, b * f32_size, &self.stream, 4, "is_weights", )?; // Actions: CudaSlice → i32 buf (async DtoD, zero CPU sync) - dtod_copy( + dtod_copy_async( self.actions_buf.raw_ptr(), gpu_batch.actions_ptr, b * std::mem::size_of::(), &self.stream, 4, "actions", )?; @@ -12038,7 +12038,7 @@ impl GpuDqnTrainer { ))); } let num_bytes = sizes[i] * std::mem::size_of::(); - dtod_copy(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten")?; + dtod_copy_async(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten")?; byte_offset += num_bytes as u64; } @@ -12085,7 +12085,7 @@ impl GpuDqnTrainer { let mut byte_offset: u64 = 0; for (i, &dst) in ptrs.iter().enumerate() { let num_bytes = sizes[i] * std::mem::size_of::(); - dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten")?; + dtod_copy_async(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten")?; byte_offset += num_bytes as u64; } @@ -12371,7 +12371,7 @@ impl GpuDqnTrainer { for (i, &src) in ptrs.iter().enumerate() { let num_bytes = sizes[i] * std::mem::size_of::(); let dst = dst_base + byte_offset; - dtod_copy(dst, src, num_bytes, &self.stream, i, "flatten_target")?; + dtod_copy_async(dst, src, num_bytes, &self.stream, i, "flatten_target")?; byte_offset += num_bytes as u64; } @@ -12412,7 +12412,7 @@ impl GpuDqnTrainer { let mut byte_offset: u64 = 0; for (i, &dst) in ptrs.iter().enumerate() { let num_bytes = sizes[i] * std::mem::size_of::(); - dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?; + dtod_copy_async(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?; byte_offset += num_bytes as u64; } @@ -13217,7 +13217,7 @@ pub(crate) fn query_max_shmem_bytes() -> usize { // raw_device_ptr wrappers removed — use CudaSlice::raw_ptr() directly (no event tracking). /// Async device-to-device memcpy with error context. -pub(crate) fn dtod_copy( +pub(crate) fn dtod_copy_async( dst: u64, src: u64, num_bytes: usize, @@ -13247,7 +13247,7 @@ fn dtod_from_slice_f32( ) -> Result<(), MLError> { let src_ptr = src.raw_ptr(); let dst_ptr = dst.raw_ptr(); - dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) + dtod_copy_async(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) } /// DtoD copy from a CudaSlice to a CudaSlice. @@ -13262,7 +13262,7 @@ fn dtod_from_u32_to_i32( ) -> Result<(), MLError> { let src_ptr = src.raw_ptr(); let dst_ptr = dst.raw_ptr(); - dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) + dtod_copy_async(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) } // ── Allocation helpers ────────────────────────────────────────────────────── @@ -13331,7 +13331,7 @@ fn dtod_from_staging( ) -> Result<(), MLError> { let src_ptr = src.raw_ptr(); let dst_ptr = dst.raw_ptr(); - dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) + dtod_copy_async(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) } /// Load ensemble aggregate + diversity kernels from precompiled cubin. diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 1b7b8c292..4a9dc9d50 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -697,15 +697,15 @@ impl GpuIqnHead { .map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?; } - use super::gpu_dqn_trainer::dtod_copy; + use super::gpu_dqn_trainer::dtod_copy_async; let f32_bytes = std::mem::size_of::(); let r_src = dqn_rewards_buf.raw_ptr(); let r_dst = self.rewards_buf.raw_ptr(); - dtod_copy(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?; + dtod_copy_async(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?; let d_src = dqn_dones_buf.raw_ptr(); let d_dst = self.dones_buf.raw_ptr(); - dtod_copy(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?; + dtod_copy_async(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?; Ok(()) } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index a484e8ed3..08dea8257 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2991,13 +2991,13 @@ fn clone_dueling_weights( stream: &Arc, head_idx: usize, ) -> Result<(gpu_weights::DuelingWeightBacking, DuelingWeightSet)> { - use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy; + use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy_async; let clone_from_ptr = |ptr: u64, len: usize| -> Result> { let dst = stream.alloc_zeros::(len) .map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?; let n_bytes = len * std::mem::size_of::(); - dtod_copy(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_dueling") + dtod_copy_async(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_dueling") .map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?; Ok(dst) }; @@ -3032,13 +3032,13 @@ fn clone_branching_weights( stream: &Arc, head_idx: usize, ) -> Result<(gpu_weights::BranchingWeightBacking, BranchingWeightSet)> { - use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy; + use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy_async; let clone_from_ptr = |ptr: u64, len: usize| -> Result> { let dst = stream.alloc_zeros::(len) .map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?; let n_bytes = len * std::mem::size_of::(); - dtod_copy(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_branching") + dtod_copy_async(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_branching") .map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?; Ok(dst) };