| """How does the packer round fp32 DST -> bfp8_b (bfp8_pack_precise)? Copy-only GELU kernel (HG_NOP) with a bfp8 output |
| vs host round-to-nearest / truncation of the same bf16 values (round 8).""" |
| import sys |
| from pathlib import Path |
| from common import * |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from hamer import tt_kernels |
|
|
| def bfp8(x, mode): |
| s = x.shape |
| g = x.reshape(-1, 16) |
| return g, s |
|
|
| dev = open_dev() |
| try: |
| torch.manual_seed(0) |
| xt = (torch.randn(1, 1, 64, 64) * torch.exp(torch.randn(1, 1, 64, 1))).bfloat16() |
| x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) |
| o = ttnn.to_torch(tt_kernels.gelu(x, defines=("HG_NOP",), out_dtype=ttnn.bfloat8_b)).float() |
| h = ttnn.to_torch(ttnn.from_torch(xt.float(), dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT)).float() |
| xf = xt.float() |
| |
| g = xf.reshape(64, 4, 16) |
| mx = g.abs().amax(-1, keepdim=True).clamp_min(1e-30) |
| step = torch.pow(2.0, torch.floor(torch.log2(mx)) - 6) |
| rn = (torch.round(g / step).clamp(-127, 127) * step).reshape(1, 1, 64, 64) |
| tr = (torch.trunc(g / step).clamp(-127, 127) * step).reshape(1, 1, 64, 64) |
| for name, ref in (("round-nearest", rn), ("truncate", tr), ("host ttnn.from_torch", h)): |
| print(f"[pack] device vs {name:22s}: equal {float((o == ref).float().mean()):.4f} max|d| {float((o - ref).abs().max()):.3e}", flush=True) |
| print(f"[pack] mean signed error device-x {float((o - xf).mean()):+.3e} rn-x {float((rn - xf).mean()):+.3e} tr-x {float((tr - xf).mean()):+.3e}") |
| finally: |
| ttnn.close_device(dev) |
|
|
