From 9347119ade9d4b82acc1c972bc7d078b7da4adea Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 01:28:32 +0000 Subject: [PATCH] diag: print expected unnorm P@V for comparison with raw kernel output --- tests/unit/test_fmha_v3_stage_c.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 899e036a..52ed0b3f 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -436,14 +436,26 @@ def test(): v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') v_kernel = v.unsqueeze(-1) c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') + debug = torch.zeros(4, dtype=torch.float32, device='cuda') # [row_sum, row_max, inv_row_sum, 0] qf = q[:, :, 0].float() kf = k[:, :, 0].float() scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) + attn_raw = qf @ kf.T * scale + attn = torch.softmax(attn_raw, dim=-1) ref = attn @ v.float() + # Expected stats for comparison + print(f' row_sum (should be 1.0): {attn.sum(dim=-1)[:4].tolist()}') + # Unnormalized softmax: exp(S - max) + S_max = attn_raw.max(dim=-1, keepdim=True).values + P_unnorm = torch.exp(attn_raw - S_max) + unnorm_pv = P_unnorm @ v.float() + unnorm_sum = P_unnorm.sum(dim=-1) + print(f' unnorm row_sum: {unnorm_sum[:4].tolist()}') + print(f' unnorm P@V[0,:4]: {unnorm_pv[0,:4].tolist()}') + print(f' kernel out[0,:4] should match unnorm P@V (no normalize)') + mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))