"""Round 8: cmm + fused add+LN epilogue (tt_kernels.cmm_linear(addln=...)) vs the separate cmm and add_layernorm ops. usage: cmm_addln_test.py shape=tiny|proj|fc2 [ws] [time] (ws: LN output width-sharded on 10 cores like qkv's input) Expect bit-identical s and h.""" 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 arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d) shape = arg("shape", "tiny") GX, GY, NT, KT, kb = {"tiny": (2, 2, 1, 8, 4), "proj": (10, 6, 4, 48, 8), "fc2": (10, 6, 4, 160, 16)}[shape] dev = open_dev() try: torch.manual_seed(0) M, K, N = 32 * GY, 32 * KT, 32 * NT * GX G = dev.compute_with_storage_grid_size() res = ttnn.create_sharded_memory_config((M, N), core_grid=ttnn.CoreGrid(y=GY, x=GX), strategy=ttnn.ShardStrategy.BLOCK, orientation=ttnn.ShardOrientation.ROW_MAJOR) at = (torch.randn(1, 1, M, K) * 0.5).to(torch.bfloat16) a = ttnn.from_torch(at, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) wt = (torch.randn(K, N) * K ** -0.5).to(torch.bfloat16) wf = tk.CmmWeight(wt, dev, GX, NT, kb, ttnn.bfloat8_b) b = ttnn.from_torch((torch.randn(N) * 0.1).reshape(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) xt = torch.randn(1, 1, M, N).to(torch.bfloat16) x1 = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) x2 = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) C = tk.addln_constants(torch.randn(N) * 0.2 + 1, torch.randn(N) * 0.1, dev) Cf = C if "noaff" in sys.argv: # round 9: reference = identity affine (gamma 1, beta 0); fused = ADDLN_NOAFF (both exact) C = tk.addln_constants(torch.ones(N), torch.zeros(N), dev) Cf = dict(tk.addln_constants(torch.ones(N), torch.zeros(N), dev), noaff=True) if "nobeta" in sys.argv: # reference = gamma with beta 0; fused = ADDLN_NOBETA (beta add skipped) gm = torch.randn(N) * 0.2 + 1 C = tk.addln_constants(gm, torch.zeros(N), dev) Cf = dict(tk.addln_constants(gm, torch.zeros(N), dev), nobeta=True) hmem = None if "ws" in sys.argv: ns = GX sg = ttnn.num_cores_to_corerangeset(ns, G, row_wise=True) hmem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(sg, [M, N // ns], ttnn.ShardOrientation.ROW_MAJOR)) if "xres" in sys.argv: # round 12: LN output straight into the resident in0 of the xmm fc1 (12 x 10 cores) hmem = "xres" defs = ("CMM_FLUSH", "CMM_LA=1") d = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True) s_ref, h_ref = tk.add_layernorm(x1, d, C, 1e-6, out_memory_config=None if hmem == "xres" else hmem) print(f"[ca] {shape} G={GX}x{GY} NT={NT} KT={KT} kb={kb} ws={hmem is not None}", flush=True) s, h = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True, addln=dict(x=x2, consts=Cf, eps=1e-6, out_mem=hmem)) if "noaff" in sys.argv or "nobeta" in sys.argv: # standalone add+LN with NOAFF / NOBETA too s3, h3 = tk.add_layernorm(x1, tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True), Cf, 1e-6, out_memory_config=None if hmem == "xres" else hmem) ttnn.synchronize_device(dev) if hmem != "xres": print(f"[ca] standalone NOAFF h equal {bool(torch.equal(ttnn.to_torch(h_ref).float(), ttnn.to_torch(h3).float()))}", flush=True) ttnn.synchronize_device(dev) sr, hr = ttnn.to_torch(s_ref).float(), ttnn.to_torch(h_ref).float() if hmem == "xres": # rebuild [M, N] from the resident shards: core (x, y), parity p, shard tile (j, r, k) = row 2 i + r, K tile (2 j + p) 4 + k rs = [ttnn.to_torch(t).float().reshape(10, 12, 5, 2, 4, 32, 32) for t in h] hf = torch.full((M, N), float("nan")) bad = 0 for y_ in range(10): for x_ in range(12): _, i_ = tk.xmm_place(x_, y_) for p_ in range(2): for j_ in range(5): for r_ in range(2): for k_ in range(4): m_, kk = 2 * i_ + r_, (2 * j_ + p_) * 4 + k_ t_ = rs[p_][y_, x_, j_, r_, k_] blk = hf[32 * m_:32 * m_ + 32, 32 * kk:32 * kk + 32] if not torch.isnan(blk).any() and not torch.equal(blk, t_): bad += 1 blk.copy_(t_) print(f"[ca] xres: copies disagreeing {bad}", flush=True) hf = hf.reshape(1, 1, M, N) sf = ttnn.to_torch(s).float() else: sf, hf = ttnn.to_torch(s).float(), ttnn.to_torch(h).float() print(f"[ca] first call OK s equal {bool(torch.equal(sr, sf))} max|ds| {float((sr - sf).abs().max()):.3e} " f"h equal {bool(torch.equal(hr, hf))} max|dh| {float((hr - hf).abs().max()):.3e} h PCC {pcc(hf, hr):.7f}", flush=True) if "noaff" in sys.argv: # vs torch LN (fp64) of the device residual s sd = sf.double() ht = torch.nn.functional.layer_norm(sd, (sd.shape[-1],), eps=1e-6) print(f"[ca] h vs torch LN(s): max|d| {float((hf.double() - ht).abs().max()):.3e} PCC {pcc(hf, ht.float()):.7f}", flush=True) if "time" in sys.argv: def sep(): d_ = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True) r = tk.add_layernorm(x1, d_, C, 1e-6, out_memory_config=hmem) ttnn.deallocate(d_) return r fus = lambda: tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True, addln=dict(x=x2, consts=Cf, eps=1e-6, out_mem=hmem)) print(f"[ca] separate {time_op(dev, sep, reps=5, iters=40):.2f} us fused {time_op(dev, fus, reps=5, iters=40):.2f} us", flush=True) finally: ttnn.close_device(dev)