P5: add reference comparison to Python multi-tile test
This commit is contained in:
@@ -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}')
|
||||
|
||||
Reference in New Issue
Block a user