diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 56f81c27a..f62b75d07 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1184,7 +1184,9 @@ pub struct GpuDqnTrainer { q_stats_buf: CudaSlice, /// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized] atom_stats_buf: CudaSlice, - /// Pinned readback buffer for q_stats [7 floats] + q_out sample 0 [12 floats] = 19 floats. + /// Pinned device-mapped readback for q_stats [7] + q_out sample 0 [12] = 19 floats. + /// GPU writes via q_readback_dev_ptr, CPU reads via q_readback_pinned — zero sync. + q_readback_dev_ptr: u64, /// DtoH uses pinned DMA-capable destination for faster async transfer. q_readback_pinned: *mut f32, @@ -4397,6 +4399,11 @@ impl GpuDqnTrainer { as *mut f32 }; unsafe { std::ptr::write_bytes(q_readback_pinned, 0, 19); } + let q_readback_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, q_readback_pinned.cast(), 0); + dp + }; info!("GpuDqnTrainer: expected_q + q_stats kernels compiled"); // ── Load mag_concat and strided_accumulate from experience_kernels cubin ── @@ -5835,6 +5842,7 @@ impl GpuDqnTrainer { q_stats_buf, atom_stats_buf, q_readback_pinned, + q_readback_dev_ptr, cql_logit_grad_kernel, cql_d_value_logits, cql_d_adv_logits, @@ -7244,13 +7252,16 @@ impl GpuDqnTrainer { let total_actions = self.total_actions() as i32; let num_atoms = self.config.num_atoms as i32; - // Stats reduction on the already-computed q_out_buf. + // Stats reduction → pinned device-mapped buffer (zero sync, one-step lag). + // q_stats_kernel writes 7 floats to q_readback_dev_ptr[0..7]. + // CPU reads from q_readback_pinned — previous step's values. + let stats_dev_ptr = self.q_readback_dev_ptr; unsafe { self.stream .launch_builder(&self.q_stats_kernel) .arg(&self.q_out_buf) .arg(&self.atom_stats_buf) - .arg(&self.q_stats_buf) + .arg(&stats_dev_ptr) .arg(&n) .arg(&total_actions) .arg(&num_atoms) @@ -7261,27 +7272,15 @@ impl GpuDqnTrainer { }) .map_err(|e| MLError::ModelError(format!("q_stats_kernel: {e}")))?; } - - // 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 + 48B DtoH. Runs every 50 training steps. - // DtoHAsync to pinned readback buffer (DMA-capable destination), then sync. + // Copy q_out sample 0 (12 floats) to pinned buffer [7..19] for per-branch Q-gap. + let q_out_dst = self.q_readback_dev_ptr + 7 * std::mem::size_of::() as u64; unsafe { - cudarc::driver::sys::cuMemcpyDtoHAsync_v2( - self.q_readback_pinned.cast(), - self.q_stats_buf.raw_ptr(), - 7 * std::mem::size_of::(), - self.stream.cu_stream(), - ); - cudarc::driver::sys::cuMemcpyDtoHAsync_v2( - self.q_readback_pinned.add(7).cast(), - self.q_out_buf.raw_ptr(), - 12 * std::mem::size_of::(), - self.stream.cu_stream(), - ); - cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); + cudarc::driver::result::memcpy_dtod_async( + q_out_dst, self.q_out_buf.raw_ptr(), + 12 * std::mem::size_of::(), self.stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("q_out sample0 DtoD: {e}")))?; } + // No sync — CPU reads previous step's values from pinned host pointer. // Q-mean EMA update — AFTER stream sync so pinned memory has the GPU-written value. self.update_q_mean_ema();