From 0fb21c970cfbe6b269f25fe935dca6de39947d76 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 21:00:27 +0200 Subject: [PATCH] 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