File size: 3,145 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 | """Custom fused attention (tt_kernels.fused_attention) vs torch fp32 and vs the library path
(head_split + SDPA 32x192 + head_concat). usage: attn_test.py [time] [soak N] [hifi2]"""
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
dev = open_dev()
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
H, D, DR, S = 16, 96, 80, 192
fid = ttnn.MathFidelity.HiFi2 if "hifi2" in sys.argv else ttnn.MathFidelity.HiFi4
q = torch.randn(1, H, S, D) * 1.5; k = torch.randn(1, H, S, D) * 1.5; v = torch.randn(1, H, S, D)
for t in (q, k, v):
t[..., DR:] = 0
qkv_t = torch.cat([q.permute(0, 2, 1, 3).reshape(1, 1, S, H * D), k.permute(0, 2, 1, 3).reshape(1, 1, S, H * D),
v.permute(0, 2, 1, 3).reshape(1, 1, S, H * D)], dim=-1).to(torch.bfloat16)
qb, kb, vb = [t.to(torch.bfloat16).float() for t in (q, k, v)]
scale = DR ** -0.5
ref = torch.softmax(qb @ kb.transpose(-1, -2) * scale, -1) @ vb # [1,H,S,D]
ref = ref.permute(0, 2, 1, 3).reshape(1, 1, S, H * D)
qkv = ttnn.from_torch(qkv_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
o = tk.fused_attention(qkv, H, D, scale, memory_config=L1, fidelity=fid)
ttnn.synchronize_device(dev)
ot = ttnn.to_torch(o).float()
print(f"custom first call OK: PCC {pcc(ot, ref):.7f} max|d| {float((ot - ref).abs().max()):.4f} mean|d| {float((ot - ref).abs().mean()):.5f}",
"pad cols zero:", bool((ot.reshape(S, H, D)[..., DR:] == 0).all()), flush=True)
pc = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=dev.compute_with_storage_grid_size(), q_chunk_size=32,
k_chunk_size=S, exp_approx_mode=False)
def lib():
qq, kk, vv = tk.head_split(qkv, H, D, memory_config=L1)
c = ttnn.transformer.scaled_dot_product_attention(qq, kk, vv, is_causal=False, scale=scale, program_config=pc, memory_config=L1)
for t in (qq, kk, vv):
ttnn.deallocate(t)
cc = tk.head_concat(c, memory_config=L1); ttnn.deallocate(c)
return cc
ol = lib(); olt = ttnn.to_torch(ol).float()
print(f"library: PCC {pcc(olt, ref):.7f} max|d| {float((olt - ref).abs().max()):.4f} mean|d| {float((olt - ref).abs().mean()):.5f}")
if "time" in sys.argv:
print(f"custom attn {time_op(dev, lambda: tk.fused_attention(qkv, H, D, scale, memory_config=L1, fidelity=fid)):.2f} us; "
f"library split+sdpa+concat {time_op(dev, lib):.2f} us", flush=True)
if "soak" in sys.argv:
n = int(sys.argv[sys.argv.index("soak") + 1])
tid = ttnn.begin_trace_capture(dev, cq_id=0)
try:
outs = [tk.fused_attention(qkv, H, D, scale, memory_config=L1, fidelity=fid) for _ in range(10)]
finally:
ttnn.end_trace_capture(dev, tid, cq_id=0)
for i in range(n // 10):
ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
if i % 50 == 0:
ttnn.synchronize_device(dev)
ttnn.synchronize_device(dev)
bad = sum(int(not torch.equal(ttnn.to_torch(x).float(), ot)) for x in outs)
print(f"soak {n} OK, differing: {bad}/10")
ttnn.release_trace(dev, tid)
ttnn.close_device(dev)
|