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

merge: HtoD→mapped-pinned migration batch 2 (HOT collector/DT/portfolio, 18 sites)

Browse Source 
Agent 2 — HOT-path migrations: bar_indices + feature_mask in
gpu_experience_collector, scratch.adam_t in decision_transformer (per-backward
4-byte HtoD eliminated), actions in gpu_portfolio simulate_batch_gpu.

3 commits, 18 sites total. Persistent MappedXxxBuffer fields for HOT-path
buffers (allocate once at ctor, host_slice_mut per call). Cargo check clean
(11 warnings — 2 unrelated baseline warnings disappeared as side-effect).

Per feedback_no_htod_htoh_only_mapped_pinned.md — second of 3 parallel batches.
This commit is contained in:
jgrusewski
2026-04-28 21:13:38 +02:00
parent 0ac7729801 1e35324bf4
commit b5d9da136a
5 changed files with 334 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
crates/ml/src/cuda_pipeline/decision_transformer.rs
View File

@@ -47,6 +47,7 @@ use cudarc::driver::{
};
use tracing::info;
use crate::MLError;
use super::mapped_pinned::{MappedF32Buffer, MappedI32Buffer};
/// Configuration for the Decision Transformer.
#[derive(Debug, Clone)]
@@ -338,8 +339,12 @@ struct DtScratch {
adam_v: CudaSlice<f32>,
/// Grad norm scratch: [1]
grad_norm: CudaSlice<f32>,
/// Adam step counter on device: [1]
adam_t: CudaSlice<i32>,
/// Adam step counter, mapped pinned: CPU writes via host_ptr each
/// `train_step_gpu`, the Adam kernel reads via dev_ptr. HOT path.
/// Pure host→device scalar — no GPU writes, so the destination type
/// can switch from CudaSlice to MappedI32Buffer without any
/// downstream changes beyond the dev_ptr extraction.
adam_t: MappedI32Buffer,
}
fn alloc_dt_scratch(
@@ -378,8 +383,11 @@ fn alloc_dt_scratch(
adam_m: alloc_f(total_params, "adam_m")?,
adam_v: alloc_f(total_params, "adam_v")?,
grad_norm: alloc_f(1, "grad_norm")?,
adam_t: stream.alloc_zeros::<i32>(1)
.map_err(|e| MLError::ModelError(format!("DT scratch adam_t: {e}")))?,
// Safety: a CUDA context is active on the constructing thread.
adam_t: unsafe {
MappedI32Buffer::new(1)
.map_err(|e| MLError::ModelError(format!("DT scratch adam_t (mapped pinned): {e}")))?
},
})
}
@@ -431,7 +439,8 @@ pub struct DecisionTransformer {
/// Adam step counter (host-side, incremented and uploaded each step).
adam_step: i32,
/// Weight decay mask [total_params]: all 1.0 (DT applies uniform decay).
wd_mask: CudaSlice<f32>,
/// Read-only by GPU — mapped pinned suffices, no GPU write-back.
wd_mask: MappedF32Buffer,
}
impl DecisionTransformer {
@@ -475,7 +484,10 @@ impl DecisionTransformer {
let mut params = stream.alloc_zeros::<f32>(total_params)
.map_err(|e| MLError::ModelError(format!("DT params alloc: {e}")))?;
super::htod_f32(&stream, &host_params, &mut params)?;
// params is mutated by the Adam kernel — destination must remain a
// CudaSlice. Upload Xavier-initialized host_params via a mapped pinned
// staging buffer + DtoD per `feedback_no_htod_htoh_only_mapped_pinned.md`.
upload_host_to_cuda_f32(&stream, &host_params, &mut params, "DT params init")?;
let context_buf = stream.alloc_zeros::<f32>(context_size)
.map_err(|e| MLError::ModelError(format!("DT context alloc: {e}")))?;
@@ -494,14 +506,17 @@ impl DecisionTransformer {
);
// Uniform weight decay mask (all 1.0) — DT applies decay to all params.
let wd_mask = {
let ones = vec![1.0_f32; total_params];
let mut buf = stream.alloc_zeros::<f32>(total_params)
.map_err(|e| MLError::ModelError(format!("DT wd_mask alloc: {e}")))?;
stream.memcpy_htod(&ones, &mut buf)
.map_err(|e| MLError::ModelError(format!("DT wd_mask htod: {e}")))?;
buf
// Read-only by the Adam kernel: mapped pinned is the cleanest fit.
// Safety: a CUDA context is active on the constructing thread.
let wd_mask = unsafe {
MappedF32Buffer::new(total_params)
.map_err(|e| MLError::ModelError(format!("DT wd_mask (mapped pinned): {e}")))?
};
{
// Direct host_ptr write — no memcpy, no staging.
let ones = vec![1.0_f32; total_params];
wd_mask.write_from_slice(&ones);
}
Ok(Self {
layout,
@@ -907,13 +922,13 @@ impl DecisionTransformer {
// ── OPTIMIZER ──────────────────────────────────────────────────
// Increment Adam step
// Increment Adam step. CPU writes the scalar straight into the mapped
// pinned buffer; the Adam kernel reads via dev_ptr — zero HtoD copy on
// the per-train_step hot path.
self.adam_step += 1;
let step_host = [self.adam_step];
stream.memcpy_htod(&step_host, &mut scratch.adam_t)
.map_err(|e| MLError::ModelError(format!("DT adam_t upload: {e}")))?;
scratch.adam_t.write_from_slice(&[self.adam_step]);
let adam_t_ptr = raw_ptr_i32(&scratch.adam_t, stream);
let adam_t_ptr = scratch.adam_t.dev_ptr;
// Zero grad_norm
{
@@ -959,7 +974,7 @@ impl DecisionTransformer {
block_dim: (256, 1, 1), shared_mem_bytes: 0,
};
let wd_mask_ptr = self.wd_mask.raw_ptr();
let wd_mask_ptr = self.wd_mask.dev_ptr;
let l1_end_dt: i32 = 0; // no L1 for DT
let l1_lambda_dt: f32 = 0.0; // disabled
unsafe {
@@ -1102,16 +1117,27 @@ impl DecisionTransformer {
.map(|i| (i * stride) as i32)
.collect();
let episode_starts_gpu = stream.clone_htod(&episode_starts)
.map_err(|e| MLError::ModelError(format!("DT episode_starts upload: {e}")))?;
// Mapped pinned: CPU writes via host_ptr, the build_trajectories kernel
// reads via dev_ptr. Replaces `clone_htod` per `feedback_no_htod_htoh_only_mapped_pinned.md`.
// Safety: CUDA context active on this thread.
let episode_starts_gpu = unsafe {
MappedI32Buffer::new(num_episodes)
.map_err(|e| MLError::ModelError(format!("DT episode_starts (mapped pinned): {e}")))?
};
episode_starts_gpu.write_from_slice(&episode_starts);
// ── Step 3: Compute per-episode return-to-go on GPU ──
// We need rewards reshaped as [num_episodes, context_len].
// Instead of copying, we'll compute RTG per-episode using the global rewards array.
// Build a temporary [num_episodes, context_len] reward slice and RTG.
let ep_total = num_episodes * context_len;
let mut ep_rewards_gpu = stream.alloc_zeros::<f32>(ep_total)
.map_err(|e| MLError::ModelError(format!("DT ep_rewards alloc: {e}")))?;
// ep_rewards is read-only by the return_to_go kernel (CPU writes once,
// kernel reads once). Mapped pinned suffices and avoids the CudaSlice +
// memcpy_htod path. Safety: CUDA context active on this thread.
let ep_rewards_gpu = unsafe {
MappedF32Buffer::new(ep_total)
.map_err(|e| MLError::ModelError(format!("DT ep_rewards (mapped pinned): {e}")))?
};
let mut rtg_gpu = stream.alloc_zeros::<f32>(ep_total)
.map_err(|e| MLError::ModelError(format!("DT rtg alloc: {e}")))?;
@@ -1144,12 +1170,14 @@ impl DecisionTransformer {
}
}
super::htod_f32(stream, &ep_rewards_host, &mut ep_rewards_gpu)?;
// Direct host_ptr write — no memcpy. The return_to_go kernel reads
// through dev_ptr.
ep_rewards_gpu.write_from_slice(&ep_rewards_host);
}
// RTG reverse cumulative sum
{
let ep_rewards_ptr = raw_ptr_f32(&ep_rewards_gpu, stream);
let ep_rewards_ptr = ep_rewards_gpu.dev_ptr;
let rtg_ptr = raw_ptr_f32_mut(&mut rtg_gpu, stream);
let gamma_f32 = gamma as f32;
let context_len_i32 = context_len as i32;
@@ -1187,11 +1215,7 @@ impl DecisionTransformer {
let _no_drop = std::mem::ManuallyDrop::new(guard);
ptr
};
let starts_ptr = {
let (ptr, guard) = episode_starts_gpu.device_ptr(stream);
let _no_drop = std::mem::ManuallyDrop::new(guard);
ptr
};
let starts_ptr = episode_starts_gpu.dev_ptr;
let traj_ptr = raw_ptr_f32_mut(&mut trajectories_gpu, stream);
let target_actions_ptr = {
let (ptr, guard) = target_actions_gpu.device_ptr_mut(stream);
@@ -1243,3 +1267,39 @@ impl DecisionTransformer {
Ok((trajectories_gpu, target_actions_gpu, num_complete_batches))
}
}
/// Upload `host` into `dst` (long-lived GPU buffer mutated by kernels) via a
/// mapped-pinned staging buffer + DtoD copy.
///
/// Used at construction-time when `dst` must remain a `CudaSlice` because a
/// kernel writes to it later (e.g. `params` mutated by Adam). Per
/// `feedback_no_htod_htoh_only_mapped_pinned.md` mapped pinned is the only
/// allowed CPU↔GPU communication path.
fn upload_host_to_cuda_f32(
stream: &Arc<CudaStream>,
host: &[f32],
dst: &mut CudaSlice<f32>,
label: &str,
) -> Result<(), MLError> {
debug_assert!(host.len() <= dst.len(), "{label}: host overflow");
// Safety: a CUDA context is active on the calling thread.
let staging = unsafe {
MappedF32Buffer::new(host.len())
.map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?
};
staging.write_from_slice(host);
let nbytes = host.len() * std::mem::size_of::<f32>();
let src_ptr = staging.dev_ptr;
let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream);
#[allow(unsafe_code)]
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
).map_err(|e| MLError::ModelError(format!("{label} DtoD: {e}")))?;
}
// Synchronise so the staging buffer can drop safely after the copy.
stream.synchronize()
.map_err(|e| MLError::ModelError(format!("{label} sync: {e}")))?;
Ok(())
}

206
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -18,7 +18,7 @@ use std::ffi::c_void;
use std::sync::Arc;
use cudarc::driver::{
CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg,
CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg,
};
use tracing::{debug, info};
@@ -29,6 +29,7 @@ use super::gpu_dqn_trainer::{
GpuDqnTrainConfig, compute_param_sizes, compute_total_params,
};
use super::gpu_weights::CuriosityWeightSet;
use super::mapped_pinned::{MappedF32Buffer, MappedI32Buffer};
use super::shared_cublas_handle::PerStreamCublasHandles;
// ---------------------------------------------------------------------------
@@ -443,7 +444,10 @@ pub struct GpuExperienceCollector {
/// [v_min, v_max, delta_z] per (episode, branch). Filled once per epoch via
/// `update_per_sample_support()` — matches the Phase 2d stride-12 contract
/// consumed by compute_expected_q / quantile_q_select / mag_concat_qdir.
per_sample_support_buf: CudaSlice<f32>,
/// Read-only by the GPU (kernels read via `dev_ptr`); CPU writes via
/// `host_ptr` from `update_per_sample_support`. Mapped pinned per
/// `feedback_no_htod_htoh_only_mapped_pinned.md`.
per_sample_support_buf: MappedF32Buffer,
/// Per-sample epsilon from IQL expectile gap (0 = use cosine schedule).
per_sample_epsilon_ptr: u64,
/// Branch sizes [4, 3, 3, 3] for 4-branch hierarchical DQN.
@@ -674,7 +678,9 @@ pub struct GpuExperienceCollector {
/// Expert demonstration actions [total_bars] — pre-computed MA crossover signals.
/// -1 = no expert opinion (use Q-network), >= 0 = expert exposure action index.
/// Uploaded once at training start, read by action selection kernel.
expert_actions_gpu: Option<CudaSlice<i32>>,
/// Mapped pinned: CPU writes once via `host_ptr`, the action kernel reads
/// via `dev_ptr` — no HtoD copy.
expert_actions_gpu: Option<MappedI32Buffer>,
/// Expert demo override probability (0.0 = disabled, decays over epochs).
expert_ratio: f32,
@@ -686,7 +692,14 @@ pub struct GpuExperienceCollector {
curiosity_trainer: Option<GpuCuriosityTrainer>,
/// #23 Causal feature mask buffer [market_dim] f32 (1.0=keep, 0.0=zero). None = no masking.
feature_mask_buf: Option<CudaSlice<f32>>,
/// Mapped pinned: CPU writes per-epoch via `host_ptr`, state_gather reads
/// via `dev_ptr` — no HtoD copy on the per-collect (t==0) path.
feature_mask_buf: Option<MappedF32Buffer>,
/// HOT path: hindsight relabel `bar_indices` upload buffer. Allocated once
/// at constructor with capacity `alloc_episodes * alloc_timesteps * 2` (cf
/// counterfactual doubles total). CPU fills via `host_ptr` per
/// `collect_experiences_gpu`; the relabel kernel reads via `dev_ptr`.
bar_indices_pinned: MappedI32Buffer,
/// #31 Bottleneck hidden buffer [alloc_episodes, bn_dim] f32. None when disabled.
exp_bn_hidden: CudaSlice<f32>,
@@ -1073,7 +1086,7 @@ impl GpuExperienceCollector {
portfolio_init[off + 9] = initial_capital; // [9] prev_equity
// [10] hold_time = 0.0, [11] realized_pnl = 0.0 (already zero)
}
super::htod_f32(&stream, &portfolio_init, &mut portfolio_states)?;
upload_host_to_cuda_f32_via_pinned(&stream, &portfolio_init, &mut portfolio_states, "portfolio_states init")?;
let episode_starts_buf = stream
.alloc_zeros::<i32>(alloc_episodes)
@@ -1090,7 +1103,12 @@ impl GpuExperienceCollector {
1.0, // dsr_var
0.0, // step_count
];
let epoch_state = super::clone_htod_f32(&stream, &epoch_state_init)?;
// GPU-mutated by env_step kernel — destination must remain a CudaSlice.
// Upload via mapped-pinned staging + DtoD per
// `feedback_no_htod_htoh_only_mapped_pinned.md`.
let mut epoch_state = stream.alloc_zeros::<f32>(epoch_state_init.len())
.map_err(|e| MLError::ModelError(format!("alloc epoch_state: {e}")))?;
upload_host_to_cuda_f32_via_pinned(&stream, &epoch_state_init, &mut epoch_state, "epoch_state init")?;
// ── Step 7: Allocate output buffers (2x for #7 counterfactual) ───
let total_output = alloc_episodes * alloc_timesteps;
@@ -1337,9 +1355,11 @@ impl GpuExperienceCollector {
.map_err(|e| MLError::ModelError(format!("alloc saboteur_params: {e}")))?;
let mut saboteur_base_buf = stream.alloc_zeros::<f32>(3)
.map_err(|e| MLError::ModelError(format!("alloc saboteur_base: {e}")))?;
// Init base params to neutral (no adversarial effect)
stream.memcpy_htod(&[1.0_f32, 0.85, 1.0], &mut saboteur_base_buf)
.map_err(|e| MLError::ModelError(format!("init saboteur_base: {e}")))?;
// Init base params to neutral (no adversarial effect). saboteur_select
// kernel writes back to this buffer, so destination stays CudaSlice;
// upload via mapped-pinned staging + DtoD.
upload_host_to_cuda_f32_via_pinned(
&stream, &[1.0_f32, 0.85, 1.0], &mut saboteur_base_buf, "saboteur_base init")?;
let saboteur_best_buf = stream.alloc_zeros::<f32>(3)
.map_err(|e| MLError::ModelError(format!("alloc saboteur_best: {e}")))?;
let saboteur_best_return_buf = stream.alloc_zeros::<f32>(1)
@@ -1493,8 +1513,20 @@ impl GpuExperienceCollector {
// stride-12 — tiled v_min/v_max/delta_z per branch. Filled once per epoch
// in update_per_sample_support(); consumers compute_expected_q /
// quantile_q_select / mag_concat_qdir index with (sample, branch, {0,1,2}).
let per_sample_support_buf = stream.alloc_zeros::<f32>(alloc_episodes * 4 * 3)
.map_err(|e| MLError::ModelError(format!("alloc per_sample_support_buf: {e}")))?;
// Mapped pinned per `feedback_no_htod_htoh_only_mapped_pinned.md`:
// CPU writes via host_ptr each epoch; kernels read via dev_ptr.
// Safety: CUDA context active on this thread (stream owns it).
let per_sample_support_buf = unsafe {
MappedF32Buffer::new(alloc_episodes * 4 * 3)
.map_err(|e| MLError::ModelError(format!("alloc per_sample_support_buf (mapped pinned): {e}")))?
};
// HOT path: hindsight relabel bar_indices buffer.
// Capacity = alloc_episodes * alloc_timesteps * 2 (cf doubles total).
let bar_indices_pinned = unsafe {
MappedI32Buffer::new(alloc_episodes * alloc_timesteps * 2)
.map_err(|e| MLError::ModelError(format!("alloc bar_indices_pinned: {e}")))?
};
// Task 8: GPU-resident step counter for experience collection loop
let step_counter_gpu = stream.alloc_zeros::<i32>(1)
@@ -1615,6 +1647,7 @@ impl GpuExperienceCollector {
ofi_dim,
curiosity_trainer,
feature_mask_buf: None,
bar_indices_pinned,
exp_h_s1_f32,
exp_h_s2_f32,
exp_h_v_f32,
@@ -2004,12 +2037,15 @@ impl GpuExperienceCollector {
/// entire training run. The action selection kernel reads this at each timestep
/// and overrides Q-network actions with probability `expert_ratio`.
pub fn upload_expert_actions(&mut self, actions: &[i32]) -> Result<(), MLError> {
let mut buf = self.stream.alloc_zeros::<i32>(actions.len())
.map_err(|e| MLError::ModelError(format!("alloc expert_actions: {e}")))?;
self.stream.memcpy_htod(actions, &mut buf)
.map_err(|e| MLError::ModelError(format!("upload expert_actions: {e}")))?;
// Mapped pinned: CPU writes via host_ptr, action-select kernel reads
// via dev_ptr — no HtoD copy. Safety: CUDA context active here.
let buf = unsafe {
MappedI32Buffer::new(actions.len())
.map_err(|e| MLError::ModelError(format!("alloc expert_actions (mapped pinned): {e}")))?
};
buf.write_from_slice(actions);
self.expert_actions_gpu = Some(buf);
info!(n_actions = actions.len(), "Expert demo actions uploaded to GPU");
info!(n_actions = actions.len(), "Expert demo actions uploaded to GPU (mapped pinned)");
Ok(())
}
@@ -2779,9 +2815,11 @@ impl GpuExperienceCollector {
let n = self.alloc_episodes;
let n_i32 = n as i32;
// Zero the output buffer before reduction
let zeros = vec![0.0_f32; TRADE_STATS_FLOATS];
super::htod_f32(&self.stream, &zeros, &mut self.trade_stats_buf)?;
// Zero the output buffer before reduction. memset_zeros stays GPU-side
// and avoids any CPU↔GPU traffic — the prior `htod_f32` of an all-zeros
// host vector was needlessly going through PCIe.
self.stream.memset_zeros(&mut self.trade_stats_buf)
.map_err(|e| MLError::ModelError(format!("zero trade_stats_buf: {e}")))?;
// Launch reduction kernel: 1 block, 256 threads
let blocks = 1_u32;
@@ -2976,7 +3014,15 @@ impl GpuExperienceCollector {
// v8: Hindsight relabeling — replace fraction of rewards with optimal exit PnL
if config.hindsight_fraction > 0.0 {
// Sanity check the persistent mapped-pinned buffer capacity.
if total > self.bar_indices_pinned.len {
return Err(MLError::ModelError(format!(
"bar_indices_pinned capacity {} < required {}",
self.bar_indices_pinned.len, total
)));
}
// Build bar_indices: bar_idx[ep * L + t] = episode_starts[ep] + t
// Write directly into the mapped-pinned host_ptr — no HtoD copy.
let mut bar_indices_cpu = vec![0_i32; total];
for ep in 0..n_episodes {
for t in 0..timesteps {
@@ -2987,10 +3033,8 @@ impl GpuExperienceCollector {
bar_indices_cpu[base_total + ep * timesteps + t] = bar;
}
}
let mut bar_indices_gpu = self.stream.alloc_zeros::<i32>(total)
.map_err(|e| MLError::ModelError(format!("alloc bar_indices: {e}")))?;
self.stream.memcpy_htod(&bar_indices_cpu, &mut bar_indices_gpu)
.map_err(|e| MLError::ModelError(format!("upload bar_indices: {e}")))?;
self.bar_indices_pinned.write_from_slice(&bar_indices_cpu);
let bar_indices_gpu_ptr = self.bar_indices_pinned.dev_ptr;
let b0 = self.branch_sizes[0] as i32;
let b1 = self.branch_sizes[1] as i32;
@@ -3008,7 +3052,7 @@ impl GpuExperienceCollector {
.launch_builder(&self.hindsight_relabel_kernel)
.arg(targets_buf) // const float* targets
.arg(&actions) // const int* actions
.arg(&bar_indices_gpu) // const int* bar_indices
.arg(&bar_indices_gpu_ptr) // const int* bar_indices (mapped pinned dev_ptr)
.arg(&rewards) // float* rewards
.arg(&b0) // int b0_size
.arg(&b1) // int b1_size
@@ -3093,9 +3137,11 @@ impl GpuExperienceCollector {
}
// ── Upload episode starts ───────────────────────────────────────
self.stream
.memcpy_htod(episode_starts, &mut self.episode_starts_buf)
.map_err(|e| MLError::ModelError(format!("upload episode_starts: {e}")))?;
// episode_starts_buf is GPU-mutated by `domain_rand_episode_starts`,
// so the destination must remain a CudaSlice. Stage via mapped-pinned
// + DtoD per `feedback_no_htod_htoh_only_mapped_pinned.md`.
upload_host_to_cuda_i32_via_pinned(
&self.stream, episode_starts, &mut self.episode_starts_buf, "episode_starts upload")?;
// ── Initialize current_timesteps to zeros ───────────────────────
self.stream
@@ -3178,14 +3224,28 @@ impl GpuExperienceCollector {
let noise_scale = config.feature_noise_scale;
let vol_norm = config.vol_normalizer;
// Upload feature mask if provided (once per epoch, not per timestep)
// Upload feature mask if provided (once per epoch, not per timestep).
// Mapped pinned: CPU writes via host_ptr, state_gather kernel
// reads via dev_ptr — no HtoD copy.
if t == 0 {
if let Some(ref mask_data) = config.feature_mask {
let mut buf = self.stream.alloc_zeros::<f32>(mask_data.len())
.map_err(|e| MLError::ModelError(format!("alloc feature_mask: {e}")))?;
self.stream.memcpy_htod(mask_data, &mut buf)
.map_err(|e| MLError::ModelError(format!("upload feature_mask: {e}")))?;
self.feature_mask_buf = Some(buf);
// Reallocate only when the mask size differs from a
// previously-cached buffer (epoch-boundary cold path).
let need_alloc = match &self.feature_mask_buf {
Some(buf) => buf.len != mask_data.len(),
None => true,
};
if need_alloc {
// Safety: CUDA context active on this thread (we're
// inside the surrounding unsafe block already).
self.feature_mask_buf = Some(
MappedF32Buffer::new(mask_data.len())
.map_err(|e| MLError::ModelError(format!("alloc feature_mask (mapped pinned): {e}")))?
);
}
// SAFETY: feature_mask_buf is Some after the alloc above.
let buf = self.feature_mask_buf.as_ref().expect("feature_mask_buf just allocated");
buf.write_from_slice(mask_data);
} else {
self.feature_mask_buf = None;
}
@@ -3193,7 +3253,7 @@ impl GpuExperienceCollector {
// #23 feature mask pointer (0 = NULL = no masking)
let mask_ptr: u64 = self.feature_mask_buf.as_ref()
.map(|b| b.device_ptr(&self.stream).0)
.map(|b| b.dev_ptr)
.unwrap_or(0u64);
// Plan params, ISV signals, and OFI features pointers
@@ -3351,6 +3411,7 @@ impl GpuExperienceCollector {
// handles num_atoms=1 correctly, and logits are now f32.
{
let na = self.num_atoms as i32;
let support_dev_ptr = self.per_sample_support_buf.dev_ptr;
unsafe {
let null_atom_stats = 0u64;
let null_atom_positions = 0u64; // NULL = use linear atom spacing
@@ -3365,7 +3426,7 @@ impl GpuExperienceCollector {
.arg(&b1)
.arg(&b2)
.arg(&b3)
.arg(&self.per_sample_support_buf) // [N, 4, 3] stride-12 per-sample per-branch support
.arg(&support_dev_ptr) // [N, 4, 3] stride-12 per-sample per-branch support (mapped pinned)
.arg(&null_atom_stats)
.arg(&mut self.q_var_buf)
.arg(&null_atom_positions) // atom_positions (NULL = linear)
@@ -3473,6 +3534,7 @@ impl GpuExperienceCollector {
let na = self.num_atoms as i32;
let iqn_r = self.iqn_readiness;
let util_e = self.util_ema;
let support_dev_ptr = self.per_sample_support_buf.dev_ptr;
unsafe {
self.stream
.launch_builder(&self.quantile_q_select_kernel)
@@ -3485,7 +3547,7 @@ impl GpuExperienceCollector {
.arg(&b1)
.arg(&b2)
.arg(&b3)
.arg(&self.per_sample_support_buf) // [N, 4, 3] stride-12 per-sample per-branch support
.arg(&support_dev_ptr) // [N, 4, 3] stride-12 per-sample per-branch support (mapped pinned)
.arg(&iqn_r)
.arg(&util_e)
.launch(launch_cfg)
@@ -3813,8 +3875,10 @@ impl GpuExperienceCollector {
let data: Vec<f32> = (0..n)
.flat_map(|_| (0..4).flat_map(|_| [v_min, v_max, delta_z]))
.collect();
super::htod_f32(&self.stream, &data, &mut self.per_sample_support_buf)?;
debug!(v_min, v_max, delta_z, n, "per_sample_support tiled [N, 4, 3] for experience collector");
// Direct host_ptr write into the mapped-pinned buffer — no HtoD copy.
// The kernel reads via dev_ptr after the next stream sync barrier.
self.per_sample_support_buf.write_from_slice(&data);
debug!(v_min, v_max, delta_z, n, "per_sample_support tiled [N, 4, 3] for experience collector (mapped pinned)");
Ok(())
}
pub fn set_per_sample_epsilon_ptr(&mut self, ptr: u64) {
@@ -3845,7 +3909,8 @@ impl GpuExperienceCollector {
portfolio_init[off + 9] = initial_capital; // [9] prev_equity
// [10] hold_time = 0.0, [11] realized_pnl = 0.0 (already zero)
}
super::htod_f32(&self.stream, &portfolio_init, &mut self.portfolio_states)?;
upload_host_to_cuda_f32_via_pinned(
&self.stream, &portfolio_init, &mut self.portfolio_states, "portfolio_states reset")?;
debug!(
initial_capital,
@@ -4523,3 +4588,66 @@ fn build_next_states_f32(
Ok(dst)
}
/// Upload `host` (f32) into `dst` via mapped-pinned staging + DtoD copy.
///
/// `dst` is a long-lived `CudaSlice` mutated by GPU kernels (so the
/// destination type cannot itself be a `MappedF32Buffer`). The staging
/// `MappedF32Buffer` provides the only allowed CPU↔GPU path per
/// `feedback_no_htod_htoh_only_mapped_pinned.md`; the DtoD copy is fully on-
/// device and trivially compliant.
fn upload_host_to_cuda_f32_via_pinned(
stream: &Arc<CudaStream>,
host: &[f32],
dst: &mut CudaSlice<f32>,
label: &str,
) -> Result<(), MLError> {
debug_assert!(host.len() <= dst.len(), "{label}: host overflow");
// Safety: a CUDA context is active on the calling thread.
let staging = unsafe {
MappedF32Buffer::new(host.len())
.map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?
};
staging.write_from_slice(host);
let nbytes = host.len() * std::mem::size_of::<f32>();
let src_ptr = staging.dev_ptr;
let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream);
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
).map_err(|e| MLError::ModelError(format!("{label} DtoD: {e}")))?;
}
// Sync so the staging buffer can drop safely after the copy completes.
stream.synchronize()
.map_err(|e| MLError::ModelError(format!("{label} sync: {e}")))?;
Ok(())
}
/// `i32` counterpart to `upload_host_to_cuda_f32_via_pinned`.
fn upload_host_to_cuda_i32_via_pinned(
stream: &Arc<CudaStream>,
host: &[i32],
dst: &mut CudaSlice<i32>,
label: &str,
) -> Result<(), MLError> {
debug_assert!(host.len() <= dst.len(), "{label}: host overflow");
// Safety: a CUDA context is active on the calling thread.
let staging = unsafe {
MappedI32Buffer::new(host.len())
.map_err(|e| MLError::ModelError(format!("{label} staging alloc: {e}")))?
};
staging.write_from_slice(host);
let nbytes = host.len() * std::mem::size_of::<i32>();
let src_ptr = staging.dev_ptr;
let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream);
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
).map_err(|e| MLError::ModelError(format!("{label} DtoD: {e}")))?;
}
stream.synchronize()
.map_err(|e| MLError::ModelError(format!("{label} sync: {e}")))?;
Ok(())
}

72
crates/ml/src/cuda_pipeline/gpu_portfolio.rs
View File

@@ -8,7 +8,7 @@
//! Uses cudarc 0.19 API:
//! - CudaContext / CudaStream (not CudaDevice — that's the older API)
//! - stream.launch_builder(&func).arg(&buf).launch(config)
//! - stream.alloc_zeros / memcpy_htod / memcpy_dtoh
//! - stream.alloc_zeros / memcpy_dtod (no host→device memcpy: mapped-pinned only)
use std::sync::Arc;
@@ -16,6 +16,7 @@ use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, DevicePtr
use tracing::{debug, info};
use crate::MLError;
use super::mapped_pinned::{MappedF32Buffer, MappedI32Buffer};
/// Maximum batch size for GPU portfolio simulation
const MAX_BATCH_SIZE: usize = 256;
@@ -34,7 +35,10 @@ pub struct GpuPortfolioSimulator {
kernel_func: CudaFunction,
// GPU buffers (persistent across batches)
actions_buf: CudaSlice<i32>,
/// Per-call action upload — mapped pinned (CPU writes via host_ptr, kernel
/// reads via dev_ptr). HOT path: every `simulate_batch_gpu` call writes
/// `batch_size` i32s into this buffer.
actions_buf: MappedI32Buffer,
portfolio_state_buf: CudaSlice<f32>,
portfolio_out_buf: CudaSlice<f32>,
rewards_out_buf: CudaSlice<f32>,
@@ -91,10 +95,15 @@ impl GpuPortfolioSimulator {
let kernel_func = compile_experience_kernels(stream.context())?;
info!("GPU portfolio simulator: CUDA kernels compiled and loaded");
// Allocate persistent buffers
let actions_buf = stream
.alloc_zeros::<i32>(MAX_BATCH_SIZE)
.map_err(|e| MLError::ModelError(format!("Failed to alloc actions buffer: {e}")))?;
// Allocate persistent buffers.
// `actions_buf` is mapped pinned: CPU writes via host_ptr, kernel reads
// via dev_ptr. Per `feedback_no_htod_htoh_only_mapped_pinned.md`, this is
// the only allowed CPU↔GPU path on the per-simulate hot path.
// Safety: a CUDA context is active on the calling thread (stream owns it).
let actions_buf = unsafe {
MappedI32Buffer::new(MAX_BATCH_SIZE)
.map_err(|e| MLError::ModelError(format!("Failed to alloc actions buffer (mapped pinned): {e}")))?
};
let mut portfolio_state_buf = stream
.alloc_zeros::<f32>(PORTFOLIO_STATE_SIZE)
.map_err(|e| MLError::ModelError(format!("Failed to alloc portfolio state: {e}")))?;
@@ -119,7 +128,7 @@ impl GpuPortfolioSimulator {
cash_reserve_pct,
0.0,
];
super::htod_f32(&stream, &init_state, &mut portfolio_state_buf)?;
upload_portfolio_state(&stream, &init_state, &mut portfolio_state_buf)?;
debug!(
"GPU portfolio sim initialized: capital={}, spread={}, reserve={}%, max_pos={}, episode_len={}, total_bars={}",
@@ -169,10 +178,10 @@ impl GpuPortfolioSimulator {
)));
}
// Upload actions to GPU
self.stream
.memcpy_htod(actions, &mut self.actions_buf)
.map_err(|e| MLError::ModelError(format!("Failed to upload actions: {e}")))?;
// Upload actions to GPU: CPU writes the host_ptr, the kernel reads the
// dev_ptr — mapped pinned coherence makes the data visible across the
// launch boundary without an HtoD memcpy.
self.actions_buf.write_from_slice(actions);
let config = LaunchConfig {
grid_dim: (1, 1, 1),
@@ -182,13 +191,14 @@ impl GpuPortfolioSimulator {
let batch_start_i32 = batch_start as i32;
let batch_size_i32 = batch_size as i32;
let actions_dev_ptr = self.actions_buf.dev_ptr;
// Safety: kernel parameters match the CUDA function signature exactly.
unsafe {
self.stream
.launch_builder(&self.kernel_func)
.arg(targets_buf)
.arg(&self.actions_buf)
.arg(&actions_dev_ptr)
.arg(&mut self.portfolio_state_buf)
.arg(&mut self.portfolio_out_buf)
.arg(&mut self.rewards_out_buf)
@@ -293,7 +303,7 @@ impl GpuPortfolioSimulator {
/// Reset portfolio state to initial capital.
pub fn reset(&mut self, initial_capital: f32, avg_spread: f32, cash_reserve_pct: f32) -> Result<(), MLError> {
let init_state = [initial_capital, 0.0, 0.0, initial_capital, avg_spread, 0.0, cash_reserve_pct, 0.0_f32];
super::htod_f32(&self.stream, &init_state, &mut self.portfolio_state_buf)?;
upload_portfolio_state(&self.stream, &init_state, &mut self.portfolio_state_buf)?;
Ok(())
}
@@ -305,3 +315,39 @@ impl GpuPortfolioSimulator {
&self.portfolio_state_buf
}
}
/// Upload `host` into `dst` via mapped-pinned + DtoD copy.
///
/// `dst` is a long-lived `CudaSlice` that the kernel later writes to (so we
/// can't replace the destination type with a `MappedF32Buffer`). The mapped
/// pinned staging buffer is the only allowed CPU↔GPU path per
/// `feedback_no_htod_htoh_only_mapped_pinned.md`; the DtoD copy is fully on-
/// device and trivially compliant.
fn upload_portfolio_state(
stream: &Arc<CudaStream>,
host: &[f32],
dst: &mut CudaSlice<f32>,
) -> Result<(), MLError> {
debug_assert!(host.len() <= dst.len(), "upload_portfolio_state overflow");
// Safety: a CUDA context is active on the calling thread.
let staging = unsafe {
MappedF32Buffer::new(host.len())
.map_err(|e| MLError::ModelError(format!("upload_portfolio_state staging alloc: {e}")))?
};
staging.write_from_slice(host);
let nbytes = host.len() * std::mem::size_of::<f32>();
let src_ptr = staging.dev_ptr;
let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream);
#[allow(unsafe_code)]
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
).map_err(|e| MLError::ModelError(format!("upload_portfolio_state DtoD: {e}")))?;
}
// The staging buffer must stay alive until the DtoD copy completes.
// Synchronize the stream so the host-side drop of `staging` is safe.
stream.synchronize()
.map_err(|e| MLError::ModelError(format!("upload_portfolio_state sync: {e}")))?;
Ok(())
}

2
crates/ml/src/cuda_pipeline/mapped_pinned.rs
View File

@@ -22,6 +22,7 @@ use std::mem::MaybeUninit;
/// CPU+GPU visible buffer of `i32`s allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. The kernel writes via `dev_ptr`
/// (with `__threadfence_system()`); the CPU reads via `host_ptr`.
#[allow(missing_debug_implementations)] // raw pointer + CUdeviceptr have no useful Debug
pub struct MappedI32Buffer {
pub host_ptr: *mut i32,
pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
@@ -202,6 +203,7 @@ impl Drop for MappedU32Buffer {
/// CPU+GPU visible buffer of `f32`s allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. The kernel writes via `dev_ptr`
/// (with `__threadfence_system()`); the CPU reads via `host_ptr`.
#[allow(missing_debug_implementations)] // raw pointer + CUdeviceptr have no useful Debug
pub struct MappedF32Buffer {
pub host_ptr: *mut f32,
pub dev_ptr: cudarc::driver::sys::CUdeviceptr,

16
docs/dqn-gpu-hot-path-audit.md
View File

@@ -72,6 +72,22 @@
| `gamma_update_kernel.cu` (Plan 1 Task 10) | single-thread kernel: reads ISV[12]; writes ISV[43] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; ISV[43] read by fill_gamma_buf on hot path via read_isv_signal_at (pinned, zero-copy) | OK |
| `kelly_cap_update_kernel.cu` (Plan 1 Task 11) | single-thread kernel: reads ISV[12] + portfolio_states[n_envs, PS_STRIDE]; aggregates mean half-Kelly × health-safety; writes ISV[44] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; portfolio_states dev ptr passed from GpuExperienceCollector without CPU roundtrip | OK |
| `atoms_update_kernel.cu` (Plan 1 Task 9) | 4-block kernel (one per branch): reads ISV[23..31) + spacing_raw params[52..56); writes atom_positions_buf[4 * num_atoms] via softmax + cumsum | COLD-PATH | Per-epoch boundary + SGD step (every 50 steps); grid=(4,1,1), block=(256,1,1); replaces CPU loop that launched adaptive_atom_positions 4× separately | OK |
| `gpu_portfolio.rs:174` | `actions_buf` upload (per simulate_batch_gpu) | **MIGRATED** | Promoted `actions_buf` from `CudaSlice<i32>` to `MappedI32Buffer(MAX_BATCH_SIZE)` allocated once at constructor. Per-call host writes go through `host_ptr`; kernel arg switched to `dev_ptr` u64. Zero HtoD copy on the simulate hot path. | FIXED |
| `gpu_portfolio.rs:122,296` | `portfolio_state_buf` init/reset | **MIGRATED** | env_step kernel writes back to `portfolio_state_buf`, so the destination must remain a `CudaSlice<f32>`. Init/reset now stage via a temporary `MappedF32Buffer` + `memcpy_dtod_async` (new helper `upload_portfolio_state`). | FIXED |
| `decision_transformer.rs:913` | `scratch.adam_t` upload per `train_step_gpu` | **MIGRATED** | HOT path: 4-byte step counter was uploaded via `stream.memcpy_htod` every backward pass. Promoted `DtScratch.adam_t` from `CudaSlice<i32>[1]` to `MappedI32Buffer(1)`. Per-step host write goes through `host_ptr`; Adam kernel reads via `dev_ptr`. | FIXED |
| `decision_transformer.rs:478` | DT params init upload (constructor) | **MIGRATED** | params is mutated by Adam, so destination must remain `CudaSlice<f32>`. Xavier-init upload now stages via `MappedF32Buffer` + `memcpy_dtod_async` through new helper `upload_host_to_cuda_f32`. | FIXED |
| `decision_transformer.rs:501` | DT wd_mask init (constructor) | **MIGRATED** | wd_mask is read-only by GPU. Promoted from `CudaSlice<f32>` to `MappedF32Buffer`. CPU writes 1.0s through `host_ptr`; kernel reads via `dev_ptr`. | FIXED |
| `decision_transformer.rs:1105` | trajectory build episode_starts | **MIGRATED** | One-time per pre-training data build (cold path). Replaced `clone_htod` with `MappedI32Buffer` + direct `host_ptr` write; kernel reads via `dev_ptr`. | FIXED |
| `decision_transformer.rs:1147` | trajectory build ep_rewards | **MIGRATED** | Same trajectory builder cold path. Replaced `htod_f32` with `MappedF32Buffer`; CPU writes via `host_ptr`, return_to_go kernel reads via `dev_ptr`. | FIXED |
| `gpu_experience_collector.rs:2992` | hindsight `bar_indices` upload (per `collect_experiences_gpu`) | **MIGRATED** | HOT path: alloc + memcpy_htod ran every collect when hindsight active. Added persistent `bar_indices_pinned: MappedI32Buffer(alloc_episodes * alloc_timesteps * 2)` allocated at constructor. Per-call host writes go through `host_ptr`; kernel arg switched to `dev_ptr` u64. | FIXED |
| `gpu_experience_collector.rs:3186` | feature_mask upload (per-collect at t==0) | **MIGRATED** | HOT path: was alloc + memcpy_htod every epoch. Promoted `feature_mask_buf` from `Option<CudaSlice<f32>>` to `Option<MappedF32Buffer>`. Reallocates only when mask size changes (still a cold-path event); CPU writes via `host_ptr`, state_gather reads via `dev_ptr`. | FIXED |
| `gpu_experience_collector.rs:3097` | episode_starts upload (per-collect) | **MIGRATED** | episode_starts_buf is GPU-mutated by `domain_rand_episode_starts`, so the destination must remain a CudaSlice. Stage via `MappedI32Buffer` + `memcpy_dtod_async` through `upload_host_to_cuda_i32_via_pinned`. | FIXED |
| `gpu_experience_collector.rs:2009` | `upload_expert_actions` (warm) | **MIGRATED** | Promoted `expert_actions_gpu` from `Option<CudaSlice<i32>>` to `Option<MappedI32Buffer>`. Direct host_ptr write replaces the alloc + memcpy_htod pair. | FIXED |
| `gpu_experience_collector.rs:3816` | `update_per_sample_support` (per-epoch) | **MIGRATED** | per_sample_support_buf is read-only by GPU; promoted from `CudaSlice<f32>` to `MappedF32Buffer`. Compute_expected_q + quantile_q_select now receive `dev_ptr` u64 kernel args. | FIXED |
| `gpu_experience_collector.rs:1076,3848` | portfolio_states init/reset (cold) | **MIGRATED** | portfolio_states is GPU-mutated by env_step. Init/reset stage via `MappedF32Buffer` + `memcpy_dtod_async` through `upload_host_to_cuda_f32_via_pinned`. | FIXED |
| `gpu_experience_collector.rs:1093` | epoch_state init (cold) | **MIGRATED** | epoch_state is GPU-mutated. Replaced `clone_htod_f32` with explicit `alloc_zeros` + `upload_host_to_cuda_f32_via_pinned`. | FIXED |
| `gpu_experience_collector.rs:1341` | saboteur_base init (cold) | **MIGRATED** | saboteur_base is GPU-mutated by `saboteur_select`. Replaced `memcpy_htod` with `upload_host_to_cuda_f32_via_pinned`. | FIXED |
| `gpu_experience_collector.rs:2784` | trade_stats_buf zero (per-epoch) | **MIGRATED** | The htod_f32 of an all-zeros vector was needlessly going through PCIe. Replaced with `stream.memset_zeros` — fully GPU-side. | FIXED |
## Summary