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

Debug: QK only, skip softmax+PV

Browse Source
This commit is contained in:
biondizzle
2026-05-28 14:15:18 +00:00
parent 79cee32125
commit fa6c124163
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/unit/test_fmha_ts_full.cu
View File

@@ -90,7 +90,8 @@ test_fmha_ts_full(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
// ===== STEP 2: Softmax =====
// ===== STEP 2: Softmax — SKIPPED FOR DEBUG =====
/*
if (wid == 0) {
float s_vals[SK], row_max = -INFINITY;
for (int n = 0; n < SK / 8; n++) {
@@ -127,6 +128,7 @@ test_fmha_ts_full(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
tmem_fence_store();
}
__syncthreads();
*/
// ===== STEP 3: PV GEMM (TS) — SKIPPED FOR DEBUG =====
/*
@@ -148,8 +150,8 @@ test_fmha_ts_full(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
__syncthreads();
*/
// ===== STEP 4: Epilogue — read O from TMEM =====
// MMA output is scaled by 0.5, so multiply by 2.0
// ===== STEP 4: Epilogue — read S from TMEM (QK output) =====
// Just verify QK works by reading first 16 values
if (wid == 0) {
float o_vals[HD];
for (int n = 0; n < HD / 8; n++) {
@@ -157,9 +159,9 @@ test_fmha_ts_full(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
: "r"(tb_o + n*8));
: "r"(tb + n*8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c] * 2.0f;
if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c];
}
if (lane == 0) for (int d=0;d<HD;d++) o_out[d] = f32_to_bf16(o_vals[d]);
}