From 274ff82494a57b4b2751ea116c3b47b954861494 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 09:12:35 +0200 Subject: [PATCH] =?UTF-8?q?perf(rl):=20eliminate=20per-step=20mapped-pinne?= =?UTF-8?q?d=20allocs=20=E2=80=94=20pre-allocated=20staging?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/src/trainer/integrated.rs | 137 +++++++++++++++++----- 1 file changed, 109 insertions(+), 28 deletions(-) diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 848565ee2..128f68e8f 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -958,6 +958,19 @@ pub struct IntegratedTrainer { pub ss_frd_grad_w2_d: CudaSlice, pub ss_frd_grad_b2_d: CudaSlice, + // ── 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, pub ss_iqn_target_q_d: CudaSlice, @@ -1952,6 +1965,18 @@ impl IntegratedTrainer { let ss_iqn_grad_b_embed_d = stream.alloc_zeros::(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::(); + 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::() / ((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::(); + 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::(); + 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, src: &CudaSlice) -> Result { +/// 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, + src: &CudaSlice, + staging: &MappedF32Buffer, +) -> Result { 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, src: &CudaSlice) -> Result std::mem::size_of::(), 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) }) }