diff --git a/dsv4/kernels/attention/fmha_6warp_multihead.cuh b/dsv4/kernels/attention/fmha_6warp_multihead.cuh index cdc46272..8d5094c7 100644 --- a/dsv4/kernels/attention/fmha_6warp_multihead.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multihead.cuh @@ -229,9 +229,16 @@ fmha_6warp_multihead_kernel(FmhaParams params) { *sRowMax = new_max; } - // Store UN-NORMALIZED P for PV (critical for multi-tile!) - // P_unnorm[j] = exp(s[j] - max) (NOT divided by row_sum) - if (lane == 0) for (int j=0;j 1) { + float inv_sum = 1.0f / running_sum; + for (int d = 0; d < HD; d++) { + o_head[d] = f32_to_bf16(sOacc[d] * inv_sum); + } + } else { + // Single tile: P was normalized, sOacc is already the correct output + for (int d = 0; d < HD; d++) { + o_head[d] = f32_to_bf16(sOacc[d]); + } } if (lse_head) { lse_head[0] = logf(running_sum) + running_max;