determinism(atom_stats): 2-phase reduction — final training-path atomicAdd removed

compute_expected_q had 2 atomicAdds accumulating per-block entropy + utilization
into atom_stats[0..1]. These feed HEALTH_DIAG's atoms=X%ent/Y%util field AND
the adaptive_atom_positions kernel (which reshapes C51 atom positions based on
utilization) — so non-determinism here propagated into training dynamics.

Replaced with standard 2-phase deterministic reduction:

  Phase 1 (inside compute_expected_q):
    Per-block warp→block shared-mem reduce unchanged, but the final write is
    atom_stats_block_sums[blockIdx.x * 2 + i] = value (unique slot per block,
    no atomic). Kernel signature: `atom_stats` param renamed to
    `atom_stats_block_sums` to reflect new semantics.

  Phase 2 (new kernel atom_stats_finalize):
    Sequentially sums block_sums[0..num_blocks*2] into atom_stats[0..1].
    Launch grid=(1,1,1) block=(1,1,1) — bit-stable across runs.

Rust glue:
  - New atom_stats_block_sums_buf [max_blocks * 2] allocated at trainer init.
    max_blocks = ceil(batch_size / 256); smoke uses 1, production 64.
  - New atom_stats_finalize_kernel loaded from experience_kernels cubin.
  - populate_q_out() now launches phase 1 + phase 2 sequentially.
  - Other launch sites (replay_forward_for_q_values, compute_denoise_target_q)
    already passed NULL atom_stats and don't need changes — the kernel skips
    accumulation on NULL.

This was the FINAL training-path atomicAdd call site. Post-commit, the entire
crates/ml/src/cuda_pipeline/ has ZERO atomicAdd call sites — full CUDA-level
determinism for the training path (cuBLAS algorithm selection remains the
main remaining nondeterminism source, mitigated via CUBLAS_WORKSPACE_CONFIG
for tests).

Smoke verification (post-commit, 5-trial multi-trial):
  q_gaps:       2.24 / 0.98 / 2.75 / 5.58 / 1.08 (all > 0.02, 5/5 pass)
  Best Sharpe:  19.41 / 38.71 / 25.33 / 35.94 / 24.22 (median 25.3)
  mean_degradation = -10.66 (NEGATIVE means sharpe_ema IMPROVED over epochs)
  All assertions pass.

Session atomicAdd tally:
  69 string occurrences → 13 real call sites (most "hits" were comments).
  Today removed: 4 (CQL barrier+IB) + 2 (monitoring_reduce) + 2 (trade_stats
  + dqn_utility) + deleted-kernel (ensemble_diversity's atomic + G6/G10's
  atomics removed with scaffolding) + 2 (atom_stats, this commit) = 13 total.
  Zero remaining.

Files touched:
  crates/ml/src/cuda_pipeline/experience_kernels.cu  (+38 / -9 net)
  crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs     (+40 / -10 net)

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

crates/ml/src/cuda_pipeline/experience_kernels.cu

@@ -2332,7 +2332,8 @@ extern "C" __global__ void compute_expected_q(
     int b2_size,
     int b3_size,
     const float* __restrict__ per_sample_support, /* [N*3]: per-sample [v_min, v_max, delta_z] from IQL */
-    float* __restrict__ atom_stats,   /* [2]: [0]=sum_entropy, [1]=sum_utilized. NULL to skip. */
+    float* __restrict__ atom_stats_block_sums, /* [num_blocks * 2]: per-block partial sums (phase 1).
+                                                 * NULL to skip atom-stat collection entirely. */
     float* __restrict__ q_variance,   /* [N, total_actions] Var[Q] per action. NULL to skip. */
     const float* __restrict__ atom_positions  /* [4, num_atoms] adaptive positions. NULL = use linear. */
 ) {
@@ -2412,15 +2413,23 @@ extern "C" __global__ void compute_expected_q(
         branch_logit_offset += N * n_d * num_atoms;
     }
 
-    /* Accumulate atom stats across all samples (monitoring only).
-     * Uses warp+block reduction to minimize atomicAdd non-determinism.
-     * One atomicAdd per BLOCK (not per thread). */
-    if (atom_stats != NULL) {
+    /* Atom-stat accumulation across samples (monitoring + adaptive_atom_positions input).
+     *
+     * 2-phase deterministic reduction (replaces cross-block atomicAdd):
+     *   Phase 1 (here): warp reduce → block reduce via shared mem → write
+     *     one (entropy, util) pair per block to atom_stats_block_sums.
+     *   Phase 2: `atom_stats_finalize` kernel sequentially sums block_sums
+     *     into atom_stats[0..1]. Bit-stable across runs.
+     *
+     * atom_stats_block_sums layout: [num_blocks * 2], block b writes
+     *   [b*2+0] = block's entropy sum, [b*2+1] = block's util sum.
+     */
+    if (atom_stats_block_sums != NULL) {
         float inv_actions = 1.0f / (float)total_actions;
         float my_entropy = sample_entropy * inv_actions;
         float my_util    = (float)sample_utilized * inv_actions;
 
-        /* Warp-level reduction */
+        /* Warp-level reduction (shfl — deterministic across lanes) */
         for (int offset = 16; offset > 0; offset >>= 1) {
             my_entropy += __shfl_xor_sync(0xFFFFFFFF, my_entropy, offset);
             my_util    += __shfl_xor_sync(0xFFFFFFFF, my_util,    offset);
@@ -2441,13 +2450,35 @@ extern "C" __global__ void compute_expected_q(
                 vu += __shfl_xor_sync(0xFFFFFFFF, vu, off);
             }
             if (lane == 0) {
-                atomicAdd(&atom_stats[0], ve);
-                atomicAdd(&atom_stats[1], vu);
+                /* One write per block to unique scratch slot — no atomic. */
+                atom_stats_block_sums[blockIdx.x * 2 + 0] = ve;
+                atom_stats_block_sums[blockIdx.x * 2 + 1] = vu;
             }
         }
     }
 }
 
+/* ══════════════════════════════════════════════════════════════════════
+ * atom_stats_finalize — phase 2 of the deterministic reduction.
+ *
+ * Sequentially sums atom_stats_block_sums[0..num_blocks*2] into
+ * atom_stats[0..1]. Launch grid=(1,1,1) block=(1,1,1); deterministic order.
+ * ══════════════════════════════════════════════════════════════════════ */
+extern "C" __global__ void atom_stats_finalize(
+    const float* __restrict__ atom_stats_block_sums,  /* [num_blocks * 2] from phase 1 */
+    float* __restrict__ atom_stats,                    /* [2] final (entropy_sum, util_sum) */
+    int num_blocks
+) {
+    float sum_entropy = 0.0f;
+    float sum_util    = 0.0f;
+    for (int b = 0; b < num_blocks; b++) {
+        sum_entropy += atom_stats_block_sums[b * 2 + 0];
+        sum_util    += atom_stats_block_sums[b * 2 + 1];
+    }
+    atom_stats[0] = sum_entropy;
+    atom_stats[1] = sum_util;
+}
+
 /* ================================================================== */
 /* Kernel 5: expert_action_override                                    */
 /* ================================================================== */

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -1310,8 +1310,15 @@ pub struct GpuDqnTrainer {
     q_stats_kernel: CudaFunction,
     /// GPU buffer for Q-value statistics [7 floats]
     q_stats_buf: CudaSlice<f32>,
-    /// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized]
+    /// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized].
+    /// Populated by `atom_stats_finalize` kernel (phase 2 of the deterministic reduction).
     atom_stats_buf: CudaSlice<f32>,
+    /// Per-block partial sums for `compute_expected_q`'s atom-stat reduction.
+    /// Phase 1 output: [max_blocks * 2] where block b writes [b*2+0]=entropy, [b*2+1]=util.
+    /// Phase 2 (`atom_stats_finalize`) reduces these deterministically into atom_stats_buf.
+    atom_stats_block_sums_buf: CudaSlice<f32>,
+    /// Phase-2 finalize kernel for atom_stats — sequential reduction, bit-stable.
+    atom_stats_finalize_kernel: CudaFunction,
     /// Pinned device-mapped readback for q_stats [7] + q_out sample 0 [total_actions] floats.
     /// GPU writes via q_readback_dev_ptr, CPU reads via q_readback_pinned — zero sync.
     q_readback_dev_ptr: u64,
@@ -5660,6 +5667,13 @@ impl GpuDqnTrainer {
             .map_err(|e| MLError::ModelError(format!("alloc q_stats_f32: {e}")))?;
         let atom_stats_buf = stream.alloc_zeros::<f32>(2)
             .map_err(|e| MLError::ModelError(format!("alloc atom_stats: {e}")))?;
+        // Per-block partial sums for the atom-stat deterministic reduction.
+        // Sized to max possible num_blocks for compute_expected_q (one thread
+        // per sample, block_dim=256). For smoke (B=16) → 1 block; for
+        // production (B=16384) → 64 blocks. Overallocate to a fixed cap.
+        let atom_stats_max_blocks: usize = ((config.batch_size + 255) / 256).max(1);
+        let atom_stats_block_sums_buf = stream.alloc_zeros::<f32>(atom_stats_max_blocks * 2)
+            .map_err(|e| MLError::ModelError(format!("alloc atom_stats_block_sums: {e}")))?;
         // q_readback — pinned host buffer for q_stats[7] + q_out sample 0 [total_actions]
         let total_actions = config.branch_0_size + config.branch_1_size + config.branch_2_size + config.branch_3_size;
         let q_readback_size = 7 + total_actions;  // 7 stats + 13 Q-values = 20
@@ -5693,6 +5707,8 @@ impl GpuDqnTrainer {
             .map_err(|e| MLError::ModelError(format!("branch_confidence_routing load: {e}")))?;
         let adaptive_atom_kernel = exp_module_for_mag.load_function("adaptive_atom_positions")
             .map_err(|e| MLError::ModelError(format!("adaptive_atom_positions load: {e}")))?;
+        let atom_stats_finalize_kernel = exp_module_for_mag.load_function("atom_stats_finalize")
+            .map_err(|e| MLError::ModelError(format!("atom_stats_finalize load: {e}")))?;
         let atom_position_grad_kernel = exp_module_for_mag.load_function("atom_position_gradient")
             .map_err(|e| MLError::ModelError(format!("atom_position_gradient load: {e}")))?;
         let q_anchor_kernel = exp_module_for_mag.load_function("q_anchor_to_flat")
@@ -7414,6 +7430,8 @@ impl GpuDqnTrainer {
             q_stats_kernel,
             q_stats_buf,
             atom_stats_buf,
+            atom_stats_block_sums_buf,
+            atom_stats_finalize_kernel,
             q_readback_pinned,
             q_readback_dev_ptr,
             cql_logit_grad_kernel,
@@ -9225,12 +9243,10 @@ impl GpuDqnTrainer {
         let grid_dim = ((batch_size as u32 + block_dim - 1) / block_dim).max(1);
         let q_out_ptr = self.q_out_buf.raw_ptr();
         let atom_stats_ptr = self.atom_stats_buf.raw_ptr();
+        let atom_stats_block_sums_ptr = self.atom_stats_block_sums_buf.raw_ptr();
         let q_var_ptr = self.q_var_buf_trainer.raw_ptr();
         let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr();
-        // Zero atom_stats before accumulation (kernel uses atomicAdd)
-        unsafe {
-            cudarc::driver::sys::cuMemsetD32Async(atom_stats_ptr, 0, 2, self.stream.cu_stream());
-        }
+        // Phase 1: per-block sums (deterministic, no atomic) → block_sums
         unsafe {
             self.stream
                 .launch_builder(&self.expected_q_kernel)
@@ -9244,7 +9260,7 @@ impl GpuDqnTrainer {
                 .arg(&b2)
                 .arg(&b3)
                 .arg(&self.per_sample_support_ptr)
-                .arg(&atom_stats_ptr)
+                .arg(&atom_stats_block_sums_ptr)
                 .arg(&q_var_ptr)
                 .arg(&atom_positions_buf_ptr)
                 .launch(LaunchConfig {
@@ -9254,6 +9270,21 @@ impl GpuDqnTrainer {
                 })
                 .map_err(|e| MLError::ModelError(format!("populate_q_out expected_q: {e}")))?;
         }
+        // Phase 2: sequential reduce block_sums → atom_stats (deterministic)
+        let num_blocks_i32 = grid_dim as i32;
+        unsafe {
+            self.stream
+                .launch_builder(&self.atom_stats_finalize_kernel)
+                .arg(&atom_stats_block_sums_ptr)
+                .arg(&atom_stats_ptr)
+                .arg(&num_blocks_i32)
+                .launch(LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (1, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("populate_q_out atom_stats_finalize: {e}")))?;
+        }
         Ok(&self.q_out_buf)
     }