From e4f838d4e3de146ae963c0f1ad5736dcfe5171eb Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 09:24:36 +0200 Subject: [PATCH] =?UTF-8?q?perf:=20GPU-native=20HER=20random=20donors=20?= =?UTF-8?q?=E2=80=94=20eliminate=20last=20CPU=20touch=20in=20hot=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Added her_sample_random_donors kernel to her_episode_kernel.cu. Uses the existing per-sample LCG RNG state (already allocated in GpuHer constructor, already used by Future strategy). The per-step training hot path now has: - Zero memcpy_htod (was 1 for HER random donors) - Zero memcpy_dtoh - Zero Vec/alloc - Zero .clone() - Zero cuStreamSynchronize - Zero format!/String Every CPU→GPU and GPU→CPU transfer has been eliminated. The only remaining sync is the 1-step-lagged async readback event check in replay_adam_and_readback (non-blocking in practice). Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_her.rs | 38 +++++++++++++++++-- .../src/cuda_pipeline/her_episode_kernel.cu | 27 +++++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 10 +---- 3 files changed, 64 insertions(+), 11 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_her.rs b/crates/ml/src/cuda_pipeline/gpu_her.rs index c1d4a47d3..04ec4c8f7 100644 --- a/crates/ml/src/cuda_pipeline/gpu_her.rs +++ b/crates/ml/src/cuda_pipeline/gpu_her.rs @@ -146,6 +146,8 @@ pub struct GpuHer { future_donors_func: CudaFunction, /// GPU kernel: find the last transition in the same episode (HER Final). episode_end_func: CudaFunction, + /// GPU kernel: sample uniformly random donors in [0, batch_size) (HER Random). + random_donors_func: CudaFunction, // Per-sample LCG RNG state for Future strategy donor selection. // Seeded once at construction; survives across relabel_batch_with_strategy calls. @@ -175,7 +177,7 @@ impl GpuHer { let relabel_func = compile_her_kernel(&stream, &config)?; // Compile HER episode boundary kernels (Future + Final strategies) - let (future_donors_func, episode_end_func) = compile_her_episode_kernels(&stream)?; + let (future_donors_func, episode_end_func, random_donors_func) = compile_her_episode_kernels(&stream)?; // Pre-allocate staging buffers let out_states = alloc_f32(&stream, her_batch * config.state_dim, "her_out_states")?; @@ -221,6 +223,7 @@ impl GpuHer { relabel_func, future_donors_func, episode_end_func, + random_donors_func, rng_states, stream, }) @@ -457,6 +460,32 @@ impl GpuHer { Ok(()) } + + /// Generate uniformly random donor indices on GPU — zero CPU involvement. + /// Writes to `self.donor_indices` using the per-sample LCG RNG. + pub fn generate_random_donors_gpu(&mut self, batch_size: usize) -> Result<(), MLError> { + let her_batch_size = self.config.her_batch_size(); + let blocks = ((her_batch_size + 255) / 256) as u32; + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let batch_i32 = batch_size as i32; + let her_i32 = her_batch_size as i32; + + unsafe { + self.stream + .launch_builder(&self.random_donors_func) + .arg(&mut self.donor_indices) + .arg(&mut self.rng_states) + .arg(&batch_i32) + .arg(&her_i32) + .launch(launch_cfg) + .map_err(|e| MLError::ModelError(format!("her_sample_random_donors: {e}")))?; + } + Ok(()) + } } // --------------------------------------------------------------------------- @@ -472,7 +501,7 @@ static HER_RELABEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/her_ fn compile_her_episode_kernels( stream: &Arc, -) -> Result<(CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction), MLError> { let context = stream.context(); let module = context.load_cubin(HER_EPISODE_CUBIN.to_vec()).map_err(|e| { MLError::ModelError(format!("her_episode_kernel module load: {e}")) @@ -483,7 +512,10 @@ fn compile_her_episode_kernels( let episode_end = module.load_function("her_find_episode_end").map_err(|e| { MLError::ModelError(format!("load her_find_episode_end: {e}")) })?; - Ok((future_donors, episode_end)) + let random_donors = module.load_function("her_sample_random_donors").map_err(|e| { + MLError::ModelError(format!("load her_sample_random_donors: {e}")) + })?; + Ok((future_donors, episode_end, random_donors)) } /// Compile the HER relabel kernel with dimension defines. diff --git a/crates/ml/src/cuda_pipeline/her_episode_kernel.cu b/crates/ml/src/cuda_pipeline/her_episode_kernel.cu index 90087cb64..762bcc283 100644 --- a/crates/ml/src/cuda_pipeline/her_episode_kernel.cu +++ b/crates/ml/src/cuda_pipeline/her_episode_kernel.cu @@ -110,3 +110,30 @@ extern "C" __global__ void her_find_episode_end( end_indices[idx] = ep_end; } + +/* ══════════════════════════════════════════════════════════════════════ + * KERNEL 4: her_sample_random_donors + * + * HER Random strategy: generate uniformly random donor indices in + * [0, batch_size) using per-thread LCG. Zero CPU involvement — + * eliminates the last memcpy_htod in the training hot path. + * + * Launch config: grid=(ceil(her_batch_size/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void her_sample_random_donors( + int* __restrict__ donor_indices, /* [her_batch_size] output */ + unsigned int* __restrict__ rng_states, /* [her_batch_size] per-thread LCG state */ + int batch_size, /* range: [0, batch_size) */ + int her_batch_size +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= her_batch_size) return; + + unsigned int rng = rng_states[idx]; + rng = rng * 1664525u + 1013904223u; + rng_states[idx] = rng; + + int donor = (int)((float)(rng >> 8) * 5.9604644775390625e-8f * (float)batch_size); + if (donor >= batch_size) donor = batch_size - 1; + donor_indices[idx] = donor; +} diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 927b63e6f..bc227be4c 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -596,14 +596,8 @@ impl FusedTrainingCtx { // Compute donor indices (GPU-native) match her.config.strategy { HerGpuStrategy::Random => { - // Random donors: generate on CPU (tiny: ~100 ints), upload - use rand::Rng; - let mut rng = rand::thread_rng(); - let donors: Vec = (0..her_batch_size) - .map(|_| rng.gen_range(0..batch_size as i32)) - .collect(); - self.stream.memcpy_htod(&donors, &mut her.donor_indices) - .map_err(|e| anyhow::anyhow!("HER random donors HtoD: {e}"))?; + her.generate_random_donors_gpu(batch_size) + .map_err(|e| anyhow::anyhow!("HER random donors GPU: {e}"))?; } HerGpuStrategy::Future | HerGpuStrategy::Final => { let episode_ids = gpu_batch.episode_ids.as_ref().ok_or_else(|| {