diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 1118550e7..51c1c6f78 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -328,6 +328,26 @@ const RL_HINDSIGHT_FORWARD_CUBIN: &[u8] = /// Phase E.3+ without touching the kernel. const RL_LR_CONTROLLER_ALPHA: f32 = 0.4; +/// Pre-extracted raw CUdeviceptr values for hot-path buffers. +/// All pointers come from pre-allocated `CudaSlice` fields whose +/// device addresses are stable for the trainer's lifetime. +/// Reading these cached u64s avoids cudarc's `device_ptr()` / +/// `device_ptr_mut()` event-tracking path which internally does +/// `cuStreamWaitEvent` + `cuEventRecord` per call — responsible +/// for ~23 `cuStreamSynchronize` per step (89% wall time per nsys). +struct HotPathPtrs { + isv: u64, + ss_pi_loss: u64, + ss_q_loss: u64, + ss_v_loss_sum: u64, + ss_frd_loss_per_b_h: u64, + /// perception.h_t_d — stable (pre-allocated at PerceptionTrainer init). + h_t_encoder: u64, + h_tp1: u64, + trade_context: u64, + multires_output: u64, +} + /// Configuration for [`IntegratedTrainer`]. Wraps a `PerceptionTrainerConfig` /// (the encoder + BCE + aux side) plus RL-specific overrides for the new /// heads. @@ -1038,6 +1058,9 @@ pub struct IntegratedTrainer { /// loss-readback syncs per step_synthetic. loss_staging_3: MappedF32Buffer, + /// Cached raw CUdeviceptr values — zero `device_ptr()` calls in hot path. + hot: HotPathPtrs, + // IQN replay-step scratch (dqn_replay_step only). pub ss_iqn_tau_target_d: CudaSlice, pub ss_iqn_target_q_d: CudaSlice, @@ -2138,6 +2161,18 @@ impl IntegratedTrainer { .new_event(None) .map_err(|e| anyhow::anyhow!("replay_done_event: {e}"))?; + let hot = HotPathPtrs { + isv: isv_d.raw_ptr(), + ss_pi_loss: ss_pi_loss_d.raw_ptr(), + ss_q_loss: ss_q_loss_d.raw_ptr(), + ss_v_loss_sum: ss_v_loss_sum_d.raw_ptr(), + ss_frd_loss_per_b_h: ss_frd_loss_per_b_h_d.raw_ptr(), + h_t_encoder: perception.h_t_view().raw_ptr(), + h_tp1: h_tp1_d.raw_ptr(), + trade_context: trade_context_d.raw_ptr(), + multires_output: multires_output_d.raw_ptr(), + }; + Ok(Self { cfg, perception, @@ -2442,6 +2477,7 @@ impl IntegratedTrainer { frd_loss_staging, scalar_staging, loss_staging_3, + hot, } .with_controllers_bootstrapped()?) } @@ -3612,9 +3648,8 @@ impl IntegratedTrainer { // Queue DtoD for NEXT step's read — async, no sync needed here. 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(), + self.isv_staging.dev_ptr, self.hot.isv, nbytes, self.stream.cu_stream(), ) .context("step_synthetic isv staging DtoD (async, for next step)")?; } @@ -4476,30 +4511,27 @@ impl IntegratedTrainer { let f32_bytes = std::mem::size_of::(); unsafe { // Slot 0: pi_loss - let (pi_src, _g0) = self.ss_pi_loss_d.device_ptr(&self.stream); cudarc::driver::result::memcpy_dtod_async( self.loss_staging_3.dev_ptr, - pi_src, + self.hot.ss_pi_loss, f32_bytes, self.stream.cu_stream(), ) .context("batched loss DtoD: pi_loss")?; // Slot 1: q_loss - let (q_src, _g1) = self.ss_q_loss_d.device_ptr(&self.stream); cudarc::driver::result::memcpy_dtod_async( self.loss_staging_3.dev_ptr + f32_bytes as u64, - q_src, + self.hot.ss_q_loss, f32_bytes, self.stream.cu_stream(), ) .context("batched loss DtoD: q_loss")?; // Slot 2: v_loss_sum - let (v_src, _g2) = self.ss_v_loss_sum_d.device_ptr(&self.stream); cudarc::driver::result::memcpy_dtod_async( self.loss_staging_3.dev_ptr + 2 * f32_bytes as u64, - v_src, + self.hot.ss_v_loss_sum, f32_bytes, self.stream.cu_stream(), ) @@ -4510,10 +4542,9 @@ impl IntegratedTrainer { let frd_n_h = crate::rl::common::FRD_N_HORIZONS; let frd_n = b_size * frd_n_h; unsafe { - let (frd_src, _g3) = self.ss_frd_loss_per_b_h_d.device_ptr(&self.stream); cudarc::driver::result::memcpy_dtod_async( self.frd_loss_staging.dev_ptr, - frd_src, + self.hot.ss_frd_loss_per_b_h, frd_n * f32_bytes, self.stream.cu_stream(), ) @@ -5152,11 +5183,10 @@ impl IntegratedTrainer { { let nbytes = b_size * HIDDEN_DIM * std::mem::size_of::(); unsafe { - let s = self.stream.cu_stream(); - let (src, _gs) = self.perception.h_t_view().device_ptr(&self.stream); - let (dst, _gd) = self.h_tp1_d.device_ptr_mut(&self.stream); - cudarc::driver::result::memcpy_dtod_async(dst, src, nbytes, s) - .context("step_with_lobsim: DtoD perception.h_t_d → self.h_tp1_d")?; + cudarc::driver::result::memcpy_dtod_async( + self.hot.h_tp1, self.hot.h_t_encoder, nbytes, self.stream.cu_stream(), + ) + .context("step_with_lobsim: DtoD perception.h_t_d → self.h_tp1_d")?; } } @@ -5165,10 +5195,8 @@ impl IntegratedTrainer { // backward in step_synthetic, AND lands h_t at slot K-1 for // the action-sampling forwards below. { - let (tc_ptr, _g1) = self.trade_context_d.device_ptr(&self.stream); - let (mr_ptr, _g2) = self.multires_output_d.device_ptr(&self.stream); - self.perception.encoder_ctx_trade_ptr = tc_ptr; - self.perception.encoder_ctx_multires_ptr = mr_ptr; + self.perception.encoder_ctx_trade_ptr = self.hot.trade_context; + self.perception.encoder_ctx_multires_ptr = self.hot.multires_output; } let _ = self .perception