biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

P5: fix single-tile backward compat (normalized P for n_kv_tiles==1)

Browse Source
This commit is contained in:
biondizzle
2026-05-30 08:47:47 +00:00
parent 2649488d13
commit 0f34f60494
1 changed files with 23 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
dsv4/kernels/attention/fmha_6warp_multihead.cuh
View File

@@ -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<SK_TILE;j++) s_p_vals[j] = s_vals[j];
// Store P for PV. For single KV tile, use NORMALIZED P (same as old kernel).
// For multi-tile, use UN-NORMALIZED P (critical for FA2 rescale).
// Single-tile: O = P_norm × V. Multi-tile: O = Σ(P_unnorm × V) / running_sum
if (lane == 0) {
if (n_kv_tiles == 1) {
// Normalized P (backward compatible with single-tile)
for (int j=0;j<SK_TILE;j++) s_vals[j] /= row_sum;
}
for (int j=0;j<SK_TILE;j++) s_p_vals[j] = s_vals[j];
}
// Update running sum
if (lane == 0) *sRowSum = old_sum + row_sum;
@@ -303,16 +310,25 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
} // end KV tile loop
// ================================================================
// Epilogue: normalize sOacc by running_sum, write to GMEM
// Epilogue: write O to GMEM
// For single KV tile (normalized P): sOacc = P_norm × V (already normalized)
// For multi KV tile (un-normalized P): O = sOacc / running_sum
// ================================================================
if (wid == 0) {
float running_max = *sRowMax;
float running_sum = *sRowSum;
if (lane == 0) {
float inv_sum = 1.0f / running_sum;
for (int d = 0; d < HD; d++) {
o_head[d] = f32_to_bf16(sOacc[d] * inv_sum);
if (n_kv_tiles > 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;