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

test: debug with (128,16) SMEM matching HD=16 exactly

Browse Source
This commit is contained in:
biondizzle
2026-05-28 12:52:19 +00:00
parent b703dc0a50
commit db4f661843
1 changed files with 8 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
tests/unit/test_umma_qk_hd64.cu
View File

@@ -49,10 +49,9 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
bf16_t* sQ = (bf16_t*)(sbuf + 64); // 64-byte aligned
bf16_t* sK = sQ + BLOCK_MN * HD; // follows Q
// Load Q (1, HD) into (128, HD) canonical layout — only row 0 has data
write_q_to_smem<HD>(sQ, q);
// Load K (SK, HD) into (128, HD) canonical layout
write_k_to_smem<SK, HD>(sK, k);
// Load Q (1, 16) into (128, 16) canonical — only first K-tile for debug
write_q_to_smem<16>(sQ, q);
write_k_to_smem<SK, 16>(sK, k);
__syncthreads();
// TMEM alloc — 128 columns for (128, 128) Layout D output
@@ -76,24 +75,12 @@ 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 < 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
// at the (2*kt)-th core column. Each core column is BLOCK_MN/8 * 128 bytes
// = BLOCK_MN * 16 bytes. So offset = 2*kt * BLOCK_MN * 16 bytes = kt * BLOCK_MN * 32 bytes.
uint32_t q_addr = sQ_smem + kt * BLOCK_MN * 32;
uint32_t k_addr = sK_smem + kt * BLOCK_MN * 32;
// Single K-tile: descriptors point to start of SMEM
uint64_t dq = make_umma_desc_kmajor_none(sQ_smem, BLOCK_MN);
uint64_t dk = make_umma_desc_kmajor_none(sK_smem, BLOCK_MN);
uint64_t dq = make_umma_desc_kmajor_none(q_addr, BLOCK_MN);
uint64_t dk = make_umma_desc_kmajor_none(k_addr, BLOCK_MN);
if (tid == 0) {
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");
__syncthreads();
if (tid == 0) {
umma_ss_f16(tb, dq, dk, idesc, false);
}
// Final fence before TMEM read