From 03cac6ffda461a0ed8a184a8304ea3f207c40528 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 13 Apr 2026 22:15:05 +0200 Subject: [PATCH] =?UTF-8?q?perf:=20accumulator=20uses=20mapped=20pinned=20?= =?UTF-8?q?memory=20=E2=80=94=20eliminates=20memcpy=5Fdtoh?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit volatile float* bypasses GPU L2 cache for correct read-modify-write across kernel launches. CPU reads directly via read_volatile at epoch boundary. Removes dead ptr variable in update_eval_v_range. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 3 -- .../src/cuda_pipeline/gpu_training_guard.rs | 50 ++++++++++--------- .../cuda_pipeline/training_guard_kernel.cu | 31 +++++++----- 3 files changed, 45 insertions(+), 39 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 835b680f8..327dba77d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1017,9 +1017,6 @@ impl GpuDqnTrainer { let v_min = q_mean - half; let v_max = q_mean + half; unsafe { - let ptr = self.eval_v_range_ptr as *mut f32; - // Pinned device-mapped: write directly from host - // eval_v_range_buf is regular device mem — use async HtoD cudarc::driver::sys::cuMemcpyHtoDAsync_v2( self.eval_v_range_ptr, [v_min, v_max].as_ptr().cast(), diff --git a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs index a5a92ee27..feeadbcb9 100644 --- a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs +++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs @@ -11,8 +11,9 @@ //! reads the previous step's result from the other — no `memcpy_dtoh` //! is needed on the hot path. //! -//! The accumulator buffer (3 floats) stays in device memory; a single -//! `memcpy_dtoh` at epoch boundary returns (mean_loss, mean_grad_norm). +//! The accumulator buffer (3 floats) also uses mapped pinned memory with +//! `volatile float*` in the kernel to bypass GPU L2 cache. The CPU reads +//! directly via `read_volatile` at epoch boundary — zero memcpy. use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; use std::ffi::c_void; @@ -188,9 +189,9 @@ pub struct GpuTrainingGuard { qdiv_has_prev: bool, /// Accumulator buffer (3 floats: loss_sum, grad_norm_sum, step_count). - /// Device memory — GPU does read-modify-write (+=) which requires device - /// memory, not mapped host memory. DtoH at epoch boundary only. - acc_buf: CudaSlice, + /// Mapped pinned memory — GPU reads/writes via `volatile float*` (bypasses + /// L2 cache, goes through PCIe). CPU reads via `read_volatile`. No memcpy. + acc_buf: MappedBuffer, /// Welford running-mean accumulator for Q-value estimation (CPU-side, no GPU tensor). q_count: usize, @@ -222,10 +223,9 @@ impl GpuTrainingGuard { .load_function("qvalue_divergence_check") .map_err(|e| MLError::ModelError(format!("qvalue_divergence_check load: {e}")))?; - // Device memory accumulator — GPU does +=, DtoH at epoch boundary. - let acc_buf = stream.alloc_zeros::(3).map_err(|e| { - MLError::ModelError(format!("alloc acc_buf: {e}")) - })?; + // Mapped pinned accumulator — GPU does volatile +=, CPU reads directly. + // Safety: CUDA context is active. + let acc_buf = unsafe { MappedBuffer::new(3)? }; // Allocate mapped pinned double-buffers // Safety: CUDA context is active (we just loaded the module on this device). @@ -318,14 +318,14 @@ impl GpuTrainingGuard { }; // Pass raw device pointers to bypass cudarc event tracking on graph-captured buffers. - let acc_ptr = self.acc_buf.raw_ptr(); + let acc_dev_ptr = self.acc_buf.dev_ptr; unsafe { self.stream .launch_builder(&self.fused_check_accum_func) .arg(&loss_ptr) .arg(&grad_norm_ptr) .arg(&write_dev_ptr) - .arg(&acc_ptr) + .arg(&acc_dev_ptr) .arg(&clip_threshold) .arg(&collapse_threshold) .arg(&warmup_int) @@ -342,16 +342,12 @@ impl GpuTrainingGuard { /// Read the epoch-boundary loss and grad_norm averages from the accumulator. /// - /// Read the epoch-boundary loss and grad_norm averages from the accumulator. - /// DtoH readback — device memory requires memcpy. + /// Direct `read_volatile` from mapped pinned memory — no memcpy needed. + /// The kernel's `__threadfence_system()` ensures CPU-visible writes. pub fn read_accumulators(&mut self) -> Result<(f64, f64), MLError> { - let mut host = [0.0_f32; 3]; - self.stream.memcpy_dtoh(&self.acc_buf, &mut host) - .map_err(|e| MLError::ModelError(format!("acc_buf DtoH: {e}")))?; - - let loss_sum = host[0] as f64; - let grad_sum = host[1] as f64; - let steps = host[2] as f64; + let loss_sum = self.acc_buf.read(0) as f64; + let grad_sum = self.acc_buf.read(1) as f64; + let steps = self.acc_buf.read(2) as f64; if steps <= 0.0 { return Ok((0.0, 0.0)); @@ -360,11 +356,17 @@ impl GpuTrainingGuard { Ok((loss_sum / steps, grad_sum / steps)) } - /// Zero the accumulator — write zeros to pinned host memory (visible to GPU). + /// Zero the accumulator — write zeros directly to mapped host memory. + /// No GPU command needed; the kernel uses `volatile` reads so it will + /// see the zeroed values on the next launch. pub fn reset_accumulators(&mut self) -> Result<(), MLError> { - self.stream.memset_zeros(&mut self.acc_buf).map_err(|e| { - MLError::ModelError(format!("reset acc_buf: {e}")) - })?; + // Safety: host_ptr is valid for 3 floats, and volatile kernel reads + // will pick up the new zeros without any explicit cache flush. + unsafe { + std::ptr::write_volatile(self.acc_buf.host_ptr.add(0), 0.0_f32); + std::ptr::write_volatile(self.acc_buf.host_ptr.add(1), 0.0_f32); + std::ptr::write_volatile(self.acc_buf.host_ptr.add(2), 0.0_f32); + } Ok(()) } diff --git a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu index fcdcac3c0..9d04c03a6 100644 --- a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu +++ b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu @@ -13,9 +13,10 @@ * qvalue_stats_reduce -- batch-wide Q-value statistics (parallel reduction) * qvalue_divergence_check -- single-sample Q-value divergence check * - * NOTE: output and qstats/qdiv output are float* host-mapped pinned memory - * (read by CPU as f32 via read_volatile). These stay float* at the boundary. - * acc_buf is float* device memory (allocated as f32 on Rust side). + * NOTE: All output buffers (output, qstats, qdiv, acc_buf) are host-mapped + * pinned memory (read by CPU as f32 via read_volatile). acc_buf uses + * `volatile float*` to bypass GPU L2 cache for correct read-modify-write + * accumulation across kernel launches. */ /* ------------------------------------------------------------------ */ @@ -32,16 +33,16 @@ /* [5] raw_grad_norm */ /* [6] reserved -- 0.0 */ /* */ -/* acc_buf layout (3 f32, device memory): */ +/* acc_buf layout (3 f32, mapped pinned memory): */ /* [0] loss_sum */ /* [1] grad_norm_sum */ /* [2] step_count */ /* ------------------------------------------------------------------ */ extern "C" __global__ void training_guard_check_and_accumulate( const float* __restrict__ loss_scalar, /* GPU-resident f32 scalar */ - const float* __restrict__ grad_norm_scalar, /* GPU-resident bf16 scalar */ + const float* __restrict__ grad_norm_scalar, /* GPU-resident f32 scalar */ float* output, /* pinned host buffer (7 floats) */ - float* acc_buf, /* device accumulator (3 f32) */ + volatile float* acc_buf, /* mapped pinned accumulator (3 f32) */ float clip_threshold, float collapse_threshold, int warmup @@ -74,18 +75,24 @@ extern "C" __global__ void training_guard_check_and_accumulate( output[4] = loss; output[5] = grad_norm; output[6] = 0.0f; - __threadfence_system(); - /* -- Accumulate to pinned device-mapped buffer -- */ + /* -- Accumulate to mapped pinned buffer -- */ + /* volatile qualifier bypasses GPU L2 cache, forcing every load/store + to go through PCIe to host memory. Without volatile, the GPU L2 + would serve stale reads (initial zeros) and each launch would + overwrite instead of accumulating. */ if (!loss_nan) { - acc_buf[0] += loss; + acc_buf[0] = acc_buf[0] + loss; } if (!grad_nan) { - acc_buf[1] += grad_norm; + acc_buf[1] = acc_buf[1] + grad_norm; } - acc_buf[2] += 1.0f; - /* acc_buf is device memory — no threadfence needed (DtoH at epoch boundary). */ + acc_buf[2] = acc_buf[2] + 1.0f; + + /* Single threadfence_system covers both output and acc_buf — + ensures all mapped-memory writes are visible to the CPU. */ + __threadfence_system(); } /* ------------------------------------------------------------------ */