| """[add + LN + s2i + qkv] vs [addln + qkv] in a trace (op-to-op effects included). usage: addln_chain.py [eth|worker]""" |
| 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 |
| res = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK, |
| orientation=ttnn.ShardOrientation.ROW_MAJOR) |
| x = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) |
| y = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) |
| gt = 1 + 0.3 * torch.randn(1280); bt = 0.2 * torch.randn(1280) |
| C = tk.addln_constants(gt, bt, dev) |
| ga = ttnn.from_torch(gt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| be = ttnn.from_torch(bt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| lnpc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(10, 6), subblock_w=4, block_h=1, block_w=4, inplace=False) |
| w = ttnn.from_torch(torch.randn(1280, 4608) / 36, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev) |
| b = ttnn.from_torch(torch.randn(1, 1, 1, 4608), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| G = dev.compute_with_storage_grid_size() |
| pc = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(compute_with_storage_grid_size=ttnn.CoreCoord(G.x, G.y), in0_block_w=10, |
| out_subblock_h=2, out_subblock_w=2, per_core_M=6, per_core_N=2, fuse_batch=True, fused_activation=None, mcast_in0=True) |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True) |
| eps = 1e-6 |
| def lib(): |
| s = ttnn.add(x, y, memory_config=res) |
| h = ttnn.layer_norm(s, weight=ga, bias=be, epsilon=eps, memory_config=res, program_config=lnpc) |
| hi = ttnn.sharded_to_interleaved(h, L1); ttnn.deallocate(h) |
| o = ttnn.linear(hi, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck); ttnn.deallocate(hi) |
| return [s, o] |
| def cus(): |
| s, h = tk.add_layernorm(x, y, C, eps) |
| o = ttnn.linear(h, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck); ttnn.deallocate(h) |
| return [s, o] |
| def cus_only(): |
| return list(tk.add_layernorm(x, y, C, eps)) |
| def mm_only(): |
| hi = ttnn.sharded_to_interleaved(x, L1) |
| return [hi] |
| hpre = ttnn.sharded_to_interleaved(x, L1) |
| def qkv_only(): |
| return [ttnn.linear(hpre, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck)] |
| def addln_then_qkv_pre(): |
| s, h = tk.add_layernorm(x, y, C, eps) |
| o = ttnn.linear(hpre, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck) |
| return [s, h, o] |
| def ln_then_qkv_pre(): |
| s = ttnn.add(x, y, memory_config=res) |
| h = ttnn.layer_norm(s, weight=ga, bias=be, epsilon=eps, memory_config=res, program_config=lnpc) |
| o = ttnn.linear(hpre, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck) |
| return [s, h, o] |
| for name, fn in [("qkv only", qkv_only), ("addln; qkv(pre)", addln_then_qkv_pre), ("add+ln; qkv(pre)", ln_then_qkv_pre), ("lib add+ln+s2i+qkv", lib), ("addln+qkv", cus), ("addln", cus_only), ("lib", lib), ("addln+qkv", cus)]: |
| print(f"{name:22s} {time_op(dev, fn, reps=32, iters=20):7.2f} us", flush=True) |
| ttnn.close_device(dev) |
|
|
