| """Round 10: MANO-head cross-attention (1 query token, 8 heads x 64, 192 keys): the library SDPA as used now |
| ([1, 8, 32, 64] query, row 0 meaningful) vs flash-decode SDPA ([1, 1, 8->32, 64] query, heads as rows). |
| usage: head_sdpa_decode.py""" |
| import sys, os |
| from pathlib import Path |
| from common import * |
| dev = open_dev() |
| try: |
| torch.manual_seed(0) |
| H, S, D = 8, 192, 64 |
| q = torch.randn(H, D) |
| k = torch.randn(H, S, D) |
| v = torch.randn(H, S, D) |
| ref = torch.softmax((q[:, None, :] @ k.transpose(-1, -2)) * D ** -0.5, -1) @ v |
| ref = ref[:, 0] |
| L1 = ttnn.L1_MEMORY_CONFIG |
| G = dev.compute_with_storage_grid_size() |
| tk_ = lambda t: ttnn.from_torch(t.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| K, V = tk_(k[None]), tk_(v[None]) |
| qh = torch.zeros(1, H, 32, D); qh[0, :, 0] = q |
| Q = tk_(qh) |
| pc = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=G, 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) |
| f0 = 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) |
| o0 = ttnn.to_torch(f0()).float()[0, :, 0] |
| print(f"[sdpa] current: PCC {pcc(o0, ref):.7f} max|d| {float((o0 - ref).abs().max()):.4f} {time_op(dev, f0):.2f} us", flush=True) |
| qd = torch.zeros(1, 1, 32, D); qd[0, 0, :H] = q |
| DR = ttnn.DRAM_MEMORY_CONFIG |
| tkd = lambda t: ttnn.from_torch(t.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR) |
| QD, K, V = tkd(qd), tkd(k[None]), tkd(v[None]) |
| for kc in (32, 64, 128): |
| pcd = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=G, q_chunk_size=32, k_chunk_size=kc, exp_approx_mode=False) |
| f1 = lambda pcd=pcd: ttnn.transformer.scaled_dot_product_attention_decode( |
| QD, K, V, is_causal=True, cur_pos=[S - 1], scale=D ** -0.5, program_config=pcd, memory_config=DR, |
| compute_kernel_config=cc) |
| try: |
| o1 = ttnn.to_torch(f1()).float()[0, 0, :H] |
| print(f"[sdpa] decode k_chunk {kc}: PCC {pcc(o1, ref):.7f} max|d| {float((o1 - ref).abs().max()):.4f} " |
| f"{time_op(dev, f1):.2f} us", flush=True) |
| except Exception as e: |
| print(f"[sdpa] decode k_chunk {kc} FAIL {str(e).splitlines()[0][:200]}", flush=True) |
| finally: |
| ttnn.close_device(dev) |
|
|
