File size: 2,204 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """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)
|