fix: reward_rank_normalize cross-block race — two-buffer read/write split

In-place ranking races across blocks: block X writes shaped rewards[i]
before block Y loads it in tile scan. __syncthreads is block-local.
Fix: read from rewards_out (raw), write to gpu_batch.rewards (clone).
The DtoD clone from collect_experiences_gpu provides the output buffer.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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<f32>, 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<f32>, 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 {

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