fix: compare un-normalized O against un-normalized reference

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"}')