merge: HtoD→mapped-pinned migration batch 3 (action selector + DQN ctor + fused training, 23 sites)
Agent 3 — set_count_bonuses HOT path (3-6 HtoD/call → 0), gpu_dqn_trainer constructor block (11 COLD sites: weight_decay_mask, branch metadata, gamma, q_quantile, spectral_norm, stochastic_depth, VSN groups, mamba2 init), upload_params/upload_target_params (2 WARM), HER + curriculum episode_starts (3 WARM sites in fused_training and training_loop). Strategy: COLD/WARM sites use mapped-pinned staging + DtoD into existing CudaSlice destinations to avoid disturbing 100+ downstream consumers; HOT set_count_bonuses uses persistent MappedF32Buffer fields (no DtoD per call). 5 commits, 23 sites total. Adds MappedU64Buffer for spectral-norm host_desc[78]. Per feedback_no_htod_htoh_only_mapped_pinned.md — third of 3 parallel batches. # Conflicts: # crates/ml/src/cuda_pipeline/mapped_pinned.rs
This commit is contained in:
@@ -19,6 +19,7 @@ use std::sync::Arc;
|
||||
use tracing::info;
|
||||
|
||||
use crate::MLError;
|
||||
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"));
|
||||
@@ -29,20 +30,26 @@ 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,
|
||||
stream: Arc<CudaStream>,
|
||||
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<CudaSlice<f32>>,
|
||||
bonus_order_buf: Option<CudaSlice<f32>>,
|
||||
bonus_urgency_buf: Option<CudaSlice<f32>>,
|
||||
bonus_exposure_buf: Option<MappedF32Buffer>,
|
||||
bonus_order_buf: Option<MappedF32Buffer>,
|
||||
bonus_urgency_buf: Option<MappedF32Buffer>,
|
||||
/// 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<f32>,
|
||||
@@ -61,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,
|
||||
@@ -85,27 +100,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(())
|
||||
}
|
||||
|
||||
@@ -116,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}")))?;
|
||||
}
|
||||
@@ -136,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)
|
||||
@@ -169,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)
|
||||
|
||||
@@ -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::<usize>() 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<T>` — 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<CudaStream>,
|
||||
n: usize,
|
||||
data: &[f32],
|
||||
label: &str,
|
||||
) -> Result<CudaSlice<f32>, MLError> {
|
||||
use cudarc::driver::DevicePtrMut;
|
||||
assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
|
||||
let mut buf = stream.alloc_zeros::<f32>(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::<f32>();
|
||||
{
|
||||
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<CudaStream>,
|
||||
n: usize,
|
||||
data: &[i32],
|
||||
label: &str,
|
||||
) -> Result<CudaSlice<i32>, MLError> {
|
||||
use cudarc::driver::DevicePtrMut;
|
||||
assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
|
||||
let mut buf = stream.alloc_zeros::<i32>(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::<i32>();
|
||||
{
|
||||
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<CudaStream>,
|
||||
n: usize,
|
||||
data: &[u32],
|
||||
label: &str,
|
||||
) -> Result<CudaSlice<u32>, MLError> {
|
||||
use cudarc::driver::DevicePtrMut;
|
||||
assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
|
||||
let mut buf = stream.alloc_zeros::<u32>(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::<u32>();
|
||||
{
|
||||
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<CudaStream>,
|
||||
n: usize,
|
||||
data: &[u64],
|
||||
label: &str,
|
||||
) -> Result<CudaSlice<u64>, MLError> {
|
||||
use cudarc::driver::DevicePtrMut;
|
||||
assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
|
||||
let mut buf = stream.alloc_zeros::<u64>(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::<u64>();
|
||||
{
|
||||
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<f32>` 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<CudaStream>,
|
||||
dst: &mut CudaSlice<f32>,
|
||||
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::<f32>();
|
||||
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::<f32>(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::<i32>(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::<i32>(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::<f32>(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::<f32>(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::<f32>(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::<f32>(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::<i32>(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::<i32>(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::<u64>(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::<f32>(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::<f32>(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::<u32>(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::<i32>(num_groups)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc vsn_group_begins: {e}")))?;
|
||||
let mut ends = stream
|
||||
.alloc_zeros::<i32>(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::<f32>(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<f32> = (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}")))?;
|
||||
@@ -12430,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).
|
||||
@@ -12449,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.
|
||||
|
||||
@@ -502,3 +502,83 @@ pub fn upload_u32_via_pinned(
|
||||
.map_err(|e| format!("upload_u32_via_pinned sync: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<T>`, 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<cudarc::driver::CudaStream>,
|
||||
n: usize,
|
||||
data: &[i32],
|
||||
label: &str,
|
||||
) -> Result<cudarc::driver::CudaSlice<i32>> {
|
||||
use cudarc::driver::DevicePtrMut;
|
||||
assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
|
||||
let mut buf = stream.alloc_zeros::<i32>(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::<i32>();
|
||||
{
|
||||
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<cudarc::driver::CudaStream>,
|
||||
n: usize,
|
||||
data: &[f32],
|
||||
label: &str,
|
||||
) -> Result<cudarc::driver::CudaSlice<f32>> {
|
||||
use cudarc::driver::DevicePtrMut;
|
||||
assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
|
||||
let mut buf = stream.alloc_zeros::<f32>(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::<f32>();
|
||||
{
|
||||
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<i32> = (normal_count..batch_size).map(|i| i as i32).collect();
|
||||
let mut src_gpu = stream.alloc_zeros::<i32>(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<f32> = vec![1.0_f32; her_batch_size];
|
||||
let mut ones_gpu = stream.alloc_zeros::<f32>(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)
|
||||
|
||||
@@ -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, Q_ABS_REF_INDEX,
|
||||
Q_DIR_ABS_REF_INDEX, TLOB_REGIME_FOCUS_EMA_INDEX,
|
||||
@@ -488,7 +489,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<i32>`; 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::<i32>();
|
||||
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!(
|
||||
|
||||
@@ -1152,7 +1152,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 | — |
|
||||
@@ -1460,6 +1460,50 @@ 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.
|
||||
|
||||
`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<T>`, 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<T>` 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).
|
||||
|
||||
`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<f32>`. 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<i32>/Vec<f32>. 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<i32>` / `CudaSlice<f32>`) 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<i32>`; 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
|
||||
|
||||
Reference in New Issue
Block a user