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

perf: accumulator uses mapped pinned memory — eliminates memcpy_dtoh

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 22:15:05 +02:00
parent 40cf48868a
commit 03cac6ffda
3 changed files with 45 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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(),

50
crates/ml/src/cuda_pipeline/gpu_training_guard.rs
View File

@@ -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<f32>,
/// 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::<f32>(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(())
}

31
crates/ml/src/cuda_pipeline/training_guard_kernel.cu
View File

@@ -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();
}
/* ------------------------------------------------------------------ */