From 8fef46ce732c5def6880069c70b02a972d1dc190 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 10:45:02 +0000 Subject: [PATCH] P5: add reference comparison to Python multi-tile test --- tests/unit/test_p5_python_minimal.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_p5_python_minimal.py b/tests/unit/test_p5_python_minimal.py index 9033be0a..4f00895a 100644 --- a/tests/unit/test_p5_python_minimal.py +++ b/tests/unit/test_p5_python_minimal.py @@ -20,6 +20,16 @@ try: o, lse = fmha_multitile_decode_raw(q, k, v, scale) print(f'Output[0,0,0,:5]: {o[0,0,0,:5].float()}') print(f'LSE: {lse[0,0,0].item():.4f}') - print('SUCCESS') + + # Reference + q_r = q[0,0].float() # (1, hd) + k_r = k[0,0].float() # (N, hd) + v_r = v[0,0].float().T # (N, hd) — V is (hd, N), transpose for reference + s = torch.matmul(q_r, k_r.T) * scale + s = torch.softmax(s, dim=-1) + o_ref = torch.matmul(s, v_r) + cos = torch.nn.functional.cosine_similarity(o[0,0].float().flatten().unsqueeze(0), o_ref.flatten().unsqueeze(0)).item() + print(f'Cosine vs reference: {cos:.6f}') + print(f'{"PASS" if cos >= 0.999990 else "FAIL"}') except Exception as e: print(f'FAILED: {e}')