From 137480c887b8fa593db18a8a3e45255a33a1166d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 20:35:20 +0200 Subject: [PATCH] fix: shared memory race in block_reduce (shmem_reduce read/write hazard) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compute-sanitizer --tool racecheck found 4000+ hazards per kernel invocation in both c51_loss_batched and mse_loss_batched. The race is between consecutive block_reduce calls sharing the same shmem_reduce buffer: 1. reduce_max returns shmem_reduce[0] → next reduce_sum writes shmem_reduce[warp] 2. reduce_sum returns shmem_reduce[0] → next reduce_max writes shmem_reduce[warp] Fixed by adding __syncthreads() after reading the return value in both block_reduce_max_f and block_reduce_sum_f (read-then-sync pattern). Also added explicit barriers between spectral decoupling reduces and subsequent softmax/log_softmax calls. This race caused non-deterministic NaN in loss values. The effect was masked by GPU printf statements (which acted as implicit memory barriers) and by CUDA Graph replay timing differences. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/c51_loss_kernel.cu | 10 ++++++++-- crates/ml/src/cuda_pipeline/mse_loss_kernel.cu | 11 +++++++++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index d485a2836..30c93bcbd 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -42,7 +42,9 @@ __device__ __forceinline__ float block_reduce_max_f( if (tid == 0) shmem_reduce[0] = result; } __syncthreads(); - return shmem_reduce[0]; + float ret = shmem_reduce[0]; + __syncthreads(); + return ret; } __device__ __forceinline__ float block_reduce_sum_f( @@ -63,7 +65,9 @@ __device__ __forceinline__ float block_reduce_sum_f( if (tid == 0) shmem_reduce[0] = result; } __syncthreads(); - return shmem_reduce[0]; + float ret = shmem_reduce[0]; + __syncthreads(); + return ret; } /* log_softmax in float — reads/writes float shmem, no BF16 overflow */ @@ -77,6 +81,7 @@ __device__ void block_log_softmax_f( for (int i = tid; i < num_atoms; i += BLOCK_THREADS) local_max = fmaxf(local_max, logits[i]); float bmax = block_reduce_max_f(local_max, shmem_reduce, tid); + __syncthreads(); /* barrier between consecutive reduce calls sharing shmem_reduce */ float local_sum = 0.0f; for (int i = tid; i < num_atoms; i += BLOCK_THREADS) @@ -349,6 +354,7 @@ extern "C" __global__ void c51_loss_batched( } float branch_sq = block_reduce_sum_f(local_sq, shmem_reduce, tid); total_logit_sq += branch_sq; + __syncthreads(); /* barrier before next reduce in log_softmax */ } block_log_softmax_f(shmem_lp, shmem_lp, shmem_reduce, tid, num_atoms); diff --git a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu index 2351cdead..e3711a61d 100644 --- a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu @@ -34,7 +34,10 @@ __device__ __forceinline__ float block_reduce_max_f( if (tid == 0) shmem_reduce[0] = result; } __syncthreads(); - return shmem_reduce[0]; + /* Read result THEN sync — prevents next reduce's write from racing with this read */ + float ret = shmem_reduce[0]; + __syncthreads(); + return ret; } __device__ __forceinline__ float block_reduce_sum_f( @@ -55,7 +58,9 @@ __device__ __forceinline__ float block_reduce_sum_f( if (tid == 0) shmem_reduce[0] = result; } __syncthreads(); - return shmem_reduce[0]; + float ret = shmem_reduce[0]; + __syncthreads(); + return ret; } /** @@ -76,6 +81,7 @@ __device__ float block_softmax_expected_q_f( for (int i = tid; i < num_atoms; i += BLOCK_THREADS) local_max = fmaxf(local_max, shmem_logits[i]); float bmax = block_reduce_max_f(local_max, shmem_reduce, tid); + __syncthreads(); /* barrier between consecutive reduce calls sharing shmem_reduce */ /* Pass 2: sum(exp(x - max)) — float exp, no BF16 overflow */ float local_sum = 0.0f; @@ -296,6 +302,7 @@ extern "C" __global__ void mse_loss_batched( } float branch_sq = block_reduce_sum_f(local_sq, shmem_reduce, tid); total_logit_sq += branch_sq; + __syncthreads(); /* barrier before next reduce in softmax */ } /* softmax + E[Q], saving probs to save_current_lp (BF16 output) */