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

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""WP-P Stage-3 backbone sweeps on the p150a (IMPLEMENTATION_PLAN.md §6 Stage 3 row, §8 session P; every lever that
lowers trace A without a new kernel).  Results feed ``tt/layers.py`` (tables), ``tt/blocks.py`` (head ops),
``tt/qwen3.py`` (residual memory), ``tt/qwen3vl_vision.py`` (``TunedVisionConfig`` defaults), ``tt/tracer.py``
(``cq1_uploads`` / ``DeviceEmbeddingGather``) and the summary ``tests/tt/results/p_summary.md``.

Sub-sweeps (``--what``; several may be given, comma separated):

| what | shapes | measures |
|---|---|---|
| ``vit`` (n17) | the DROID tower on the golden ``pixel_values`` | ``TunedQwen3VLVisionTT`` variants (``TunedVisionConfig``: SDPA chunks / fidelity, RoPE composite vs ``rotary_embedding_hf``, L1 vs DRAM activations, LoFi FF, tuned vs tt_dit mergers) **and** the verbatim oracle on the same weights: every §7.2 vision tap gated with ``harness.check_taps`` (multi gates), fork-vs-verbatim PCC, traced tower ms (5 replays, min), traced == untraced |
| ``rope`` | ``[4, 16, 448, 64]`` random, HF tables | composite 7-op rotate-half vs ``rotary_embedding_hf`` on a batch of 4: PCC / max abs vs torch, µs |
| ``matmul`` | the M = 1792 Qwen-ViT rows (qkv / out / fc1 (+GELU_TANH) / fc2 / patch embed) and the M = 448 merger rows | ``MatmulMultiCoreReuseMultiCastProgramConfig`` rows x ``in0_block_w`` grid, ``minimal_matmul``, L1 vs DRAM activations, fused vs separate GELU: µs + PCC |
| ``sdpa`` | the 8 backbone attention shapes (SigLIP 256 / 352+mask, Qwen-ViT B=4 448, LLM causal 128 / 384 / 512, VL-SA 384 / 512 + mask) | ``SDPAProgramConfig`` chunk pairs x HiFi2 / HiFi4 (+fp32 acc), full 11x10 vs the tt_dit 10x10 grid: µs + PCC vs torch |
| ``heads`` | fused-QKV split / heads concat at S = 128 / 384 / 512 (LLM, VL-SA) and M = 256 / 352 / 1792 (ViTs) | ``nlp_create_qkv_heads`` (L1 / DRAM out), ``split_query_key_value_and_split_heads``, slice+reshape+permute; ``nlp_concat_heads`` vs ``concatenate_heads`` vs permute+reshape: µs + bit-exactness |
| ``residual`` (n15 / n17) | ``Qwen3LayerTT`` at S = 384 / 512 with random weights | residual stream in DRAM (plan §3.3) vs L1: layer µs + output equality |
| ``embed`` | the version's canonical prompt | ``DeviceEmbeddingGather`` (resident table + uint32 ids) vs the host ``embed_rows`` + tilize + upload path: bit-exactness, µs, upload bytes |
| ``cq`` | ``Gr00tTT`` (needs ``--num-cqs 2``) | ``TraceExecutor.replay`` wall ms with ``cq1_uploads`` off / on; results bit-equal |
| ``e2e`` | ``Gr00tTT`` (``bench_e2e.run`` method) | the tuned per-stage JSON ``benchmarks/results/backbone_tuned_<version>.json`` (fork tower for n17 via ``GR00T_VIT_FORK``; ``--refresh-baseline`` also writes a ``bench_e2e_<version>_<ts>.json`` for ``test_tt_perf``) |

Method: one Metal trace of ``--iters`` back-to-back ops (whole towers: 1), replayed ``--reps`` times, **min** over
replays (mb1 §0 / bench_stage_ops); PCC through ``common.pcc``.  Every variant records its error instead of aborting
the sweep.  Output: ``benchmarks/results/sweep_backbone/sweep_backbone_<what>_<version>_<YYYYmmdd-HHMMSS>.json`` (``PERF_SCHEMA``
via ``bench_stage_ops.perf_payload``; ``stages`` = the winners, ``rows`` = every variant).

Run (device rule; from the tt-metal root)::

    DEVICE_LOCK_TIMEOUT=14400 /home/deepgadget/experiments/gr00t/bin/with-device.sh \\
        python -m models.experimental.gr00t.benchmarks.sweep_backbone --what vit,rope --version n17
    python -m models.experimental.gr00t.benchmarks.sweep_backbone --what all --dry-run     # no ttnn, no device
"""

from __future__ import annotations

import argparse
import math
import os
import statistics
import sys
import time
from dataclasses import asdict, dataclass, field, replace
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

from models.experimental.gr00t.benchmarks import bench_e2e as E2E
from models.experimental.gr00t.benchmarks import bench_stage_ops as SO
from models.experimental.gr00t.benchmarks import mb1_common as MB
from models.experimental.gr00t.tt import device as D
from models.experimental.gr00t.tt import layers as L

HERE = Path(__file__).resolve().parent
RESULTS_DIR = SO.RESULTS_DIR
#: sweep JSONs live in a sub-directory: ``tests/tt/test_tt_perf.py::test_existing_results_documents_load`` requires every
#: JSON directly under ``benchmarks/results`` to be a bench_stage_ops / bench_e2e / bench_denoise document (the tuned
#: ``backbone_tuned_<version>.json`` files are bench_e2e payloads and stay at the top level).
SWEEP_RESULTS_DIR = RESULTS_DIR / "sweep_backbone"
BENCH_NAME = "sweep_backbone"
SWEEPS: Tuple[str, ...] = ("vit", "rope", "matmul", "sdpa", "heads", "mask", "residual", "embed", "cq", "e2e")
VERSIONS: Tuple[str, ...] = ("n15", "n16", "n17")
DEFAULT_ITERS = 8
DEFAULT_REPS = 3
DEFAULT_TOWER_REPS = 5
PCC_ADMIT = 0.999  # random-data admission bar of a program config (D0 convention)
KNOWN_VIT_MISSES: Dict[str, float] = {"vit_deepstack_mergers[0]": 0.9995}  # D0/D1/D2 device fact (xfail floor)


def tuned_json_path(version: str) -> Path:
    """``benchmarks/results/backbone_tuned_<version>.json`` (WP-P deliverable)."""
    if version not in VERSIONS:
        raise ValueError(f"version must be one of {VERSIONS}, got {version!r}")
    return RESULTS_DIR / f"backbone_tuned_{version}.json"


# --------------------------------------------------------------------------------------------------------------------
# shared helpers
# --------------------------------------------------------------------------------------------------------------------
def _ttnn() -> Any:
    import ttnn

    return ttnn


def _stats(xs: Sequence[float]) -> Dict[str, float]:
    return {"min": min(xs), "median": statistics.median(xs), "mean": statistics.fmean(xs), "max": max(xs), "n": len(xs)}


def _pcc(got: Any, ref: Any) -> Dict[str, Any]:
    from models.experimental.gr00t.common.pcc import max_abs, pcc

    return {"pcc": pcc(got, ref), "max_abs": max_abs(got, ref)}


def _rand(shape: Sequence[int], seed: int, scale: float = 1.0) -> Any:
    import torch

    g = torch.Generator().manual_seed(int(seed))
    return (torch.randn(tuple(shape), generator=g) * scale).to(torch.bfloat16)


def _dev(t: Any, device: Any, mem: str = "L1", dtype: str = "bf16", layout: str = "TILE") -> Any:
    return L.to_device(t, device, dtype=dtype, mem=mem, layout=layout)


def _time(device: Any, fn: Callable[[], Any], iters: int, reps: int) -> Dict[str, Any]:
    t = MB.time_traced(_ttnn(), device, fn, iters, reps=reps)
    return {
        "us": round(float(t["min"]), 3),
        "us_mean": round(float(t["mean"]), 3),
        "reps_us": [round(x, 3) for x in t["reps"]],
    }


def _try(row: Dict[str, Any], fn: Callable[[], Dict[str, Any]], errors: List[str], label: str) -> Dict[str, Any]:
    try:
        row.update(fn())
        row.setdefault("error", None)
    except Exception as e:  # noqa: BLE001 - record and continue (mb1 convention)
        row["error"] = MB.short_error(e)
        errors.append(f"{label}: {row['error']}")
    return row


def _grid(cols: int, rows: int) -> Any:
    return D.core_grid((cols, rows))


def _mcast2d_config(m: int, k: int, n: int, rows: int, ibw: int, act: Optional[str]) -> Any:
    """A full-width (11 columns) 2-D multicast program config with ``rows`` grid rows and ``in0_block_w=ibw``."""
    ttnn = _ttnn()
    mt, kt, nt = L.tiles(m), L.tiles(k), L.tiles(n)
    if kt % ibw:
        raise ValueError(f"in0_block_w={ibw} does not divide Kt={kt}")
    per_core_M = -(-mt // rows)
    per_core_N = -(-nt // L.GRID_COLS)
    sub = L.subblock_2d(per_core_M, per_core_N)
    return ttnn.MatmulMultiCoreReuseMultiCastProgramConfig(
        compute_with_storage_grid_size=_grid(L.GRID_COLS, rows),
        in0_block_w=ibw,
        out_subblock_h=sub[0],
        out_subblock_w=sub[1],
        per_core_M=per_core_M,
        per_core_N=per_core_N,
        transpose_mcast=False,
        fused_activation=L.fused_activation_param(act),
        fuse_batch=True,
    )


def _payload(
    what: str,
    version: str,
    args: Mapping[str, Any],
    stages: Mapping[str, Any],
    rows: Sequence[Mapping[str, Any]],
    errors: Sequence[str],
    meta: Mapping[str, Any],
    device: Any = None,
    method: str = "",
) -> Dict[str, Any]:
    facts = None
    if device is not None:
        facts = D.device_facts(device)
        facts.update(MB.device_report(_ttnn(), device))
    return SO.perf_payload(
        f"{BENCH_NAME}_{what}",
        version,
        dry_run=device is None,
        args=dict(args),
        method={
            "timing": "Metal trace of --iters back-to-back ops (towers: 1), --reps replays, min (mb1 §0)",
            "notes": method,
        },
        stages=stages,
        rows=list(rows),
        device=facts,
        meta=dict(meta),
        errors=list(errors),
    )


def _write(payload: Mapping[str, Any], what: str, version: str, out_dir: Path) -> Path:
    sweep_dir = out_dir / "sweep_backbone" if out_dir == RESULTS_DIR else out_dir
    path = SO.perf_json_path(f"{BENCH_NAME}_{what}", version, sweep_dir)
    return SO.write_perf_json(payload, sweep_dir, path=path)


def _golden(version: str, sample: str) -> Dict[str, Any]:
    ok, reason = E2E.golden_available(version, sample)
    if not ok:
        raise FileNotFoundError(f"{version}/{sample}: {reason}")
    setup = E2E.resolve_setup(version)
    inputs = E2E.load_golden_inputs(version, sample, setup["embodiment_tag"], setup["layout"].name)
    inputs["setup"] = setup
    return inputs


# --------------------------------------------------------------------------------------------------------------------
# vit: the fork vs the verbatim tower on the golden pixels
# --------------------------------------------------------------------------------------------------------------------
def vit_variants() -> Dict[str, Any]:
    """``name -> TunedVisionConfig`` (the default first; every other variant changes one knob)."""
    from models.experimental.gr00t.tt.qwen3vl_vision import TunedVisionConfig

    base = TunedVisionConfig()
    return {
        "default": base,
        "rope_composite": replace(base, rope_mode="composite"),
        "sdpa_hifi4_nofp32_256": replace(base, sdpa_ckc="sdpa_hifi4"),
        "sdpa_hifi4_nofp32_448": replace(base, sdpa_ckc="sdpa_hifi4", sdpa_chunks=(448, 448)),
        "sdpa_hifi2_256": replace(base, sdpa_ckc="sdpa_hifi2"),
        "sdpa_hifi2_448": replace(base, sdpa_ckc="sdpa_hifi2", sdpa_chunks=(448, 448)),
        "merger_gelu_fused": replace(base, merger_gelu_fused=True),
        "ff_l1": replace(base, ff_mem="L1"),
        "mergers_ttdit": replace(base, tuned_mergers=False),
        "embed_ttdit": replace(base, tuned_embed=False),
        "lofi_ff": replace(base, ff_ckc="lofi_fp32acc"),
    }


def _vit_taps_host(QV: Any, taps: Dict[str, Any]) -> Dict[str, Any]:
    out = {k: QV.to_torch(v) for k, v in taps.items()}
    return out


def _gate_rows(
    version: str, gs: Any, host_taps: Mapping[str, Any], label: str
) -> Tuple[List[Dict[str, Any]], bool, List[str]]:
    """``harness.check_taps`` rows (multi gates); ``ok`` honours the known merger[0] floor as an xfail, not a pass."""
    from models.experimental.gr00t.tests.tt.harness import check_taps

    avail = {gs.threshold_key(t, **kw) for t, kw in gs.available("fp32")}
    gated = {k: v for k, v in host_taps.items() if k in avail}
    report = check_taps(version, gated, gs, multi=True)
    rows = []
    known: List[str] = []
    ok = True
    for r in report.rows:
        row = {
            "variant": label,
            "tap": r.get("tap"),
            "key": r.get("key"),
            "pcc": r.get("pcc"),
            "gate_pcc": r.get("gate_pcc"),
            "max_abs": r.get("max_abs"),
            "ok": bool(r.get("ok")),
        }
        if not row["ok"]:
            floor = KNOWN_VIT_MISSES.get(str(r.get("key")))
            if floor is not None and isinstance(r.get("pcc"), float) and r["pcc"] >= floor:
                row["known_miss_floor"] = floor
                known.append(str(r.get("key")))
            else:
                ok = False
        rows.append(row)
    return rows, ok, known


def sweep_vit(
    device: Any, *, version: str, sample: str, variants: Optional[Sequence[str]], reps: int, blocks_timing: bool
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    from models.experimental.gr00t.common.checkpoint import LazyCheckpoint
    from models.experimental.gr00t.tt import qwen3vl_vision as QV
    from models.experimental.gr00t.tt.layout import DeviceLayout

    if version != "n17":
        raise ValueError("--what vit is the N1.7 Qwen3-VL tower (n17)")
    ttnn = _ttnn()
    g = _golden(version, sample)
    gs, mi, cfg = g["gs"], g["mi"], g["setup"]["cfg"]
    dl = DeviceLayout.for_inputs(cfg, mi)
    ck = LazyCheckpoint(version)
    all_variants = vit_variants()
    names = list(all_variants) if not variants else list(variants)
    for n in names:
        if n not in all_variants:
            raise KeyError(f"unknown vit variant {n!r}; known: {sorted(all_variants)}")
    t0 = time.perf_counter()
    tower = QV.TunedQwen3VLVisionTT(cfg, dl, ck, device, tuned=all_variants[names[0]])
    load_s = time.perf_counter() - t0
    host = tower.prepare_inputs(mi)
    dev = tower.upload(host)
    rows: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    meta: Dict[str, Any] = {
        "shape_key": str(mi.shape_key),
        "layout": dl.layout_name,
        "load_s": round(load_s, 2),
        "tower": tower.describe(),
        "variants": {n: all_variants[n].to_dict() for n in names},
    }

    def run_path(
        label: str, forward: Callable[[Dict[str, Any]], Tuple[Any, List[Any]]], is_fork: bool
    ) -> Dict[str, Any]:
        rec: Dict[str, Any] = {"variant": label, "fork": is_fork}
        taps: Dict[str, Any] = {}
        ttnn.synchronize_device(device)
        t1 = time.perf_counter()
        merged, deep = forward(taps)
        ttnn.synchronize_device(device)
        rec["untraced_forward_s"] = round(time.perf_counter() - t1, 4)
        host_taps = _vit_taps_host(QV, taps)
        merged_h = QV.to_torch(merged)
        for t in [merged, *deep]:
            if not any(t is v for v in taps.values()):
                ttnn.deallocate(t)
        for t in taps.values():
            ttnn.deallocate(t)
        gate_rows, ok, known = _gate_rows(version, gs, host_taps, label)
        rec["gates_ok"] = ok
        rec["known_misses"] = known
        rec["gate_rows"] = gate_rows
        rec["min_pcc"] = min((r["pcc"] for r in gate_rows if isinstance(r.get("pcc"), float)), default=None)
        rec["host_taps"] = host_taps  # stripped before the JSON
        # traced tower (untraced == traced bit-equality on the projector output)
        traced = QV.traced_seconds(device, lambda: forward(None), reps=reps)
        m_t, d_t = traced["output"]
        rec["traced_equals_untraced"] = bool(torch.equal(QV.to_torch(m_t), merged_h))
        for t in [m_t, *d_t]:
            ttnn.deallocate(t)
        rec["tower_ms"] = round(traced["min_s"] * 1e3, 3)
        rec["tower_ms_mean"] = round(traced["mean_s"] * 1e3, 3)
        rec["tower_capture_s"] = round(traced["capture_s"], 3)
        return rec

    # --- verbatim oracle (same weights)
    ver = _try(
        {"variant": "verbatim", "fork": False},
        lambda: run_path("verbatim", lambda taps: tower.forward_verbatim(dev, taps), False),
        errors,
        "vit:verbatim",
    )
    ver_taps = ver.pop("host_taps", {})
    rows.append(ver)
    if "tower_ms" in ver:
        stages["tower_verbatim"] = {
            "metric": "ms",
            "value": ver["tower_ms"],
            "gates_ok": ver.get("gates_ok"),
            "known_misses": ver.get("known_misses"),
        }

    # --- fork variants
    best: Optional[Tuple[float, str]] = None
    for name in names:
        tuned = all_variants[name]
        try:
            tower.reconfigure(tuned)
        except Exception as e:  # noqa: BLE001
            rows.append({"variant": name, "fork": True, "error": MB.short_error(e)})
            errors.append(f"vit:{name}: {rows[-1]['error']}")
            continue
        rec = _try(
            {"variant": name, "fork": True, "tuned": tuned.to_dict()},
            lambda: run_path(name, lambda taps: tower.forward(dev, taps), True),
            errors,
            f"vit:{name}",
        )
        taps_h = rec.pop("host_taps", {})
        if taps_h and ver_taps:
            rec["pcc_vs_verbatim"] = {k: round(_pcc(taps_h[k], ver_taps[k])["pcc"], 6) for k in taps_h if k in ver_taps}
        if blocks_timing and "tower_ms" in rec and name == names[0]:
            try:
                x0 = tower.add_pos(tower.embed(dev.patches), dev.pos)
                tb = QV.traced_seconds(
                    device, lambda: tower.forward_blocks(x0, tower.depth, dev.rope, dev.cu_seqlens)[0], reps=reps
                )
                ttnn.deallocate(tb["output"])
                ttnn.deallocate(x0)
                rec["blocks24_ms"] = round(tb["min_s"] * 1e3, 3)
                rec["per_block_us"] = round(tb["min_s"] * 1e6 / tower.depth, 1)
            except Exception as e:  # noqa: BLE001
                rec["blocks_error"] = MB.short_error(e)
                errors.append(f"vit:{name}:blocks: {rec['blocks_error']}")
        rows.append(rec)
        if "tower_ms" in rec and rec.get("gates_ok"):
            if best is None or rec["tower_ms"] < best[0]:
                best = (rec["tower_ms"], name)
        stages[f"tower_fork/{name}"] = {
            "metric": "ms",
            "value": rec.get("tower_ms"),
            "gates_ok": rec.get("gates_ok"),
            "known_misses": rec.get("known_misses"),
            "min_pcc": rec.get("min_pcc"),
            "traced_equals_untraced": rec.get("traced_equals_untraced"),
            "error": rec.get("error"),
        }
    meta["best_fork_variant"] = None if best is None else {"variant": best[1], "tower_ms": best[0]}
    tower.release()
    for t in dev.tensors().values():
        ttnn.deallocate(t)
    return stages, rows, errors, meta


# --------------------------------------------------------------------------------------------------------------------
# rope: composite vs rotary_embedding_hf on a batch of 4 images
# --------------------------------------------------------------------------------------------------------------------
def _rope_torch(x: Any, cos: Any, sin: Any) -> Any:
    import torch

    xf = x.float()
    half = xf.shape[-1] // 2
    rot = torch.cat([-xf[..., half:], xf[..., :half]], dim=-1)
    return xf * cos.float() + rot * sin.float()


def sweep_rope(
    device: Any,
    *,
    iters: int,
    reps: int,
    shapes: Sequence[Tuple[int, int, int, int]] = ((4, 16, 448, 64), (1, 16, 1792, 64)),
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    ttnn = _ttnn()
    rows: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    for b, h, s, d in shapes:
        x_h = _rand((b, h, s, d), 1)
        ang = torch.rand(s, d // 2, generator=torch.Generator().manual_seed(2)) * 6.28
        emb = torch.cat([ang, ang], dim=-1)
        cos_h, sin_h = emb.cos().to(torch.bfloat16), emb.sin().to(torch.bfloat16)
        ref = _rope_torch(x_h, cos_h, sin_h)
        for mem in ("L1", "DRAM"):
            x = _dev(x_h, device, mem)
            cos = _dev(cos_h.reshape(1, 1, s, d), device, mem)
            sin = _dev(sin_h.reshape(1, 1, s, d), device, mem)
            mc = D.resolve_mem(mem)

            def composite() -> Any:
                hi = ttnn.slice(x, [0, 0, 0, d // 2], [b, h, s, d], memory_config=mc)
                lo = ttnn.slice(x, [0, 0, 0, 0], [b, h, s, d // 2], memory_config=mc)
                neg = ttnn.neg(hi, memory_config=mc)
                ttnn.deallocate(hi)
                rot = ttnn.concat([neg, lo], dim=-1, memory_config=mc)
                ttnn.deallocate(neg)
                ttnn.deallocate(lo)
                xc = ttnn.mul(x, cos, memory_config=mc)
                rs = ttnn.mul(rot, sin, memory_config=mc)
                ttnn.deallocate(rot)
                out = ttnn.add(xc, rs, memory_config=mc)
                ttnn.deallocate(xc)
                ttnn.deallocate(rs)
                return out

            def hf_op() -> Any:
                return L.rope(x, cos, sin, layout="hf", mem=mem)

            for name, fn in (("composite", composite), ("hf_op", hf_op)):
                key = f"rope/{name}/{b}x{h}x{s}x{d}/{mem}"
                row = {"shape": [b, h, s, d], "variant": name, "mem": mem, "ops": 7 if name == "composite" else 1}

                def run() -> Dict[str, Any]:
                    out = fn()
                    ttnn.synchronize_device(device)
                    got = L.to_torch(out)
                    ttnn.deallocate(out)
                    m = _pcc(got, ref)
                    t = _time(device, fn, iters, reps)
                    return {**m, **t, "admit": m["pcc"] >= PCC_ADMIT}

                rows.append(_try(row, run, errors, key))
                stages[key] = {
                    "metric": "us",
                    "value": rows[-1].get("us"),
                    "pcc": rows[-1].get("pcc"),
                    "error": rows[-1].get("error"),
                }
            for t in (x, cos, sin):
                ttnn.deallocate(t)
    return stages, rows, errors, {"shapes": [list(s) for s in shapes]}


# --------------------------------------------------------------------------------------------------------------------
# matmul: the M = 1792 / 448 Qwen-ViT rows (L1 vs DRAM, ibw x rows, minimal, fused GELU)
# --------------------------------------------------------------------------------------------------------------------
MATMUL_SHAPES: Tuple[Tuple[str, int, int, int, Optional[str]], ...] = (
    ("vit_qkv", 1792, 1024, 3072, None),
    ("vit_out", 1792, 1024, 1024, None),
    ("vit_fc1", 1792, 1024, 4096, "gelu_tanh"),
    ("vit_fc2", 1792, 4096, 1024, None),
    ("vit_embed", 1792, 1536, 1024, None),
    ("merger_fc1", 448, 4096, 4096, "gelu"),
    ("merger_fc2", 448, 4096, 2048, None),
)


def matmul_candidates(m: int, k: int, n: int, dtype: str) -> List[Dict[str, Any]]:
    """The table row plus a grid of 2-D configs (``rows`` x ``in0_block_w``) and ``minimal`` (CPU-side descriptions)."""
    kt = L.tiles(k)
    out: List[Dict[str, Any]] = [{"name": "table", "kind": "table"}]
    for rows in (8, 10):
        for ibw in (2, 4, 8, 16):
            if kt % ibw == 0:
                out.append({"name": f"mcast2d_r{rows}_ibw{ibw}", "kind": "mcast2d", "rows": rows, "ibw": ibw})
    out.append({"name": "minimal", "kind": "minimal"})
    L.matmul_spec(m, k, n, dtype)  # the table row must exist
    return out


def sweep_matmul(
    device: Any,
    *,
    iters: int,
    reps: int,
    dtypes: Sequence[str] = ("bf16",),
    shapes: Sequence[Tuple[str, int, int, int, Optional[str]]] = MATMUL_SHAPES,
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    from models.tt_dit.utils.matmul import get_matmul_config

    ttnn = _ttnn()
    rows_out: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    ckc = D.resolve_ckc(D.CKC_HIFI2_FP32ACC)
    for name, m, k, n, act in shapes:
        x_h = _rand((m, k), 11, 0.5)
        for dtype in dtypes:
            w_h = _rand((k, n), 12, 1.0 / math.sqrt(k))
            b_h = _rand((1, n), 13, 0.1)
            ref = x_h.float() @ w_h.float() + b_h.float()
            if act == "gelu_tanh":
                ref = torch.nn.functional.gelu(ref, approximate="tanh")
            elif act == "gelu":
                ref = torch.nn.functional.gelu(ref)
            w = _dev(w_h, device, "DRAM", dtype)
            bias = _dev(b_h, device, "DRAM")
            for mem in ("L1", "DRAM"):
                x = _dev(x_h, device, mem)
                mc = D.resolve_mem(mem)
                for cand in matmul_candidates(m, k, n, dtype):
                    for fused in ((True, False) if act else (True,)):
                        label = f"{name}/{dtype}/{mem}/{cand['name']}" + (
                            "" if act is None else ("/fused" if fused else "/separate_gelu")
                        )
                        row = {
                            "shape": name,
                            "m": m,
                            "k": k,
                            "n": n,
                            "dtype": dtype,
                            "act": act,
                            "mem": mem,
                            "candidate": cand["name"],
                            "fused": fused if act else None,
                        }

                        def build_fn() -> Callable[[], Any]:
                            a = act if (act and fused) else None
                            if cand["kind"] == "table":
                                plan = L.matmul_plan(m, k, n, dtype, a)
                                row["config"] = (
                                    f"{plan.kind} grid={plan.grid} pcM={plan.per_core_M} pcN={plan.per_core_N} ibw={plan.in0_block_w} sub={plan.out_subblock}"
                                )
                                pc = plan.to_ttnn()
                                kind = plan.kind
                            elif cand["kind"] == "mcast2d":
                                pc = _mcast2d_config(m, k, n, cand["rows"], cand["ibw"], a)
                                row["config"] = f"mcast2d rows={cand['rows']} ibw={cand['ibw']}"
                                kind = "mcast2d"
                            else:
                                pc = get_matmul_config(m, k, n, _grid(*L.EXPECTED_COMPUTE_GRID), use_heuristic=True)
                                row["config"] = "minimal_matmul heuristic"
                                kind = "minimal"

                            def mm() -> Any:
                                if kind == "minimal":
                                    out = ttnn.experimental.minimal_matmul(
                                        x,
                                        w,
                                        bias_tensor=bias,
                                        fused_activation=L.fused_activation_param(a),
                                        config=pc,
                                        memory_config=mc,
                                        dtype=ttnn.bfloat16,
                                        compute_kernel_config=ckc,
                                    )
                                else:
                                    out = ttnn.linear(
                                        x,
                                        w,
                                        bias=bias,
                                        program_config=pc,
                                        memory_config=mc,
                                        dtype=ttnn.bfloat16,
                                        compute_kernel_config=ckc,
                                    )
                                if act and not fused:
                                    g = L.gelu(out, "tanh" if act == "gelu_tanh" else "erf", mem=mem)
                                    ttnn.deallocate(out)
                                    return g
                                return out

                            return mm

                        def run() -> Dict[str, Any]:
                            fn = build_fn()
                            out = fn()
                            ttnn.synchronize_device(device)
                            got = L.to_torch(out)
                            ttnn.deallocate(out)
                            met = _pcc(got, ref)
                            t = _time(device, fn, iters, reps)
                            return {**met, **t, "admit": met["pcc"] >= PCC_ADMIT}

                        rows_out.append(_try(row, run, errors, f"matmul:{label}"))
                        r = rows_out[-1]
                        stages[f"mm/{label}"] = {
                            "metric": "us",
                            "value": r.get("us"),
                            "pcc": r.get("pcc"),
                            "config": r.get("config"),
                            "error": r.get("error"),
                        }
                ttnn.deallocate(x)
            ttnn.deallocate(w)
            ttnn.deallocate(bias)
    return stages, rows_out, errors, {"shapes": [list(s) for s in shapes], "pcc_admit": PCC_ADMIT}


# --------------------------------------------------------------------------------------------------------------------
# sdpa: chunk pairs x fidelity at the 8 backbone shapes
# --------------------------------------------------------------------------------------------------------------------
@dataclass(frozen=True)
class SdpaShape:
    name: str
    cls: str
    b: int
    h: int
    hkv: int
    s: int
    d: int
    real_d: int
    mask: Optional[str]  # None | "key_valid" (last `pad` keys masked) ; causal handled by is_causal
    causal: bool
    valid: int
    pairs: Tuple[Tuple[int, int], ...]


SDPA_SHAPES: Tuple[SdpaShape, ...] = (
    SdpaShape(
        "siglip_256",
        "vit_siglip",
        1,
        16,
        16,
        256,
        96,
        72,
        None,
        False,
        256,
        ((64, 128), (64, 256), (64, 64), (128, 128), (128, 256), (256, 256)),
    ),
    SdpaShape(
        "siglip_352",
        "vit_siglip",
        1,
        16,
        16,
        352,
        96,
        72,
        "key_valid",
        False,
        324,
        ((64, 64), (64, 128), (32, 128), (128, 128), (32, 352), (64, 352), (352, 352)),
    ),
    SdpaShape(
        "qwenvit_448x4",
        "vit_qwen3vl",
        4,
        16,
        16,
        448,
        64,
        64,
        None,
        False,
        448,
        ((256, 256), (128, 256), (128, 128), (64, 128), (64, 64), (448, 448), (256, 448)),
    ),
    SdpaShape(
        "llm_128", "llm_causal", 1, 16, 8, 128, 128, 128, None, True, 128, ((32, 32), (64, 64), (128, 128), (32, 128))
    ),
    SdpaShape(
        "llm_384",
        "llm_causal",
        1,
        16,
        8,
        384,
        128,
        128,
        None,
        True,
        384,
        ((128, 128), (128, 384), (64, 64), (64, 128), (192, 192), (256, 256)),
    ),
    SdpaShape(
        "llm_512",
        "llm_causal",
        1,
        16,
        8,
        512,
        128,
        128,
        None,
        True,
        512,
        ((128, 128), (256, 256), (128, 256), (256, 512), (64, 128)),
    ),
    SdpaShape(
        "vlsa_384",
        "vlsa",
        1,
        32,
        32,
        384,
        64,
        64,
        "key_valid",
        False,
        296,
        ((128, 128), (192, 192), (128, 192), (128, 384), (64, 128), (256, 256), (384, 384)),
    ),
    SdpaShape(
        "vlsa_512",
        "vlsa",
        1,
        32,
        32,
        512,
        64,
        64,
        "key_valid",
        False,
        469,
        ((256, 256), (128, 256), (128, 128), (128, 512), (64, 256), (256, 512), (512, 512)),
    ),
)
SDPA_CKCS: Tuple[str, ...] = ("sdpa_hifi2", "sdpa_hifi4", "hifi4_fp32acc")


def _sdpa_ref(q: Any, k: Any, v: Any, scale: float, mask: Any, causal: bool) -> Any:
    import torch

    qf, kf, vf = q.float(), k.float(), v.float()
    b, h, s, d = qf.shape
    hkv = kf.shape[1]
    if hkv != h:
        rep = h // hkv
        kf = kf.repeat_interleave(rep, dim=1)
        vf = vf.repeat_interleave(rep, dim=1)
    scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale
    if causal:
        causal_mask = torch.triu(torch.ones(s, s, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(causal_mask, float("-inf"))
    if mask is not None:
        scores = scores + mask.float()
    p = torch.softmax(scores, dim=-1)
    return torch.matmul(p, vf)


def sweep_sdpa(
    device: Any,
    *,
    iters: int,
    reps: int,
    shapes: Sequence[SdpaShape] = SDPA_SHAPES,
    ckcs: Sequence[str] = SDPA_CKCS,
    reserved_column: bool = True,
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    ttnn = _ttnn()
    rows_out: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    for sh in shapes:
        q_h = _rand((sh.b, sh.h, sh.s, sh.d), 21, 0.5)
        k_h = _rand((sh.b, sh.hkv, sh.s, sh.d), 22, 0.5)
        v_h = _rand((sh.b, sh.hkv, sh.s, sh.d), 23, 0.5)
        if sh.real_d != sh.d:  # padded heads carry zeros (SigLIP 72 -> 96)
            q_h[..., sh.real_d :] = 0
            k_h[..., sh.real_d :] = 0
            v_h[..., sh.real_d :] = 0
        scale = float(sh.real_d) ** -0.5
        mask_h = None
        if sh.mask == "key_valid":
            mask_h = torch.zeros(1, 1, sh.s, sh.s, dtype=torch.bfloat16)
            mask_h[..., sh.valid :] = -1e9
        ref = _sdpa_ref(q_h, k_h, v_h, scale, mask_h, sh.causal)
        q, k, v = (_dev(t, device, "L1") for t in (q_h, k_h, v_h))
        mask = _dev(mask_h, device, "DRAM") if mask_h is not None else None
        grids = [("11x10", L.EXPECTED_COMPUTE_GRID)]
        if reserved_column and sh.name.startswith("qwenvit"):
            grids.append(("10x10", (L.EXPECTED_COMPUTE_GRID[0] - 1, L.EXPECTED_COMPUTE_GRID[1])))
        for gname, grid in grids:
            for ckc_name in ckcs:
                ck = D.resolve_ckc(ckc_name)
                for qc, kc in sh.pairs:
                    if qc > sh.s or kc > sh.s:
                        continue
                    label = f"{sh.name}/{gname}/{ckc_name}/q{qc}k{kc}"
                    row = {
                        "shape": sh.name,
                        "cls": sh.cls,
                        "b": sh.b,
                        "h": sh.h,
                        "hkv": sh.hkv,
                        "s": sh.s,
                        "d": sh.d,
                        "mask": sh.mask,
                        "causal": sh.causal,
                        "grid": gname,
                        "ckc": ckc_name,
                        "q_chunk": qc,
                        "k_chunk": kc,
                        "forbidden": (qc, kc) in L.FORBIDDEN_SDPA_CHUNKS,
                        "in_allowlist": L.sdpa_chunk_ok(sh.s, sh.s, sh.d, (qc, kc)),
                    }
                    pc = ttnn.SDPAProgramConfig(
                        compute_with_storage_grid_size=_grid(*grid),
                        q_chunk_size=qc,
                        k_chunk_size=kc,
                        exp_approx_mode=False,
                    )

                    def fn() -> Any:
                        return ttnn.transformer.scaled_dot_product_attention(
                            q,
                            k,
                            v,
                            attn_mask=mask,
                            is_causal=sh.causal,
                            scale=scale,
                            program_config=pc,
                            compute_kernel_config=ck,
                            memory_config=ttnn.L1_MEMORY_CONFIG,
                        )

                    def run() -> Dict[str, Any]:
                        out = fn()
                        ttnn.synchronize_device(device)
                        got = L.to_torch(out)
                        ttnn.deallocate(out)
                        gv, rv = got[..., : sh.real_d], ref[..., : sh.real_d]
                        if sh.mask == "key_valid" or sh.causal:
                            gv, rv = gv[:, :, : sh.valid], rv[:, :, : sh.valid]
                        met = _pcc(gv, rv)
                        t = _time(device, fn, iters, reps)
                        return {**met, **t, "admit": met["pcc"] >= PCC_ADMIT}

                    rows_out.append(_try(row, run, errors, f"sdpa:{label}"))
                    r = rows_out[-1]
                    stages[f"sdpa/{label}"] = {
                        "metric": "us",
                        "value": r.get("us"),
                        "pcc": r.get("pcc"),
                        "admit": r.get("admit"),
                        "error": r.get("error"),
                    }
        for t in (q, k, v, mask):
            if t is not None:
                ttnn.deallocate(t)
    return stages, rows_out, errors, {"shapes": [asdict(s) for s in shapes], "pcc_admit": PCC_ADMIT}


# --------------------------------------------------------------------------------------------------------------------
# heads: fused-QKV split and heads concat variants
# --------------------------------------------------------------------------------------------------------------------
HEAD_SHAPES: Tuple[Tuple[str, int, int, int, int, int], ...] = (
    # name, B, S, H, Hkv, d
    ("llm_128", 1, 128, 16, 8, 128),
    ("llm_384", 1, 384, 16, 8, 128),
    ("llm_512", 1, 512, 16, 8, 128),
    ("vlsa_384", 1, 384, 32, 32, 64),
    ("vlsa_512", 1, 512, 32, 32, 64),
    ("siglip_256", 1, 256, 16, 16, 96),
    ("siglip_352", 1, 352, 16, 16, 96),
    ("qwenvit_448x4", 4, 448, 16, 16, 64),
)


def sweep_heads(
    device: Any, *, iters: int, reps: int, shapes: Sequence[Tuple[str, int, int, int, int, int]] = HEAD_SHAPES
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    ttnn = _ttnn()
    rows_out: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    for name, b, s, h, hkv, d in shapes:
        width = (h + 2 * hkv) * d
        qkv_h = _rand((b, 1, s, width), 31)
        for in_mem in ("L1", "DRAM"):
            qkv = _dev(qkv_h, device, in_mem)
            base_q = base_k = base_v = None

            def split_nlp(out_mem: Optional[str]) -> Callable[[], Any]:
                def fn() -> Any:
                    return L.split_heads(qkv, h, hkv, mem=out_mem)

                return fn

            def split_slices() -> Any:
                mc = ttnn.L1_MEMORY_CONFIG
                outs = []
                off = 0
                for nh in (h, hkv, hkv):
                    part = ttnn.slice(qkv, [0, 0, 0, off], [b, 1, s, off + nh * d], memory_config=mc)
                    off += nh * d
                    r = ttnn.reshape(part, (b, s, nh, d))
                    p = ttnn.permute(r, (0, 2, 1, 3), memory_config=mc)
                    ttnn.deallocate(part)
                    outs.append(p)
                return tuple(outs)

            variants = [
                ("nlp_create_qkv_heads/out_default", split_nlp(None)),
                ("nlp_create_qkv_heads/out_L1", split_nlp("L1")),
                ("nlp_create_qkv_heads/out_DRAM", split_nlp("DRAM")),
                ("slice+reshape+permute", split_slices),
            ]
            for vname, fn in variants:
                label = f"split/{name}/in_{in_mem}/{vname}"
                row = {
                    "op": "split",
                    "shape": name,
                    "b": b,
                    "s": s,
                    "h": h,
                    "hkv": hkv,
                    "d": d,
                    "in_mem": in_mem,
                    "variant": vname,
                }

                def run() -> Dict[str, Any]:
                    nonlocal base_q, base_k, base_v
                    outs = fn()
                    ttnn.synchronize_device(device)
                    got = [L.to_torch(t) for t in outs]
                    for t in outs:
                        ttnn.deallocate(t)
                    if base_q is None:
                        base_q, base_k, base_v = got
                        exact = True
                    else:
                        exact = all(torch.equal(a, r) for a, r in zip(got, (base_q, base_k, base_v)))
                    t = _time(device, fn, iters, reps)
                    return {**t, "exact_vs_nlp": exact, "shapes": [list(x.shape) for x in got]}

                rows_out.append(_try(row, run, errors, f"heads:{label}"))
                r = rows_out[-1]
                stages[f"heads/{label}"] = {
                    "metric": "us",
                    "value": r.get("us"),
                    "exact": r.get("exact_vs_nlp"),
                    "error": r.get("error"),
                }
            ttnn.deallocate(qkv)
        # concat variants on [B, H, S, d] -> [B, 1, S, H d]
        o_h = _rand((b, h, s, d), 32)
        for in_mem in ("L1",):
            o = _dev(o_h, device, in_mem)
            base_c = None

            def concat_nlp() -> Any:
                return L.concat_heads(o, mem="L1")

            def concat_perm() -> Any:
                p = ttnn.permute(o, (0, 2, 1, 3), memory_config=ttnn.L1_MEMORY_CONFIG)
                r = ttnn.reshape(p, (b, 1, s, h * d))
                return r

            def concat_tr() -> Any:
                return ttnn.transformer.concatenate_heads(o, memory_config=ttnn.L1_MEMORY_CONFIG)

            for vname, fn in (
                ("nlp_concat_heads", concat_nlp),
                ("permute+reshape", concat_perm),
                ("transformer.concatenate_heads", concat_tr),
            ):
                label = f"concat/{name}/in_{in_mem}/{vname}"
                row = {
                    "op": "concat",
                    "shape": name,
                    "b": b,
                    "s": s,
                    "h": h,
                    "d": d,
                    "in_mem": in_mem,
                    "variant": vname,
                }

                def run() -> Dict[str, Any]:
                    nonlocal base_c
                    out = fn()
                    ttnn.synchronize_device(device)
                    got = L.to_torch(out).reshape(b, 1, s, h * d)
                    ttnn.deallocate(out)
                    if base_c is None:
                        base_c = got
                        exact = True
                    else:
                        exact = bool(torch.equal(got, base_c))
                    t = _time(device, fn, iters, reps)
                    return {**t, "exact_vs_nlp": exact}

                rows_out.append(_try(row, run, errors, f"heads:{label}"))
                r = rows_out[-1]
                stages[f"heads/{label}"] = {
                    "metric": "us",
                    "value": r.get("us"),
                    "exact": r.get("exact_vs_nlp"),
                    "error": r.get("error"),
                }
            ttnn.deallocate(o)
    return stages, rows_out, errors, {"shapes": [list(s) for s in shapes]}


# --------------------------------------------------------------------------------------------------------------------
# mask: N1.6 SigLIP block with the additive key mask vs the pad-channel key bias (tt/blocks.py MASK_IMPLS)
# --------------------------------------------------------------------------------------------------------------------
def sweep_mask(
    device: Any, *, iters: int, reps: int, m_rows: int = 352, valid: int = 324
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    """``EncoderBlockTT`` (SigLIP family, 16 x 72 -> 96, random weights with zero pad channels) at the N1.6 shape:
    ``mask_impl="additive"`` vs ``"key_bias"`` -- block µs (traced x iters), output PCC / max abs on the valid rows."""
    import torch

    from models.experimental.gr00t.tt import blocks as B
    from models.experimental.gr00t.tt.layout import additive_key_mask

    ttnn = _ttnn()
    dim, heads, real_d, d_pad, ff = 1152, 16, 72, 96, 4304
    inner = heads * d_pad
    g = torch.Generator().manual_seed(77)

    def rnd(*shape: int, scale: float = 1.0) -> Any:
        return (torch.randn(*shape, generator=g) * scale).to(torch.bfloat16)

    wqkv = rnd(dim, 3 * inner, scale=1 / math.sqrt(dim)).float()
    bqkv = rnd(1, 3 * inner, scale=0.05).float()
    for part in range(3):  # zero the pad channels 72..95 of every head (pad_heads semantics)
        for h in range(heads):
            c0 = part * inner + h * d_pad + real_d
            wqkv[:, c0 : c0 + (d_pad - real_d)] = 0
            bqkv[:, c0 : c0 + (d_pad - real_d)] = 0
    wo = rnd(inner, dim, scale=1 / math.sqrt(inner)).float()
    for h in range(heads):
        wo[h * d_pad + real_d : (h + 1) * d_pad, :] = 0
    weights_h = {
        "ln1.w": torch.ones(1, dim),
        "ln1.b": torch.zeros(1, dim),
        "Wqkv": wqkv,
        "bqkv": bqkv,
        "Wo": wo,
        "bo": rnd(1, dim, scale=0.05).float(),
        "ln2.w": torch.ones(1, dim),
        "ln2.b": torch.zeros(1, dim),
        "fc1.W": rnd(dim, ff, scale=1 / math.sqrt(dim)).float(),
        "fc1.b": rnd(1, ff, scale=0.05).float(),
        "fc2.W": rnd(ff, dim, scale=1 / math.sqrt(ff)).float(),
        "fc2.b": rnd(1, dim, scale=0.05).float(),
    }
    weights = {k: _dev(v.to(torch.bfloat16), device, "DRAM") for k, v in weights_h.items()}
    mask_h = additive_key_mask(torch.arange(m_rows) < valid, m_rows)
    mask = _dev(mask_h, device, "DRAM")
    x_h = rnd(1, 1, m_rows, dim)
    x_h[:, :, valid:] = 0
    x = _dev(x_h, device, "L1")
    rows_out: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    base = None
    for impl in ("additive", "key_bias"):
        label = f"mask/siglip_{m_rows}/{impl}"
        row: Dict[str, Any] = {"impl": impl, "m_rows": m_rows, "valid": valid}
        try:
            blk = B.EncoderBlockTT(
                dim,
                heads,
                d_pad,
                ff,
                1e-6,
                real_d**-0.5,
                mask_mode="key_valid",
                sdpa_chunks=(64, 128),  # tt/siglip.py SDPA_CHUNKS_SIGLIP
                m_rows=m_rows,
                weights=weights,
                mask=mask,
                mask_impl=impl,
                device=device,
            )
        except Exception as e:  # noqa: BLE001
            row["error"] = MB.short_error(e)
            errors.append(f"{label}: {row['error']}")
            rows_out.append(row)
            continue

        def run(blk: Any = blk) -> Dict[str, Any]:
            nonlocal base
            out = blk.forward(x)
            ttnn.synchronize_device(device)
            got = L.to_torch(out)[:, :, :valid]
            ttnn.deallocate(out)
            if base is None:
                base = got
                met = {"pcc": 1.0, "max_abs": 0.0, "exact": True}
            else:
                met = {**_pcc(got, base), "exact": bool(torch.equal(got, base))}
            t = _time(device, lambda: blk.forward(x), iters, reps)
            return {**met, **t, "block": blk.describe()}

        rows_out.append(_try(row, run, errors, label))
        r = rows_out[-1]
        stages[label] = {
            "metric": "us",
            "value": r.get("us"),
            "pcc_vs_additive": r.get("pcc"),
            "exact": r.get("exact"),
            "error": r.get("error"),
        }
    for t in (x, mask, *weights.values()):
        ttnn.deallocate(t)
    return stages, rows_out, errors, {"m_rows": m_rows, "valid": valid}


# --------------------------------------------------------------------------------------------------------------------
# residual: Qwen3LayerTT residual stream DRAM vs L1 at S = 384 / 512
# --------------------------------------------------------------------------------------------------------------------
class _DictW:
    """Minimal ``TTWeights`` stand-in for ``Qwen3LayerTT`` (``get(name)`` of random device tensors, no options)."""

    def __init__(self, tensors: Mapping[str, Any]) -> None:
        self._t = dict(tensors)

    def get(self, name: str) -> Any:
        return self._t[name]


def sweep_residual(
    device: Any, *, version: str, iters: int, reps: int
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    from models.experimental.gr00t.common.configs import get_config
    from models.experimental.gr00t.tt import qwen3 as Q
    from models.experimental.gr00t.tt.layout import DeviceLayout
    from models.experimental.gr00t.tt.policy import TTPolicy
    from models.experimental.gr00t.tt.weights import rope_transformation_matrix

    ttnn = _ttnn()
    cfg = get_config(version)
    dl = DeviceLayout.from_layout(cfg, cfg.canonical_layout)
    s_pad = int(dl.s_pad)
    if dl.text_padding != "right":
        raise ValueError(
            f"{version}: residual sweep runs the causal layer (right padding), layout is {dl.text_padding!r}"
        )
    pol = TTPolicy()
    llm = cfg.llm
    names = Q.layer_weight_names(0)
    shapes = {
        "Wqkv": (llm.hidden, (llm.n_heads + 2 * llm.n_kv_heads) * llm.head_dim),
        "Wo": (llm.n_heads * llm.head_dim, llm.hidden),
        "W_gate_up": (llm.hidden, 2 * llm.intermediate),
        "W_down": (llm.intermediate, llm.hidden),
        "input_ln.w": (1, llm.hidden),
        "post_attn_ln.w": (1, llm.hidden),
        "q_norm.w": (1, llm.head_dim),
        "k_norm.w": (1, llm.head_dim),
    }
    tensors: Dict[str, Any] = {}
    for i, (stem, shp) in enumerate(shapes.items()):
        scale = 1.0 / math.sqrt(shp[0]) if shp[0] > 1 else 1.0
        t = _rand(shp, 40 + i, scale)
        if stem.endswith(".w"):
            t = (torch.ones(shp) + 0.05 * torch.randn(shp, generator=torch.Generator().manual_seed(50 + i))).to(
                torch.bfloat16
            )
        tensors[names[stem]] = _dev(t, device, "DRAM")
    W = _DictW(tensors)
    x = _dev(_rand((1, 1, s_pad, llm.hidden), 60), device, "DRAM")
    cos = _dev(_rand((1, 1, s_pad, llm.head_dim), 61), device, "DRAM")
    sin = _dev(_rand((1, 1, s_pad, llm.head_dim), 62), device, "DRAM")
    trans = _dev(rope_transformation_matrix(), device, "DRAM")
    rows_out: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    base = None
    for mem_name in ("DRAM", "L1"):
        layer = Q.Qwen3LayerTT(cfg, 0, W, pol, dl, residual_mem=D.MemPreset(mem_name))
        label = f"residual/{version}/S{s_pad}/{mem_name}"
        row = {
            "version": version,
            "s_pad": s_pad,
            "residual_mem": mem_name,
            "swiglu": layer.swiglu,
            "layer": layer.describe(),
        }

        def fn() -> Any:
            return layer.forward(x, cos, sin, trans, None, None)

        def run() -> Dict[str, Any]:
            nonlocal base
            out = fn()
            ttnn.synchronize_device(device)
            got = L.to_torch(out)
            ttnn.deallocate(out)
            if base is None:
                base = got
                exact = True
            else:
                exact = bool(torch.equal(got, base))
            t = _time(device, fn, iters, reps)
            return {**t, "exact_vs_dram": exact, "finite": bool(torch.isfinite(got).all())}

        rows_out.append(_try(row, run, errors, label))
        r = rows_out[-1]
        stages[label] = {"metric": "us", "value": r.get("us"), "exact": r.get("exact_vs_dram"), "error": r.get("error")}
    for t in (x, cos, sin, trans, *tensors.values()):
        ttnn.deallocate(t)
    return stages, rows_out, errors, {"s_pad": s_pad}


# --------------------------------------------------------------------------------------------------------------------
# embed: resident-table gather vs the host path
# --------------------------------------------------------------------------------------------------------------------
def sweep_embed(
    device: Any, *, version: str, sample: str, iters: int, reps: int, host_reps: int = 20
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    from models.experimental.gr00t.common import weights as CW
    from models.experimental.gr00t.common.checkpoint import LazyCheckpoint
    from models.experimental.gr00t.tt.layout import DeviceLayout
    from models.experimental.gr00t.tt.tracer import DeviceEmbeddingGather, InputSpec, TTNNTraceBackend

    ttnn = _ttnn()
    g = _golden(version, sample)
    mi, cfg = g["mi"], g["setup"]["cfg"]
    dl = DeviceLayout.for_inputs(cfg, mi)
    s_pad = int(dl.s_pad)
    ck = LazyCheckpoint(version)
    table = ck.get(cfg.checkpoint.llm_prefix + "embed_tokens.weight")
    if table.dtype != torch.bfloat16:
        raise ValueError(f"embed table dtype {table.dtype} != bf16")
    ids = mi.input_ids_padded
    slots = tuple(tuple(int(x) for x in run) for run in mi.slots.image_slots)
    # host reference == TTWeights.embed_rows semantics (image slots zeroed)
    ref = CW.embed_rows(table, torch.as_tensor(ids).long()).clone()
    for lo, hi in slots:
        ref[:, lo:hi] = 0
    ref = ref.reshape(1, 1, s_pad, table.shape[1]).to(torch.bfloat16)
    rows_out: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    meta: Dict[str, Any] = {
        "s_pad": s_pad,
        "vocab": int(table.shape[0]),
        "n_image_slots": len(slots),
        "shape_key": str(mi.shape_key),
    }
    gather = DeviceEmbeddingGather(table, device, s_pad=s_pad, zero_slots=slots, mem="L1")
    meta["gather"] = gather.describe()
    be = TTNNTraceBackend()
    ids_dev = be.allocate(gather.spec, device)
    ids_host = gather.host_ids(ids)

    def write_ids() -> None:
        be.write(ids_host, ids_dev, gather.spec, 0)

    # exactness (with and without the row mask)
    for label, zero in (("device_gather/masked", True), ("device_gather/unmasked", False)):
        row = {"variant": label}

        def run(zero: bool = zero) -> Dict[str, Any]:
            write_ids()
            ttnn.synchronize_device(device)
            if not zero:
                gather.row_mask, saved = None, gather.row_mask
            try:
                out = gather.gather(ids_dev)
                ttnn.synchronize_device(device)
                got = L.to_torch(out).to(torch.bfloat16)
                ttnn.deallocate(out)
                t = _time(device, lambda: gather.gather(ids_dev), iters, reps)
            finally:
                if not zero:
                    gather.row_mask = saved
            exact = bool(torch.equal(got, ref))
            if not zero:  # the image slots differ by construction; compare the other rows
                keep = torch.ones(s_pad, dtype=torch.bool)
                for lo, hi in slots:
                    keep[lo:hi] = False
                exact = bool(torch.equal(got[:, :, keep], ref[:, :, keep]))
            return {**t, "exact_vs_host_embed_rows": exact}

        rows_out.append(_try(row, run, errors, f"embed:{label}"))
        r = rows_out[-1]
        stages[label] = {
            "metric": "us",
            "value": r.get("us"),
            "exact": r.get("exact_vs_host_embed_rows"),
            "error": r.get("error"),
        }
    # ids upload (host -> uint32 ROW_MAJOR buffer) wall time
    ttnn.synchronize_device(device)
    tw = []
    for _ in range(host_reps):
        t0 = time.perf_counter()
        write_ids()
        ttnn.synchronize_device(device)
        tw.append((time.perf_counter() - t0) * 1e6)
    stages["ids_upload"] = {
        "metric": "us",
        "value": round(statistics.median(tw), 1),
        "bytes": 4 * s_pad,
        **{k: round(v, 1) for k, v in _stats(tw).items() if k != "n"},
    }
    # host path: embed_rows + tilize + copy into a persistent [1,1,S_pad,E] bf16 buffer (the D2 path)
    spec = InputSpec("text_embeds", (1, 1, s_pad, int(table.shape[1])), "bf16", True)
    buf = be.allocate(spec, device)
    th, tg, tt_ = [], [], []
    for _ in range(host_reps):
        t0 = time.perf_counter()
        rows = CW.embed_rows(table, torch.as_tensor(ids).long()).clone()
        for lo, hi in slots:
            rows[:, lo:hi] = 0
        host = rows.reshape(1, 1, s_pad, table.shape[1]).to(torch.bfloat16)
        t1 = time.perf_counter()
        h = ttnn.from_torch(host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
        t2 = time.perf_counter()
        ttnn.copy_host_to_device_tensor(h, buf)
        ttnn.synchronize_device(device)
        t3 = time.perf_counter()
        tg.append((t1 - t0) * 1e6)
        th.append((t2 - t1) * 1e6)
        tt_.append((t3 - t2) * 1e6)
    stages["host_gather"] = {"metric": "us", "value": round(statistics.median(tg), 1)}
    stages["host_tilize"] = {"metric": "us", "value": round(statistics.median(th), 1)}
    stages["host_copy"] = {
        "metric": "us",
        "value": round(statistics.median(tt_), 1),
        "bytes": 2 * s_pad * int(table.shape[1]),
    }
    stages["host_path_total"] = {
        "metric": "us",
        "value": round(statistics.median([a + b + c for a, b, c in zip(tg, th, tt_)]), 1),
    }
    ttnn.deallocate(buf)
    ttnn.deallocate(ids_dev)
    gather.release()
    return stages, rows_out, errors, meta


# --------------------------------------------------------------------------------------------------------------------
# cq: Gr00tTT replay with cq1_uploads off / on (2-CQ device)
# --------------------------------------------------------------------------------------------------------------------
def sweep_cq(
    device: Any, *, version: str, sample: str, reps: int, warmup: int = 5
) -> Tuple[Dict[str, Any], List[Dict[str, Any]], List[str], Dict[str, Any]]:
    import torch

    from models.experimental.gr00t.tt.model import Gr00tTT
    from models.experimental.gr00t.tt.policy import TTPolicy

    ttnn = _ttnn()
    g = _golden(version, sample)
    mi = g["mi"]
    noise = E2E.make_noise(g["setup"]["layout"], 0)
    rows_out: List[Dict[str, Any]] = []
    errors: List[str] = []
    stages: Dict[str, Any] = {}
    preds: Dict[str, Any] = {}
    for flag in (False, True):
        label = f"replay/cq1_uploads={int(flag)}"
        row: Dict[str, Any] = {"cq1_uploads": flag}
        model = None
        try:
            # the Stage-1 model on the default-L1 device (the megakernel default needs open_model_device)
            model = Gr00tTT.from_pretrained(version, policy=TTPolicy(dit_backend="ttnn"), device=device)
            model.executor.set_cq1_uploads(flag)
            model.warm_and_capture(mi, noise)
            times = []
            for i in range(warmup + reps):
                t0 = time.perf_counter()
                out = model.executor.replay(mi, noise=noise)
                ttnn.synchronize_device(device)
                dt = (time.perf_counter() - t0) * 1e3
                if i >= warmup:
                    times.append(dt)
            preds[label] = ttnn.to_torch(out).clone()
            row.update(
                {k: round(v, 3) for k, v in _stats(times).items()},
                upload_s=model.executor.timing.get("last_replay_upload_s"),
                executor=model.executor.describe()["cq1_uploads"],
            )
            stages[label] = {
                "metric": "ms",
                "value": round(min(times), 3),
                "median_ms": round(statistics.median(times), 3),
            }
        except Exception as e:  # noqa: BLE001
            row["error"] = MB.short_error(e)
            errors.append(f"cq:{label}: {row['error']}")
        finally:
            if model is not None:
                model.release()
        rows_out.append(row)
    if len(preds) == 2:
        a, b = preds.values()
        eq = bool(torch.equal(a, b))
        stages["bit_equal"] = {
            "metric": "us",
            "value": 1.0 if eq else 0.0,
            "note": "1 = cq1 result bit-equal to the 1-CQ result",
        }
    return stages, rows_out, errors, {"shape_key": str(mi.shape_key), "warmup": warmup, "reps": reps}


# --------------------------------------------------------------------------------------------------------------------
# e2e: the tuned per-stage JSON (bench_e2e method)
# --------------------------------------------------------------------------------------------------------------------
def run_e2e(
    device: Any,
    *,
    version: str,
    sample: str,
    fork: Optional[bool],
    warmup: int,
    reps: int,
    refresh_baseline: bool,
    out_dir: Path,
) -> Tuple[Dict[str, Any], Path, Optional[Path]]:
    from models.experimental.gr00t.tt.qwen3vl_vision import VIT_FORK_ENV

    use_fork = (version == "n17") if fork is None else bool(fork)
    prev = os.environ.get(VIT_FORK_ENV)
    os.environ[VIT_FORK_ENV] = (
        "1" if use_fork else "0"
    )  # adapter: BackboneTT.build constructs the tower without a policy
    try:
        payload = E2E.run(device, version=version, sample=sample, warmup=warmup, reps=reps, verbose=True)
    finally:
        if prev is None:
            os.environ.pop(VIT_FORK_ENV, None)
        else:
            os.environ[VIT_FORK_ENV] = prev
    payload["meta"]["wp"] = "WP-P"
    payload["meta"]["vit_fork"] = use_fork
    payload["meta"]["levers"] = {
        "vit_fork": use_fork,
        "layers_tables": "tt/layers.py MATMUL_CONFIGS / SDPA_CHUNKS as of this run",
        "residual_mem": "tt/qwen3.py residual_mem_for",
        "head_ops": "tt/blocks.py / tt/layers.py split_heads / concat_heads",
    }
    tuned_path = SO.write_perf_json(
        payload,
        out_dir,
        path=tuned_json_path(version) if out_dir == RESULTS_DIR else out_dir / tuned_json_path(version).name,
    )
    base_path = SO.write_perf_json(payload, out_dir) if refresh_baseline else None
    return payload, tuned_path, base_path


# --------------------------------------------------------------------------------------------------------------------
# report: winners vs the current tables (no ttnn)
# --------------------------------------------------------------------------------------------------------------------
def report_matmul(path: Path) -> List[Dict[str, Any]]:
    """Per (shape, dtype, mem, fused): the table row's µs vs the fastest admitted candidate (PCC >= PCC_ADMIT)."""
    d = SO.load_perf_json(path)
    groups: Dict[Tuple[Any, ...], List[Dict[str, Any]]] = {}
    for r in d["rows"]:
        groups.setdefault((r["shape"], r["dtype"], r["mem"], r.get("fused")), []).append(r)
    out = []
    for key, rows in groups.items():
        ok = [r for r in rows if r.get("us") is not None and r.get("admit")]
        table = next((r for r in rows if r["candidate"] == "table"), None)
        best = min(ok, key=lambda r: r["us"]) if ok else None
        out.append(
            {
                "shape": key[0],
                "dtype": key[1],
                "mem": key[2],
                "fused": key[3],
                "table_us": None if table is None else table.get("us"),
                "table_config": None if table is None else table.get("config"),
                "table_error": None if table is None else table.get("error"),
                "best_us": None if best is None else best["us"],
                "best": None if best is None else best["candidate"],
                "best_config": None if best is None else best.get("config"),
                "gain_us": (
                    None
                    if (best is None or table is None or table.get("us") is None)
                    else round(table["us"] - best["us"], 2)
                ),
                "n_failed": sum(1 for r in rows if r.get("error")),
            }
        )
    return out


def report_sdpa(path: Path) -> List[Dict[str, Any]]:
    """Per (shape, grid, ckc): the current allowlist default's µs vs the fastest admitted pair."""
    d = SO.load_perf_json(path)
    groups: Dict[Tuple[Any, ...], List[Dict[str, Any]]] = {}
    for r in d["rows"]:
        groups.setdefault((r["shape"], r["grid"], r["ckc"]), []).append(r)
    out = []
    for (shape, grid, ckc), rows in groups.items():
        sh = next(s for s in SDPA_SHAPES if s.name == shape)
        default = L.default_sdpa_chunks(sh.cls, sh.s, sh.s, sh.d)
        cur = next((r for r in rows if (r["q_chunk"], r["k_chunk"]) == default), None)
        ok = [r for r in rows if r.get("us") is not None and r.get("admit") and not r["forbidden"]]
        best = min(ok, key=lambda r: r["us"]) if ok else None
        out.append(
            {
                "shape": shape,
                "grid": grid,
                "ckc": ckc,
                "default": list(default),
                "default_us": None if cur is None else cur.get("us"),
                "default_pcc": None if cur is None else cur.get("pcc"),
                "best": None if best is None else [best["q_chunk"], best["k_chunk"]],
                "best_us": None if best is None else best["us"],
                "best_pcc": None if best is None else best.get("pcc"),
                "wrong_pairs": [
                    [r["q_chunk"], r["k_chunk"]] for r in rows if r.get("pcc") is not None and r["pcc"] < PCC_ADMIT
                ],
                "failed_pairs": [[r["q_chunk"], r["k_chunk"]] for r in rows if r.get("error")],
            }
        )
    return out


def print_report(path: Path) -> None:
    d = SO.load_perf_json(path)
    bench = d["bench"]
    if bench.endswith("_matmul"):
        for r in report_matmul(path):
            print(
                f"{r['shape']:12s} {r['dtype']:6s} {r['mem']:5s} fused={str(r['fused']):5s} table={r['table_us']} "
                f"best={r['best_us']} ({r['best']}) gain={r['gain_us']} failed={r['n_failed']} | {r['best_config']}"
            )
    elif bench.endswith("_sdpa"):
        for r in report_sdpa(path):
            print(
                f"{r['shape']:14s} {r['grid']:5s} {r['ckc']:14s} default={r['default']}={r['default_us']} "
                f"(pcc {r['default_pcc']}) best={r['best']}={r['best_us']} (pcc {r['best_pcc']}) wrong={r['wrong_pairs']} "
                f"failed={r['failed_pairs']}"
            )
    else:
        for k, st in d["stages"].items():
            print(
                f"{k:60s} {st.get('value')} {st.get('metric')} {({kk: vv for kk, vv in st.items() if kk not in ('metric', 'value')})}"
            )


# --------------------------------------------------------------------------------------------------------------------
# dry run (no ttnn)
# --------------------------------------------------------------------------------------------------------------------
def dry_run(whats: Sequence[str], version: str) -> Dict[str, Any]:
    """Resolve every table lookup the sweeps rely on without ttnn; returns a summary dict."""
    out: Dict[str, Any] = {"version": version, "whats": list(whats)}
    if "vit" in whats:
        out["vit_variants"] = {k: v.to_dict() for k, v in vit_variants().items()}
    if "matmul" in whats:
        out["matmul_candidates"] = {
            name: [c["name"] for c in matmul_candidates(m, k, n, "bf16")] for name, m, k, n, _ in MATMUL_SHAPES
        }
    if "sdpa" in whats:
        out["sdpa_shapes"] = {
            s.name: {
                "pairs": [list(p) for p in s.pairs],
                "allowlisted": [list(p) for p in s.pairs if L.sdpa_chunk_ok(s.s, s.s, s.d, p)],
            }
            for s in SDPA_SHAPES
        }
    if "heads" in whats:
        out["head_shapes"] = [list(s) for s in HEAD_SHAPES]
    if "e2e" in whats:
        out["tuned_json"] = str(tuned_json_path(version))
    return out


# --------------------------------------------------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------------------------------------------------
def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--what", default="vit", help="comma-separated sub-sweeps or 'all': " + ",".join(SWEEPS))
    ap.add_argument("--version", default="n17", choices=VERSIONS)
    ap.add_argument("--sample", default="canonical")
    ap.add_argument("--iters", type=int, default=DEFAULT_ITERS)
    ap.add_argument("--reps", type=int, default=DEFAULT_REPS)
    ap.add_argument("--tower-reps", type=int, default=DEFAULT_TOWER_REPS)
    ap.add_argument("--variants", default="", help="vit: comma-separated TunedVisionConfig variant names (default all)")
    ap.add_argument("--no-blocks-timing", action="store_true", help="vit: skip the 24-block trace timing")
    ap.add_argument("--dtypes", default="bf16", help="matmul: comma-separated weight dtypes")
    ap.add_argument("--num-cqs", type=int, default=1, help="command queues to open (2 for --what cq)")
    ap.add_argument(
        "--fork", default="auto", choices=("auto", "on", "off"), help="e2e: N1.7 tower (auto = fork for n17)"
    )
    ap.add_argument("--warmup", type=int, default=E2E.DEFAULT_WARMUP)
    ap.add_argument("--e2e-reps", type=int, default=E2E.DEFAULT_REPS)
    ap.add_argument("--refresh-baseline", action="store_true", help="e2e: also write bench_e2e_<version>_<ts>.json")
    ap.add_argument("--out", default=str(RESULTS_DIR))
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument(
        "--report", default=None, help="print the winners-vs-table report of a sweep JSON and exit (no ttnn)"
    )
    return ap.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)
    whats = list(SWEEPS) if args.what == "all" else [w.strip() for w in args.what.split(",") if w.strip()]
    for w in whats:
        if w not in SWEEPS:
            raise SystemExit(f"unknown sweep {w!r}; choose from {SWEEPS}")
    out_dir = Path(args.out)
    if args.report:
        print_report(Path(args.report))
        return 0
    if args.dry_run:
        import json

        print(json.dumps(dry_run(whats, args.version), indent=1, default=str))
        return 0
    os.environ.setdefault("TT_METAL_CACHE", str(SO.DEFAULT_TT_METAL_CACHE))
    device = D.open_gr00t_device(num_command_queues=args.num_cqs)
    rc = 0
    try:
        for what in whats:
            t0 = time.perf_counter()
            common = {
                "version": args.version,
                "sample": args.sample,
                "iters": args.iters,
                "reps": args.reps,
                "num_cqs": args.num_cqs,
            }
            if what == "e2e":
                fork = None if args.fork == "auto" else (args.fork == "on")
                payload, tuned_path, base_path = run_e2e(
                    device,
                    version=args.version,
                    sample=args.sample,
                    fork=fork,
                    warmup=args.warmup,
                    reps=args.e2e_reps,
                    refresh_baseline=args.refresh_baseline,
                    out_dir=out_dir,
                )
                print(f"[{BENCH_NAME}] e2e -> {tuned_path}" + (f" (+ baseline {base_path})" if base_path else ""))
                continue
            if what == "vit":
                variants = [v for v in args.variants.split(",") if v] or None
                stages, rows, errors, meta = sweep_vit(
                    device,
                    version=args.version,
                    sample=args.sample,
                    variants=variants,
                    reps=args.tower_reps,
                    blocks_timing=not args.no_blocks_timing,
                )
            elif what == "rope":
                stages, rows, errors, meta = sweep_rope(device, iters=args.iters, reps=args.reps)
            elif what == "matmul":
                stages, rows, errors, meta = sweep_matmul(
                    device, iters=args.iters, reps=args.reps, dtypes=[d for d in args.dtypes.split(",") if d]
                )
            elif what == "sdpa":
                stages, rows, errors, meta = sweep_sdpa(device, iters=args.iters, reps=args.reps)
            elif what == "heads":
                stages, rows, errors, meta = sweep_heads(device, iters=args.iters, reps=args.reps)
            elif what == "mask":
                stages, rows, errors, meta = sweep_mask(device, iters=args.iters, reps=args.reps)
            elif what == "residual":
                stages, rows, errors, meta = sweep_residual(
                    device, version=args.version, iters=max(1, args.iters // 2), reps=args.reps
                )
            elif what == "embed":
                stages, rows, errors, meta = sweep_embed(
                    device, version=args.version, sample=args.sample, iters=args.iters, reps=args.reps
                )
            elif what == "cq":
                stages, rows, errors, meta = sweep_cq(
                    device, version=args.version, sample=args.sample, reps=args.e2e_reps
                )
            else:
                raise AssertionError(what)
            meta = dict(meta, wall_s=round(time.perf_counter() - t0, 1))
            payload = _payload(what, args.version, common, stages, rows, errors, meta, device)
            path = _write(payload, what, args.version, out_dir)
            print(f"[{BENCH_NAME}] {what}: {len(rows)} rows, {len(errors)} errors, {meta['wall_s']} s -> {path}")
            for k, st in stages.items():
                v = st.get("value")
                extra = {kk: vv for kk, vv in st.items() if kk not in ("metric", "value") and vv is not None}
                print(
                    f"   {k:60s} {v if v is not None else '-':>10} {st.get('metric')}  {extra if extra else ''}"[:220]
                )
            if errors:
                rc = 1
                for e in errors[:40]:
                    print("   ERROR", e[:300], file=sys.stderr)
    finally:
        D.close_gr00t_device(device)
    return rc


__all__ = [
    "BENCH_NAME",
    "HEAD_SHAPES",
    "KNOWN_VIT_MISSES",
    "MATMUL_SHAPES",
    "PCC_ADMIT",
    "SDPA_SHAPES",
    "SWEEPS",
    "SdpaShape",
    "dry_run",
    "main",
    "matmul_candidates",
    "print_report",
    "report_matmul",
    "report_sdpa",
    "run_e2e",
    "sweep_cq",
    "sweep_embed",
    "sweep_heads",
    "sweep_mask",
    "sweep_matmul",
    "sweep_residual",
    "sweep_rope",
    "sweep_sdpa",
    "sweep_vit",
    "tuned_json_path",
    "vit_variants",
]

if __name__ == "__main__":
    sys.exit(main())