File size: 3,387 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 51 52 53 54 55 56 57 58 59 60 61 62 | """Round 15 (D5): tt_kernels.hattn2 (dense-score 1-query attention, one core per head) vs the library SDPA as the MANO head
uses it (8x1 grid, q chunk 32, k chunk 192, HiFi4 + fp32) and vs torch fp32. usage: hattn2_test.py [time] [soak=N] [big]
big: score magnitudes up to ~60 (stress the max subtraction / exp range)"""
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 = 8, 192, 64
sc = 3.0 if "big" in sys.argv else 1.0
q = (torch.randn(H, D) * sc).to(torch.bfloat16)
k = (torch.randn(H, S, D) * sc).to(torch.bfloat16)
v = torch.randn(H, S, D).to(torch.bfloat16)
ref = torch.softmax((q.float()[:, None, :] @ k.float().transpose(-1, -2)) * D ** -0.5, -1) @ v.float() # [H, 1, D]
ref = ref[:, 0]
L1 = ttnn.L1_MEMORY_CONFIG
tt_ = lambda t: ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
qh = torch.zeros(1, H, 32, D, dtype=torch.bfloat16)
qh[0, :, 0] = q
Q, K, V = tt_(qh), tt_(k[None]), tt_(v[None])
pc = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=ttnn.CoreCoord(8, 1), q_chunk_size=32, k_chunk_size=192,
exp_approx_mode=False)
cc = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
math_approx_mode=False)
lib = lambda: ttnn.transformer.scaled_dot_product_attention(Q, K, V, is_causal=False, scale=D ** -0.5, program_config=pc,
memory_config=L1, compute_kernel_config=cc)
new = lambda: tk.hattn2(Q, K, V, D ** -0.5)
o1 = new()
ttnn.synchronize_device(dev)
O1 = ttnn.to_torch(o1).float()[0, :, 0]
o0 = lib()
O0 = ttnn.to_torch(o0).float()[0, :, 0]
print(f"[ha2] first call OK; finite {bool(torch.isfinite(O1).all())}", flush=True)
print(f"[ha2] vs fp32 ref: hattn2 PCC {pcc(O1, ref):.7f} max|d| {float((O1 - ref).abs().max()):.3e} mean|d| "
f"{float((O1 - ref).abs().mean()):.3e} | lib PCC {pcc(O0, ref):.7f} max|d| {float((O0 - ref).abs().max()):.3e} "
f"mean|d| {float((O0 - ref).abs().mean()):.3e}", flush=True)
print(f"[ha2] hattn2 vs lib: PCC {pcc(O1, O0):.7f} max|d| {float((O1 - O0).abs().max()):.3e}", flush=True)
if "time" in sys.argv:
print(f"[ha2] lib SDPA {time_op(dev, lib):.2f} us", flush=True)
print(f"[ha2] hattn2 {time_op(dev, new):.2f} us", flush=True)
if "prof" in sys.argv: # HA2_PROF=1 TT_METAL_DEVICE_PROFILER=1: zone timeline of the last hattn2 call
for _ in range(4):
ttnn.deallocate(new())
ttnn.synchronize_device(dev)
ttnn.ReadDeviceProfiler(dev)
n = int(arg("soak", "0"))
if n:
bad = 0
for _ in range(n):
o = new()
bad += int(not torch.equal(ttnn.to_torch(o).float()[0, :, 0], O1))
ttnn.deallocate(o)
print(f"[ha2] soak {n}: differing {bad}", flush=True)
finally:
ttnn.close_device(dev)
if "prof" in sys.argv:
tl_report(risks=("NCRISC", "BRISC", "TRISC_0", "TRISC_1", "TRISC_2"))
|