# 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_.json`` (fork tower for n17 via ``GR00T_VIT_FORK``; ``--refresh-baseline`` also writes a ``bench_e2e__.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___.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_.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_.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__.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())