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

fix(gpu_dqn_trainer): migrate 11 COLD ctor HtoD sites to mapped pinned

Browse Source 
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<T>`, then stream-syncs so the staging buffer is safe to
drop. Destination buffer types remain `CudaSlice<T>` 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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 21:06:13 +02:00
parent 0fb21c970c
commit 87ea9fa416
2 changed files with 188 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

240
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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}")))?;

15
docs/dqn-wire-up-audit.md
View File

@@ -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<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).
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