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

test: debug single K-tile with offset descriptor

Browse Source
This commit is contained in:
biondizzle
2026-05-28 12:51:33 +00:00
parent 435ca037cf
commit b703dc0a50
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/unit/test_umma_qk_hd64.cu
View File

@@ -76,7 +76,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
uint32_t sK_smem = __cvta_generic_to_shared(sK);
uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN);
for (int kt = 0; kt < NKT; kt++) {
for (int kt = 0; kt < 1; kt++) { // DEBUG: only first K-tile
// Offset to the kt-th K-tile: kt * BLOCK_MN * 32 bytes
// (each K-tile is 16 BF16 = 2 core columns, BLOCK_MN * 16B apart, total 32B per row)
// Wait, that's not right. In canonical layout, the kt-th 16-column slice starts
@@ -89,7 +89,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
uint64_t dk = make_umma_desc_kmajor_none(k_addr, BLOCK_MN);
if (tid == 0) {
umma_ss_f16(tb, dq, dk, idesc, kt > 0); // first tile: no accumulate
umma_ss_f16(tb, dq, dk, idesc, false); // no accumulate, single tile
}
// Fence after each K-tile MMA to ensure TMEM is updated
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
@@ -125,7 +125,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
if (tid == 0) {
for (int j = 0; j < SK; j++) {
float dot = 0.0f;
for (int d = 0; d < HD; d++)
for (int d = 0; d < 16; d++) // DEBUG: only first K-tile
dot += bf16_to_f32(q[d]) * bf16_to_f32(k[j * HD + d]);
s_scalar[j] = dot * scale;
}