diff --git a/tests/unit/test_fmha_hd64_smem_p.cu b/tests/unit/test_fmha_hd64_smem_p.cu index c19882bc..934a40ab 100644 --- a/tests/unit/test_fmha_hd64_smem_p.cu +++ b/tests/unit/test_fmha_hd64_smem_p.cu @@ -190,6 +190,21 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v, } } + // ===== Register-math PV for verification ===== + float* s_ref_o = s_p_vals; // Reuse s_p_vals (already consumed by PV) + if (tid == 0) { + float s_ref[SK]; + for (int j=0;j