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

Add QK GEMM to minimal PV test

Browse Source
This commit is contained in:
biondizzle
2026-05-28 14:39:51 +00:00
parent e9f8f9e6e3
commit 38b0ff0bf8
1 changed files with 18 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
tests/unit/test_minimal_pv.cu
View File

@@ -25,10 +25,16 @@ test_minimal_pv()
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
bf16_t* sPk = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
bf16_t* sK0 = sQ0 + TILE_SZ;
bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + TILE_SZ) + 127) & ~(uintptr_t)127);
bf16_t* sV = (bf16_t*)(((uintptr_t)(sPk + TILE_SZ) + 127) & ~(uintptr_t)127);
float* s_p_vals = (float*)(sV + 256); // 16 P values for 1 K-tile
// Fill sQ0, sK0 (will be used by QK)
for (int i = tid; i < TILE_SZ; i += 128) { sQ0[i] = 0; sK0[i] = 0; }
__syncthreads();
// Set s_p_vals = all 0.5
for (int j = tid; j < 16; j += 128) s_p_vals[j] = 0.5f;
__syncthreads();
@@ -59,6 +65,16 @@ test_minimal_pv()
__syncthreads();
uint32_t tb = *sTmemBase;
// QK GEMM (SS) — uses sQ0 and sK0
{
uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ0), BLOCK_MN);
uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK0), BLOCK_MN);
uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN);
if (tid == 0) umma_ss_f16(tb, dq, dk, idesc, false);
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
}
// PV SS MMA
uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), BLOCK_MN);
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
@@ -83,7 +99,7 @@ test_minimal_pv()
int main() {
printf("=== Minimal PV with s_p_vals ===\n");
int smem = (4+16 + TILE_SZ*2 + 256*2 + 16*4 + 256 + 127) & ~127;
int smem = (4+16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + 256*2 + 16*4 + 256 + 127) & ~127;
printf("SMEM: %d bytes\n", smem);
test_minimal_pv<<<1, 128, smem>>>();
cudaError_t err = cudaDeviceSynchronize();