From 9b476d87f94db1e4871fcc9a9eb09acc599e3cdd Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 22:44:11 +0000 Subject: [PATCH] fix: compare un-normalized O against un-normalized reference --- tests/unit/test_d2_regression.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_d2_regression.py b/tests/unit/test_d2_regression.py index acd67453..09f78706 100644 --- a/tests/unit/test_d2_regression.py +++ b/tests/unit/test_d2_regression.py @@ -26,6 +26,7 @@ def test(): attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] attn_exp = torch.exp(qf @ kf.T * scale - attn_max) attn_sum = attn_exp.sum(dim=-1, keepdim=True) + ref_unnorm = attn_exp @ v.float() ref_norm = (attn_exp / attn_sum) @ v.float() lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') @@ -51,7 +52,7 @@ def test(): out = c_tile[:, :, 0].float() cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) + out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) ).item() print(f'hd={hd}, s_k={s_k}: cos_norm {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}')