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

test: re-enable TMEM alloc + zero

Browse Source
This commit is contained in:
biondizzle
2026-05-28 09:23:10 +00:00
parent 09e8ea5933
commit ca419c52f3
1 changed files with 16 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
tests/unit/test_umma_qk.cu
View File

@@ -62,23 +62,23 @@ test_umma_qk_hd16(
}
// ================================================================
// TMEM allocation — SKIP for now
// TMEM allocation
// ================================================================
// if (wid == 0) {
// uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
// tmem_alloc(smem_ptr, 128);
// }
// __syncthreads();
// uint32_t tmem_base = *sTmemBase;
//
// // 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();
if (wid == 0) {
uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
tmem_alloc(smem_ptr, 128);
}
__syncthreads();
uint32_t tmem_base = *sTmemBase;
// 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