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

fix: reference should be raw dot product (MMA is unscaled)

Browse Source
This commit is contained in:
biondizzle
2026-05-29 18:48:39 +00:00
parent 2b945f255b
commit d5e20b2d42
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_tma_kqk.cu
View File

@@ -164,13 +164,13 @@ int main() {
float* h_out = (float*)malloc(SK * sizeof(float));
cudaMemcpy(h_out, d_out, SK * sizeof(float), cudaMemcpyDeviceToHost);
float scale = 1.0f / sqrtf((float)HD);
// Reference: raw dot product (MMA is unscaled, scale applied during softmax)
int fail = 0; float max_rel = 0;
for (int j = 0; j < SK; j++) {
float dot = 0;
for (int d = 0; d < HD; d++)
dot += bf16_to_f32_host(h_q[d]) * bf16_to_f32_host(h_k[j * HD + d]);
float ref = dot * scale;
float ref = dot; // NO scale — MMA output is unscaled
float got = h_out[j];
float rel = fabsf(ref) > 1e-4f ? fabsf(got - ref) / fabsf(ref) : fabsf(got - ref);
if (rel > max_rel) max_rel = rel;