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

fix: proper early return for SMEM load test

Browse Source
This commit is contained in:
biondizzle
2026-05-28 08:41:30 +00:00
parent 82bc2c4a49
commit 58b610c96c
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -74,13 +74,15 @@ fmha_qk_verify(
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
s_out[0] = dot * scale;
s_out[1] = (float)(sQ_smem);
s_out[2] = (float)(sK_smem);
s_out[3] = (float)(tmem_base);
}
__syncthreads();
return; // TEMPORARILY SKIP MMA — just check SMEM loads
// TEMPORARILY SKIP MMA — just verify SMEM loads
if (wid == 0) tmem_dealloc(tmem_base, 128);
return; // ALL threads return — we're just testing 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();