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

Fix: remove ×2 QK scale correction (MMA scale is 1.0, not 0.5)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 15:09:57 +00:00
parent 6ea7356fdd
commit 40c83c769a
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/unit/test_fmha_hd64_smem_p.cu
View File

@@ -114,8 +114,8 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
: "r"(tb + n*8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) {
s_vals[n*8+c] = tmp[c] * scale * 2.0f; // ×2 to undo QK 0.5 scale
row_max = fmaxf(row_max, tmp[c] * scale * 2.0f);
s_vals[n*8+c] = tmp[c] * scale; // No ×2 correction needed for QK MMA
row_max = fmaxf(row_max, tmp[c] * scale);
}
}
row_max = wmax(row_max);