File size: 7,608 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """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"))
|