| """Round 15 (D5 fused): the MANO-head q GEMV with the attention in the same program (hgemv(attn=...): 16 GEMV cores signal |
| 8 hattn2 cores through a semaphore) vs the two separate programs. Expect bit-identical q and context. |
| usage: hgattn_test.py [time] [soak=N]""" |
| import sys |
| from pathlib import Path |
| from common import * |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from hamer import tt_kernels as tk |
| arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d) |
| dev = open_dev() |
| try: |
| torch.manual_seed(0) |
| H, S, D, K, N = 8, 192, 64, 1024, 512 |
| L1 = ttnn.L1_MEMORY_CONFIG |
| x = ttnn.from_torch(torch.randn(1, 1, 32, K), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| w = tk.CmmWeight((torch.randn(K, N) * K ** -0.5).to(torch.bfloat16), dev, N // 32, 1, K // 32, ttnn.bfloat16) |
| b = ttnn.from_torch(torch.zeros(1, 1, 32, N), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| tt_ = lambda t: ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| Kt, Vt = tt_(torch.randn(1, H, S, D).to(torch.bfloat16)), tt_(torch.randn(1, H, S, D).to(torch.bfloat16)) |
| gcores = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(7, 3))}) |
| acores = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(8, 8), ttnn.CoreCoord(11, 9))}) |
|
|
| def sep(): |
| q = tk.hgemv(x, w, b, ln_eps=1e-5, out_dtype=ttnn.bfloat16, out_heads=H, cores=gcores) |
| return q, tk.hattn2(q, Kt, Vt, D ** -0.5, cores=acores) |
|
|
| def fused(): |
| return tk.hgemv(x, w, b, ln_eps=1e-5, out_dtype=ttnn.bfloat16, out_heads=H, cores=gcores, |
| attn=(Kt, Vt, D ** -0.5, acores)) |
| q1, c1 = fused() |
| ttnn.synchronize_device(dev) |
| print("[hga] fused first call OK", flush=True) |
| q0, c0 = sep() |
| Q0, C0 = ttnn.to_torch(q0).float()[0, :, 0], ttnn.to_torch(c0).float()[0, :, 0] |
| Q1, C1 = ttnn.to_torch(q1).float()[0, :, 0], ttnn.to_torch(c1).float()[0, :, 0] |
| print(f"[hga] q equal {bool(torch.equal(Q0, Q1))} ctx equal {bool(torch.equal(C0, C1))} (max|d| {float((C0 - C1).abs().max()):.3e})", flush=True) |
| if "time" in sys.argv: |
| print(f"[hga] separate {time_op(dev, sep):.2f} us", flush=True) |
| print(f"[hga] fused {time_op(dev, fused):.2f} us", flush=True) |
| n = int(arg("soak", "0")) |
| if n: |
| bad = 0 |
| for _ in range(n): |
| q_, c_ = fused() |
| bad += int(not torch.equal(ttnn.to_torch(c_).float()[0, :, 0], C1)) |
| ttnn.deallocate(q_); ttnn.deallocate(c_) |
| print(f"[hga] soak {n}: differing {bad}", flush=True) |
| finally: |
| ttnn.close_device(dev) |
|
|
