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

perf: GPU-native HER random donors — eliminate last CPU touch in hot path

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-02 09:24:36 +02:00
parent 4c7bbf3a12
commit e4f838d4e3
3 changed files with 64 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
crates/ml/src/cuda_pipeline/gpu_her.rs
View File

@@ -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<CudaStream>,
) -> 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.

27
crates/ml/src/cuda_pipeline/her_episode_kernel.cu
View File

@@ -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;
}

10
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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<i32> = (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(|| {