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

test: apply 1/sqrt(HD) scale to MMA output — 4x was the scale factor, not a bug!

Browse Source
This commit is contained in:
biondizzle
2026-05-28 11:34:45 +00:00
parent 013f370046
commit 3c7d9d9303
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/unit/test_umma_qk.cu
View File

@@ -86,7 +86,7 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k,
if (lane == 0 && n < 1) { // Only first 8 cols for debug
for (int c = 0; c < 8; c++) {
if (out_row < 128) {
s_out[out_row * 8 + c] = tmp[c];
s_out[out_row * 8 + c] = tmp[c] * scale; // Apply 1/sqrt(HD) scale
}
}
}
@@ -118,11 +118,11 @@ int main() {
float* h_s_out = (float*)calloc(128*8, sizeof(float));
float* h_s_scalar = (float*)calloc(SK, sizeof(float));
// Test with ALL-ONES data: Q[0,d]=1.0, K[i,d]=1.0
// Expected: S[0,j] = sum(1*1, d=0..15) = 16.0 for all j
// Test with random data — MMA output is UNSCALED (no 1/sqrt(HD))
// We apply the scale factor in the read
srand(42);
for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16_host(1.0f);
for (int i = 0; i < SK*HD; i++) h_k[i] = f32_to_bf16_host(1.0f);
for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
for (int i = 0; i < SK*HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
bf16_t *d_q, *d_k; float *d_s_out, *d_s_scalar;
cudaMalloc(&d_q, HD*sizeof(bf16_t)); cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));