File size: 5,166 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Block-tail emulation: addln1(prefetch proj) -> proj = cmm(SDPA-layout ctx, head_concat folded) -> addln2(prefetch fc2)
-> gelu -> fc2 = cmm(prefetched pairs) vs the model's library path (head_concat + 2D linear, gelu + 2D linear).
usage: cmm_block_test.py [pf2=P] [time] [soak=N]"""
import sys
from pathlib import Path
from common import *
from open_device_12x10 import open_device
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
from hamer.vit import tuned_matmul_configs
arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d)
P2 = int(arg("pf2", "3")); S2 = int(arg("s2", "3")); S1 = int(arg("s1", "3"))
dev = open_device(grid="12x10", dispatch="eth", trace_region_size=64 << 20)
try:
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
res = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK,
orientation=ttnn.ShardOrientation.ROW_MAJOR)
T = lambda t, mem=L1, dt=ttnn.bfloat16: ttnn.from_torch(t, dtype=dt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mem)
x = T(torch.randn(1, 1, 192, 1280), res); y = T(torch.randn(1, 1, 192, 1280), res)
ctx_t = torch.randn(1, 16, 192, 96) * 0.5
ctx = T(ctx_t)
h1 = T(torch.randn(1, 1, 192, 5120) * 1.5)
pw_t = (torch.randn(1536, 1280) * 1536 ** -0.5).bfloat16(); pb_t = (torch.randn(1280) * 0.1).bfloat16()
fw_t = (torch.randn(5120, 1280) * 5120 ** -0.5).bfloat16(); fb_t = (torch.randn(1280) * 0.1).bfloat16()
pw = T(pw_t, ttnn.DRAM_MEMORY_CONFIG, ttnn.bfloat8_b); fw = T(fw_t, ttnn.DRAM_MEMORY_CONFIG, ttnn.bfloat8_b)
pb = T(pb_t.reshape(1, 1, 1, -1)); fb = T(fb_t.reshape(1, 1, 1, -1))
C1 = tk.addln_constants(torch.ones(1280), torch.zeros(1280), dev); C2 = tk.addln_constants(torch.ones(1280) * 0.9, torch.zeros(1280), dev)
pwk = tk.CmmWeight(pw_t, dev, 10, 4, 8, ttnn.bfloat8_b); fwk = tk.CmmWeight(fw_t, dev, 10, 4, 16, ttnn.bfloat8_b)
V1 = "v1" in sys.argv
P1 = int(arg("pf1", "3"))
pfp = tk.CmmPrefetch(dev, 10, 6, pwk, P1, single=V1) if P1 else None
pff = tk.CmmPrefetch(dev, 10, 6, fwk, P2, single=V1) if P2 else None
defs = ("CMM_FLUSH", f"CMM_LA={arg('la', '1')}") + (() if V1 else ("CMM_DRAM0",))
ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True)
pcs = tuned_matmul_configs(dev.compute_with_storage_grid_size(), fp32_acc=True)
def cmm_path():
s, h = tk.add_layernorm(x, y, C1, 1e-6, out_memory_config=L1, prefetch=pfp)
ttnn.deallocate(h)
d = tk.cmm_linear(ctx, pwk, pb, res, kb=8, slots=S1, defines=defs, prefetch=pfp, a_heads=(16, 96), v1=V1)
s2, h2 = tk.add_layernorm(s, d, C2, 1e-6, out_memory_config=L1, prefetch=pff)
ttnn.deallocate(h2)
g = tk.gelu(h1, memory_config=L1)
d2 = tk.cmm_linear(g, fwk, fb, res, kb=16, slots=S2, defines=defs, prefetch=pff, v1=V1)
for t in (s, d, g):
ttnn.deallocate(t)
return s2, d2
def lib_path():
s, h = tk.add_layernorm(x, y, C1, 1e-6, out_memory_config=L1)
ttnn.deallocate(h)
cm = tk.head_concat(ctx, memory_config=L1)
d = ttnn.linear(cm, pw, bias=pb, program_config=pcs["proj"], memory_config=res, compute_kernel_config=ck)
s2, h2 = tk.add_layernorm(s, d, C2, 1e-6, out_memory_config=L1)
ttnn.deallocate(h2)
g = tk.gelu(h1, memory_config=L1)
d2 = ttnn.linear(g, fw, bias=fb, program_config=pcs["fc2"], memory_config=res, compute_kernel_config=ck)
for t in (s, cm, d, g):
ttnn.deallocate(t)
return s2, d2
a2, ad = cmm_path(); ttnn.synchronize_device(dev)
b2, bd = lib_path(); ttnn.synchronize_device(dev)
A2, AD, B2, BD = (ttnn.to_torch(t).float() for t in (a2, ad, b2, bd))
print(f"[blk] first calls OK s2 pcc {pcc(A2, B2):.7f} max|d| {float((A2 - B2).abs().max()):.4f} fc2 pcc {pcc(AD, BD):.7f} max|d| {float((AD - BD).abs().max()):.4f}", flush=True)
if "time" in sys.argv:
print(f"[blk] cmm path {time_op(dev, cmm_path):.2f} us library path {time_op(dev, lib_path):.2f} us", flush=True)
n = int(arg("soak", "0"))
if n:
reps = 20
tid = ttnn.begin_trace_capture(dev, cq_id=0)
outs = []
try:
for _ in range(reps):
outs.append(cmm_path())
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(p).float(), A2) and torch.equal(ttnn.to_torch(q).float(), AD))) for p, q in outs)
print(f"[blk] soak {n // reps * reps} block-tails OK, outputs differing from the first: {bad}/{reps}", flush=True)
finally:
ttnn.release_trace(dev, tid)
finally:
ttnn.close_device(dev)
|