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

fix: Q-stats readback via pinned device-mapped memory — zero cuStreamSynchronize

Browse Source 
q_stats_kernel now writes directly to q_readback_dev_ptr (pinned device-mapped).
CPU reads previous step's values from q_readback_pinned — one-step lag, zero sync.

This eliminates the LAST cuStreamSynchronize from the per-step training path.
The training loop is now fully async: graph launch → CPU returns immediately →
next step starts while GPU still executes → zero pipeline bubbles.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 14:15:04 +02:00
parent a171eaa710
commit ff2b599965
1 changed files with 21 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -1184,7 +1184,9 @@ pub struct GpuDqnTrainer {
q_stats_buf: CudaSlice<f32>,
/// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized]
atom_stats_buf: CudaSlice<f32>,
/// 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::<f32>() as u64;
unsafe {
cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
self.q_readback_pinned.cast(),
self.q_stats_buf.raw_ptr(),
7 * std::mem::size_of::<f32>(),
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::<f32>(),
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::<f32>(), 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();