diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 659bfa221..1cd1d6404 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -642,11 +642,14 @@ pub struct GpuDqnTrainer { sel_norm_partials: CudaSlice, // [1] /// [1] max grad norm for sel Adam clipping (fixed 1.0). sel_clip_buf: CudaSlice, // [1] - /// [1] sel Adam step counter on device (for CUDA kernel compatibility). - sel_t_buf: CudaSlice, // [1] + /// Selectivity Adam step counter — pinned device-mapped (no per-step HtoD copy). + sel_t_pinned: *mut i32, + sel_t_dev_ptr: u64, // ── Liquid tau ── per_branch_q_gap_ema: [f32; 4], + /// Last computed per-branch Q-gaps from reduce_current_q_stats. + last_per_branch_q_gaps: [f32; 4], liquid_mod_pinned: *mut f32, // 4 floats pinned device-mapped liquid_mod_dev_ptr: u64, @@ -1069,6 +1072,9 @@ impl Drop for GpuDqnTrainer { if !self.liquid_mod_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.liquid_mod_pinned.cast()) }; } + if !self.sel_t_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.sel_t_pinned.cast()) }; + } if !self.t_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.t_pinned.cast()) }; } @@ -2728,8 +2734,23 @@ impl GpuDqnTrainer { let mut sel_clip_buf = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc sel_clip_buf: {e}")))?; super::htod_f32(&stream, &[1.0_f32], &mut sel_clip_buf)?; - // sel_t_buf: Adam step counter for selectivity gate, initialized to 0. - let sel_t_buf = alloc_i32(&stream, 1, "sel_t_buf")?; + // Selectivity Adam step counter — pinned device-mapped (no per-step HtoD copy) + let sel_t_pinned: *mut i32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned sel_t alloc: {e}")))? + as *mut i32 + }; + unsafe { *sel_t_pinned = 0; } + let sel_t_dev_ptr = unsafe { + let mut dev_ptr: u64 = 0; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dev_ptr as *mut u64, + sel_t_pinned.cast(), + 0, + ); + dev_ptr + }; // ── Liquid tau: 4 pinned device-mapped floats ───────────────── let liquid_mod_pinned: *mut f32 = unsafe { @@ -3312,8 +3333,10 @@ impl GpuDqnTrainer { sel_norm_buf, sel_norm_partials, sel_clip_buf, - sel_t_buf, + sel_t_pinned, + sel_t_dev_ptr, per_branch_q_gap_ema: [0.0; 4], + last_per_branch_q_gaps: [0.0; 4], liquid_mod_pinned, liquid_mod_dev_ptr, vsn_masked_buf, @@ -4756,8 +4779,7 @@ impl GpuDqnTrainer { // Synchronous readback — Q-stats drive eval_v_range which drives // action selection. Stale values cause 2-state oscillation where // the system flips between cached results from different steps. - // Cost: ~5µs sync + 28B DtoH. Acceptable since this runs every - // 50 training steps, not per-step. + // Cost: ~5µs sync + 28B + 48B DtoH. Runs every 50 training steps. unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } let mut host = [0.0_f32; 7]; unsafe { @@ -4768,6 +4790,30 @@ impl GpuDqnTrainer { ); } + // Per-branch Q-gap: read first sample's 12 Q-values from q_out_buf. + // Stream already synced. Branch layout: [dir(3), mag(3), ord(3), urg(3)]. + let mut q12 = [0.0_f32; 12]; + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + q12.as_mut_ptr().cast(), + self.q_out_buf.raw_ptr(), + 12 * std::mem::size_of::(), + ); + } + let branch_sizes = [ + self.config.branch_0_size, self.config.branch_1_size, + self.config.branch_2_size, self.config.branch_3_size, + ]; + let mut offset = 0; + for d in 0..4 { + let bs = branch_sizes[d]; + let slice = &q12[offset..offset + bs]; + let max_q = slice.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let mean_q = slice.iter().sum::() / bs as f32; + self.last_per_branch_q_gaps[d] = (max_q - mean_q).max(0.0); + offset += bs; + } + Ok(QValueStatsResult { avg_max_q: host[0] as f64, q_min: host[1], @@ -4779,6 +4825,9 @@ impl GpuDqnTrainer { }) } + /// Last computed per-branch Q-gaps from reduce_current_q_stats. + pub fn get_per_branch_q_gaps(&self) -> [f32; 4] { self.last_per_branch_q_gaps } + /// Flush the last in-flight Q-stats readback at epoch end — non-blocking. /// /// Reads directly from the pinned host buffer without synchronizing. @@ -6802,16 +6851,15 @@ impl GpuDqnTrainer { /// Run Adam update on selectivity gate parameters. /// /// Computes L2 grad norm from sel_grad, then runs dqn_adam_update_kernel at LR=1e-4. - /// Increments sel_adam_step and writes the updated step counter to sel_t_buf. + /// Increments sel_adam_step and writes to pinned device-mapped sel_t_pinned. pub(crate) fn step_selectivity_adam(&mut self) -> Result<(), MLError> { self.sel_adam_step += 1; let step_val = self.sel_adam_step; let sel_dim = self.config.shared_h2 + 1; let sel_n = sel_dim as i32; - // Write step counter to device buffer (memcpy_htod — small 4-byte copy, not on hot path). - self.stream.memcpy_htod(&[step_val], &mut self.sel_t_buf) - .map_err(|e| MLError::ModelError(format!("sel_t_buf htod: {e}")))?; + // Write step counter to pinned device-mapped memory (no HtoD copy needed) + unsafe { *self.sel_t_pinned = step_val; } // Phase 1: grad_norm_kernel on sel_grad → sel_norm_partials [1 block]. let grad_ptr = self.sel_grad.raw_ptr(); @@ -6858,7 +6906,7 @@ impl GpuDqnTrainer { let m_ptr = self.sel_adam_m.raw_ptr(); let v_ptr = self.sel_adam_v.raw_ptr(); let clip_ptr = self.sel_clip_buf.raw_ptr(); - let t_ptr = self.sel_t_buf.raw_ptr(); + let t_ptr = self.sel_t_dev_ptr; let blocks = ((sel_dim as u32 + 255) / 256).max(1); unsafe { self.stream diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 7e3d5356b..a8fe31713 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2153,6 +2153,7 @@ impl FusedTrainingCtx { /// Read per_branch_q_gap_ema. pub(crate) fn per_branch_q_gap_ema(&self) -> [f32; 4] { self.trainer.per_branch_q_gap_ema() } + pub(crate) fn get_per_branch_q_gaps(&self) -> [f32; 4] { self.trainer.get_per_branch_q_gaps() } /// Set eval_q_mean_ema (trajectory backtracking restore). pub(crate) fn set_eval_q_mean_ema(&mut self, v: f32) { self.trainer.set_eval_q_mean_ema(v); } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index ec1fa9064..21511f99f 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1404,9 +1404,8 @@ impl DQNTrainer { // Update eval v_range from observed Q-stats + Q-gap let q_std = stats.q_variance.max(0.0).sqrt(); let q_gap = (stats.avg_max_q as f32 - stats.q_mean).max(0.0); - // Simple approximation: use global q_gap for all branches initially. - // Per-branch extraction requires reading 12 Q-values from GPU. - let per_branch_q_gaps = [q_gap; 4]; // uniform until per-branch extraction is added + // Per-branch Q-gaps computed in reduce_current_q_stats (48B DtoH readback). + let per_branch_q_gaps = fused.get_per_branch_q_gaps(); fused.update_eval_v_range(stats.q_mean, q_std, q_gap, per_branch_q_gaps); } } @@ -2304,9 +2303,7 @@ impl DQNTrainer { fn is_q_gap_frozen(&self) -> bool { if let Some(ref fused) = self.fused_ctx { let emas = fused.per_branch_q_gap_ema(); - // Current q_gap approximated as uniform across branches (same as update_eval_v_range) - let q_gap = self.epoch_q_gap; - let current = [q_gap; 4]; + let current = fused.get_per_branch_q_gaps(); let threshold = 0.001_f32; current.iter().zip(emas.iter()).all(|(gap, ema)| (gap - ema).abs() < threshold) } else {