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

test: match isolated TS test exactly (V=all-1, BLOCK_MN=16)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 13:44:36 +00:00
parent 482328160a
commit a944f90040
1 changed files with 10 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
tests/unit/test_fmha_ts_hd16.cu
View File

@@ -57,16 +57,14 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
bf16_t* sQ_pad = sQ + 128 * 16;
for (int i = tid; i < 4096; i += 128) sQ_pad[i] = 0;
// Load V as 8 K-tiles of (16, 16) canonical
for (int i = tid; i < VKT * V_TILE_SZ; i += 128) sV_base[i] = 0;
for (int kt = 0; kt < VKT; kt++) {
bf16_t* sv = sV_base + kt * V_TILE_SZ;
for (int i = tid; i < MMA_K_BF16 * HD; i += 128) {
int r = i / HD, d = i % HD;
int ck = d / 8, lc = d % 8;
int tmn = r / 8, lr = r % 8;
sv[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = v[d * SK + kt * MMA_K_BF16 + r];
}
// Load V = all 1.0 into (16, 16) canonical (same as isolated test)
for (int i = tid; i < 16 * 16; i += 128) sV_base[i] = 0;
__syncthreads();
for (int i = tid; i < 16 * 16; i += 128) {
int r = i / 16, c = i % 16;
int ck = c / 8, lc = c % 8;
int tmn = r / 8, lr = r % 8;
sV_base[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = f32_to_bf16(1.0f);
}
__syncthreads();
@@ -138,9 +136,8 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
uint32_t idesc_pv = make_idesc(BLOCK_MN, HD);
for (int kt = 0; kt < 1; kt++) { // DEBUG: only first PV K-tile
bf16_t* sv = sV_base + kt * V_TILE_SZ;
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sv), MMA_K_BF16);
uint32_t tmem_a = tb + kt * MMA_K_BF16; // P's K-tile columns [16*kt, 16*kt+15]
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV_base), 16);
uint32_t tmem_a = tb; // P's first 16 columns
if (tid == 0) umma_ts_f16(tb_o, tmem_a, dv, idesc_pv, false); // no accumulate for first tile
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");