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

test: read 8 TMEM columns individually (no loop)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 09:44:30 +00:00
parent bcc6ed114d
commit 6c38c6e442
1 changed files with 23 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
tests/unit/test_umma_qk.cu
View File

@@ -86,13 +86,32 @@ test_umma_qk_hd16(
if (wid == 0 && lane == 0) tmem_fence_store();
__syncthreads();
// Read S[0,0] from TMEM
// Read S from TMEM — read 128 columns one at a time
// Use separate tmem_load calls with __syncthreads between them
// to avoid the loop crash
if (wid == 0) {
uint32_t u0, u1, u2, u3;
tmem_load(tmem_base, u0, u1, u2, u3);
if (lane == 0) s_out[0] = u32_to_f32(u0);
#define READ_COL(c) do { \
uint32_t u0, u1, u2, u3; \
tmem_load(tmem_base + c, u0, u1, u2, u3); \
if (lane == 0) s_out[c] = u32_to_f32(u0); \
} while(0)
READ_COL(0);
}
__syncthreads();
if (wid == 0) { READ_COL(1); }
__syncthreads();
if (wid == 0) { READ_COL(2); }
__syncthreads();
if (wid == 0) { READ_COL(3); }
__syncthreads();
if (wid == 0) { READ_COL(4); }
__syncthreads();
if (wid == 0) { READ_COL(5); }
__syncthreads();
if (wid == 0) { READ_COL(6); }
__syncthreads();
if (wid == 0) { READ_COL(7); }
__syncthreads();
// Scalar reference
if (tid == 0) {