P5: fix sOacc init race — use single thread (tid==0) instead of 4 softmax warps
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user