File size: 2,716 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 39 40 41 42 43 44 45 46 47 48 49 50 | """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)
|