| """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) |
|
|
