File size: 1,754 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 | """add_layernorm on the MANO-head token layout: [1,1,32,1024] WIDTH-sharded on 8 cores (one tile row)."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
dev = open_dev()
try:
torch.manual_seed(0)
cg = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(7, 0))})
wm = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1,
ttnn.ShardSpec(cg, [32, 128], ttnn.ShardOrientation.ROW_MAJOR))
xt = torch.randn(1, 1, 32, 1024) * 2; yt = torch.randn(1, 1, 32, 1024)
g = torch.randn(1024) * 0.2 + 1; b = torch.randn(1024) * 0.1
x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wm)
y = ttnn.from_torch(yt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wm)
C = tk.addln_constants(g, b, dev)
s, h = tk.add_layernorm(x, y, C, 1e-5, out_memory_config=wm)
ttnn.synchronize_device(dev)
sr = xt.bfloat16().float() + yt.bfloat16().float()
hr = torch.nn.functional.layer_norm(sr.bfloat16().float(), (1024,), g.bfloat16().float(), b.bfloat16().float(), 1e-5)
st, ht = ttnn.to_torch(s).float(), ttnn.to_torch(h).float()
print(f"[ahead] first call OK s pcc {pcc(st, sr):.6f} h pcc {pcc(ht, hr):.6f} max|dh| {float((ht - hr).abs().max()):.4f}", flush=True)
ln = lambda: ttnn.layer_norm(ttnn.add(x, y, memory_config=wm), weight=None, bias=None, epsilon=1e-5, memory_config=wm)
if "time" in sys.argv:
print(f"[ahead] addln {time_op(dev, lambda: tk.add_layernorm(x, y, C, 1e-5, out_memory_config=wm)):.2f} us", flush=True)
finally:
ttnn.close_device(dev)
|