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

test: revert volatile, fix wid==0, full 4 K-tiles

Browse Source
This commit is contained in:
biondizzle
2026-05-28 12:16:09 +00:00
parent de879342dd
commit 526fafb808
1 changed files with 4 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/unit/test_umma_qk_hd64.cu
View File

@@ -31,7 +31,7 @@ test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k,
bf16_t* sK = sQ + 128 * hd;
// TMEM alloc
if (wid == 1) {
if (wid == 0) {
tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128);
}
__syncthreads();
@@ -42,18 +42,15 @@ test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k,
__syncthreads();
// Write Q (1, hd) to sQ row 0 in canonical layout
// Use volatile to prevent compiler optimization eliminating writes
volatile bf16_t* vsQ = (volatile bf16_t*)sQ;
volatile bf16_t* vsK = (volatile bf16_t*)sK;
for (int d = tid; d < hd; d += 128) {
int ck = d / 8, lc = d % 8;
vsQ[ck * 16 * 64 + lc] = q[d];
sQ[ck * 16 * 64 + lc] = q[d];
}
// Write K (sk, hd) to sK in canonical layout
for (int i = tid; i < sk * hd; i += 128) {
int r = i / hd, c = i % hd;
int tmn = r / 8, ck = c / 8, lr = r % 8, lc = c % 8;
vsK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[i];
sK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[i];
}
__syncthreads();
@@ -70,7 +67,7 @@ test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k,
uint32_t idesc = make_idesc(128, 128);
// K-tile loop with accumulate
for (int kt = 0; kt < 1; kt++) { // 1 K-tile only
for (int kt = 0; kt < hd / 16; kt++) { // Full K-tile loop
// K-tile kt: columns [16*kt, 16*kt+16)
// In canonical layout, columns start at core_k = 2*kt and 2*kt+1
// Offset = 2*kt * 2048 bytes from matrix base