From 22f30fe37fbd637b856f0cc3a89928080b945f60 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 6 Apr 2026 12:18:40 +0200 Subject: [PATCH] fix: atomicExch for seg_tree leaf writes prevents duplicate-index race MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PER samples with replacement, so a batch can contain duplicate indices. With non-atomic leaf writes, two threads processing the same index both read the same old_pa, but only one write survives — internal nodes get double the delta while the leaf changed once, permanently corrupting the tree sum. atomicExch returns the exact replaced value so each thread propagates the correct delta regardless of concurrent duplicate writes. Also fix readback_pinned layout doc: slot 9 is used for diversity loss. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/seg_tree_kernel.cu | 6 ++---- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 3 ++- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/crates/ml-dqn/src/seg_tree_kernel.cu b/crates/ml-dqn/src/seg_tree_kernel.cu index d49a3789d..8649fb3a5 100644 --- a/crates/ml-dqn/src/seg_tree_kernel.cu +++ b/crates/ml-dqn/src/seg_tree_kernel.cu @@ -56,8 +56,7 @@ extern "C" __global__ void seg_tree_update( // ancestor. atomicAdd is commutative+associative, so concurrent threads // produce correct sums without synchronization barriers. int leaf = capacity + (int)idx; - float old_pa = tree[leaf]; - tree[leaf] = pa; + float old_pa = atomicExch(&tree[leaf], pa); float delta = pa - old_pa; int node = leaf >> 1; @@ -90,8 +89,7 @@ extern "C" __global__ void seg_tree_insert( // Write leaf and propagate delta to root via atomicAdd (race-free). int leaf = capacity + (int)idx; - float old_pa = tree[leaf]; - tree[leaf] = pa; + float old_pa = atomicExch(&tree[leaf], pa); float delta = pa - old_pa; int node = leaf >> 1; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 454220567..a93d57c87 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -636,7 +636,8 @@ pub struct GpuDqnTrainer { /// [0]=loss, [1]=mse_loss, [2]=grad_norm_sq (per-step, from graph_adam) /// [3]=avg_max_q, [4]=q_min, [5]=q_max, [6]=q_mean, [7]=q_var (every 50 steps) /// [8]=causal_mean_sens (every N steps) - /// [9..16]=reserved (nan_flags, future use) + /// [9]=ensemble_diversity_loss (epoch boundary) + /// [10..16]=reserved (future use) /// Pinned memory enables true async cuMemcpyDtoHAsync without CPU blocking. readback_pinned: *mut f32, /// Whether there's an in-flight readback to collect