| """Block-tail emulation: addln1(prefetch proj) -> proj = cmm(SDPA-layout ctx, head_concat folded) -> addln2(prefetch fc2) |
| -> gelu -> fc2 = cmm(prefetched pairs) vs the model's library path (head_concat + 2D linear, gelu + 2D linear). |
| usage: cmm_block_test.py [pf2=P] [time] [soak=N]""" |
| import sys |
| from pathlib import Path |
| from common import * |
| from open_device_12x10 import open_device |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from hamer import tt_kernels as tk |
| from hamer.vit import tuned_matmul_configs |
| arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d) |
| P2 = int(arg("pf2", "3")); S2 = int(arg("s2", "3")); S1 = int(arg("s1", "3")) |
| dev = open_device(grid="12x10", dispatch="eth", trace_region_size=64 << 20) |
| try: |
| 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) |
| T = lambda t, mem=L1, dt=ttnn.bfloat16: ttnn.from_torch(t, dtype=dt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mem) |
| x = T(torch.randn(1, 1, 192, 1280), res); y = T(torch.randn(1, 1, 192, 1280), res) |
| ctx_t = torch.randn(1, 16, 192, 96) * 0.5 |
| ctx = T(ctx_t) |
| h1 = T(torch.randn(1, 1, 192, 5120) * 1.5) |
| pw_t = (torch.randn(1536, 1280) * 1536 ** -0.5).bfloat16(); pb_t = (torch.randn(1280) * 0.1).bfloat16() |
| fw_t = (torch.randn(5120, 1280) * 5120 ** -0.5).bfloat16(); fb_t = (torch.randn(1280) * 0.1).bfloat16() |
| pw = T(pw_t, ttnn.DRAM_MEMORY_CONFIG, ttnn.bfloat8_b); fw = T(fw_t, ttnn.DRAM_MEMORY_CONFIG, ttnn.bfloat8_b) |
| pb = T(pb_t.reshape(1, 1, 1, -1)); fb = T(fb_t.reshape(1, 1, 1, -1)) |
| C1 = tk.addln_constants(torch.ones(1280), torch.zeros(1280), dev); C2 = tk.addln_constants(torch.ones(1280) * 0.9, torch.zeros(1280), dev) |
| pwk = tk.CmmWeight(pw_t, dev, 10, 4, 8, ttnn.bfloat8_b); fwk = tk.CmmWeight(fw_t, dev, 10, 4, 16, ttnn.bfloat8_b) |
| V1 = "v1" in sys.argv |
| P1 = int(arg("pf1", "3")) |
| pfp = tk.CmmPrefetch(dev, 10, 6, pwk, P1, single=V1) if P1 else None |
| pff = tk.CmmPrefetch(dev, 10, 6, fwk, P2, single=V1) if P2 else None |
| defs = ("CMM_FLUSH", f"CMM_LA={arg('la', '1')}") + (() if V1 else ("CMM_DRAM0",)) |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True) |
| pcs = tuned_matmul_configs(dev.compute_with_storage_grid_size(), fp32_acc=True) |
|
|
| def cmm_path(): |
| s, h = tk.add_layernorm(x, y, C1, 1e-6, out_memory_config=L1, prefetch=pfp) |
| ttnn.deallocate(h) |
| d = tk.cmm_linear(ctx, pwk, pb, res, kb=8, slots=S1, defines=defs, prefetch=pfp, a_heads=(16, 96), v1=V1) |
| s2, h2 = tk.add_layernorm(s, d, C2, 1e-6, out_memory_config=L1, prefetch=pff) |
| ttnn.deallocate(h2) |
| g = tk.gelu(h1, memory_config=L1) |
| d2 = tk.cmm_linear(g, fwk, fb, res, kb=16, slots=S2, defines=defs, prefetch=pff, v1=V1) |
| for t in (s, d, g): |
| ttnn.deallocate(t) |
| return s2, d2 |
|
|
| def lib_path(): |
| s, h = tk.add_layernorm(x, y, C1, 1e-6, out_memory_config=L1) |
| ttnn.deallocate(h) |
| cm = tk.head_concat(ctx, memory_config=L1) |
| d = ttnn.linear(cm, pw, bias=pb, program_config=pcs["proj"], memory_config=res, compute_kernel_config=ck) |
| s2, h2 = tk.add_layernorm(s, d, C2, 1e-6, out_memory_config=L1) |
| ttnn.deallocate(h2) |
| g = tk.gelu(h1, memory_config=L1) |
| d2 = ttnn.linear(g, fw, bias=fb, program_config=pcs["fc2"], memory_config=res, compute_kernel_config=ck) |
| for t in (s, cm, d, g): |
| ttnn.deallocate(t) |
| return s2, d2 |
|
|
| a2, ad = cmm_path(); ttnn.synchronize_device(dev) |
| b2, bd = lib_path(); ttnn.synchronize_device(dev) |
| A2, AD, B2, BD = (ttnn.to_torch(t).float() for t in (a2, ad, b2, bd)) |
| print(f"[blk] first calls OK s2 pcc {pcc(A2, B2):.7f} max|d| {float((A2 - B2).abs().max()):.4f} fc2 pcc {pcc(AD, BD):.7f} max|d| {float((AD - BD).abs().max()):.4f}", flush=True) |
| if "time" in sys.argv: |
| print(f"[blk] cmm path {time_op(dev, cmm_path):.2f} us library path {time_op(dev, lib_path):.2f} us", flush=True) |
| n = int(arg("soak", "0")) |
| if n: |
| reps = 20 |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| outs = [] |
| try: |
| for _ in range(reps): |
| outs.append(cmm_path()) |
| finally: |
| ttnn.end_trace_capture(dev, tid, cq_id=0) |
| try: |
| for i in range(n // reps): |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=False) |
| if i % 40 == 39: |
| ttnn.synchronize_device(dev) |
| ttnn.synchronize_device(dev) |
| bad = sum(int(not (torch.equal(ttnn.to_torch(p).float(), A2) and torch.equal(ttnn.to_torch(q).float(), AD))) for p, q in outs) |
| print(f"[blk] soak {n // reps * reps} block-tails OK, outputs differing from the first: {bad}/{reps}", flush=True) |
| finally: |
| ttnn.release_trace(dev, tid) |
| finally: |
| ttnn.close_device(dev) |
|
|
