File size: 2,706 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 | """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 # [H, 1, D]
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: # noqa: BLE001
print(f"[sdpa] decode k_chunk {kc} FAIL {str(e).splitlines()[0][:200]}", flush=True)
finally:
ttnn.close_device(dev)
|