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

test: HD=64 all-ones, expected S[0,j]=64 (unscaled) or 8.0 scaled

Browse Source
This commit is contained in:
biondizzle
2026-05-28 11:44:31 +00:00
parent 73f9ff98c9
commit abe1870429
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
tests/unit/test_umma_qk_hd64.cu
View File

@@ -130,6 +130,14 @@ test_umma_qk_hd64(const bf16_t* q, const bf16_t* k,
bool accumulate = (kt > 0);
// Debug: write descriptor info
if (tid == 0 && kt == 0) {
s_out[240] = (float)q_ktile_addr;
s_out[241] = (float)k_ktile_addr;
s_out[242] = (float)desc_q;
s_out[243] = (float)desc_k;
}
// 4 warp leaders call MMA
if (lane == 0) {
umma_ss_f16(tb, desc_q, desc_k, idesc, accumulate);
@@ -186,9 +194,8 @@ int main() {
float* h_s_out = (float*)calloc(128 * 16, sizeof(float));
float* h_s_scalar = (float*)calloc(SK, sizeof(float));
srand(42);
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);
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);
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));