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

test: re-enable all TMEM ops (alloc, zero, dealloc)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 09:24:28 +00:00
parent acc7424a48
commit 2b1c8ce7df
1 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tests/unit/test_umma_qk.cu
View File

@@ -71,14 +71,14 @@ test_umma_qk_hd16(
__syncthreads();
uint32_t tmem_base = *sTmemBase;
// SKIP zeroing TMEM — just test alloc
// if (wid == 0) {
// for (int col = 0; col < 128; col++) {
// tmem_store(tmem_base + col, 0, 0, 0, 0);
// }
// tmem_fence_store();
// }
// __syncthreads();
// Zero TMEM
if (wid == 0) {
for (int col = 0; col < 128; col++) {
tmem_store(tmem_base + col, 0, 0, 0, 0);
}
tmem_fence_store();
}
__syncthreads();
// ================================================================
// Load Q and K into SMEM in canonical layout
@@ -149,10 +149,10 @@ test_umma_qk_hd16(
}
__syncthreads();
// TMEM dealloc — skipped
// if (wid == 0) {
// tmem_dealloc(tmem_base, 128);
// }
// TMEM dealloc
if (wid == 0) {
tmem_dealloc(tmem_base, 128);
}
}
// ==================================================================