File size: 1,810 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 | """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) # tile rows of 16 inside a face: TILE layout handled by comparing in torch layout
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() # host conversion
xf = xt.float()
# host reference: groups of 16 consecutive columns per row (face rows)
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)
|