P5: fix sOacc init race — use single thread (tid==0) instead of 4 softmax warps

This commit is contained in:
2026-05-30 08:53:50 +00:00
parent 66b126ded8
commit 5f4856d771

View File

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