gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: shared memory race in block_reduce (shmem_reduce read/write hazard)

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-04 20:35:20 +02:00
parent ca7201a1ec
commit 137480c887
2 changed files with 17 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

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

11
crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
View File

@@ -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) */