fix(action-selector): migrate rng_states to mapped pinned (COLD ctor)

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) <noreply@anthropic.com>

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<u32>,
+    /// 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<u32>,
     fill_mask_buf: CudaSlice<i32>,
     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::<u32>(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc actions_buf: {e}")))?;
         let fill_mask_buf = stream.alloc_zeros::<i32>(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::<u32>(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::<f32>(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)

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<Self, String> {
+        let num_bytes = len * std::mem::size_of::<u32>();
+        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<u32> {
+        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<Self, String> {
+        let num_bytes = len * std::mem::size_of::<u64>();
+        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<u64> {
+        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);
+        }
+    }
+}

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