| """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) |
|
|
