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

fix(action-selector): migrate set_count_bonuses to mapped pinned (HOT path)

Browse Source 
Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned
(cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path.
`set_count_bonuses` is called per action-selection — converting the
three bonus buffers from CudaSlice<f32> (HtoD memcpy) to
MappedF32Buffer eliminates 3-6 HtoD memcpys per call.

Before: 3 reused-buffer htod_f32 (lines 92/96/100) + 3 first-call
clone_htod_f32 (lines 93/97/101) on every call = 3 HtoD steady-state.
After: zero HtoD — direct host_ptr writes via write_from_slice.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 20:56:37 +02:00
parent 3c0d262927
commit 08ddf10a1a
2 changed files with 39 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
crates/ml/src/cuda_pipeline/gpu_action_selector.rs
View File

@@ -19,6 +19,7 @@ use std::sync::Arc;
use tracing::info;
use crate::MLError;
use super::mapped_pinned::MappedF32Buffer;
/// 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"));
@@ -35,14 +36,17 @@ pub struct GpuActionSelector {
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>,
@@ -85,27 +89,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(())
}

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

@@ -1100,7 +1100,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 | — |