From f782ec5a12df1a511cc118619efc09ca326ea4d0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 18 Mar 2026 17:48:43 +0100 Subject: [PATCH] =?UTF-8?q?perf(cuda):=20NoisyLinear=20forward=20+=20copy?= =?UTF-8?q?=20fully=20GPU-native=20=E2=80=94=20zero=20CPU?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit noisy_layers.rs: - forward(): cuBLAS sgemm for W=mu+sigma*epsilon matmul (was CPU loops) - copy_params_from(): 6x DtoD async memcpy (was 12 PCIe roundtrips) - Added noisy_add_bias CUDA kernel for bias addition branching.rs: - copy_weights_from(): DtoD memcpy for VarStore + NoisyLinear params - register_mu_in_varstore(): DtoD clone (was host roundtrip) Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/branching.rs | 110 ++++++++----- crates/ml-dqn/src/noisy_layers.rs | 265 +++++++++++++++++------------- 2 files changed, 219 insertions(+), 156 deletions(-) diff --git a/crates/ml-dqn/src/branching.rs b/crates/ml-dqn/src/branching.rs index 641357c1b..69557a694 100644 --- a/crates/ml-dqn/src/branching.rs +++ b/crates/ml-dqn/src/branching.rs @@ -36,15 +36,44 @@ //! - **`NoisyNet`**: Factorized Gaussian noise in value/branch heads for learned exploration. //! - **State Dim Alignment**: Auto-pad `state_dim` to multiples of 8 for tensor core HMMA. +use std::mem::ManuallyDrop; use std::sync::Arc; -use cudarc::driver::{CudaSlice, CudaStream}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; use ml_core::cuda_autograd::{GpuLinear, GpuTensor, GpuVarStore}; use serde::{Deserialize, Serialize}; use crate::noisy_layers::NoisyLinear; use ml_core::MLError; +/// Async device-to-device memcpy for a single CudaSlice buffer. Zero CPU involvement. +fn dtod_copy_slice( + src: &CudaSlice, + dst: &mut CudaSlice, + stream: &Arc, + label: &str, +) -> Result<(), MLError> { + let num_bytes = src.len() * std::mem::size_of::(); + let src_ptr = { + let (ptr, guard) = src.device_ptr(stream); + let _no_drop = ManuallyDrop::new(guard); + ptr + }; + let dst_ptr = { + let (ptr, guard) = dst.device_ptr_mut(stream); + let _no_drop = ManuallyDrop::new(guard); + ptr + }; + // SAFETY: src and dst are valid device allocations on the same context. + // num_bytes = len * sizeof(f32) does not exceed either allocation. + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("DtoD {label}: {e}")))?; + } + Ok(()) +} + /// Output of the branching network's forward pass. /// /// Provides the decomposed value + per-branch advantages @@ -223,6 +252,7 @@ impl MaybeNoisyLinear { /// /// Clones the mu data from the `NoisyLinear` into the `GpuVarStore` so that the /// GPU experience collector and serialization can find them by name. + /// Uses async DtoD memcpy — zero CPU involvement. fn register_mu_in_varstore(&self, vars: &mut GpuVarStore, name: &str) -> Result<(), MLError> { let Self::Noisy(n) = self; let [w_mu, b_mu] = n.mu_slices(); @@ -230,35 +260,20 @@ impl MaybeNoisyLinear { let in_f = n.in_features(); let stream = vars.stream().clone(); - // Clone weight_mu -> "{name}.weight" + // Clone weight_mu -> "{name}.weight" via DtoD memcpy let w_len = w_mu.len(); - let w_clone = stream.alloc_zeros::(w_len).map_err(|e| { + let mut w_clone = stream.alloc_zeros::(w_len).map_err(|e| { MLError::ModelError(format!("register_mu alloc {name}.weight: {e}")) })?; - // DtoD copy via host - let mut host_buf = vec![0.0_f32; w_len]; - stream.memcpy_dtoh(w_mu, &mut host_buf).map_err(|e| { - MLError::ModelError(format!("register_mu DtoH {name}.weight: {e}")) - })?; - let mut w_clone = w_clone; - stream.memcpy_htod(&host_buf, &mut w_clone).map_err(|e| { - MLError::ModelError(format!("register_mu HtoD {name}.weight: {e}")) - })?; + dtod_copy_slice(w_mu, &mut w_clone, &stream, &format!("register_mu {name}.weight"))?; vars.register(format!("{name}.weight"), w_clone, vec![out_f, in_f])?; - // Clone bias_mu -> "{name}.bias" + // Clone bias_mu -> "{name}.bias" via DtoD memcpy let b_len = b_mu.len(); - let b_clone = stream.alloc_zeros::(b_len).map_err(|e| { + let mut b_clone = stream.alloc_zeros::(b_len).map_err(|e| { MLError::ModelError(format!("register_mu alloc {name}.bias: {e}")) })?; - let mut host_buf = vec![0.0_f32; b_len]; - stream.memcpy_dtoh(b_mu, &mut host_buf).map_err(|e| { - MLError::ModelError(format!("register_mu DtoH {name}.bias: {e}")) - })?; - let mut b_clone = b_clone; - stream.memcpy_htod(&host_buf, &mut b_clone).map_err(|e| { - MLError::ModelError(format!("register_mu HtoD {name}.bias: {e}")) - })?; + dtod_copy_slice(b_mu, &mut b_clone, &stream, &format!("register_mu {name}.bias"))?; vars.register(format!("{name}.bias"), b_clone, vec![out_f])?; Ok(()) @@ -1046,42 +1061,51 @@ impl BranchingDuelingQNetwork { /// Copy weights from another branching network (target network sync). /// /// Copies both `GpuVarStore` vars (shared encoder) AND `NoisyLinear` head vars. + /// Zero CPU downloads — all copies are async device-to-device memcpy. pub fn copy_weights_from(&mut self, other: &BranchingDuelingQNetwork) -> Result<(), MLError> { - // 1. Copy GpuVarStore vars (shared encoder layers) + // 1. Copy GpuVarStore vars (shared encoder layers) via DtoD memcpy for (name, other_param) in other.vars.iter() { if let Some(self_param) = self.vars.get_mut(name) { - let mut host = vec![0.0_f32; other_param.data.len()]; - other.stream.memcpy_dtoh(&other_param.data, &mut host).map_err(|e| { - MLError::ModelError(format!("Copy weight DtoH {}: {}", name, e)) - })?; - self.stream.memcpy_htod(&host, &mut self_param.data).map_err(|e| { - MLError::ModelError(format!("Copy weight HtoD {}: {}", name, e)) - })?; + let num_bytes = other_param.data.len() * std::mem::size_of::(); + let src_ptr = { + let (ptr, guard) = other_param.data.device_ptr(&self.stream); + let _no_drop = ManuallyDrop::new(guard); + ptr + }; + let dst_ptr = { + let (ptr, guard) = self_param.data.device_ptr_mut(&self.stream); + let _no_drop = ManuallyDrop::new(guard); + ptr + }; + // SAFETY: src and dst are valid device allocations on the same context. + // num_bytes = len * sizeof(f32) does not exceed either allocation. + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(), + ).map_err(|e| { + MLError::ModelError(format!("Copy weight DtoD {name}: {e}")) + })?; + } } } - // 2. Copy NoisyLinear head vars (not in GpuVarStore -- standalone CudaSlice) - // NoisyLinear copy is done via GpuTensor roundtrip (host intermediary). - // This is the cold path (target network sync, not hot training). - Self::copy_noisy_layer(&mut self.value_fc, &other.value_fc, "value_fc", &self.stream, &other.stream)?; - Self::copy_noisy_layer(&mut self.value_out, &other.value_out, "value_out", &self.stream, &other.stream)?; - for (d, (self_fc, other_fc)) in self.branch_fcs.iter_mut().zip(other.branch_fcs.iter()).enumerate() { - Self::copy_noisy_layer(self_fc, other_fc, &format!("branch_{}_fc", d), &self.stream, &other.stream)?; + // 2. Copy NoisyLinear head vars via DtoD memcpy (delegated to NoisyLinear::copy_params_from) + Self::copy_noisy_layer(&mut self.value_fc, &other.value_fc)?; + Self::copy_noisy_layer(&mut self.value_out, &other.value_out)?; + for (self_fc, other_fc) in self.branch_fcs.iter_mut().zip(other.branch_fcs.iter()) { + Self::copy_noisy_layer(self_fc, other_fc)?; } - for (d, (self_out, other_out)) in self.branch_outs.iter_mut().zip(other.branch_outs.iter()).enumerate() { - Self::copy_noisy_layer(self_out, other_out, &format!("branch_{}_out", d), &self.stream, &other.stream)?; + for (self_out, other_out) in self.branch_outs.iter_mut().zip(other.branch_outs.iter()) { + Self::copy_noisy_layer(self_out, other_out)?; } Ok(()) } - /// Copy `NoisyLinear` weights between matching layers via host roundtrip. + /// Copy `NoisyLinear` weights between matching layers via DtoD memcpy. fn copy_noisy_layer( dst: &mut MaybeNoisyLinear, src: &MaybeNoisyLinear, - _label: &str, - _dst_stream: &Arc, - _src_stream: &Arc, ) -> Result<(), MLError> { let MaybeNoisyLinear::Noisy(src_noisy) = src; let MaybeNoisyLinear::Noisy(dst_noisy) = dst; diff --git a/crates/ml-dqn/src/noisy_layers.rs b/crates/ml-dqn/src/noisy_layers.rs index ecfae54d6..bba453613 100644 --- a/crates/ml-dqn/src/noisy_layers.rs +++ b/crates/ml-dqn/src/noisy_layers.rs @@ -10,12 +10,57 @@ //! - Factorized Gaussian noise: `ε_ij` = `f(ε_i)` × `f(ε_j)` where f(x) = sign(x) × √|x| //! - Reduces parameter count by ~70% vs independent noise while maintaining exploration quality +use std::mem::ManuallyDrop; use std::sync::Arc; -use cudarc::driver::{CudaSlice, CudaStream}; +use cudarc::cublas::CudaBlas; +use cudarc::cublas::sys::cublasOperation_t; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; +use cudarc::nvrtc::Ptx; use ml_core::cuda_autograd::GpuTensor; use ml_core::MLError; +// ── Pointer helpers (same pattern as gpu_tensor.rs) ────────────────────── + +/// Extract raw CUDA device pointer from a CudaSlice (read-only). +fn raw_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { + let (ptr, guard) = slice.device_ptr(stream); + let _no_drop = ManuallyDrop::new(guard); + ptr +} + +/// Extract raw mutable CUDA device pointer from a CudaSlice. +fn raw_ptr_mut(slice: &mut CudaSlice, stream: &CudaStream) -> u64 { + let (ptr, guard) = slice.device_ptr_mut(stream); + let _no_drop = ManuallyDrop::new(guard); + ptr +} + +/// Compile the bias-add kernel for NoisyLinear forward. +fn compile_bias_add_kernel(stream: &Arc) -> Result { + let src = r#" +extern "C" __global__ +void noisy_add_bias_kernel(float* __restrict__ y, + const float* __restrict__ bias, + int rows, int cols) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < rows * cols) { + int col = idx % cols; + y[idx] += bias[col]; + } +} +"#; + let context = stream.context(); + let ptx: Ptx = ml_core::cuda_compile::compile_ptx_for_device(src, &context) + .map_err(|e| MLError::ModelError(format!("noisy bias kernel compilation: {e}")))?; + let module = context.load_module(ptx).map_err(|e| { + MLError::ModelError(format!("noisy bias module load: {e}")) + })?; + module.load_function("noisy_add_bias_kernel").map_err(|e| { + MLError::ModelError(format!("noisy_add_bias_kernel load: {e}")) + }) +} + /// Noisy linear layer with factorized Gaussian noise (Rainbow DQN standard) /// /// Key features: @@ -185,74 +230,95 @@ impl NoisyLinear { /// uses the fused CUDA kernel `dqn_forward_only_kernel` with BF16 tensor core /// matmul.** /// - /// Computes: y = (mu_w + sigma_w * epsilon_w) x input + (mu_b + sigma_b * epsilon_b) + /// Computes entirely on GPU: y = (mu_w + sigma_w * epsilon_w) x input + (mu_b + sigma_b * epsilon_b) + /// Zero CPU downloads — all arithmetic via CUDA elementwise kernels + cuBLAS sgemm. #[cold] pub fn forward(&self, x: &GpuTensor) -> Result { - // Compute: y = (mu_w + sigma_w * epsilon_w) @ x^T + (mu_b + sigma_b * epsilon_b) - // This is the cold-path implementation using host-side computation. - // Hot path uses fused CUDA kernels in GpuDqnTrainer. let batch = if x.ndim() == 1 { 1 } else { x.shape().first().copied().unwrap_or(1) }; + let in_dim = self.in_features; + let out_dim = self.out_features; + let w_shape = vec![out_dim, in_dim]; + let b_shape = vec![out_dim]; - // Download all needed data to host - let x_host = x.to_host(&self.stream)?; - let mut mu_w_host = vec![0.0_f32; self.out_features * self.in_features]; - self.stream.memcpy_dtoh(&self.weight_mu, &mut mu_w_host).map_err(|e| { - MLError::ModelError(format!("NoisyLinear forward weight_mu DtoH: {e}")) - })?; - let mut sigma_w_host = vec![0.0_f32; self.out_features * self.in_features]; - self.stream.memcpy_dtoh(&self.weight_sigma, &mut sigma_w_host).map_err(|e| { - MLError::ModelError(format!("NoisyLinear forward weight_sigma DtoH: {e}")) - })?; - let mut eps_w_host = vec![0.0_f32; self.out_features * self.in_features]; - self.stream.memcpy_dtoh(&self.weight_epsilon, &mut eps_w_host).map_err(|e| { - MLError::ModelError(format!("NoisyLinear forward weight_epsilon DtoH: {e}")) - })?; - let mut mu_b_host = vec![0.0_f32; self.out_features]; - self.stream.memcpy_dtoh(&self.bias_mu, &mut mu_b_host).map_err(|e| { - MLError::ModelError(format!("NoisyLinear forward bias_mu DtoH: {e}")) - })?; - let mut sigma_b_host = vec![0.0_f32; self.out_features]; - self.stream.memcpy_dtoh(&self.bias_sigma, &mut sigma_b_host).map_err(|e| { - MLError::ModelError(format!("NoisyLinear forward bias_sigma DtoH: {e}")) - })?; - let mut eps_b_host = vec![0.0_f32; self.out_features]; - self.stream.memcpy_dtoh(&self.bias_epsilon, &mut eps_b_host).map_err(|e| { - MLError::ModelError(format!("NoisyLinear forward bias_epsilon DtoH: {e}")) + // Wrap CudaSlice refs as GpuTensor (cheap clone — cudarc 0.19 ref-counted) + let mu_w = GpuTensor::new(self.weight_mu.clone(), w_shape.clone())?; + let sigma_w = GpuTensor::new(self.weight_sigma.clone(), w_shape)?; + let eps_w = GpuTensor::new(self.weight_epsilon.clone(), vec![out_dim, in_dim])?; + + let mu_b = GpuTensor::new(self.bias_mu.clone(), b_shape.clone())?; + let sigma_b = GpuTensor::new(self.bias_sigma.clone(), b_shape)?; + let eps_b = GpuTensor::new(self.bias_epsilon.clone(), vec![out_dim])?; + + // GPU elementwise: effective_weight = mu_w + sigma_w * epsilon_w + let sigma_eps_w = sigma_w.mul(&eps_w, &self.stream)?; + let effective_w = mu_w.add(&sigma_eps_w, &self.stream)?; + + // GPU elementwise: effective_bias = mu_b + sigma_b * epsilon_b + let sigma_eps_b = sigma_b.mul(&eps_b, &self.stream)?; + let effective_b = mu_b.add(&sigma_eps_b, &self.stream)?; + + // cuBLAS sgemm: y = x @ W_eff^T (x:[batch, in], W_eff:[out, in] -> y:[batch, out]) + // Column-major: Y_col[out, B] = W_col^T[out, in] @ X_col[in, B] + // transA=T, transB=N, m=out, n=B, k=in + let cublas = CudaBlas::new(self.stream.clone()).map_err(|e| { + MLError::ModelError(format!("NoisyLinear forward cublas init: {e}")) })?; - // Compute effective weight: W = mu_w + sigma_w * epsilon_w - let effective_w: Vec = mu_w_host.iter() - .zip(sigma_w_host.iter()) - .zip(eps_w_host.iter()) - .map(|((mu, sigma), eps)| mu + sigma * eps) - .collect(); + let mut y = GpuTensor::zeros(&[batch, out_dim], &self.stream)?; + let w_ptr = raw_ptr(effective_w.data(), &self.stream); + let x_ptr = raw_ptr(x.data(), &self.stream); + let y_ptr = raw_ptr_mut(y.data_mut(), &self.stream); - // Compute effective bias: b = mu_b + sigma_b * epsilon_b - let effective_b: Vec = mu_b_host.iter() - .zip(sigma_b_host.iter()) - .zip(eps_b_host.iter()) - .map(|((mu, sigma), eps)| mu + sigma * eps) - .collect(); - - // y = x @ W^T + b (x: [batch, in], W: [out, in] -> y: [batch, out]) - let mut output = vec![0.0_f32; batch * self.out_features]; - for b in 0..batch { - for o in 0..self.out_features { - let mut sum = effective_b.get(o).copied().unwrap_or(0.0); - for i in 0..self.in_features { - let x_val = x_host.get(b * self.in_features + i).copied().unwrap_or(0.0); - let w_val = effective_w.get(o * self.in_features + i).copied().unwrap_or(0.0); - sum += x_val * w_val; - } - if let Some(slot) = output.get_mut(b * self.out_features + o) { - *slot = sum; - } - } + // SAFETY: w_ptr, x_ptr, y_ptr are valid device allocations on same context. + // Dimensions are consistent: W[out, in] x X^T[in, batch] -> Y[out, batch]. + unsafe { + cudarc::cublas::result::sgemm( + *cublas.handle(), + cublasOperation_t::CUBLAS_OP_T, // transA: W stored [out, in], need W^T + cublasOperation_t::CUBLAS_OP_N, // transB: X stored [batch, in] = X_col[in, batch] + out_dim as i32, // m + batch as i32, // n + in_dim as i32, // k + &1.0_f32 as *const f32, // alpha + w_ptr as *const f32, // A = W_eff + in_dim as i32, // lda + x_ptr as *const f32, // B = X + in_dim as i32, // ldb + &0.0_f32 as *const f32, // beta + y_ptr as *mut f32, // C = Y + out_dim as i32, // ldc + ).map_err(|e| MLError::ModelError(format!("NoisyLinear cuBLAS sgemm: {e:?}")))?; } - GpuTensor::from_host(&output, vec![batch, self.out_features], &self.stream) + // Add bias: Y[b, j] += effective_bias[j] via CUDA kernel + let bias_fn = compile_bias_add_kernel(&self.stream)?; + let total = batch * out_dim; + let threads = 256_u32; + let blocks = ((total as u32) + threads - 1) / threads; + let rows_i32 = batch as i32; + let cols_i32 = out_dim as i32; + + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (threads, 1, 1), + shared_mem_bytes: 0, + }; + + // SAFETY: y and effective_b are valid device allocations, kernel indices are bounds-checked. + unsafe { + self.stream + .launch_builder(&bias_fn) + .arg(y.data()) + .arg(effective_b.data()) + .arg(&rows_i32) + .arg(&cols_i32) + .launch(launch_cfg) + .map_err(|e| MLError::ModelError(format!("NoisyLinear bias kernel: {e}")))?; + } + + Ok(y) } /// Get all learnable parameters as CudaSlice references (for optimizer) @@ -295,64 +361,37 @@ impl NoisyLinear { Ok(()) } - /// Copy all learnable parameters from `src` to `self` via host roundtrip. + /// Copy all learnable parameters from `src` to `self` via GPU DtoD memcpy. /// - /// Copies weight_mu, bias_mu, weight_sigma, bias_sigma. Epsilon buffers - /// are NOT copied (they are resampled independently each forward pass). + /// 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> { - // weight_mu - let mut buf = vec![0.0_f32; src.weight_mu.len()]; - src.stream.memcpy_dtoh(&src.weight_mu, &mut buf).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy weight_mu DtoH: {e}")) - })?; - self.stream.memcpy_htod(&buf, &mut self.weight_mu).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy weight_mu HtoD: {e}")) - })?; - - // bias_mu - let mut buf = vec![0.0_f32; src.bias_mu.len()]; - src.stream.memcpy_dtoh(&src.bias_mu, &mut buf).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy bias_mu DtoH: {e}")) - })?; - self.stream.memcpy_htod(&buf, &mut self.bias_mu).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy bias_mu HtoD: {e}")) - })?; - - // weight_sigma - let mut buf = vec![0.0_f32; src.weight_sigma.len()]; - src.stream.memcpy_dtoh(&src.weight_sigma, &mut buf).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy weight_sigma DtoH: {e}")) - })?; - self.stream.memcpy_htod(&buf, &mut self.weight_sigma).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy weight_sigma HtoD: {e}")) - })?; - - // bias_sigma - let mut buf = vec![0.0_f32; src.bias_sigma.len()]; - src.stream.memcpy_dtoh(&src.bias_sigma, &mut buf).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy bias_sigma DtoH: {e}")) - })?; - self.stream.memcpy_htod(&buf, &mut self.bias_sigma).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy bias_sigma HtoD: {e}")) - })?; - - // Also copy epsilon buffers for deterministic eval - let mut buf = vec![0.0_f32; src.weight_epsilon.len()]; - src.stream.memcpy_dtoh(&src.weight_epsilon, &mut buf).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy weight_epsilon DtoH: {e}")) - })?; - self.stream.memcpy_htod(&buf, &mut self.weight_epsilon).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy weight_epsilon HtoD: {e}")) - })?; - - let mut buf = vec![0.0_f32; src.bias_epsilon.len()]; - src.stream.memcpy_dtoh(&src.bias_epsilon, &mut buf).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy bias_epsilon DtoH: {e}")) - })?; - self.stream.memcpy_htod(&buf, &mut self.bias_epsilon).map_err(|e| { - MLError::ModelError(format!("NoisyLinear copy bias_epsilon HtoD: {e}")) - })?; + 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)?; + Ok(()) + } + /// Async device-to-device memcpy for a single CudaSlice buffer. + fn dtod_copy( + src: &CudaSlice, + dst: &mut CudaSlice, + label: &str, + stream: &Arc, + ) -> Result<(), MLError> { + let num_bytes = src.len() * std::mem::size_of::(); + let src_ptr = raw_ptr(src, stream); + let dst_ptr = raw_ptr_mut(dst, stream); + // SAFETY: src and dst are valid device allocations on the same context. + // num_bytes = len * sizeof(f32) does not exceed either allocation. + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("NoisyLinear DtoD {label}: {e}")))?; + } Ok(()) } }