diff --git a/dsv4/kernels/attention/fmha_qk_verify.cuh b/dsv4/kernels/attention/fmha_qk_verify.cuh index e420814b..8b58c0f1 100644 --- a/dsv4/kernels/attention/fmha_qk_verify.cuh +++ b/dsv4/kernels/attention/fmha_qk_verify.cuh @@ -74,13 +74,15 @@ fmha_qk_verify( 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 + s_out[0] = dot * scale; + s_out[1] = (float)(sQ_smem); + s_out[2] = (float)(sK_smem); + s_out[3] = (float)(tmem_base); } __syncthreads(); - return; // TEMPORARILY SKIP MMA — just check SMEM loads + // TEMPORARILY SKIP MMA — just verify SMEM loads + if (wid == 0) tmem_dealloc(tmem_base, 128); + return; // ALL threads return — we're just testing 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();