53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
"""Test FMHA with pv_n_tile=16 (N=16 sub-tiles for PV GEMM)."""
|
|
import torch
|
|
import math
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
|
|
|
|
from dsv4.kernels.attention.production import dsv4_attention_per_head
|
|
|
|
|
|
def test_fmha_pv16(hd):
|
|
sk = 128
|
|
scale = 1.0 / math.sqrt(hd)
|
|
|
|
torch.manual_seed(42)
|
|
q = torch.randn(1, 1, hd, dtype=torch.bfloat16, device='cuda')
|
|
k = torch.randn(sk, hd, dtype=torch.bfloat16, device='cuda')
|
|
v = torch.randn(hd, sk, dtype=torch.bfloat16, device='cuda')
|
|
|
|
o = dsv4_attention_per_head(q, k, v, scale=scale, swa_len=sk)
|
|
|
|
# Reference
|
|
q_ref = q[0, 0].float()
|
|
k_ref = k.float()
|
|
v_ref = v.float()
|
|
|
|
s = (q_ref @ k_ref.T) * scale
|
|
p = torch.softmax(s, dim=-1)
|
|
o_ref = (p @ v_ref.T).to(torch.bfloat16)
|
|
|
|
o_f = o[0, 0].float()
|
|
o_ref_f = o_ref.float()
|
|
|
|
cs = torch.nn.functional.cosine_similarity(o_f.unsqueeze(0), o_ref_f.unsqueeze(0)).item()
|
|
print(f"HD={hd} pv_n_tile=16: cosine={cs:.8f}")
|
|
|
|
if cs < 0.999:
|
|
print(f" FAILED")
|
|
print(f" o[0:4] = {o_f[0:4].tolist()}")
|
|
print(f" o_ref[0:4] = {o_ref_f[0:4].tolist()}")
|
|
return False
|
|
print(f" PASSED")
|
|
return True
|
|
|
|
|
|
if __name__ == '__main__':
|
|
all_pass = True
|
|
for hd in [16, 64, 128]:
|
|
if not test_fmha_pv16(hd):
|
|
all_pass = False
|
|
sys.exit(0 if all_pass else 1)
|