diff --git a/tests/unit/test_fmha_ts_hd16.cu b/tests/unit/test_fmha_ts_hd16.cu index b14edf79..d88fd029 100644 --- a/tests/unit/test_fmha_ts_hd16.cu +++ b/tests/unit/test_fmha_ts_hd16.cu @@ -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