File size: 1,744 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | """hattn (kernels/hattn) vs torch fp64 and vs head_split + SDPA + head_concat timing. usage: hattn_test.py [layer=0] [time=1]"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
a = dict(s.split("=") for s in sys.argv[1:])
l = int(a.get("layer", 0))
dev = open_dev()
try:
torch.manual_seed(0)
H, L = 8, 6
qv = torch.randn(H * 64) * 2
kvt = torch.randn(192, L * 1024) * 1.5
rowf = lambda v: v.reshape(1, 1, 1, -1).expand(1, 1, 32, v.numel()).contiguous()
q = ttnn.from_torch(rowf(qv * 0.125).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
memory_config=ttnn.L1_MEMORY_CONFIG)
kv = ttnn.from_torch(kvt.reshape(1, 1, 192, -1).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT,
device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
kb = kvt.to(torch.bfloat16).double()
qb = (qv * 0.125).to(torch.bfloat16).double()
ref = []
for h in range(H):
K = kb[:, l * 1024 + h * 64: l * 1024 + h * 64 + 64]
V = kb[:, l * 1024 + 512 + h * 64: l * 1024 + 512 + h * 64 + 64]
p = torch.softmax(K @ qb[h * 64:(h + 1) * 64], 0)
ref.append(p @ V)
ref = torch.cat(ref)
o = tk.hattn(q, kv, l)
ttnn.synchronize_device(dev)
got = ttnn.to_torch(o).double()[0, 0, 0]
print(f"[ha] layer {l}: pcc {pcc(got, ref):.7f} max|d| {float((got - ref).abs().max()):.5f}", flush=True)
if a.get("time", "1") == "1":
print(f"[ha] hattn {time_op(dev, lambda: tk.hattn(q, kv, l)):.2f} us/call", flush=True)
finally:
from open_device_12x10 import close_device
close_device(dev)
|