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