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

test: properly aligned V SMEM buffer

Browse Source
This commit is contained in:
biondizzle
2026-05-28 13:47:47 +00:00
parent 1c14ada386
commit 0dcaa648b3
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/unit/test_fmha_ts_hd16.cu
View File

@@ -40,8 +40,9 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
uint32_t* sTmemBase = (uint32_t*)sbuf;
bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
bf16_t* sK = sQ + 128 * 16 + 4096;
bf16_t* sV_base = sK + 128 * 16;
float* sQ_row = (float*)(sV_base + VKT * V_TILE_SZ);
// Fresh aligned V buffer (ensure 128-byte alignment for descriptor)
bf16_t* sV = (bf16_t*)(((uintptr_t)(sK + 128 * 16) + 127) & ~(uintptr_t)127);
float* sQ_row = (float*)(sV + 16 * 16);
for (int d = tid; d < HD; d += 128) sQ_row[d] = bf16_to_f32(q[d]);
@@ -58,13 +59,13 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
for (int i = tid; i < 4096; i += 128) sQ_pad[i] = 0;
// 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;
for (int i = tid; i < 16 * 16; i += 128) sV[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);
sV[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = f32_to_bf16(1.0f);
}
__syncthreads();
@@ -136,7 +137,7 @@ 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
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV_base), 16);
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
uint32_t tmem_a = tb; // P's first 16 columns
if (tid == 0) umma_ts_f16(tb, tmem_a, dv, idesc_pv, false); // Write C to tb (overwriting P)