diff --git a/dsv4/kernels/attention/fmha_6warp_multihead.cuh b/dsv4/kernels/attention/fmha_6warp_multihead.cuh index 8d5094c7..9d19e4b8 100644 --- a/dsv4/kernels/attention/fmha_6warp_multihead.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multihead.cuh @@ -128,11 +128,11 @@ fmha_6warp_multihead_kernel(FmhaParams params) { float* s_p_vals = (float*)(sV + V_SUB_SZ); float* sOacc = (float*)(((uintptr_t)(s_p_vals + SK_TILE) + 15) & ~(uintptr_t)15); - // Initialize accumulator - if (is_softmax_warp && lane == 0) { + // Initialize accumulator (single thread to avoid race) + if (tid == 0) { *sRowMax = -INFINITY; *sRowSum = 0.0f; - for (int d = lane; d < HD; d += 32) sOacc[d] = 0.0f; + for (int d = 0; d < HD; d++) sOacc[d] = 0.0f; } // TMEM allocation