P5: add reference comparison to Python multi-tile test

This commit is contained in:
2026-05-30 10:45:02 +00:00
parent 897a70a491
commit 8fef46ce73

View File

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