From ae65740a3db917e4650b8b4329f42201ef7c281d Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Mon, 6 Apr 2026 11:15:17 +0200
Subject: [PATCH] fix: race-free segment tree propagation via atomicAdd deltas

seg_tree_update and seg_tree_insert used non-atomic
tree[node] = tree[2*node] + tree[2*node+1] with 8192 threads racing
on shared internal nodes. On H100 (132 SMs), all threads execute
simultaneously causing data races and hangs.

Replace with atomicAdd delta propagation: each thread computes
delta = new_leaf - old_leaf, then atomicAdd to every ancestor.
Commutative + associative = race-free. O(log n) per thread,
hardware-accelerated on SM 9.0.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml-dqn/src/seg_tree_kernel.cu | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/crates/ml-dqn/src/seg_tree_kernel.cu b/crates/ml-dqn/src/seg_tree_kernel.cu
index 7e6274c59..d49a3789d 100644
--- a/crates/ml-dqn/src/seg_tree_kernel.cu
+++ b/crates/ml-dqn/src/seg_tree_kernel.cu
@@ -7,10 +7,11 @@
 // Three kernels:
 //   1. seg_tree_update    — priority update: compute (|td|^alpha + eps),
 //                           write to priorities[], compute priority^alpha,
-//                           write to tree leaf, propagate sums to root.
+//                           write to tree leaf, propagate DELTA to root
+//                           via atomicAdd (race-free, O(log n) per thread).
 //   2. seg_tree_insert    — insert: take raw priorities (max_priority fill),
 //                           compute priority^alpha, write to tree leaf,
-//                           propagate sums to root.
+//                           propagate DELTA to root via atomicAdd.
 //   3. seg_tree_sample    — proportional sampling: parallel root-to-leaf
 //                           traversal with Philox RNG. Output i64 indices.
 
@@ -50,13 +51,18 @@ extern "C" __global__ void seg_tree_update(
     // Compute priority^alpha for tree leaf (same as pow_alpha_f32)
     float pa = powf(new_prio, alpha);
 
-    // Write leaf and propagate sums to root
+    // Write leaf and propagate delta to root via atomicAdd.
+    // Each thread computes delta = new_leaf - old_leaf and adds it to every
+    // 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 delta = pa - old_pa;
 
     int node = leaf >> 1;
     while (node >= 1) {
-        tree[node] = tree[2 * node] + tree[2 * node + 1];
+        atomicAdd(&tree[node], delta);
         node >>= 1;
     }
 }
@@ -82,13 +88,15 @@ extern "C" __global__ void seg_tree_insert(
     // Compute priority^alpha for tree leaf
     float pa = powf(priorities[i], alpha);
 
-    // Write leaf and propagate sums to root
+    // 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 delta = pa - old_pa;
 
     int node = leaf >> 1;
     while (node >= 1) {
-        tree[node] = tree[2 * node] + tree[2 * node + 1];
+        atomicAdd(&tree[node], delta);
         node >>= 1;
     }
 }