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

test: skip QK+softmax, write P directly to TMEM for PV debug

Browse Source
This commit is contained in:
biondizzle
2026-05-28 13:41:50 +00:00
parent f24bc583dc
commit 3a40ed6d69
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_fmha_ts_hd16.cu
View File

@@ -70,7 +70,23 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
}
__syncthreads();
// STEP 1: QK GEMM
// DEBUG: Skip QK and softmax, write P directly to TMEM
// Write P = all 1.0 (128, 128) — only row 0 has 1.0, rest 0
if (wid == 0) {
for (int n = 0; n < SK / 8; n++) {
// Lane 0 writes 1.0 for row 0, other lanes write 0
float p0=(lane==0)?0.0078125f:0, p1=(lane==0)?0.0078125f:0;
float p2=(lane==0)?0.0078125f:0, p3=(lane==0)?0.0078125f:0;
float p4=(lane==0)?0.0078125f:0, p5=(lane==0)?0.0078125f:0;
float p6=(lane==0)?0.0078125f:0, p7=(lane==0)?0.0078125f:0;
asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0],{%1,%2,%3,%4,%5,%6,%7,%8};" :: "r"(tb+n*8),"f"(p0),"f"(p1),"f"(p2),"f"(p3),"f"(p4),"f"(p5),"f"(p6),"f"(p7));
}
tmem_fence_store();
}
__syncthreads();
/*
// STEP 1: QK GEMM — DISABLED for debug
uint64_t desc_q = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ), BLOCK_MN);
uint64_t desc_k = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK), BLOCK_MN);
uint32_t idesc_qk = make_idesc(BLOCK_MN, BLOCK_MN);
@@ -78,7 +94,7 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
// STEP 2: Softmax + write P to TMEM
// STEP 2: Softmax — DISABLED for debug
if (wid == 0) {
float s_vals[SK], row_max = -INFINITY;
for (int n = 0; n < SK / 8; n++) {
@@ -93,7 +109,6 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
row_sum = wsum(row_sum);
if (lane == 0) for (int j=0;j<SK;j++) s_vals[j] /= row_sum;
// Write P to TMEM using 32x32b.x8 stores
for (int n = 0; n < SK / 8; n++) {
float p0=(lane==0)?s_vals[n*8+0]:0, p1=(lane==0)?s_vals[n*8+1]:0;
float p2=(lane==0)?s_vals[n*8+2]:0, p3=(lane==0)?s_vals[n*8+3]:0;
@@ -104,6 +119,7 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
tmem_fence_store();
}
__syncthreads();
*/
// Zero O TMEM region
if (wid == 0) {