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