File size: 1,443 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 | """Face-granular head split (unpadded 80-wide heads -> 96 padded) / concat (96 -> 80): exactness + time."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
dev = open_dev()
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
S, H, Dh, Dp = 192, 16, 80, 96
for trial in range(2):
qkv_t = torch.randn(1, 1, S, 3 * H * Dh)
qkv = ttnn.from_torch(qkv_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
ref = qkv_t.reshape(1, S, 3, H, Dh).permute(2, 0, 3, 1, 4) # 3,1,H,S,Dh
ref = torch.nn.functional.pad(ref, (0, Dp - Dh))
outs = tk.head_split_faces(qkv, H, Dh, Dp)
ok = all(torch.equal(ttnn.to_torch(o).float(), ref[i].to(torch.bfloat16).float()) for i, o in enumerate(outs))
print(f"trial{trial} split exact={ok}", flush=True)
ctx_t = torch.randn(1, H, S, Dp)
ctx = ttnn.from_torch(ctx_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
c = tk.head_concat_faces(ctx, Dh)
okc = torch.equal(ttnn.to_torch(c).float(), ctx_t[..., :Dh].permute(0, 2, 1, 3).reshape(1, 1, S, H * Dh).to(torch.bfloat16).float())
print(f"trial{trial} concat exact={okc}", flush=True)
print("face split us", time_op(dev, lambda: tk.head_split_faces(qkv, H, Dh, Dp)))
print("face concat us", time_op(dev, lambda: tk.head_concat_faces(ctx, Dh)))
ttnn.close_device(dev)
|