| """Round 11: proj-shaped cmm (+ add+LN epilogue) whose idle cores and epilogue prefetch a qkv XmmPrefetch; checks the |
| buffer contents. usage: idle_pf_test.py""" |
| 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 |
| dev = open_dev() |
| try: |
| torch.manual_seed(0) |
| GX, GY, NT, KT, kb = 10, 6, 4, 48, 8 |
| 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) |
| a = ttnn.from_torch(torch.randn(1, 1, M, K), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) |
| wf = tk.CmmWeight((torch.randn(K, N) * K ** -0.5).to(torch.bfloat16), dev, GX, NT, kb, ttnn.bfloat8_b) |
| b = ttnn.from_torch(torch.zeros(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| x = ttnn.from_torch(torch.randn(1, 1, M, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) |
| C = dict(tk.addln_constants(torch.ones(N), torch.zeros(N), dev), nobeta=True) |
| wq_t = (torch.randn(1280, 4608) * 0.05).to(torch.bfloat16) |
| XPF = tk.XmmPrefetch(dev, tk.CmmWeight(wq_t, dev, 36, 4, 4, ttnn.bfloat8_b), 9, 2) |
| tl = wq_t.reshape(40, 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(2)], 0) |
| expq = ttnn.to_torch(ttnn.from_torch(exp, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT)).float() |
| s, h = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=("CMM_FLUSH", "CMM_LA=1", "CMM_NOSTAGE", "CMM_BIAS_LATE"), |
| v1=True, ss=3, addln=dict(x=x, consts=C, eps=1e-6), idle_pf=XPF) |
| ttnn.synchronize_device(dev) |
| got = ttnn.to_torch(XPF.buf).float() |
| print(f"[ipf] first call OK, prefetch buffer equal: {bool(torch.equal(got, expq))} max|d| {float((got - expq).abs().max()):.3e}", flush=True) |
| finally: |
| ttnn.close_device(dev) |
|
|
