code/models/experimental/gr00t/benchmarks/mb2_bfp8_dit_block.py

# 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())