diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 5517b9b1e..846b0e56b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -2363,16 +2363,19 @@ impl GpuExperienceCollector { &self.rewards_out } - /// Rank-normalize rewards in-place: r_shaped = sign(r) * rank(|r|) * std_ema. - /// Call after experience collection, before replay buffer insertion. - pub fn shape_rewards(&self, rewards: &mut CudaSlice, n: usize, reward_std: f32) -> Result<(), MLError> { + /// Rank-normalize rewards: r_shaped = sign(r) * rank(|r|) * std_ema. + /// Uses rewards_out as read-only source, writes shaped values to `rewards_dst`. + /// Two-buffer design avoids cross-block race condition (in-place would race + /// because __syncthreads is block-local but ranking reads ALL rewards). + pub fn shape_rewards(&self, rewards_dst: &mut CudaSlice, n: usize, reward_std: f32) -> Result<(), MLError> { if n == 0 { return Ok(()); } let n_i32 = n as i32; let blocks = ((n + 255) / 256) as u32; unsafe { self.stream .launch_builder(&self.reward_rank_kernel) - .arg(rewards) + .arg(&self.rewards_out) // read-only raw rewards (source) + .arg(rewards_dst) // output shaped rewards (separate buffer) .arg(&n_i32) .arg(&reward_std) .launch(LaunchConfig { diff --git a/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu b/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu index 07a14ea41..0ce5a7dbb 100644 --- a/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu +++ b/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu @@ -5,8 +5,10 @@ * * rank(|r[i]|) = count(|r[j]| <= |r[i]| for all j in batch) / N * - * Naive counting-sort does N global reads per thread (256M at B=16384). - * Tiled version: load TILE_SIZE rewards into shared memory, count within + * Two-buffer design: reads from `rewards_in`, writes to `rewards_out`. + * In-place would race across blocks (__syncthreads is block-local). + * + * Tiled: load TILE_SIZE rewards into shared memory, count within * tile, sum counts across tiles. Reduces global memory reads by TILE_SIZE×. * * Grid: ceil(N/256), Block: 256. One thread per reward. @@ -16,14 +18,15 @@ #define TILE_SIZE 256 extern "C" __global__ void reward_rank_normalize( - float* __restrict__ rewards, /* [N] in-place */ + const float* __restrict__ rewards_in, /* [N] read-only raw rewards */ + float* __restrict__ rewards_out, /* [N] output shaped rewards */ int N, - float std_ema /* running std of raw rewards */ + float std_ema /* running std of raw rewards */ ) { int i = blockIdx.x * blockDim.x + threadIdx.x; /* Each thread loads its own |reward| once from global memory */ - float r_i = (i < N) ? rewards[i] : 0.0f; + float r_i = (i < N) ? rewards_in[i] : 0.0f; float abs_i = fabsf(r_i); __shared__ float tile[TILE_SIZE]; @@ -34,7 +37,7 @@ extern "C" __global__ void reward_rank_normalize( for (int t = 0; t < num_tiles; t++) { /* Cooperative tile load: each thread loads one element */ int tile_idx = t * TILE_SIZE + threadIdx.x; - tile[threadIdx.x] = (tile_idx < N) ? fabsf(rewards[tile_idx]) : 1e30f; + tile[threadIdx.x] = (tile_idx < N) ? fabsf(rewards_in[tile_idx]) : 1e30f; __syncthreads(); /* Count within tile — shared memory scan, no global reads */ @@ -52,5 +55,5 @@ extern "C" __global__ void reward_rank_normalize( /* Rank in [0, 1], sign preserved, scaled by std_ema */ float rank = (float)count / (float)N; float sign = (r_i > 0.0f) ? 1.0f : (r_i < 0.0f) ? -1.0f : 0.0f; - rewards[i] = sign * rank * fmaxf(std_ema, 1e-6f); + rewards_out[i] = sign * rank * fmaxf(std_ema, 1e-6f); }