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

test: fix K source stride mismatch in manual SMEM write

Browse Source
This commit is contained in:
biondizzle
2026-05-28 12:54:03 +00:00
parent 7d16a30cb6
commit 926ae5d7bf
1 changed files with 19 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tests/unit/test_umma_qk_hd64.cu
View File

@@ -41,9 +41,25 @@ test_umma_hd64_debug(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
__syncthreads();
uint32_t tb = *sTmemBase;
// Load Q and K into SMEM in canonical layout — ONLY first 16 dims
write_q_to_smem<16>(sQ, q);
write_k_to_smem<128, 16>(sK, k); // k is (128, 64), but we only read first 16 cols
// Load Q (first 16 dims) into (128, 16) canonical — manual write
// Zero sQ
for (int i = tid; i < 128 * 16; i += 128) sQ[i] = 0;
// Q row 0, first 16 dims
for (int c = tid; c < 16; c += 128) {
int ck = c / 8, lc = c % 8;
sQ[ck * 16 * 64 + lc] = q[c];
}
// Load K (first 16 cols of (128, 64)) into (128, 16) canonical
// K source stride is 64, but we only write 16 columns to SMEM
for (int i = tid; i < 128 * 16; i += 128) sK[i] = 0;
for (int r = 0; r < 128; r++) {
for (int c = tid; c < 16; c += 128) {
int ck = c / 8, lc = c % 8;
int tmn = r / 8, lr = r % 8;
sK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * 64 + c];
}
}
bf16_t* sQ_pad = sQ + 128 * 16;
for (int i = tid; i < 4096; i += 128) sQ_pad[i] = 0;
__syncthreads();