From 08ddf10a1ae2bf0405d27a77e98d5624546b1ad5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 20:56:37 +0200 Subject: [PATCH 1/5] fix(action-selector): migrate set_count_bonuses to mapped pinned (HOT path) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned (cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path. `set_count_bonuses` is called per action-selection — converting the three bonus buffers from CudaSlice (HtoD memcpy) to MappedF32Buffer eliminates 3-6 HtoD memcpys per call. Before: 3 reused-buffer htod_f32 (lines 92/96/100) + 3 first-call clone_htod_f32 (lines 93/97/101) on every call = 3 HtoD steady-state. After: zero HtoD — direct host_ptr writes via write_from_slice. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/gpu_action_selector.rs | 62 ++++++++++++------- docs/dqn-wire-up-audit.md | 2 +- 2 files changed, 39 insertions(+), 25 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index b830007b8..367c85588 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use tracing::info; use crate::MLError; +use super::mapped_pinned::MappedF32Buffer; /// Precompiled epsilon_greedy_kernel cubin, embedded at compile time by build.rs. static EPSILON_GREEDY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_greedy_kernel.cubin")); @@ -35,14 +36,17 @@ pub struct GpuActionSelector { max_batch_size: usize, stream: Arc, q_gap_threshold: f32, - /// UCB count bonus for branching heads: [5], [3], [3] on GPU. + /// UCB count bonus for branching heads: [4], [3], [3] in mapped pinned memory. /// Device pointer = 0 means disabled (kernel treats NULL as no bonus). + /// Mapped pinned per `feedback_no_htod_htoh_only_mapped_pinned.md` — + /// `set_count_bonuses` is hot (per action-selection call). Direct host_ptr + /// writes via `write_from_slice` eliminate the per-call HtoD memcpy. bonus_exposure_ptr: u64, bonus_order_ptr: u64, bonus_urgency_ptr: u64, - bonus_exposure_buf: Option>, - bonus_order_buf: Option>, - bonus_urgency_buf: Option>, + bonus_exposure_buf: Option, + bonus_order_buf: Option, + bonus_urgency_buf: Option, /// Per-sample epsilon buffer for branching kernel: [max_batch_size] on GPU. /// Filled via GPU-side fill_f32 kernel before each branching launch. epsilon_buf: CudaSlice, @@ -85,27 +89,37 @@ impl GpuActionSelector { /// Upload per-branch UCB count bonuses for branching action selection. /// Bonuses are added to Q-values before argmax in the greedy path. - /// Pass empty slices or zeros to disable. + /// Pass zeros to disable (kernel sees zero contribution; same effect as NULL). + /// + /// HOT path: called per action-selection. Mapped pinned buffers (allocated + /// lazily on first call) make subsequent updates pure host_ptr writes — + /// no HtoD memcpy. Per `feedback_no_htod_htoh_only_mapped_pinned.md`. pub fn set_count_bonuses(&mut self, exposure: &[f32; 4], order: &[f32; 3], urgency: &[f32; 3]) -> Result<(), MLError> { - // Allocate or reuse GPU buffers - let be = match self.bonus_exposure_buf.take() { - Some(mut buf) => { super::htod_f32(&self.stream, exposure, &mut buf)?; buf } - None => super::clone_htod_f32(&self.stream, exposure)?, - }; - let bo = match self.bonus_order_buf.take() { - Some(mut buf) => { super::htod_f32(&self.stream, order, &mut buf)?; buf } - None => super::clone_htod_f32(&self.stream, order)?, - }; - let bu = match self.bonus_urgency_buf.take() { - Some(mut buf) => { super::htod_f32(&self.stream, urgency, &mut buf)?; buf } - None => super::clone_htod_f32(&self.stream, urgency)?, - }; - self.bonus_exposure_ptr = be.device_ptr(&self.stream).0; - self.bonus_order_ptr = bo.device_ptr(&self.stream).0; - self.bonus_urgency_ptr = bu.device_ptr(&self.stream).0; - self.bonus_exposure_buf = Some(be); - self.bonus_order_buf = Some(bo); - self.bonus_urgency_buf = Some(bu); + // Lazy-allocate mapped pinned buffers on first call. Safety: a CUDA + // context is active because `self.stream` was constructed against it. + if self.bonus_exposure_buf.is_none() { + let buf = unsafe { MappedF32Buffer::new(4) } + .map_err(|e| MLError::ModelError(format!("alloc bonus_exposure mapped pinned: {e}")))?; + self.bonus_exposure_ptr = buf.dev_ptr; + self.bonus_exposure_buf = Some(buf); + } + if self.bonus_order_buf.is_none() { + let buf = unsafe { MappedF32Buffer::new(3) } + .map_err(|e| MLError::ModelError(format!("alloc bonus_order mapped pinned: {e}")))?; + self.bonus_order_ptr = buf.dev_ptr; + self.bonus_order_buf = Some(buf); + } + if self.bonus_urgency_buf.is_none() { + let buf = unsafe { MappedF32Buffer::new(3) } + .map_err(|e| MLError::ModelError(format!("alloc bonus_urgency mapped pinned: {e}")))?; + self.bonus_urgency_ptr = buf.dev_ptr; + self.bonus_urgency_buf = Some(buf); + } + // Direct host_ptr writes — no memcpy. Kernel reads via dev_ptr after + // stream sync barrier (mapped pinned coherence). + self.bonus_exposure_buf.as_ref().unwrap().write_from_slice(exposure); + self.bonus_order_buf.as_ref().unwrap().write_from_slice(order); + self.bonus_urgency_buf.as_ref().unwrap().write_from_slice(urgency); Ok(()) } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index ababd2efa..ef022574b 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -1100,7 +1100,7 @@ P5T5 Phase E (2026-04-26): per-fold reset for `MetricBandsRegistry` regression-d | `cuda_pipeline/shared_cublas_handle.rs` | `fused_training.rs`, `gpu_iqn_head.rs`, `gpu_iql_trainer.rs`, `gpu_attention.rs`, `gpu_curiosity_trainer.rs` (10 consumers) | Wired | Shared cuBLAS/cuBLASLt handle | — | | `cuda_pipeline/cublas_algo_deterministic.rs` | `gpu_dqn_trainer.rs`, `gpu_curiosity_trainer.rs`, `gpu_iqn_head.rs` (7 consumers) | Wired | Deterministic cuBLASLt algo selection | — | | `cuda_pipeline/gpu_weights.rs` | `trainer/mod.rs`, `gpu_dqn_trainer.rs`, `hyperopt/adapters/dqn.rs` (11 consumers) | Wired | Weight tensor layout constants | — | -| `cuda_pipeline/gpu_action_selector.rs` (`GpuActionSelector`) | `trainer/mod.rs`, `fused_training.rs`, `gpu_backtest_evaluator.rs` | Wired | Epsilon-greedy + routed action selection | — | +| `cuda_pipeline/gpu_action_selector.rs` (`GpuActionSelector`) | `trainer/mod.rs`, `fused_training.rs`, `gpu_backtest_evaluator.rs` | Wired | Epsilon-greedy + routed action selection. UCB count-bonus buffers (exposure/order/urgency) are mapped pinned per `feedback_no_htod_htoh_only_mapped_pinned.md` — eliminates per-call HtoD memcpys in `set_count_bonuses` (HOT). | — | | `cuda_pipeline/gpu_monitoring.rs` (`GpuMonitor`) | `fused_training.rs`, `trainer/metrics.rs`, `trainer/training_loop.rs` | Wired | GPU-side monitoring_reduce launch | — | | `cuda_pipeline/gpu_training_guard.rs` | `gpu_dqn_trainer.rs`, `trainer/training_loop.rs` | Wired | NaN / gradient anomaly guard | — | | `cuda_pipeline/gpu_backtest_evaluator.rs` (`GpuBacktestEvaluator`) | `trainer/metrics.rs` (eval path) | Wired | Per-epoch GPU backtest evaluation | — | From 0fb21c970cfbe6b269f25fe935dca6de39947d76 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 21:00:27 +0200 Subject: [PATCH 2/5] fix(action-selector): migrate rng_states to mapped pinned (COLD ctor) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned (cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path. `GpuActionSelector::new` previously did one HtoD memcpy to upload RNG seeds. Now allocates `MappedU32Buffer`, writes seeds via host_ptr, and passes `dev_ptr` (CUdeviceptr) to the three kernels that consume rng_states (epsilon_greedy, epsilon_greedy_routed, branching_action_select). Adds `MappedU32Buffer` and `MappedU64Buffer` to `mapped_pinned.rs` mirroring the existing `MappedF32Buffer`/`MappedI32Buffer` API (new/write_from_slice/read_all/Drop). The U64 variant is staged for the upcoming spectral-norm host_desc[78] descriptor table migration in `gpu_dqn_trainer::new`. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/gpu_action_selector.rs | 38 +++-- crates/ml/src/cuda_pipeline/mapped_pinned.rs | 160 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 9 + 3 files changed, 195 insertions(+), 12 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 367c85588..30f8743fe 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use tracing::info; use crate::MLError; -use super::mapped_pinned::MappedF32Buffer; +use super::mapped_pinned::{MappedF32Buffer, MappedU32Buffer}; /// Precompiled epsilon_greedy_kernel cubin, embedded at compile time by build.rs. static EPSILON_GREEDY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_greedy_kernel.cubin")); @@ -30,7 +30,10 @@ pub struct GpuActionSelector { routed_kernel_func: CudaFunction, branching_kernel_func: CudaFunction, route_func: CudaFunction, - rng_states: CudaSlice, + /// RNG state, mapped pinned. Kernel reads & writes via `rng_states_ptr`; + /// the seed array is initialised via direct host_ptr writes (no HtoD). + rng_states: MappedU32Buffer, + rng_states_ptr: u64, actions_buf: CudaSlice, fill_mask_buf: CudaSlice, max_batch_size: usize, @@ -65,16 +68,24 @@ impl GpuActionSelector { let fill_f32_kernel = module.load_function("fill_f32").map_err(|e| MLError::ModelError(format!("fill_f32 function load: {e}")))?; let actions_buf = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc actions_buf: {e}")))?; let fill_mask_buf = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc fill_mask_buf: {e}")))?; - let mut rng_seeds = Vec::with_capacity(max_batch_size); - for i in 0..max_batch_size { - let s = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(i as u64); - rng_seeds.push((s >> 32) as u32 | 1); + // RNG state — mapped pinned per `feedback_no_htod_htoh_only_mapped_pinned.md`. + // Seeds written via direct host_ptr writes; kernel reads & writes + // back via rng_states_ptr (CUdeviceptr). + // Safety: a CUDA context is active (we just resolved kernels through it). + let rng_states = unsafe { MappedU32Buffer::new(max_batch_size) } + .map_err(|e| MLError::ModelError(format!("alloc rng_states mapped pinned: {e}")))?; + let rng_states_ptr = rng_states.dev_ptr; + { + let mut rng_seeds = Vec::with_capacity(max_batch_size); + for i in 0..max_batch_size { + let s = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(i as u64); + rng_seeds.push((s >> 32) as u32 | 1); + } + rng_states.write_from_slice(&rng_seeds); } - let mut rng_states = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?; - stream.memcpy_htod(&rng_seeds, &mut rng_states).map_err(|e| MLError::ModelError(format!("upload rng_states: {e}")))?; let epsilon_buf = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc epsilon_buf: {e}")))?; info!("GpuActionSelector initialized: max_batch_size={max_batch_size}, precompiled cubin loaded"); - Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, actions_buf, fill_mask_buf, max_batch_size, stream, q_gap_threshold: 0.0, + Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, rng_states_ptr, actions_buf, fill_mask_buf, max_batch_size, stream, q_gap_threshold: 0.0, bonus_exposure_ptr: 0, bonus_order_ptr: 0, bonus_urgency_ptr: 0, bonus_exposure_buf: None, bonus_order_buf: None, bonus_urgency_buf: None, epsilon_buf, @@ -130,9 +141,10 @@ impl GpuActionSelector { let bs_i32 = batch_size as i32; let na_i32 = num_actions as i32; let q_gap = self.q_gap_threshold; + let rng_ptr = self.rng_states_ptr; unsafe { self.stream.launch_builder(&self.kernel_func) - .arg(q_values).arg(&mut self.rng_states).arg(&mut self.actions_buf) + .arg(q_values).arg(&rng_ptr).arg(&mut self.actions_buf) .arg(&epsilon).arg(&bs_i32).arg(&na_i32).arg(&q_gap) .launch(config).map_err(|e| MLError::ModelError(format!("epsilon_greedy kernel launch: {e}")))?; } @@ -150,9 +162,10 @@ impl GpuActionSelector { let config = launch_config_1d(batch_size); let bs_i32 = batch_size as i32; let na_i32 = num_actions as i32; + let rng_ptr = self.rng_states_ptr; unsafe { self.stream.launch_builder(&self.routed_kernel_func) - .arg(q_values).arg(&mut self.rng_states).arg(&mut self.actions_buf).arg(&mut self.fill_mask_buf) + .arg(q_values).arg(&rng_ptr).arg(&mut self.actions_buf).arg(&mut self.fill_mask_buf) .arg(&epsilon).arg(&bs_i32).arg(&na_i32).arg(&step_offset) .arg(&spread).arg(&median_spread).arg(&volatility).arg(&median_vol) .arg(&spread_bps).arg(&ioc_fill_prob).arg(&limit_fill_min).arg(&limit_fill_max) @@ -183,10 +196,11 @@ impl GpuActionSelector { let be_ptr = self.bonus_exposure_ptr; let bo_ptr = self.bonus_order_ptr; let bu_ptr = self.bonus_urgency_ptr; + let rng_ptr = self.rng_states_ptr; unsafe { self.stream.launch_builder(&self.branching_kernel_func) .arg(q_exposure).arg(q_order).arg(q_urgency) - .arg(&mut self.rng_states).arg(&mut self.actions_buf) + .arg(&rng_ptr).arg(&mut self.actions_buf) .arg(&self.epsilon_buf) .arg(&bs_i32).arg(&self.q_gap_threshold) .arg(&be_ptr).arg(&bo_ptr).arg(&bu_ptr) // UCB count bonuses (0 = NULL = disabled) diff --git a/crates/ml/src/cuda_pipeline/mapped_pinned.rs b/crates/ml/src/cuda_pipeline/mapped_pinned.rs index 039398868..1ac89e4e7 100644 --- a/crates/ml/src/cuda_pipeline/mapped_pinned.rs +++ b/crates/ml/src/cuda_pipeline/mapped_pinned.rs @@ -191,3 +191,163 @@ impl Drop for MappedF32Buffer { } } } + +// ── MappedU32Buffer ─────────────────────────────────────────────────────────── + +/// CPU+GPU visible buffer of `u32`s allocated via +/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. Used for RNG seed/state arrays where +/// the kernel reads the seeds via `dev_ptr` (no HtoD upload required). +pub struct MappedU32Buffer { + pub host_ptr: *mut u32, + pub dev_ptr: cudarc::driver::sys::CUdeviceptr, + pub len: usize, +} + +unsafe impl Send for MappedU32Buffer {} +unsafe impl Sync for MappedU32Buffer {} + +impl MappedU32Buffer { + /// Allocate `len` u32s of mapped pinned memory. + /// + /// # Safety + /// Caller must ensure a CUDA context is active on the current thread. + pub unsafe fn new(len: usize) -> Result { + let num_bytes = len * std::mem::size_of::(); + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP + | cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE; + + let host_ptr = cudarc::driver::result::malloc_host(num_bytes, flags) + .map_err(|e| format!("MappedU32Buffer alloc ({len} u32): {e}"))? + as *mut u32; + + std::ptr::write_bytes(host_ptr, 0, len); + + let mut dev_ptr_raw = MaybeUninit::uninit(); + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + dev_ptr_raw.as_mut_ptr(), + host_ptr as *mut c_void, + 0, + ) + .result() + .map_err(|e| format!("cuMemHostGetDevicePointer (u32 buf): {e}"))?; + + Ok(Self { + host_ptr, + dev_ptr: dev_ptr_raw.assume_init(), + len, + }) + } + + /// Read all `len` entries via `read_volatile`. Caller must have + /// synchronised the producing stream first. + pub fn read_all(&self) -> Vec { + let mut out = Vec::with_capacity(self.len); + unsafe { + for i in 0..self.len { + out.push(std::ptr::read_volatile(self.host_ptr.add(i))); + } + } + out + } + + /// Write CPU-side data into the buffer via host_ptr. + /// Direct memory write — no memcpy. Caller must ensure + /// `slice.len() <= self.len`. + pub fn write_from_slice(&self, slice: &[u32]) { + assert!(slice.len() <= self.len, "MappedU32Buffer write overflow"); + unsafe { + for (i, &v) in slice.iter().enumerate() { + std::ptr::write_volatile(self.host_ptr.add(i), v); + } + } + } +} + +impl Drop for MappedU32Buffer { + fn drop(&mut self) { + unsafe { + #[allow(clippy::let_underscore_must_use)] + let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void); + } + } +} + +// ── MappedU64Buffer ─────────────────────────────────────────────────────────── + +/// CPU+GPU visible buffer of `u64`s allocated via +/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. Used for descriptor tables (e.g. +/// spectral-norm pointer/dim arrays) where each entry is a 64-bit value. +pub struct MappedU64Buffer { + pub host_ptr: *mut u64, + pub dev_ptr: cudarc::driver::sys::CUdeviceptr, + pub len: usize, +} + +unsafe impl Send for MappedU64Buffer {} +unsafe impl Sync for MappedU64Buffer {} + +impl MappedU64Buffer { + /// Allocate `len` u64s of mapped pinned memory. + /// + /// # Safety + /// Caller must ensure a CUDA context is active on the current thread. + pub unsafe fn new(len: usize) -> Result { + let num_bytes = len * std::mem::size_of::(); + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP + | cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE; + + let host_ptr = cudarc::driver::result::malloc_host(num_bytes, flags) + .map_err(|e| format!("MappedU64Buffer alloc ({len} u64): {e}"))? + as *mut u64; + + std::ptr::write_bytes(host_ptr, 0, len); + + let mut dev_ptr_raw = MaybeUninit::uninit(); + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + dev_ptr_raw.as_mut_ptr(), + host_ptr as *mut c_void, + 0, + ) + .result() + .map_err(|e| format!("cuMemHostGetDevicePointer (u64 buf): {e}"))?; + + Ok(Self { + host_ptr, + dev_ptr: dev_ptr_raw.assume_init(), + len, + }) + } + + /// Read all `len` entries via `read_volatile`. Caller must have + /// synchronised the producing stream first. + pub fn read_all(&self) -> Vec { + let mut out = Vec::with_capacity(self.len); + unsafe { + for i in 0..self.len { + out.push(std::ptr::read_volatile(self.host_ptr.add(i))); + } + } + out + } + + /// Write CPU-side data into the buffer via host_ptr. + /// Direct memory write — no memcpy. Caller must ensure + /// `slice.len() <= self.len`. + pub fn write_from_slice(&self, slice: &[u64]) { + assert!(slice.len() <= self.len, "MappedU64Buffer write overflow"); + unsafe { + for (i, &v) in slice.iter().enumerate() { + std::ptr::write_volatile(self.host_ptr.add(i), v); + } + } + } +} + +impl Drop for MappedU64Buffer { + fn drop(&mut self) { + unsafe { + #[allow(clippy::let_underscore_must_use)] + let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void); + } + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index ef022574b..9e1221770 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -1408,6 +1408,15 @@ one implementation. Adds `write_from_slice` helper for direct host_ptr writes (no memcpy). Test 0.F bit-identical post-move. Per `feedback_no_htod_htoh_only_mapped_pinned.md`. +Mapped pinned types extended (2026-04-28): `MappedU32Buffer` and +`MappedU64Buffer` added to `cuda_pipeline/mapped_pinned.rs` to cover RNG +state arrays (u32) and descriptor/pointer tables (u64). First consumers: +`gpu_action_selector::rng_states` migration (eliminates HtoD seed upload +in constructor; kernel reads/writes via dev_ptr) and the upcoming +`gpu_dqn_trainer::new` constructor block (spectral-norm host_desc[78] +u64 table). All four mapped pinned variants share the same +`new`/`write_from_slice`/`read_all` API. + MoE moe_mixture_forward kernel + Rust wrapper (2026-04-27): first MoE CUDA kernel landed. Single-thread-per-(b,c) kernel computes h_s2[b,c] = Σ_k g[b,k]·expert_outputs[k,b,c]. No atomicAdd, capture-friendly. Rust From 87ea9fa4168b30ce87bf47a10a983fd2ff485bf9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 21:06:13 +0200 Subject: [PATCH 3/5] fix(gpu_dqn_trainer): migrate 11 COLD ctor HtoD sites to mapped pinned MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned (cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path. Adds module-level `upload_via_mapped_{f32,i32,u32,u64}` helpers and an in-place `update_via_mapped_f32`. Each stages CPU data through a transient mapped pinned buffer and DtoD-copies into the destination `CudaSlice`, then stream-syncs so the staging buffer is safe to drop. Destination buffer types remain `CudaSlice` so every existing consumer (raw_ptr / kernel arg) is untouched, satisfying `feedback_no_partial_refactor.md` for these one-shot init paths. Migrated COLD sites in `GpuDqnTrainer::new`: - weight_decay_mask (TOTAL_PARAMS f32) - branch_slice_starts_dev, branch_slice_lens_dev ([i32; 4]) - branch_grad_scales_dev ([f32; 4]) - per_branch_gamma_base/max_dev ([f32; 4] each) - q_quantile_branch_offsets/sizes_dev ([i32; 4] each) - spectral_norm_descriptors_dev ([u64; 78]) - stochastic_depth_scale_buf ([f32; 3]) - stochastic_depth_rng_state ([u32; 1]) - vsn_group_begins_buf, vsn_group_ends_buf ([i32; num_groups] each) - mamba2_params Xavier init (mamba2_param_count f32) 11 COLD HtoD sites eliminated. cargo check clean (15 warnings; +2 over baseline 13 are the new MappedU32/U64 struct visibility warnings, matching the existing MappedF32/I32 pattern). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 240 +++++++++++++----- docs/dqn-wire-up-audit.md | 15 ++ 2 files changed, 188 insertions(+), 67 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index f27ad0611..77d7eff05 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -64,6 +64,7 @@ use super::gpu_aux_heads::{ AuxHeadsBackwardOps, AuxHeadsForwardOps, AUX_HIDDEN_DIM, AUX_NEXT_BAR_K, AUX_REGIME_K, }; use super::gpu_moe_head::GpuMoeHead; +use super::mapped_pinned::{MappedF32Buffer, MappedI32Buffer, MappedU32Buffer, MappedU64Buffer}; // ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ────── pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); @@ -1719,6 +1720,157 @@ pub(crate) fn padded_byte_offset(param_sizes: &[usize], idx: usize) -> u64 { .sum::() as u64 } +// ── Mapped-pinned upload helpers ───────────────────────────────────────────── +// Per `feedback_no_htod_htoh_only_mapped_pinned.md`, HtoD memcpys are +// forbidden. These helpers stage CPU data through a transient mapped pinned +// buffer and DtoD-copy into a regular GPU-resident `CudaSlice` — preserving +// the existing kernel-arg surface while eliminating the HtoD step. +// COLD path only: each helper allocates + frees a mapped buffer per call and +// stream-syncs to keep the staging buffer alive until the DtoD completes. + +fn upload_via_mapped_f32( + stream: &Arc, + n: usize, + data: &[f32], + label: &str, +) -> Result, MLError> { + use cudarc::driver::DevicePtrMut; + assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len()); + let mut buf = stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("{label} alloc: {e}")))?; + // Safety: a CUDA context is active on this thread (the stream was built + // against it; the caller is mid-construction and holds it live). + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?; + staging.write_from_slice(data); + let n_bytes = n * std::mem::size_of::(); + { + let (dst_ptr, _g) = buf.device_ptr_mut(stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("{label} dtod: {e}")))?; + } + } + // Sync so staging is safe to drop (DtoD completes before drop frees it). + stream.synchronize() + .map_err(|e| MLError::ModelError(format!("{label} dtod sync: {e}")))?; + Ok(buf) +} + +fn upload_via_mapped_i32( + stream: &Arc, + n: usize, + data: &[i32], + label: &str, +) -> Result, MLError> { + use cudarc::driver::DevicePtrMut; + assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len()); + let mut buf = stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("{label} alloc: {e}")))?; + let staging = unsafe { MappedI32Buffer::new(n) } + .map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?; + staging.write_from_slice(data); + let n_bytes = n * std::mem::size_of::(); + { + let (dst_ptr, _g) = buf.device_ptr_mut(stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("{label} dtod: {e}")))?; + } + } + stream.synchronize() + .map_err(|e| MLError::ModelError(format!("{label} dtod sync: {e}")))?; + Ok(buf) +} + +fn upload_via_mapped_u32( + stream: &Arc, + n: usize, + data: &[u32], + label: &str, +) -> Result, MLError> { + use cudarc::driver::DevicePtrMut; + assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len()); + let mut buf = stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("{label} alloc: {e}")))?; + let staging = unsafe { MappedU32Buffer::new(n) } + .map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?; + staging.write_from_slice(data); + let n_bytes = n * std::mem::size_of::(); + { + let (dst_ptr, _g) = buf.device_ptr_mut(stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("{label} dtod: {e}")))?; + } + } + stream.synchronize() + .map_err(|e| MLError::ModelError(format!("{label} dtod sync: {e}")))?; + Ok(buf) +} + +fn upload_via_mapped_u64( + stream: &Arc, + n: usize, + data: &[u64], + label: &str, +) -> Result, MLError> { + use cudarc::driver::DevicePtrMut; + assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len()); + let mut buf = stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("{label} alloc: {e}")))?; + let staging = unsafe { MappedU64Buffer::new(n) } + .map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?; + staging.write_from_slice(data); + let n_bytes = n * std::mem::size_of::(); + { + let (dst_ptr, _g) = buf.device_ptr_mut(stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("{label} dtod: {e}")))?; + } + } + stream.synchronize() + .map_err(|e| MLError::ModelError(format!("{label} dtod sync: {e}")))?; + Ok(buf) +} + +/// Update an existing GPU-resident `CudaSlice` from a host slice via +/// mapped-pinned staging + DtoD copy. WARM path equivalent of +/// `upload_via_mapped_f32` for in-place updates of existing buffers. +fn update_via_mapped_f32( + stream: &Arc, + dst: &mut CudaSlice, + data: &[f32], + label: &str, +) -> Result<(), MLError> { + use cudarc::driver::DevicePtrMut; + let n = data.len(); + assert!(dst.len() >= n, "{label}: dst.len {} < data.len {n}", dst.len()); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?; + staging.write_from_slice(data); + let n_bytes = n * std::mem::size_of::(); + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("{label} dtod: {e}")))?; + } + stream.synchronize() + .map_err(|e| MLError::ModelError(format!("{label} dtod sync: {e}")))?; + Ok(()) +} + /// Pre-resolved raw u64 CUDA device pointers for all GPU buffers. /// Computed once at construction. Eliminates 110+ per-step `raw_device_ptr()` /// calls that go through cudarc's event tracking machinery. @@ -9081,19 +9233,16 @@ impl GpuDqnTrainer { // G1: Weight decay mask — 1.0 for trunk+value (indices 0-7), 0.0 for branch heads. // Allocated before CUDA Graph capture so pointer is stable across replays. + // Mapped pinned staging eliminates HtoD per + // `feedback_no_htod_htoh_only_mapped_pinned.md` — GPU reads via DtoD. let weight_decay_mask = { let param_sizes = compute_param_sizes(&config); let mut mask = vec![0.0_f32; total_params]; - // Indices 0-7 = trunk (w_s1, b_s1, w_s2, b_s2) + value head (w_v1, b_v1, w_v2, b_v2) let trunk_end: usize = (0..8).map(|i| align4(param_sizes[i])).sum(); for i in 0..trunk_end { mask[i] = 1.0; } - let mut buf = stream.alloc_zeros::(total_params) - .map_err(|e| MLError::ModelError(format!("wd_mask alloc: {e}")))?; - stream.memcpy_htod(&mask, &mut buf) - .map_err(|e| MLError::ModelError(format!("wd_mask htod: {e}")))?; - buf + upload_via_mapped_f32(&stream, total_params, &mask, "wd_mask")? }; // Task 0.4 — pinned host buffer for grad readback. @@ -9244,14 +9393,8 @@ impl GpuDqnTrainer { lens[b] = len as i32; if len > max_len { max_len = len; } } - let mut starts_dev = stream.alloc_zeros::(4) - .map_err(|e| MLError::ModelError(format!("alloc branch_slice_starts_dev: {e}")))?; - stream.memcpy_htod(&starts, &mut starts_dev) - .map_err(|e| MLError::ModelError(format!("htod branch_slice_starts_dev: {e}")))?; - let mut lens_dev = stream.alloc_zeros::(4) - .map_err(|e| MLError::ModelError(format!("alloc branch_slice_lens_dev: {e}")))?; - stream.memcpy_htod(&lens, &mut lens_dev) - .map_err(|e| MLError::ModelError(format!("htod branch_slice_lens_dev: {e}")))?; + let starts_dev = upload_via_mapped_i32(&stream, 4, &starts, "branch_slice_starts_dev")?; + let lens_dev = upload_via_mapped_i32(&stream, 4, &lens, "branch_slice_lens_dev")?; (starts_dev, lens_dev, max_len) }; let branch_grad_norms_dev = stream.alloc_zeros::(4) @@ -9260,11 +9403,7 @@ impl GpuDqnTrainer { // Initialised to 1.0 so HEALTH_DIAG reads a sensible "no-op" // value before the first rescale launch writes real scales. let ones = [1.0_f32; 4]; - let mut buf = stream.alloc_zeros::(4) - .map_err(|e| MLError::ModelError(format!("alloc branch_grad_scales_dev: {e}")))?; - stream.memcpy_htod(&ones, &mut buf) - .map_err(|e| MLError::ModelError(format!("htod branch_grad_scales_dev: {e}")))?; - buf + upload_via_mapped_f32(&stream, 4, &ones, "branch_grad_scales_dev")? }; // Load reduce + rescale kernels from a dedicated CUmodule so their // CUfunction handles are isolated from all other graphs (Hopper @@ -9312,20 +9451,10 @@ impl GpuDqnTrainer { // D.2: per-branch gamma base/max device arrays — allocated once at construction. let gamma_base_host: [f32; 4] = [0.92, 0.88, 0.85, 0.80]; // DIR, MAG, ORD, URG let gamma_max_host: [f32; 4] = [0.99, 0.95, 0.93, 0.90]; // DIR, MAG, ORD, URG - let per_branch_gamma_base_dev = { - let mut dev = stream.alloc_zeros::(4) - .map_err(|e| MLError::ModelError(format!("alloc per_branch_gamma_base: {e}")))?; - stream.memcpy_htod(&gamma_base_host, &mut dev) - .map_err(|e| MLError::ModelError(format!("htod per_branch_gamma_base: {e}")))?; - dev - }; - let per_branch_gamma_max_dev = { - let mut dev = stream.alloc_zeros::(4) - .map_err(|e| MLError::ModelError(format!("alloc per_branch_gamma_max: {e}")))?; - stream.memcpy_htod(&gamma_max_host, &mut dev) - .map_err(|e| MLError::ModelError(format!("htod per_branch_gamma_max: {e}")))?; - dev - }; + let per_branch_gamma_base_dev = + upload_via_mapped_f32(&stream, 4, &gamma_base_host, "per_branch_gamma_base")?; + let per_branch_gamma_max_dev = + upload_via_mapped_f32(&stream, 4, &gamma_max_host, "per_branch_gamma_max")?; // Plan 1 Task 11: load kelly_cap_update kernel (cold-path, per-epoch). let kelly_cap_update_kernel = { @@ -9392,14 +9521,8 @@ impl GpuDqnTrainer { let b3 = config.branch_3_size as i32; let offsets: [i32; 4] = [0, b0, b0 + b1, b0 + b1 + b2]; let sizes: [i32; 4] = [b0, b1, b2, b3]; - let mut off_dev = stream.alloc_zeros::(4) - .map_err(|e| MLError::ModelError(format!("alloc q_quantile branch_offsets: {e}")))?; - stream.memcpy_htod(&offsets, &mut off_dev) - .map_err(|e| MLError::ModelError(format!("htod q_quantile branch_offsets: {e}")))?; - let mut sz_dev = stream.alloc_zeros::(4) - .map_err(|e| MLError::ModelError(format!("alloc q_quantile branch_sizes: {e}")))?; - stream.memcpy_htod(&sizes, &mut sz_dev) - .map_err(|e| MLError::ModelError(format!("htod q_quantile branch_sizes: {e}")))?; + let off_dev = upload_via_mapped_i32(&stream, 4, &offsets, "q_quantile_branch_offsets")?; + let sz_dev = upload_via_mapped_i32(&stream, 4, &sizes, "q_quantile_branch_sizes")?; (off_dev, sz_dev) }; @@ -10055,11 +10178,7 @@ impl GpuDqnTrainer { // [12] W_bn [bottleneck_dim, market_dim] — GOFF 33 (was 24) w_ptrs[33], spec_u_bn.raw_ptr(), spec_v_bn.raw_ptr(), bn_dim_u64, md_u64, 0, ]; - let mut desc_buf = stream.alloc_zeros::(78) - .map_err(|e| MLError::ModelError(format!("alloc spectral_norm_descriptors: {e}")))?; - stream.memcpy_htod(&host_desc, &mut desc_buf) - .map_err(|e| MLError::ModelError(format!("upload spectral_norm_descriptors: {e}")))?; - desc_buf + upload_via_mapped_u64(&stream, 78, &host_desc, "spectral_norm_descriptors")? }; // ── Initialize shared cuBLAS handle (one handle for all forward/backward) ── @@ -10231,16 +10350,12 @@ impl GpuDqnTrainer { let ensemble_std_buf = stream.alloc_zeros::(b) .map_err(|e| MLError::ModelError(format!("alloc ensemble_std_buf: {e}")))?; // #21 Stochastic depth: 3 layer scales + GPU RNG state - let mut stochastic_depth_scale_buf = stream.alloc_zeros::(3) - .map_err(|e| MLError::ModelError(format!("alloc stochastic_depth_scale: {e}")))?; - stream.memcpy_htod(&[1.0_f32, 1.0, 1.0], &mut stochastic_depth_scale_buf) - .map_err(|e| MLError::ModelError(format!("init stochastic_depth_scale: {e}")))?; - let mut stochastic_depth_rng_state = stream.alloc_zeros::(1) - .map_err(|e| MLError::ModelError(format!("alloc sd_rng_state: {e}")))?; + let stochastic_depth_scale_buf = + upload_via_mapped_f32(&stream, 3, &[1.0_f32, 1.0, 1.0], "stochastic_depth_scale")?; // Deterministic seed for reproducible stochastic depth masks let sd_seed: u32 = 0x5D5E_ED00; - stream.memcpy_htod(&[sd_seed], &mut stochastic_depth_rng_state) - .map_err(|e| MLError::ModelError(format!("seed sd_rng: {e}")))?; + let stochastic_depth_rng_state = + upload_via_mapped_u32(&stream, 1, &[sd_seed], "sd_rng_state")?; // Curiosity cuBLAS GEMM pipeline buffers: // CUR_INPUT = market_dim + 3, CUR_HIDDEN = 128, CUR_OUTPUT = market_dim let cur_input = config.market_dim + 3; @@ -10494,16 +10609,8 @@ impl GpuDqnTrainer { host_begins[g] = gb as i32; host_ends[g] = ge as i32; } - let mut begins = stream - .alloc_zeros::(num_groups) - .map_err(|e| MLError::ModelError(format!("alloc vsn_group_begins: {e}")))?; - let mut ends = stream - .alloc_zeros::(num_groups) - .map_err(|e| MLError::ModelError(format!("alloc vsn_group_ends: {e}")))?; - stream.memcpy_htod(&host_begins, &mut begins) - .map_err(|e| MLError::ModelError(format!("htod vsn_group_begins: {e}")))?; - stream.memcpy_htod(&host_ends, &mut ends) - .map_err(|e| MLError::ModelError(format!("htod vsn_group_ends: {e}")))?; + let begins = upload_via_mapped_i32(&stream, num_groups, &host_begins, "vsn_group_begins")?; + let ends = upload_via_mapped_i32(&stream, num_groups, &host_ends, "vsn_group_ends")?; (begins, ends) }; // Plan 4 Task 1B-iv: VSN backward scratch buffers. The backward chain @@ -11430,7 +11537,7 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc mamba2_adam_m: {e}")))?; let mamba2_adam_v = stream.alloc_zeros::(mamba2_param_count) .map_err(|e| MLError::ModelError(format!("alloc mamba2_adam_v: {e}")))?; - // Xavier init for mamba2_params + // Xavier init for mamba2_params (mapped pinned upload, no HtoD) { let scale = (2.0_f32 / (h_width + MAMBA2_STATE_DIM) as f32).sqrt(); let init_data: Vec = (0..mamba2_param_count).map(|j| { @@ -11438,8 +11545,7 @@ impl GpuDqnTrainer { let u = (hash as f32) / (u32::MAX as f32) * 2.0 - 1.0; u * scale }).collect(); - stream.memcpy_htod(&init_data, &mut mamba2_params) - .map_err(|e| MLError::ModelError(format!("mamba2_params xavier init: {e}")))?; + update_via_mapped_f32(&stream, &mut mamba2_params, &init_data, "mamba2_params_xavier")?; } let mamba2_module = stream.context().load_cubin(MAMBA2_CUBIN.to_vec()) .map_err(|e| MLError::ModelError(format!("mamba2 cubin load: {e}")))?; diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 9e1221770..e20758729 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -1417,6 +1417,21 @@ in constructor; kernel reads/writes via dev_ptr) and the upcoming u64 table). All four mapped pinned variants share the same `new`/`write_from_slice`/`read_all` API. +`gpu_dqn_trainer::new` HtoD elimination (2026-04-28): added module-level +`upload_via_mapped_{f32,i32,u32,u64}` helpers and an in-place +`update_via_mapped_f32`. Each stages CPU bytes through a transient +mapped pinned buffer and DtoD-copies into the destination +`CudaSlice`, then stream-syncs so the staging buffer is safe to +drop. Migrated 11 COLD ctor sites: weight_decay_mask, +branch_slice_starts/lens, branch_grad_scales, per_branch_gamma_base/max, +q_quantile_branch_offsets/sizes, spectral_norm_descriptors (78 u64), +stochastic_depth_scale (3 f32), sd_rng_state (1 u32), vsn_group_begins, +vsn_group_ends, mamba2_params Xavier init. All consumer surfaces +unchanged — destination remains `CudaSlice` with all subsequent +kernel-arg call sites intact. Per +`feedback_no_htod_htoh_only_mapped_pinned.md` and +`feedback_no_partial_refactor.md` (no consumer migration needed). + MoE moe_mixture_forward kernel + Rust wrapper (2026-04-27): first MoE CUDA kernel landed. Single-thread-per-(b,c) kernel computes h_s2[b,c] = Σ_k g[b,k]·expert_outputs[k,b,c]. No atomicAdd, capture-friendly. Rust From c86318ad5711c420bc0242b2da26b935ce6cfaa8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 21:07:45 +0200 Subject: [PATCH 4/5] fix(gpu_dqn_trainer): migrate upload_params/upload_target_params (WARM) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned (cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path. Both `upload_params` and `upload_target_params` previously did `stream.memcpy_htod` of TOTAL_PARAMS f32 (~MB) at fold boundaries / external weight loads. Now route through the helper `update_via_mapped_f32` introduced in the previous commit — stages CPU bytes through a transient mapped pinned buffer and DtoD-copies into the existing `params_buf` / `target_params_buf` `CudaSlice`. No public API change. The destination buffer types remain `CudaSlice` so every downstream kernel-arg consumer is untouched. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 12 ++++++------ docs/dqn-wire-up-audit.md | 7 +++++++ 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 77d7eff05..adf99f783 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -12536,10 +12536,10 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("download_params post-dtoh sync: {e}"))) } - /// Upload online params from host to GPU (HtoD). + /// Upload online params from host to GPU via mapped-pinned staging. + /// No HtoD memcpy per `feedback_no_htod_htoh_only_mapped_pinned.md`. pub fn upload_params(&mut self, src: &[f32]) -> Result<(), MLError> { - self.stream.memcpy_htod(src, &mut self.params_buf) - .map_err(|e| MLError::ModelError(format!("upload params: {e}"))) + update_via_mapped_f32(&self.stream, &mut self.params_buf, src, "upload_params") } /// Download target params from GPU to host (synchronous DtoH). @@ -12555,10 +12555,10 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("download_target_params post-dtoh sync: {e}"))) } - /// Upload target params from host to GPU (HtoD). + /// Upload target params from host to GPU via mapped-pinned staging. + /// No HtoD memcpy per `feedback_no_htod_htoh_only_mapped_pinned.md`. pub fn upload_target_params(&mut self, src: &[f32]) -> Result<(), MLError> { - self.stream.memcpy_htod(src, &mut self.target_params_buf) - .map_err(|e| MLError::ModelError(format!("upload target params: {e}"))) + update_via_mapped_f32(&self.stream, &mut self.target_params_buf, src, "upload_target_params") } /// Speculative forward: pre-compute trunk from intermediate features. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index e20758729..7b7736289 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -1432,6 +1432,13 @@ kernel-arg call sites intact. Per `feedback_no_htod_htoh_only_mapped_pinned.md` and `feedback_no_partial_refactor.md` (no consumer migration needed). +`upload_params` / `upload_target_params` WARM migration (2026-04-28): +both methods are called at fold boundaries / external weight loads. +Previously did `memcpy_htod` of TOTAL_PARAMS f32 (~MB). Now route +through `update_via_mapped_f32` — mapped pinned staging + DtoD copy +into the existing `params_buf` / `target_params_buf` +`CudaSlice`. No API change for callers, no HtoD. + MoE moe_mixture_forward kernel + Rust wrapper (2026-04-27): first MoE CUDA kernel landed. Single-thread-per-(b,c) kernel computes h_s2[b,c] = Σ_k g[b,k]·expert_outputs[k,b,c]. No atomicAdd, capture-friendly. Rust From af18ae39c88096bbdc61f13779f9d462f72c99d8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 21:10:03 +0200 Subject: [PATCH 5/5] fix(fused-training,training-loop): migrate WARM HtoD sites to mapped pinned MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned (cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path. Migrates 3 WARM HtoD sites: - `fused_training.rs:505` HER source_indices Vec upload - `fused_training.rs:511` HER reward_ones Vec upload - `training_loop.rs:470` curriculum episode_starts Vec upload For `fused_training.rs` adds module-level `staging_upload_{i32,f32}` helpers. For `training_loop.rs` uses an inline staging+DtoD block because the destination is reached via the out-of-scope accessor `collector.episode_starts_buf_mut() -> &mut CudaSlice`. All destination buffer types remain `CudaSlice` so no downstream consumer (incl. `gpu_her::relabel_batch_with_strategy` and any `episode_starts_buf` reader) needs to change. Note: if Agent 2's `gpu_experience_collector` merge changes `episode_starts_buf_mut`'s return type to a mapped pinned buffer, the inline staging block in `training_loop.rs:470` can be simplified to a direct `write_from_slice` (follow-up). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/src/trainers/dqn/fused_training.rs | 79 ++++++++++++++++--- .../src/trainers/dqn/trainer/training_loop.rs | 22 +++++- docs/dqn-wire-up-audit.md | 13 +++ 3 files changed, 104 insertions(+), 10 deletions(-) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 32be3016d..57d0af028 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -33,6 +33,7 @@ use tracing::info; use cudarc::driver::DevicePtr; use cudarc::driver::sys as cuda_sys; +use crate::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer}; use crate::cuda_pipeline::gpu_attention::{GpuAttention, GpuAttentionConfig}; use crate::cuda_pipeline::gpu_tlob::GpuTlob; use crate::cuda_pipeline::gpu_dqn_trainer::{ @@ -57,6 +58,67 @@ pub(crate) struct FusedStepResult { pub grad_norm: f32, } +// ── Mapped-pinned staging helpers (HtoD elimination) ───────────────────────── +// Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned +// (cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path. These helpers +// stage CPU data through a transient mapped pinned buffer and DtoD-copy into +// a GPU-resident `CudaSlice`, preserving the existing kernel-arg surface +// (e.g., `her_source_indices_gpu` consumed by gpu_her.rs out of scope). + +fn staging_upload_i32( + stream: &Arc, + n: usize, + data: &[i32], + label: &str, +) -> Result> { + use cudarc::driver::DevicePtrMut; + assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len()); + let mut buf = stream.alloc_zeros::(n) + .map_err(|e| anyhow::anyhow!("{label} alloc: {e}"))?; + let staging = unsafe { MappedI32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("{label} staging alloc: {e}"))?; + staging.write_from_slice(data); + let n_bytes = n * std::mem::size_of::(); + { + let (dst_ptr, _g) = buf.device_ptr_mut(stream); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ).map_err(|e| anyhow::anyhow!("{label} dtod: {e}"))?; + } + } + stream.synchronize() + .map_err(|e| anyhow::anyhow!("{label} dtod sync: {e}"))?; + Ok(buf) +} + +fn staging_upload_f32( + stream: &Arc, + n: usize, + data: &[f32], + label: &str, +) -> Result> { + use cudarc::driver::DevicePtrMut; + assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len()); + let mut buf = stream.alloc_zeros::(n) + .map_err(|e| anyhow::anyhow!("{label} alloc: {e}"))?; + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("{label} staging alloc: {e}"))?; + staging.write_from_slice(data); + let n_bytes = n * std::mem::size_of::(); + { + let (dst_ptr, _g) = buf.device_ptr_mut(stream); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ).map_err(|e| anyhow::anyhow!("{label} dtod: {e}"))?; + } + } + stream.synchronize() + .map_err(|e| anyhow::anyhow!("{label} dtod sync: {e}"))?; + Ok(buf) +} + /// Child sub-graph: owns both the CUgraph topology and an instantiated CUgraphExec. /// /// The exec comes directly from cudarc's `end_capture` instantiation — we do NOT @@ -494,22 +556,21 @@ impl FusedTrainingCtx { None }; - // Pre-compute HER source indices and reward ones on GPU (avoid per-step Vec + HtoD) + // Pre-compute HER source indices and reward ones on GPU. + // Mapped pinned staging + DtoD copy — no HtoD per + // `feedback_no_htod_htoh_only_mapped_pinned.md`. + // Destination types remain `CudaSlice<{i32,f32}>` so the downstream + // `relabel_batch_with_strategy` kernel-arg surface in gpu_her.rs is + // unchanged (out of scope per task). let (her_source_indices_gpu, her_reward_ones_gpu) = if gpu_her.is_some() { let her_batch_size = (batch_size as f64 * f64::from(hyperparams.her_ratio)) as usize; let normal_count = batch_size.saturating_sub(her_batch_size); // Source indices: [normal_count, normal_count+1, ..., batch_size-1] let source_idx: Vec = (normal_count..batch_size).map(|i| i as i32).collect(); - let mut src_gpu = stream.alloc_zeros::(her_batch_size) - .map_err(|e| anyhow::anyhow!("HER source indices alloc: {e}"))?; - stream.memcpy_htod(&source_idx, &mut src_gpu) - .map_err(|e| anyhow::anyhow!("HER source indices HtoD: {e}"))?; + let src_gpu = staging_upload_i32(&stream, her_batch_size, &source_idx, "HER_source_indices")?; // Reward ones: [1.0; her_batch_size] let ones: Vec = vec![1.0_f32; her_batch_size]; - let mut ones_gpu = stream.alloc_zeros::(her_batch_size) - .map_err(|e| anyhow::anyhow!("HER reward ones alloc: {e}"))?; - stream.memcpy_htod(&ones, &mut ones_gpu) - .map_err(|e| anyhow::anyhow!("HER reward ones HtoD: {e}"))?; + let ones_gpu = staging_upload_f32(&stream, her_batch_size, &ones, "HER_reward_ones")?; (Some(src_gpu), Some(ones_gpu)) } else { (None, None) diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index db500d8b5..20ceb6677 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -21,6 +21,7 @@ use cudarc::driver::CudaSlice; use common::metrics::{questdb_sink, training_metrics}; use tracing::{debug, info, warn}; +use crate::cuda_pipeline::mapped_pinned::MappedI32Buffer; use crate::cuda_pipeline::gpu_dqn_trainer::{ EPOCH_IDX_INDEX, EPSILON_EFF_INDEX, GAMMA_DIR_EFF_INDEX, TLOB_REGIME_FOCUS_EMA_INDEX, }; @@ -467,7 +468,26 @@ impl DQNTrainer { .map(|i| restricted[(i * stride).min(restricted.len() - 1)]) .collect(); if let Some(ref stream) = self.cuda_stream { - let _ = stream.memcpy_htod(&episode_starts, collector.episode_starts_buf_mut()); + // Mapped pinned staging + DtoD — no HtoD per + // `feedback_no_htod_htoh_only_mapped_pinned.md`. + // The collector's `episode_starts_buf` is currently a + // `CudaSlice`; if Agent 2's merge reshapes it to + // a mapped pinned buffer this site simplifies to a + // direct `write_from_slice` (follow-up). + use cudarc::driver::DevicePtrMut; + let n = episode_starts.len(); + if let Ok(staging) = unsafe { MappedI32Buffer::new(n) } { + staging.write_from_slice(&episode_starts); + let dst_buf = collector.episode_starts_buf_mut(); + let n_bytes = n * std::mem::size_of::(); + let (dst_ptr, _g) = dst_buf.device_ptr_mut(stream); + unsafe { + let _ = cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(), + ); + } + let _ = stream.synchronize(); + } } if epoch % 5 == 0 { debug!( diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 7b7736289..a1f68aa43 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -1439,6 +1439,19 @@ through `update_via_mapped_f32` — mapped pinned staging + DtoD copy into the existing `params_buf` / `target_params_buf` `CudaSlice`. No API change for callers, no HtoD. +`fused_training.rs` HER block + `training_loop.rs` curriculum WARM +migration (2026-04-28): both sites previously did `stream.memcpy_htod` +on small Vec/Vec. Migrated to mapped pinned staging + DtoD +copy via local helpers in `fused_training.rs` +(`staging_upload_{i32,f32}`) and an inline pattern in +`training_loop.rs`. Destination buffer types unchanged +(`CudaSlice` / `CudaSlice`) so the `relabel_batch_with_strategy` +kernel-arg surface in gpu_her.rs (out of scope) is untouched. The +`training_loop.rs:470` site goes through +`collector.episode_starts_buf_mut() -> &mut CudaSlice`; if Agent 2's +merge changes that accessor type to a mapped pinned buffer, the inline +staging block can be simplified to a direct `write_from_slice`. + MoE moe_mixture_forward kernel + Rust wrapper (2026-04-27): first MoE CUDA kernel landed. Single-thread-per-(b,c) kernel computes h_s2[b,c] = Σ_k g[b,k]·expert_outputs[k,b,c]. No atomicAdd, capture-friendly. Rust