diff --git a/dsv4/kernels/attention/fmha_qk_verify.cuh b/dsv4/kernels/attention/fmha_qk_verify.cuh index 0803950e..e420814b 100644 --- a/dsv4/kernels/attention/fmha_qk_verify.cuh +++ b/dsv4/kernels/attention/fmha_qk_verify.cuh @@ -67,28 +67,20 @@ fmha_qk_verify( uint64_t desc_q = make_umma_desc_mn_none(sQ_smem, HD); uint64_t desc_k = make_umma_desc_k_none(sK_smem, HD); - // Quick test: write known data to TMEM and read it back - // This verifies TMEM read/write works in this kernel - if (wid == 0) { - tmem_store(tmem_base + 0, f32_to_u32(42.0f), f32_to_u32(43.0f), f32_to_u32(44.0f), f32_to_u32(45.0f)); - tmem_fence_store(); - } - __syncthreads(); + // Quick test: verify SMEM data was loaded correctly + // Write Q[0,0..3] * K[0,0..3] dot product (scalar) to s_out[0] as sanity check if (tid == 0) { - uint32_t u0, u1, u2, u3; - tmem_load(tmem_base + 0, u0, u1, u2, u3); - printf("[qk] TMEM verify: %f %f %f %f\n", u32_to_f32(u0), u32_to_f32(u1), u32_to_f32(u2), u32_to_f32(u3)); + float dot = 0; + for (int d = 0; d < HD; d++) { + dot += bf16_to_f32(sQ[d]) * bf16_to_f32(sK[d]); + } + s_out[0] = dot * scale; // sanity: scalar QK result + s_out[1] = (float)(sQ_smem); // debug: SMEM address + s_out[2] = (float)(sK_smem); // debug: SMEM address + s_out[3] = (float)(tmem_base); // debug: TMEM base } __syncthreads(); - - // Now zero it for the actual MMA - if (wid == 0) { - tmem_store(tmem_base + 0, 0, 0, 0, 0); - tmem_fence_store(); - } - __syncthreads(); - - // QK GEMM: S = Q @ K^T (SS, both SMEM → TMEM) + return; // TEMPORARILY SKIP MMA — just check SMEM loads // MMA: called by ONE lane (elect_one_sync pattern) if (tid == 0) printf("[qk] tmem_base=%u sQ_smem=0x%x sK_smem=0x%x\n", tmem_base, sQ_smem, sK_smem); __syncthreads();