# SPDX-License-Identifier: Apache-2.0 """mb2 part B (8) — bfp8 / math-fidelity / fp32-accumulate fidelity of a REAL N1.6 DiT block on the golden input. 01-architecture-synthesis §6 item 3 ("bfp8 / LoFi accuracy on device ... device math fidelity was not emulated"). * Weights: ``action_head.model.transformer_blocks.{0,1}.*`` of ``nvidia/GR00T-N1.6-3B`` @ d0814e7e, read with ``safetensors.safe_open`` from the shard named by ``model.safetensors.index.json`` (weight-loader-spec §1.1). Block 0 = cross-attention to the 35 NON-image tokens (``AlternateVLDiT.forward``, ``dit.py:365-387``: ``non_image_attention_mask = ~image_mask & backbone_attention_mask`` for ``idx % 4 == 0``); block 1 = self-attention. * Golden (``golden/n16/fp32_math/tensors``, golden-tap-map.json ``dit_block0_out`` / ``dit_block1_out`` k=0): input ``step0_dit_in_sa_embs`` [1,51,1536], ``step0_dit_temb`` [1,1536], ``dit_in_encoder_hidden_states`` [1,116,2048] (= ``ah_vlln_out``), ``dit_in_{image_mask,backbone_attention_mask}``; outputs ``step0_dit_block_00_out``, ``step0_dit_block_01_out`` (block 1 is fed the golden block-0 output so each block's PCC is its own error); intermediate taps ``adanorm1_out / crossattn_out / norm3_out / ff_out / selfattn_out`` pin the reference semantics. * Reference semantics (verified on CPU first, ``--cpu-only``): AdaLN = ``LN(x, eps 1e-5, no affine) * (1 + scale) + shift`` with ``scale, shift = chunk(norm1.linear(SiLU(temb)))`` (weight-loader-spec §4.6, scale first); ``attn1`` with q/k/v/out biases, 32 heads x 48, torch SDPA default scale 48^-0.5, bool key mask; ``norm3`` non-affine LN; FF = ``Linear(1536->6144) -> GELU(tanh) -> Linear(6144->1536)`` (diffusers ``FeedForward("gelu-approximate")``); ungated residuals. * Device path: heads padded 48->64 (Q/K/V output columns + biases, ``to_out`` input rows; static-shape-plan §7), ``scale=48**-0.5`` passed explicitly, M padded 51->64 with zero rows, cross K/V hoisted (116 -> 128 rows, zero pad) with an additive key mask [1,1,64,128] (-1e9 on image / pad keys); self block: fused QKV linear -> ``nlp_create_qkv_heads(transpose_k_heads=False)`` -> SDPA with pad-key mask. 12 ttnn ops per block (+ 4 hoisting ops for the cross K/V, outside the per-step path). * Sweep: weight dtype {bf16, bfp8_b} x math fidelity {HiFi4, HiFi2, LoFi} x fp32_dest_acc_en {on, off} for the six projection matmuls; extra variants: hoisted K/V in bfp8, FF weights in bfp4, SDPA fp32 accumulate. LN always HiFi4/fp32acc (layernorm default), SDPA HiFi2/no fp32acc (tt_dit convention) unless stated. Reported: PCC / max|Δ| vs the fp32 golden block output on the 51 valid rows, plus the official bf16 run's PCC (golden-tap-map thresholds: 0.9999948 / 0.9999924) for context. """ from __future__ import annotations import argparse import json import os import sys import time from typing import Any, Dict, List, Optional import torch import torch.nn.functional as F sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import mb2_common as C # noqa: E402 SNAP = "/home/deepgadget/.cache/huggingface/hub/models--nvidia--GR00T-N1.6-3B/snapshots/d0814e7ecb19202e7c8468b46098b0b7ef3a6d61" GOLD = "/home/deepgadget/experiments/gr00t/golden/n16/fp32_math/tensors" GOLD_BF16 = "/home/deepgadget/experiments/gr00t/golden/n16/bf16_sdpa/tensors" D = 1536 NH = 32 HD = 48 HDP = 64 FF = 6144 ENC = 2048 M_LOG = 51 M_PAD = 64 L_LOG = 116 L_PAD = 128 EPS = 1e-5 NEG = -1e9 # -------------------------------------------------------------------------------------------------------------------- # checkpoint + golden loading # -------------------------------------------------------------------------------------------------------------------- def load_block_weights(block: int) -> Dict[str, torch.Tensor]: """Read the 14 tensors of one DiT block (fp32 copies) from the right shard; raise on any missing key.""" from safetensors import safe_open with open(os.path.join(SNAP, "model.safetensors.index.json")) as f: wm = json.load(f)["weight_map"] prefix = f"action_head.model.transformer_blocks.{block}." names = [ "attn1.to_q.weight", "attn1.to_q.bias", "attn1.to_k.weight", "attn1.to_k.bias", "attn1.to_v.weight", "attn1.to_v.bias", "attn1.to_out.0.weight", "attn1.to_out.0.bias", "ff.net.0.proj.weight", "ff.net.0.proj.bias", "ff.net.2.weight", "ff.net.2.bias", "norm1.linear.weight", "norm1.linear.bias", ] # fmt: skip out: Dict[str, torch.Tensor] = {} by_shard: Dict[str, List[str]] = {} for n in names: key = prefix + n if key not in wm: raise KeyError(f"missing checkpoint key {key}") by_shard.setdefault(wm[key], []).append(n) for shard, ns in by_shard.items(): with safe_open(os.path.join(SNAP, shard), framework="pt", device="cpu") as f: for n in ns: out[n] = f.get_tensor(prefix + n).float() exp = { "attn1.to_q.weight": (D, D), "attn1.to_k.weight": (D, ENC if block % 2 == 0 else D), "attn1.to_v.weight": (D, ENC if block % 2 == 0 else D), "attn1.to_out.0.weight": (D, D), "ff.net.0.proj.weight": (FF, D), "ff.net.2.weight": (D, FF), "norm1.linear.weight": (2 * D, D), } for n, s in exp.items(): if tuple(out[n].shape) != s: raise ValueError(f"{prefix}{n}: shape {tuple(out[n].shape)} != expected {s}") return out def load_golden() -> Dict[str, torch.Tensor]: names = { "x0": "step0_dit_in_sa_embs", "temb": "step0_dit_temb", "enc": "dit_in_encoder_hidden_states", "valid": "dit_in_backbone_attention_mask", "image": "dit_in_image_mask", "b0_adanorm1": "step0_dit_block_00_adanorm1_out", "b0_crossattn": "step0_dit_block_00_crossattn_out", "b0_norm3": "step0_dit_block_00_norm3_out", "b0_ff": "step0_dit_block_00_ff_out", "b0_out": "step0_dit_block_00_out", "b1_selfattn": "step0_dit_block_01_selfattn_out", "b1_out": "step0_dit_block_01_out", } g = {k: torch.load(os.path.join(GOLD, v + ".pt"), map_location="cpu") for k, v in names.items()} for k in ("b0_out", "b1_out"): p = os.path.join(GOLD_BF16, names[k] + ".pt") if os.path.exists(p): g[k + "_official_bf16"] = torch.load(p, map_location="cpu").float() if tuple(g["x0"].shape) != (1, M_LOG, D) or tuple(g["enc"].shape) != (1, L_LOG, ENC): raise ValueError(f"unexpected golden shapes x0 {tuple(g['x0'].shape)} enc {tuple(g['enc'].shape)}") return g # -------------------------------------------------------------------------------------------------------------------- # torch fp32 reference (dit.py BasicTransformerBlock.forward + AdaLayerNorm + diffusers Attention / FeedForward) # -------------------------------------------------------------------------------------------------------------------- def adaln_vectors(w: Dict[str, torch.Tensor], temb: torch.Tensor): ada = F.silu(temb.double()) @ w["norm1.linear.weight"].double().T + w["norm1.linear.bias"].double() scale, shift = ada.chunk(2, dim=1) # scale FIRST (weight-loader-spec §4.6, verified) return (1.0 + scale).float(), shift.float() def ref_block( x: torch.Tensor, w: Dict[str, torch.Tensor], temb: torch.Tensor, enc: Optional[torch.Tensor], key_valid: Optional[torch.Tensor], ) -> Dict[str, torch.Tensor]: """fp64 reference of one block; returns intermediates named like the golden taps.""" x = x.double() one_plus_scale, shift = adaln_vectors(w, temb) h = F.layer_norm(x, (D,), eps=EPS) * one_plus_scale.double()[:, None] + shift.double()[:, None] ctx = h if enc is None else enc.double() q = h @ w["attn1.to_q.weight"].double().T + w["attn1.to_q.bias"].double() k = ctx @ w["attn1.to_k.weight"].double().T + w["attn1.to_k.bias"].double() v = ctx @ w["attn1.to_v.weight"].double().T + w["attn1.to_v.bias"].double() B, S, _ = q.shape Lk = k.shape[1] q = q.reshape(B, S, NH, HD).permute(0, 2, 1, 3) k = k.reshape(B, Lk, NH, HD).permute(0, 2, 1, 3) v = v.reshape(B, Lk, NH, HD).permute(0, 2, 1, 3) mask = None if key_valid is not None: mask = key_valid.reshape(1, 1, 1, Lk).bool() a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) # default scale = 48^-0.5 a = a.permute(0, 2, 1, 3).reshape(B, S, D) attn = a @ w["attn1.to_out.0.weight"].double().T + w["attn1.to_out.0.bias"].double() x = x + attn h2 = F.layer_norm(x, (D,), eps=EPS) f = F.gelu(h2 @ w["ff.net.0.proj.weight"].double().T + w["ff.net.0.proj.bias"].double(), approximate="tanh") ff = f @ w["ff.net.2.weight"].double().T + w["ff.net.2.bias"].double() x = x + ff return {"adanorm1": h.float(), "attn": attn.float(), "norm3": h2.float(), "ff": ff.float(), "out": x.float()} def cpu_reference_check( g: Dict[str, torch.Tensor], w0: Dict[str, torch.Tensor], w1: Dict[str, torch.Tensor] ) -> Dict[str, Any]: key_valid_b0 = (~g["image"]) & g["valid"] # block 0: non-image tokens (idx % 4 == 0) r0 = ref_block(g["x0"], w0, g["temb"], g["enc"], key_valid_b0) r1 = ref_block(g["b0_out"], w1, g["temb"], None, None) res = { "block0_num_valid_keys": int(key_valid_b0.sum()), "block0_adanorm1_vs_golden": C.compare(g["b0_adanorm1"], r0["adanorm1"]), "block0_crossattn_vs_golden": C.compare(g["b0_crossattn"], r0["attn"]), "block0_norm3_vs_golden": C.compare(g["b0_norm3"], r0["norm3"]), "block0_ff_vs_golden": C.compare(g["b0_ff"], r0["ff"]), "block0_out_vs_golden": C.compare(g["b0_out"], r0["out"]), "block1_selfattn_vs_golden": C.compare(g["b1_selfattn"], r1["attn"]), "block1_out_vs_golden": C.compare(g["b1_out"], r1["out"]), } # negative controls that must be worse (pin the mask semantics) r0_all = ref_block(g["x0"], w0, g["temb"], g["enc"], None) r0_img = ref_block(g["x0"], w0, g["temb"], g["enc"], g["image"] & g["valid"]) res["control_block0_all_keys_pcc"] = C.pcc(g["b0_out"], r0_all["out"]) res["control_block0_image_keys_pcc"] = C.pcc(g["b0_out"], r0_img["out"]) for k in ("b0_out_official_bf16", "b1_out_official_bf16"): if k in g: res[f"{k}_vs_fp32_golden"] = C.compare(g[k[:6]], g[k]) return res # -------------------------------------------------------------------------------------------------------------------- # device block # -------------------------------------------------------------------------------------------------------------------- def pad_heads_out(wt: torch.Tensor, bias: torch.Tensor): """Linear weight [out=NH*HD, in] -> [in, NH*HDP] (x @ W layout) with each head's 48 columns zero-padded to 64.""" out_dim, in_dim = wt.shape w = wt.reshape(NH, HD, in_dim) w = F.pad(w, (0, 0, 0, HDP - HD)).reshape(NH * HDP, in_dim) b = F.pad(bias.reshape(NH, HD), (0, HDP - HD)).reshape(1, 1, 1, NH * HDP) return w.T.contiguous(), b def pad_heads_in(wt: torch.Tensor, bias: torch.Tensor): """to_out weight [out=1536, in=NH*HD] -> [NH*HDP, 1536] with zero rows for the padded head dims.""" out_dim, in_dim = wt.shape w = wt.reshape(out_dim, NH, HD) w = F.pad(w, (0, HDP - HD)).reshape(out_dim, NH * HDP) return w.T.contiguous(), bias.reshape(1, 1, 1, out_dim) class DeviceBlock: def __init__( self, ttnn, device, w: Dict[str, torch.Tensor], cross: bool, temb: torch.Tensor, variant: Dict[str, Any] ): self.ttnn = ttnn self.device = device self.cross = cross self.variant = variant wd = C.tt_dtype(ttnn, variant["w_dtype"]) ffd = C.tt_dtype(ttnn, variant.get("ff_dtype", variant["w_dtype"])) self.ck_mm = C.compute_config( ttnn, variant["fidelity"], fp32_acc=variant["fp32_acc"], packer_l1_acc=variant.get("packer_l1_acc", False) ) self.ck_ln = C.compute_config(ttnn, "HiFi4", fp32_acc=True) self.ck_sdpa = C.compute_config( ttnn, variant.get("sdpa_fidelity", "HiFi2"), fp32_acc=variant.get("sdpa_fp32_acc", False) ) grid = device.compute_with_storage_grid_size() self.sdpa_cfg = ttnn.SDPAProgramConfig( compute_with_storage_grid_size=grid, q_chunk_size=64, k_chunk_size=64, exp_approx_mode=False ) one_plus_scale, shift = adaln_vectors(w, temb) self.ln_w = C.to_tt(ttnn, device, one_plus_scale.to(torch.bfloat16)) self.ln_b = C.to_tt(ttnn, device, shift.to(torch.bfloat16)) wq, bq = pad_heads_out(w["attn1.to_q.weight"], w["attn1.to_q.bias"]) wk, bk = pad_heads_out(w["attn1.to_k.weight"], w["attn1.to_k.bias"]) wv, bv = pad_heads_out(w["attn1.to_v.weight"], w["attn1.to_v.bias"]) if cross: self.w_q = C.to_tt(ttnn, device, wq.to(torch.bfloat16), dtype=wd) self.b_q = C.to_tt(ttnn, device, bq.to(torch.bfloat16)) self.w_k = C.to_tt(ttnn, device, wk.to(torch.bfloat16), dtype=wd) self.b_k = C.to_tt(ttnn, device, bk.to(torch.bfloat16)) self.w_v = C.to_tt(ttnn, device, wv.to(torch.bfloat16), dtype=wd) self.b_v = C.to_tt(ttnn, device, bv.to(torch.bfloat16)) else: self.w_qkv = C.to_tt(ttnn, device, torch.cat([wq, wk, wv], dim=1).to(torch.bfloat16), dtype=wd) self.b_qkv = C.to_tt(ttnn, device, torch.cat([bq, bk, bv], dim=-1).to(torch.bfloat16)) wo, bo = pad_heads_in(w["attn1.to_out.0.weight"], w["attn1.to_out.0.bias"]) self.w_o = C.to_tt(ttnn, device, wo.to(torch.bfloat16), dtype=wd) self.b_o = C.to_tt(ttnn, device, bo.to(torch.bfloat16)) self.w_1 = C.to_tt(ttnn, device, w["ff.net.0.proj.weight"].T.contiguous().to(torch.bfloat16), dtype=ffd) self.b_1 = C.to_tt(ttnn, device, w["ff.net.0.proj.bias"].reshape(1, 1, 1, FF).to(torch.bfloat16)) self.w_2 = C.to_tt(ttnn, device, w["ff.net.2.weight"].T.contiguous().to(torch.bfloat16), dtype=ffd) self.b_2 = C.to_tt(ttnn, device, w["ff.net.2.bias"].reshape(1, 1, 1, D).to(torch.bfloat16)) self.k_h = None self.v_h = None def split_heads(self, flat, S: int): ttnn = self.ttnn r = ttnn.reshape(flat, (1, S, NH, HDP)) return ttnn.permute(r, (0, 2, 1, 3)) def hoist_kv(self, enc_dev): """K/V projections of the VL tokens (once per inference in the real model).""" ttnn = self.ttnn k = ttnn.linear(enc_dev, self.w_k, bias=self.b_k, compute_kernel_config=self.ck_mm) v = ttnn.linear(enc_dev, self.w_v, bias=self.b_v, compute_kernel_config=self.ck_mm) k = self.split_heads(k, L_PAD) v = self.split_heads(v, L_PAD) kvd = self.variant.get("kv_dtype") if kvd is not None and kvd != "bf16": k = ttnn.typecast(k, C.tt_dtype(ttnn, kvd)) v = ttnn.typecast(v, C.tt_dtype(ttnn, kvd)) self.k_h, self.v_h = k, v def forward(self, x, mask): ttnn = self.ttnn h = ttnn.layer_norm(x, epsilon=EPS, weight=self.ln_w, bias=self.ln_b, compute_kernel_config=self.ck_ln) if self.cross: q = ttnn.linear(h, self.w_q, bias=self.b_q, compute_kernel_config=self.ck_mm) q = self.split_heads(q, M_PAD) k, v = self.k_h, self.v_h else: qkv = ttnn.linear(h, self.w_qkv, bias=self.b_qkv, compute_kernel_config=self.ck_mm) q, k, v = ttnn.experimental.nlp_create_qkv_heads( qkv, num_heads=NH, num_kv_heads=NH, transpose_k_heads=False ) a = ttnn.transformer.scaled_dot_product_attention( q, k, v, attn_mask=mask, is_causal=False, scale=HD**-0.5, program_config=self.sdpa_cfg, compute_kernel_config=self.ck_sdpa, ) a = ttnn.experimental.nlp_concat_heads(a) o = ttnn.linear(a, self.w_o, bias=self.b_o, compute_kernel_config=self.ck_mm) x = ttnn.add(x, o) h2 = ttnn.layer_norm(x, epsilon=EPS, compute_kernel_config=self.ck_ln) f = ttnn.linear(h2, self.w_1, bias=self.b_1, compute_kernel_config=self.ck_mm) f = ttnn.gelu(f, variant=ttnn.GeluVariant.Tanh) f = ttnn.linear(f, self.w_2, bias=self.b_2, compute_kernel_config=self.ck_mm) x2 = ttnn.add(x, f) return {"adanorm1": h, "attn": o, "ff": f, "out": x2} def deallocate(self): for t in list(vars(self).values()): if hasattr(t, "buffer_address") or type(t).__name__ == "Tensor": try: self.ttnn.deallocate(t) except Exception: # noqa: BLE001 pass VARIANTS: List[Dict[str, Any]] = [ dict(name="bf16_HiFi4_fp32acc", w_dtype="bf16", fidelity="HiFi4", fp32_acc=True), dict(name="bf16_HiFi2_fp32acc", w_dtype="bf16", fidelity="HiFi2", fp32_acc=True), dict(name="bf16_HiFi2_nofp32acc", w_dtype="bf16", fidelity="HiFi2", fp32_acc=False), dict(name="bf16_LoFi_fp32acc", w_dtype="bf16", fidelity="LoFi", fp32_acc=True), dict(name="bf16_LoFi_nofp32acc", w_dtype="bf16", fidelity="LoFi", fp32_acc=False), dict(name="bfp8_HiFi4_fp32acc", w_dtype="bfp8", fidelity="HiFi4", fp32_acc=True), dict(name="bfp8_HiFi2_fp32acc", w_dtype="bfp8", fidelity="HiFi2", fp32_acc=True), dict(name="bfp8_HiFi2_nofp32acc", w_dtype="bfp8", fidelity="HiFi2", fp32_acc=False), dict(name="bfp8_LoFi_fp32acc", w_dtype="bfp8", fidelity="LoFi", fp32_acc=True), dict(name="bfp8_LoFi_nofp32acc", w_dtype="bfp8", fidelity="LoFi", fp32_acc=False), dict(name="bfp8_HiFi2_fp32acc_kv_bfp8", w_dtype="bfp8", fidelity="HiFi2", fp32_acc=True, kv_dtype="bfp8"), dict(name="bfp8_HiFi2_fp32acc_ff_bfp4", w_dtype="bfp8", fidelity="HiFi2", fp32_acc=True, ff_dtype="bfp4"), dict(name="bf16_HiFi2_fp32acc_sdpa_fp32acc", w_dtype="bf16", fidelity="HiFi2", fp32_acc=True, sdpa_fp32_acc=True), dict(name="bf16_HiFi2_fp32acc_sdpa_HiFi4", w_dtype="bf16", fidelity="HiFi2", fp32_acc=True, sdpa_fidelity="HiFi4"), dict(name="bf16_HiFi2_fp32acc_packer_l1", w_dtype="bf16", fidelity="HiFi2", fp32_acc=True, packer_l1_acc=True), ] def device_sweep( ttnn, device, rep: C.Report, g: Dict[str, torch.Tensor], w0, w1, variants: List[Dict[str, Any]] ) -> None: bf = torch.bfloat16 x0_pad = F.pad(g["x0"], (0, 0, 0, M_PAD - M_LOG)).reshape(1, 1, M_PAD, D) x1_pad = F.pad(g["b0_out"], (0, 0, 0, M_PAD - M_LOG)).reshape(1, 1, M_PAD, D) enc_pad = F.pad(g["enc"], (0, 0, 0, L_PAD - L_LOG)).reshape(1, 1, L_PAD, ENC) key_valid_b0 = ((~g["image"]) & g["valid"]).reshape(-1) m0 = torch.full((1, 1, M_PAD, L_PAD), NEG) m0[..., :L_LOG][..., key_valid_b0] = 0.0 m1 = torch.zeros(1, 1, M_PAD, M_PAD) m1[..., M_LOG:] = NEG x0_dev = C.to_tt(ttnn, device, x0_pad.to(bf)) x1_dev = C.to_tt(ttnn, device, x1_pad.to(bf)) enc_dev = C.to_tt(ttnn, device, enc_pad.to(bf)) m0_dev = C.to_tt(ttnn, device, m0.to(bf)) m1_dev = C.to_tt(ttnn, device, m1.to(bf)) rep.results["input_bf16_rounding_pcc_x0"] = C.pcc(g["x0"], C.bf16_round(g["x0"])) rep.results["input_bf16_rounding_pcc_enc"] = C.pcc(g["enc"], C.bf16_round(g["enc"])) # reference of the block computed from the bf16-rounded inputs (what the device actually sees) ref0_b = ref_block(C.bf16_round(g["x0"]), w0, g["temb"], C.bf16_round(g["enc"]), (~g["image"]) & g["valid"])["out"] ref1_b = ref_block(C.bf16_round(g["b0_out"]), w1, g["temb"], None, None)["out"] rep.results["fp32_ref_on_bf16_inputs_vs_golden"] = { "block0_pcc": C.pcc(g["b0_out"], ref0_b), "block1_pcc": C.pcc(g["b1_out"], ref1_b), } for var in variants: def run(var=var): res: Dict[str, Any] = {"variant": var} for bidx, (w, cross, x_dev, mask, gold_out, gold_attn, gold_ff) in enumerate( ( (w0, True, x0_dev, m0_dev, g["b0_out"], g["b0_crossattn"], g["b0_ff"]), (w1, False, x1_dev, m1_dev, g["b1_out"], g["b1_selfattn"], None), ) ): blk = DeviceBlock(ttnn, device, w, cross, g["temb"], var) if cross: blk.hoist_kv(enc_dev) outs = blk.forward(x_dev, mask) ttnn.synchronize_device(device) o = C.to_torch(ttnn, outs["out"]).reshape(M_PAD, D)[:M_LOG].reshape(1, M_LOG, D) a = C.to_torch(ttnn, outs["attn"]).reshape(M_PAD, D)[:M_LOG].reshape(1, M_LOG, D) r: Dict[str, Any] = { "out_vs_fp32_golden": C.compare(gold_out, o), "attn_out_vs_fp32_golden": C.compare(gold_attn, a), } if gold_ff is not None: f = C.to_torch(ttnn, outs["ff"]).reshape(M_PAD, D)[:M_LOG].reshape(1, M_LOG, D) r["ff_out_vs_fp32_golden"] = C.compare(gold_ff, f) key = f"b{bidx}_out_official_bf16" if key in g: r["out_vs_official_bf16_run_pcc"] = C.pcc(g[key], o) pad_rows = C.to_torch(ttnn, outs["out"]).reshape(M_PAD, D)[M_LOG:] r["pad_rows_finite"] = bool(torch.isfinite(pad_rows).all()) r["eager_block_us"] = C.time_calls(ttnn, device, lambda: blk.forward(x_dev, mask), reps=5)["min_us"] res[f"block{bidx}_{'cross' if cross else 'self'}"] = r blk.deallocate() return res rep.run(f"variant_{var['name']}", run) def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--out", type=str, default=C.default_out("mb2_bfp8_dit_block")) ap.add_argument("--cpu-only", action="store_true", help="only the torch reference vs golden check (no device)") ap.add_argument("--variants", type=str, default="all") args = ap.parse_args() rep = C.Report("mb2_bfp8_dit_block") rep.env = C.env_info() t0 = time.perf_counter() g = load_golden() w0 = load_block_weights(0) w1 = load_block_weights(1) rep.results["load_s"] = round(time.perf_counter() - t0, 2) rep.run("cpu_reference_vs_golden", lambda: cpu_reference_check(g, w0, w1)) ok = rep.results["cpu_reference_vs_golden"].get("block0_out_vs_golden", {}).get("pcc", 0) > 0.99999 if not ok: rep.note( "CPU reference does NOT match the golden block-0 output at PCC > 0.99999 — device PCCs below are relative to a suspect reference" ) if args.cpu_only: rep.write(args.out.replace(".json", "_cpu_only.json")) return 0 import ttnn # lazy rep.env.update(C.env_info(ttnn)) device = C.open_device(ttnn) rep.env["grid"] = str(device.compute_with_storage_grid_size()) try: variants = ( VARIANTS if args.variants == "all" else [v for v in VARIANTS if v["name"] in args.variants.split(",")] ) device_sweep(ttnn, device, rep, g, w0, w1, variants) finally: C.close_device(ttnn, device) # compact table print("\nvariant | block0(cross) PCC / max|d| | block1(self) PCC / max|d| | eager us b0/b1") for var in variants: r = rep.results.get(f"variant_{var['name']}", {}) b0 = r.get("block0_cross", {}).get("out_vs_fp32_golden", {}) b1 = r.get("block1_self", {}).get("out_vs_fp32_golden", {}) print( f"{var['name']:34s} | {b0.get('pcc', float('nan')):.7f} / {b0.get('max_abs', float('nan')):7.3f} | " f"{b1.get('pcc', float('nan')):.7f} / {b1.get('max_abs', float('nan')):7.3f} | " f"{r.get('block0_cross', {}).get('eager_block_us', float('nan')):7.0f} / {r.get('block1_self', {}).get('eager_block_us', float('nan')):7.0f}" ) rep.write(args.out) return 0 if __name__ == "__main__": sys.exit(main())