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

perf(rl): eliminate per-step mapped-pinned allocs — pre-allocated staging

Browse Source 
Replace 3 read_slice_d calls (each allocs+frees MappedF32Buffer) and
4 read_scalar_d calls with pre-allocated persistent isv_staging,
frd_loss_staging, and scalar_staging buffers. Deletes the now-unused
read_scalar_d function.

Eliminates ~20 cuMemHostAlloc + ~20 cuMemFree per step from the hot
path. GPU completes all kernels in 160us/step — host overhead was 14ms.
Target: <1ms host overhead at b=256.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 09:12:35 +02:00
parent e0ea71e90e
commit 274ff82494
1 changed files with 109 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -958,6 +958,19 @@ pub struct IntegratedTrainer {
pub ss_frd_grad_w2_d: CudaSlice<f32>,
pub ss_frd_grad_b2_d: CudaSlice<f32>,
// ── Pre-allocated mapped-pinned staging for hot-path device reads ──
// Eliminates per-step cuMemHostAlloc/Free from read_slice_d /
// read_scalar_d. Per `feedback_no_htod_htoh_only_mapped_pinned`:
// mapped-pinned is the only CPU↔GPU path.
/// ISV mirror staging `[RL_SLOTS_END]` — reused by both step_synthetic
/// and step_with_lobsim ISV refresh paths.
isv_staging: MappedF32Buffer,
/// FRD loss readback staging `[B × FRD_N_HORIZONS]`.
frd_loss_staging: MappedF32Buffer,
/// Scalar loss readback staging `[1]` — reused across all read_scalar_d
/// call sites within a single step (serialised by stream order).
scalar_staging: MappedF32Buffer,
// IQN replay-step scratch (dqn_replay_step only).
pub ss_iqn_tau_target_d: CudaSlice<f32>,
pub ss_iqn_target_q_d: CudaSlice<f32>,
@@ -1952,6 +1965,18 @@ impl IntegratedTrainer {
let ss_iqn_grad_b_embed_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)
.context("alloc ss_iqn_grad_b_embed_d")?;
// Pre-allocated mapped-pinned staging buffers — eliminates per-step
// cuMemHostAlloc/Free from the hot path. Allocated once at init,
// reused every step. Per `feedback_no_htod_htoh_only_mapped_pinned`.
let isv_staging = unsafe { MappedF32Buffer::new(RL_SLOTS_END) }
.map_err(|e| anyhow::anyhow!("isv_staging: {e}"))?;
let frd_loss_staging = unsafe {
MappedF32Buffer::new(b_size * crate::rl::common::FRD_N_HORIZONS)
}
.map_err(|e| anyhow::anyhow!("frd_loss_staging: {e}"))?;
let scalar_staging = unsafe { MappedF32Buffer::new(1) }
.map_err(|e| anyhow::anyhow!("scalar_staging: {e}"))?;
Ok(Self {
cfg,
perception,
@@ -2228,6 +2253,9 @@ impl IntegratedTrainer {
ss_iqn_grad_b_out_d,
ss_iqn_grad_w_embed_d,
ss_iqn_grad_b_embed_d,
isv_staging,
frd_loss_staging,
scalar_staging,
}
.with_controllers_bootstrapped()?)
}
@@ -3356,14 +3384,26 @@ impl IntegratedTrainer {
self.last_v_loss,
)
.context("rl_lr_controller launch")?;
// Mapped-pinned ISV mirror refresh per
// `feedback_no_htod_htoh_only_mapped_pinned` — was a raw
// `stream.memcpy_dtoh` (forbidden by
// `feedback_no_htod_htoh_only_mapped_pinned`).
let isv_host_fresh =
read_slice_d(&self.stream, &self.isv_d, self.isv_host.len())
.context("step_synthetic: read isv")?;
self.isv_host.copy_from_slice(&isv_host_fresh);
// Mapped-pinned ISV mirror refresh via pre-allocated staging per
// `feedback_no_htod_htoh_only_mapped_pinned`. Reuses
// `self.isv_staging` — no per-step cuMemHostAlloc/Free.
{
let n = self.isv_host.len();
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (src, _g) = self.isv_d.device_ptr(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
self.isv_staging.dev_ptr, src, nbytes, self.stream.cu_stream(),
)
.context("step_synthetic isv staging DtoD")?;
}
self.stream.synchronize().context("step_synthetic isv staging sync")?;
for i in 0..n {
self.isv_host[i] = unsafe {
std::ptr::read_volatile(self.isv_staging.host_ptr.add(i))
};
}
}
let lambdas = read_loss_lambdas_from_isv(&self.isv_host);
// Mutate each per-head Adam's lr field from the ISV mirror. The
// controller has just emitted into ISV[412..417], so this read
@@ -4160,20 +4200,40 @@ impl IntegratedTrainer {
// to synchronize before reading mapped-pinned loss scalars.
self.stream.synchronize().context("step_synthetic sync")?;
let l_pi_host = read_scalar_d(&self.stream, &self.ss_pi_loss_d)?;
let l_pi_host = read_scalar_via_staging(
&self.stream, &self.ss_pi_loss_d, &self.scalar_staging,
)?;
self.last_pi_loss = l_pi_host;
let l_q_host = read_scalar_d(&self.stream, &self.ss_q_loss_d)?;
let l_q_host = read_scalar_via_staging(
&self.stream, &self.ss_q_loss_d, &self.scalar_staging,
)?;
self.last_q_loss = l_q_host / (b_size as f32);
let l_v_sum_host = read_scalar_d(&self.stream, &self.ss_v_loss_sum_d)?;
let l_v_sum_host = read_scalar_via_staging(
&self.stream, &self.ss_v_loss_sum_d, &self.scalar_staging,
)?;
let l_v_host = l_v_sum_host / (b_size as f32);
self.last_v_loss = l_v_host;
let frd_n_h = crate::rl::common::FRD_N_HORIZONS;
let loss_pb_h =
read_slice_d(&self.stream, &self.ss_frd_loss_per_b_h_d, b_size * frd_n_h)?;
let l_frd_host = loss_pb_h.iter().sum::<f32>() / ((b_size * frd_n_h) as f32);
let l_frd_host = {
let n = b_size * frd_n_h;
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (src, _g) = self.ss_frd_loss_per_b_h_d.device_ptr(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
self.frd_loss_staging.dev_ptr, src, nbytes, self.stream.cu_stream(),
)
.context("frd loss staging DtoD")?;
}
self.stream.synchronize().context("frd loss staging sync")?;
let mut sum = 0.0_f32;
for i in 0..n {
sum += unsafe { std::ptr::read_volatile(self.frd_loss_staging.host_ptr.add(i)) };
}
sum / (n as f32)
};
// ── Step 12: compose stats ───────────────────────────────────
// BCE / aux losses are NOT read this phase — perception is driven
@@ -4497,7 +4557,9 @@ impl IntegratedTrainer {
&mut self.ss_q_grad_logits_d,
)
.context("dqn_replay_step: dqn_head.backward_logits")?;
let l_q_host = read_scalar_d(&self.stream, &self.ss_q_loss_d)?;
let l_q_host = read_scalar_via_staging(
&self.stream, &self.ss_q_loss_d, &self.scalar_staging,
)?;
let l_q = l_q_host / (b_size as f32);
self.dqn_head
@@ -5744,14 +5806,26 @@ impl IntegratedTrainer {
// is later (Step 2 of step_synthetic) — too late for this
// pre-step_synthetic call. The cost is one DtoH of the full
// ISV slice (424 floats), bounded.
// Mapped-pinned staging for ISV mirror refresh per
// `feedback_no_htod_htoh_only_mapped_pinned`. The full ISV
// slice is 424 floats; one mapped-pinned alloc + DtoD per
// step is bounded.
let isv_host_fresh =
read_slice_d(&self.stream, &self.isv_d, self.isv_host.len())
.context("step_with_lobsim: read isv (pre-PER)")?;
self.isv_host.copy_from_slice(&isv_host_fresh);
// Mapped-pinned ISV mirror refresh via pre-allocated staging per
// `feedback_no_htod_htoh_only_mapped_pinned`. Reuses
// `self.isv_staging` — no per-step cuMemHostAlloc/Free.
{
let n = self.isv_host.len();
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (src, _g) = self.isv_d.device_ptr(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
self.isv_staging.dev_ptr, src, nbytes, self.stream.cu_stream(),
)
.context("step_with_lobsim isv staging DtoD")?;
}
self.stream.synchronize().context("step_with_lobsim isv staging sync")?;
for i in 0..n {
self.isv_host[i] = unsafe {
std::ptr::read_volatile(self.isv_staging.host_ptr.add(i))
};
}
}
// GPU PER push — split into two coalesced kernels:
// Phase 1 (ring): n-step accumulation + flush decision per batch.
@@ -6623,10 +6697,17 @@ fn write_slice_i32_d(
Ok(())
}
fn read_scalar_d(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<f32> {
/// Hot-path scalar read reusing a pre-allocated `MappedF32Buffer` staging
/// slot. Same DtoD + sync + volatile-read pattern as `read_scalar_d` but
/// without per-call `cuMemHostAlloc` / `cuMemFreeHost`. The staging buffer
/// is `[1]` and serialised by stream order across consecutive calls.
fn read_scalar_via_staging(
stream: &Arc<CudaStream>,
src: &CudaSlice<f32>,
staging: &MappedF32Buffer,
) -> Result<f32> {
debug_assert!(src.len() >= 1);
let staging = unsafe { MappedF32Buffer::new(1) }
.map_err(|e| anyhow::anyhow!("read_scalar_d staging: {e}"))?;
debug_assert!(staging.len >= 1);
unsafe {
let (src_ptr, _g) = src.device_ptr(stream);
cudarc::driver::result::memcpy_dtod_async(
@@ -6635,9 +6716,9 @@ fn read_scalar_d(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<f32>
std::mem::size_of::<f32>(),
stream.cu_stream(),
)
.context("read_scalar_d DtoD")?;
.context("read_scalar_via_staging DtoD")?;
}
stream.synchronize().context("read_scalar_d sync")?;
stream.synchronize().context("read_scalar_via_staging sync")?;
Ok(unsafe { std::ptr::read_volatile(staging.host_ptr) })
}