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

fix: warp-collective TMEM read/dealloc in minimal QK test

Browse Source
This commit is contained in:
biondizzle
2026-05-29 18:42:03 +00:00
parent c755e6fdde
commit 9fc2d549e4
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/unit/test_qk_minimal.cu
View File

@@ -90,8 +90,9 @@ test_qk_minimal_kernel(float* __restrict__ out_s,
asm volatile("fence.sc.gpu;" ::: "memory");
__syncthreads();
// Read from TMEM
if (lane < T) {
// TMEM read — ALL 32 lanes must participate (warp-collective)
// For T=1, only lane 0's data is meaningful
{
for (int n = 0; n < NUM_READS; n++) {
float tmp[8];
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
@@ -108,7 +109,8 @@ test_qk_minimal_kernel(float* __restrict__ out_s,
}
}
__syncthreads();
tmem_dealloc(tb, TMEM_N);
// TMEM dealloc — warp 0 does it (must be warp-collective)
if (tid < 32) tmem_dealloc(tb, TMEM_N);
}
int main() {