CUDA: Fix data-races when reusing SMEM in block_reduce (#26385)

* CUDA: Fix data-races when reusing block_reduce

block_reduce currently doesn't resync after reading from SMEM, causing
potential data-races when reusing SMEM for multiple reductions.

One may consider simply always adding this in block_reduce, but this
comes at a potential perf cost

* double-buffering for single-row softmax

* double-buffering for norm as well

* Add comment

* Add explanatory comment to block_reduce

* Specify need for + do memory barrier only in multi-warp scenario

* Implement review-suggestion from @gaugarg-nv
This commit is contained in:
Oliver Simons
2026-08-03 14:22:44 +02:00
committed by GitHub
parent 0b14b87d7c
commit 9bd4c09ea5
3 changed files with 18 additions and 9 deletions
+2 -1
View File
@@ -627,7 +627,8 @@ template <typename T> struct block_reduce_policy<block_reduce_method::MAX, T> {
};
template <block_reduce_method reduce_method_t, const unsigned int block_size_template = 0, typename T>
static __device__ T block_reduce(T val, T * shared_vals) {
static __device__ T block_reduce(T val, [[maybe_unused]] T * shared_vals) {
// for multi-warp reductions, callers must not reuse shared_vals until all reads from this invocation have completed
val = block_reduce_policy<reduce_method_t, T>::reduce(val);
const unsigned int block_size = block_size_template == 0 ? blockDim.x : block_size_template;
if (block_size > WARP_SIZE) {