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