gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(cuda): NoisyLinear forward + copy fully GPU-native — zero CPU

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 17:48:43 +01:00
parent c4f3b643e2
commit f782ec5a12
2 changed files with 219 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

110
crates/ml-dqn/src/branching.rs
View File

@@ -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<f32>,
dst: &mut CudaSlice<f32>,
stream: &Arc<CudaStream>,
label: &str,
) -> Result<(), MLError> {
let num_bytes = src.len() * std::mem::size_of::<f32>();
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::<f32>(w_len).map_err(|e| {
let mut w_clone = stream.alloc_zeros::<f32>(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::<f32>(b_len).map_err(|e| {
let mut b_clone = stream.alloc_zeros::<f32>(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::<f32>();
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<CudaStream>,
_src_stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
let MaybeNoisyLinear::Noisy(src_noisy) = src;
let MaybeNoisyLinear::Noisy(dst_noisy) = dst;

265
crates/ml-dqn/src/noisy_layers.rs
View File

@@ -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<f32>, 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<f32>, 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<CudaStream>) -> Result<cudarc::driver::CudaFunction, MLError> {
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<GpuTensor, MLError> {
// 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<f32> = 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<f32> = 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<f32>,
dst: &mut CudaSlice<f32>,
label: &str,
stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
let num_bytes = src.len() * std::mem::size_of::<f32>();
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(())
}
}