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