"""Round 11: fc2-style cmm (GELU helpers + fused add+LN epilogue) in the k2 K-split layout vs the round-8 "out" layout. usage: cmm_k2_test.py shape=s|fc2 [time] [soak=N] (expects s / h equal up to the summation order of the two K halves)""" 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", "s") GX, GY, NT, KT, kb_out, kb_k2 = {"s": (4, 6, 2, 32, 8, 4), "fc2": (10, 6, 4, 160, 16, 8)}[shape] dev = open_dev() try: torch.manual_seed(0) M, K, N = 32 * GY, 32 * KT, 32 * NT * GX 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) pre_t = (torch.randn(1, 1, M, K) * 1.5).to(torch.bfloat16) pre = ttnn.from_torch(pre_t, 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) 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) C = dict(tk.addln_constants(torch.randn(N) * 0.2 + 1, torch.zeros(N), dev), nobeta=True) defs = ("CMM_FLUSH", "CMM_LA=1", "CMM_NOSTAGE", "CMM_BIAS_LATE") + tuple(d for d in __import__("os").environ.get("K2_DEFS", "").split(",") if d) def amem(kb, mpc): return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.BLOCK_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec( ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(GX - 1, GY - 1))}), [32, 32 * kb * mpc], ttnn.ShardOrientation.ROW_MAJOR)) w_out = tk.CmmWeight(wt, dev, GX, NT, kb_out, ttnn.bfloat8_b) w_k2 = tk.CmmWeight(wt, dev, GX, NT, kb_k2, ttnn.bfloat8_b) X = {L: ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) for L in ("out", "k2")} # persistent (no host writes inside a trace) PF = {L: tk.CmmPrefetch(dev, GX, GY, w, int(arg("pf", "0")), single=True) if int(arg("pf", "0")) else None for L, w in (("out", w_out), ("k2", w_k2))} if PF["k2"] is not None: PF["k2"] = PF["k2"].for_k2() XPF = None if "xpf" in sys.argv and shape == "fc2": # round 11: the next qkv's XmmPrefetch filled by the epilogue + GELU helpers qkt = 40 wq_t = (torch.randn(32 * qkt, 4608) * 0.05).to(torch.bfloat16) wq_x = tk.CmmWeight(wq_t, dev, 36, 4, 4, ttnn.bfloat8_b) XPF = tk.XmmPrefetch(dev, wq_x, 9, 4) tl = wq_t.reshape(qkt, 32, 144, 32).permute(0, 2, 1, 3) exp = torch.cat([tl[b_ * 4:(b_ + 1) * 4, ng * 4:(ng + 1) * 4].reshape(16 * 32, 32) for (_, ng) in sorted(XPF.senders.items(), key=lambda kv: kv[1]) for b_ in range(4)], 0) expq = ttnn.to_torch(ttnn.from_torch(exp, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT)).float() def run(layout, use_pf=False, fill_pf=False): """use_pf: read the first blocks from the prefetch buffer; fill_pf: this call's add+LN reader fills it""" kb, mpc, wf = (kb_out, 1, w_out) if layout == "out" else (kb_k2, 2, w_k2) g = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, M, K]), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev, amem(kb, mpc)) s, h = tk.cmm_linear(g, wf, b, res, kb=kb, slots=3, defines=defs, v1=True, ss=3, gelu_from=pre, layout=layout, prefetch=PF[layout] if use_pf else None, gelu_pf=XPF, addln=dict(x=X[layout], consts=C, eps=1e-6, prefetch=PF[layout] if fill_pf else None, prefetch2=XPF)) ttnn.deallocate(g) return s, h print(f"[k2] shape={shape} G={GX}x{GY} NT={NT} KT={KT} kb out/k2 {kb_out}/{kb_k2}", flush=True) s0, h0 = run("out") s1, h1 = run("k2") ttnn.synchronize_device(dev) if XPF is not None: got = ttnn.to_torch(XPF.buf).float() print(f"[k2] qkv prefetch buffer equal to the packed blocks: {bool(torch.equal(got, expq))} " f"(max|d| {float((got - expq).abs().max()):.3e})", flush=True) S0, H0, S1, H1 = (ttnn.to_torch(t).float() for t in (s0, h0, s1, h1)) print(f"[k2] first call OK s PCC {pcc(S1, S0):.7f} max|ds| {float((S1 - S0).abs().max()):.3e} eq {float((S1 == S0).float().mean()):.4f} " f"h PCC {pcc(H1, H0):.7f} max|dh| {float((H1 - H0).abs().max()):.3e} eq {float((H1 == H0).float().mean()):.4f}", flush=True) # fp32 torch reference of s (GELU exact erf, bfp8-quantized weight) wq = ttnn.to_torch(ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT)).float() act = torch.nn.functional.gelu(pre_t.float()).to(torch.bfloat16).float().reshape(M, K) sref = xt.float().reshape(M, N) + (act @ wq + ttnn.to_torch(b).float().reshape(1, N)) print(f"[k2] vs fp32 ref: s PCC out {pcc(S0, sref):.7f} k2 {pcc(S1, sref):.7f} mean|ds| out {float((S0.reshape(M, N) - sref).abs().mean()):.3e} " f"k2 {float((S1.reshape(M, N) - sref).abs().mean()):.3e} signed mean out {float((S0.reshape(M, N) - sref).mean()):.3e} " f"k2 {float((S1.reshape(M, N) - sref).mean()):.3e}", flush=True) if PF["k2"] is not None: for L in ("out", "k2"): run(L, fill_pf=True) sp, hp = run(L, use_pf=True) ttnn.synchronize_device(dev) ref = (S0, H0) if L == "out" else (S1, H1) print(f"[k2] {L} with prefetch pf={arg('pf', '0')}: s equal {bool(torch.equal(ttnn.to_torch(sp).float(), ref[0]))} " f"h equal {bool(torch.equal(ttnn.to_torch(hp).float(), ref[1]))}", flush=True) if "time" in sys.argv: for L in ("out", "k2"): print(f"[k2] {L} {time_op(dev, lambda: run(L)):.2f} us", flush=True) if PF[L] is not None: print(f"[k2] {L} with prefetch (fill + use per call) {time_op(dev, lambda: run(L, use_pf=True, fill_pf=True)):.2f} us", flush=True) if "prof" in sys.argv: # TT_METAL_DEVICE_PROFILER=1 K2_DEFS=CMM_PROF [ADDLN_EXP=prof,wprof]: zone timeline of k2 (+pf) L = arg("proflayout", "k2") for _ in range(4): if PF[L] is not None: run(L, fill_pf=True) o = run(L, use_pf=PF[L] is not None) ttnn.synchronize_device(dev) ttnn.ReadDeviceProfiler(dev) n = int(arg("soak", "0")) if n: reps = 20 outs = [] tid = ttnn.begin_trace_capture(dev, cq_id=0) try: for _ in range(reps): outs.append(run("k2")) 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(s).float(), S1) and torch.equal(ttnn.to_torch(h).float(), H1))) for s, h in outs) print(f"[k2] soak {n // reps * reps} replays, outputs differing from the first k2 call: {bad}/{reps}", flush=True) finally: ttnn.release_trace(dev, tid) finally: ttnn.close_device(dev) if "prof" in sys.argv: tl_report(risks=("NCRISC", "BRISC", "TRISC_0", "TRISC_1", "TRISC_2"))