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

test: re-enable TMEM zeroing with tmem_base debug

Browse Source
This commit is contained in:
biondizzle
2026-05-28 09:26:16 +00:00
parent a7f81331f8
commit f1aaa50326
1 changed files with 19 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
tests/unit/test_umma_qk.cu
View File

@@ -71,14 +71,25 @@ test_umma_qk_hd16(
__syncthreads();
uint32_t tmem_base = *sTmemBase;
// Zero TMEM — skip for now
// if (wid == 0) {
// for (int col = 0; col < 32; col++) {
// tmem_store(tmem_base + col, 0, 0, 0, 0);
// }
// tmem_fence_store();
// }
// __syncthreads();
// Debug: write tmem_base to output
if (tid == 0) {
s_out[138] = 999.0f; // sentinel before tmem ops
}
__syncthreads();
// Zero TMEM — re-enable with debug
if (wid == 0) {
for (int col = 0; col < 128; col++) {
tmem_store(tmem_base + col, 0, 0, 0, 0);
}
tmem_fence_store();
}
__syncthreads();
if (tid == 0) {
s_out[138] = (float)tmem_base; // write tmem_base after tmem ops
}
__syncthreads();
// ================================================================
// Load Q and K into SMEM in canonical layout