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

debug: verify SMEM loads + scalar QK sanity check

Browse Source
This commit is contained in:
biondizzle
2026-05-28 08:40:16 +00:00
parent 53139d24bf
commit 82bc2c4a49
1 changed files with 11 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
dsv4/kernels/attention/fmha_qk_verify.cuh
View File

@@ -67,28 +67,20 @@ fmha_qk_verify(
uint64_t desc_q = make_umma_desc_mn_none(sQ_smem, HD);
uint64_t desc_k = make_umma_desc_k_none(sK_smem, HD);
// Quick test: write known data to TMEM and read it back
// This verifies TMEM read/write works in this kernel
if (wid == 0) {
tmem_store(tmem_base + 0, f32_to_u32(42.0f), f32_to_u32(43.0f), f32_to_u32(44.0f), f32_to_u32(45.0f));
tmem_fence_store();
}
__syncthreads();
// Quick test: verify SMEM data was loaded correctly
// Write Q[0,0..3] * K[0,0..3] dot product (scalar) to s_out[0] as sanity check
if (tid == 0) {
uint32_t u0, u1, u2, u3;
tmem_load(tmem_base + 0, u0, u1, u2, u3);
printf("[qk] TMEM verify: %f %f %f %f\n", u32_to_f32(u0), u32_to_f32(u1), u32_to_f32(u2), u32_to_f32(u3));
float dot = 0;
for (int d = 0; d < HD; d++) {
dot += bf16_to_f32(sQ[d]) * bf16_to_f32(sK[d]);
}
s_out[0] = dot * scale; // sanity: scalar QK result
s_out[1] = (float)(sQ_smem); // debug: SMEM address
s_out[2] = (float)(sK_smem); // debug: SMEM address
s_out[3] = (float)(tmem_base); // debug: TMEM base
}
__syncthreads();
// Now zero it for the actual MMA
if (wid == 0) {
tmem_store(tmem_base + 0, 0, 0, 0, 0);
tmem_fence_store();
}
__syncthreads();
// QK GEMM: S = Q @ K^T (SS, both SMEM → TMEM)
return; // TEMPORARILY SKIP MMA — just check SMEM loads
// MMA: called by ONE lane (elect_one_sync pattern)
if (tid == 0) printf("[qk] tmem_base=%u sQ_smem=0x%x sK_smem=0x%x\n", tmem_base, sQ_smem, sK_smem);
__syncthreads();