code/models/experimental/gr00t/tt/model.py

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

"""``Gr00tTT``: the assembled GR00T N1.5 / N1.6 / N1.7 device policy -- ``BackboneTT`` -> ``VLAdapterTT`` ->
``ActionHeadTT`` under one ``TraceExecutor`` (IMPLEMENTATION_PLAN.md §2.2 row ``tt/model.py``, §2.3 "Version
parametrisation", §3.1 "Host per call", §4 "Trace plan", §7.2 last row, §7.4 / §7.5; WP-14).

This is the **only** module that switches on the version, and it does so once, at load, through the recipe flags of
plan §2.3 (:class:`ModelRecipe`; the backbone half is ``tt.backbone.BackboneRecipe``).  No forward contains an
``if version``: every difference is a constructor choice (tower / projector / LLM mask / RoPE table / final norm /
deepstack / VL-SA / future tokens / cross key sets / Euler dtype) or a fixed mask-presence pattern of the layout.

Per call (plan §3.1)::

    obs --common.preprocessing.encode(layout_name fixed)--> mi
    host: text_embeds = TTWeights.embed_rows(ids, slots)    patches / pixels   N1.6 cos/sin rows   N1.6 LLM mask
          VL-SA mask (n15 / n17)   DiT cross masks (text | all)   state_rows [1,1,32,D_pad]   x_t = noise_to_xt(noise)
    TraceExecutor.replay: write every per-call input -> execute_trace(A_vis, A_llm, A_adp, B) | (A, B)
    action_pred [1,1,64,A_pad] --xt_to_action_pred--> [1,H,A_max] --common.normalization.decode--> {group: (1,H,D)}

Stages registered on the executor (plan §4; names are the trace names of the ``per_stage`` layout, the ``two``
layout groups them into ``A`` / ``B``)::

    vision   (A)  BackboneTT.forward_vision   inputs: patches | vis_patches (+ constant vis_pos / vis_cos / vis_sin)
    llm      (A)  BackboneTT.forward_llm      inputs: text_embeds, rope_cos/sin (per call for N1.6), llm_mask (N1.6)
    adapter  (A)  VLAdapterTT.forward         inputs: vl_mask (n15 / n17), state_rows      -> state_features
                                              side effect: hoisted K_i / V_i in the adapter's persistent buffers
    denoise  (B)  ActionHeadTT.denoise        inputs: x_t, cross_text_mask | cross_all_mask; constants self_mask,
                                              cross_image_mask (when L_img_pad != L_img_real)  -> action_pred (a clone
                                              of the head's persistent buffer, so the executor owns the trace output).
                                              ``policy.dit_backend``: the Stage-1 ttnn ops or the ``steps4``
                                              megakernel ``generic_op`` (WP-K5) -- same inputs, same output buffer.

The megakernel backend -- the ``TTPolicy`` **default** since 2026-09-18 (``tt.policy.DEFAULT_DIT_BACKEND``; the ttnn
Stage-1 path stays selectable through ``TTPolicy(dit_backend="ttnn")`` / ``GR00T_DIT_BACKEND=ttnn``) -- needs the
device opened with a smaller ``worker_l1_size`` (the kernel-config ring buffer of ``mk_k4_summary.md`` §7.2 fact 2;
:data:`MK_L1_CUT_BYTES_DEFAULT`): :func:`open_model_device` does that from the policy and
``from_pretrained(device=None)`` uses it; a device that was opened another way is refused at build
(:func:`device_worker_l1_size`).  Under the megakernel the ttnn DiT weights are not uploaded
(:func:`megakernel_weight_names`: the arena replaces the ``dit`` / ``dit_precompute`` / ``encoders`` tensors except
the adapter's ``dit.block.{i}.Wkv`` / ``bkv``, ``enc.state.*`` and ``enc.action.b3_pos`` that ``TTWeights.constants``
reads); ``from_pretrained(load_ttnn_dit_weights=True)`` keeps the full Stage-1 upload.

Every per-call tensor -- the masks included -- is an executor buffer (static-shape §9 / §10: written unconditionally
before every replay, address-checked, presence fixed per layout), so ``ActionHeadTT`` is built on the executor's
mask buffers and ``update_masks`` is never used here.

Per-call host cost (``docs/plan/n17-perf.md``; every item bit-exact by construction):

* ``embed_gather="device"`` (default): the text rows come from :class:`~tt.tracer.DeviceEmbeddingGather` inside the
  LLM stage (resident 622 MB ROW_MAJOR table, ``ttnn.embedding`` + image-slot row mask, bit-equal to
  ``TTWeights.embed_rows``); the per-call input ``text_embeds`` (2 MB tilize + upload) becomes the uint32 token ids
  ``input_ids`` (2 KB).  ``"host"`` keeps the WP-12 host gather.
* :meth:`Gr00tTT.host_inputs` builds the per-call inputs itself (``BackboneTT.check_call`` once, the patch rows as
  ``pixel_values.to(bf16)``, ``text_embeds_host`` / ``rope_host`` / ``llm_mask_host``) instead of
  ``BackboneTT.host_inputs``, whose N1.7 path re-derives the constant ``vis_pos`` / ``vis_cos`` / ``vis_sin`` tables
  through ``tower.prepare_inputs`` on every call (1.9 ms) although they are bound once in ``constant_host_inputs``.
* the executor runs with ``prepare_asserts=True`` (``check_call`` is the static-shape pass) and checks finiteness once
  per call; the per-tensor upload accounting is ``executor.upload_stats``.
* ``cq1_uploads`` (default: on when the device was opened with two command queues) writes the inputs of trace ``i + 1``
  on CQ 1 while trace ``i`` runs (WP-P); ``from_pretrained(device=None)`` opens the device with ``num_command_queues``.

Untraced path (:meth:`Gr00tTT.run_untraced`, plan §7.2 last row / §7.5 step 1): the same stage functions run eagerly
on the same buffers with every module's tap collection on, and the taps come back as **host** tensors in the goldens'
shapes under the canonical golden names (``GoldenSet.threshold_key`` form: ``vit_block_last``, ``llm_layer_first``,
``sa_embs[k=2]``, ``vit_deepstack_mergers[1]`` ...; :func:`expected_tap_keys` lists them per version).  Device taps are
freed before the method returns, so an untraced run between two replays leaves the captured allocation state intact.

``ttnn`` is imported lazily (constructor / forwards); :class:`ModelRecipe`, :func:`model_input_specs`,
:func:`expected_tap_keys` and the constants below are CPU-testable.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple

import numpy as np
import torch

from models.experimental.gr00t.common.configs import DIT_QUERY_ROWS_PAD, GR00TConfig, LayoutConfig, get_config
from models.experimental.gr00t.tt import layers as L
from models.experimental.gr00t.tt.action_head import ActionHeadTT, head_tap_names
from models.experimental.gr00t.tt.adapter import STATE_ROWS, HoistOut, VLAdapterTT, slots_to_rows, state_rows_host
from models.experimental.gr00t.tt.backbone import (
    IN_COS,
    IN_LLM_MASK,
    IN_PATCHES,
    IN_SIN,
    IN_TEXT,
    IN_VIS_PATCHES,
    OUT_FEATURES,
    STAGE_LLM,
    STAGE_VISION,
    TAP_INPUTS_EMBEDS,
    TAP_PROJECTOR_OUT,
    BackboneRecipe,
    BackboneTT,
)
from models.experimental.gr00t.tt.device import (
    DEFAULT_L1_SMALL_SIZE,
    DEFAULT_TRACE_REGION_SIZE,
    DRAM,
    L1,
    close_gr00t_device,
    open_gr00t_device,
    resolve_mem,
)
from models.experimental.gr00t.tt.encoders import noise_to_xt, xt_dtype, xt_to_action_pred
from models.experimental.gr00t.tt.layout import MASK_NAMES, DeviceLayout, StaticShapeError
from models.experimental.gr00t.tt.policy import TTPolicy
from models.experimental.gr00t.tt.qwen3 import padded_to_rows
from models.experimental.gr00t.tt.qwen3vl_vision import TAP_KEYS as QV_TAP_KEYS
from models.experimental.gr00t.tt.siglip import TAP_EMBEDDINGS as SIGLIP_TAP_EMBEDDINGS
from models.experimental.gr00t.tt.siglip import TAP_POST_LN as SIGLIP_TAP_POST_LN
from models.experimental.gr00t.tt.siglip import tap_block as siglip_tap_block
from models.experimental.gr00t.tt.tracer import DeviceEmbeddingGather, InputSpec, TraceExecutor, TraceStateError

# executor input / constant / output names added by this module (the backbone's live in tt/backbone.py)
IN_TEXT_IDS = "input_ids"  # uint32 token ids of the device embedding gather (replaces the backbone's text_embeds)
EMBED_GATHERS: Tuple[str, ...] = ("host", "device")
DEFAULT_EMBED_GATHER = "device"
IN_VL_MASK = "vl_mask"
IN_STATE_ROWS = "state_rows"
IN_XT = "x_t"
IN_CROSS_TEXT_MASK = "cross_text_mask"
IN_CROSS_ALL_MASK = "cross_all_mask"
CONST_SELF_MASK = "self_mask"
CONST_CROSS_IMAGE_MASK = "cross_image_mask"
OUT_STATE_FEATURES = "state_features"
OUT_ACTION_PRED = "action_pred"
STAGE_ADAPTER = "adapter"
STAGE_DENOISE = "denoise"
#: trace names per layout (plan §4; ``benchmarks/bench_e2e.py`` ``STAGE_NAMES`` expects exactly these)
STAGE_NAMES: Dict[str, Tuple[str, ...]] = {
    "per_stage": (STAGE_VISION, STAGE_LLM, STAGE_ADAPTER, STAGE_DENOISE),
    "two": ("A", "B"),
}
#: mask name of tt.layout -> executor buffer name (the LLM mask belongs to the backbone's input specs)
MASK_BUFFER_NAMES: Dict[str, str] = {
    "self": CONST_SELF_MASK,
    "cross_image": CONST_CROSS_IMAGE_MASK,
    "cross_text": IN_CROSS_TEXT_MASK,
    "cross_all": IN_CROSS_ALL_MASK,
    "llm": IN_LLM_MASK,
    "vl": IN_VL_MASK,
}
#: Flow-matching noise seed of the deployed policies (reference.model.DEFAULT_NOISE_SEED; foundation-f4 §2 item 5:
#: N1.5 ``manual_seed(0)``, N1.6 / N1.7 seed 42) -- ``predict(obs, noise=None)`` mirrors it with a CPU generator.
NOISE_SEED: Mapping[str, int] = {"n15": 0, "n16": 42, "n17": 42}
#: Canonical golden taps the device path does not produce (host inputs, host tables, reference-only intermediates);
#: everything else in ``GoldenSet.available()`` must come out of ``run_untraced`` (plan §7.2 last row).
HOST_OR_UNPRODUCED_TAPS: Tuple[str, ...] = (
    "input_ids",
    "attention_mask",
    "pixel_values",
    "image_geometry",
    "state",
    "embodiment_id",
    "initial_noise",
    "llm_embed_tokens",
    "llm_rotary_cos",
    "llm_rotary_sin",
    "llm_position_ids",
    "backbone_attention_mask",
    "image_mask",
    "position_embedding",
    "noisy_actions_in",
    "timestep",
    "temb",
)
#: produced by ``decode`` (needs the observation), not by ``run_untraced``
DECODE_TAP = "unnormalized_actions"


def unproduced_taps(cfg: GR00TConfig) -> Tuple[str, ...]:
    """:data:`HOST_OR_UNPRODUCED_TAPS` plus ``llm_final_norm`` when the version does not apply the final norm (N1.7:
    the golden is reference-only, ``backbone_features`` is the pre-norm residual; plan §7.2 "llm_final_norm (n15/n16)").
    """
    out = HOST_OR_UNPRODUCED_TAPS
    if not bool(cfg.llm.final_norm_applied):
        out = out + ("llm_final_norm",)
    return out


#: prefix of the raw device-layout host copies ``run_untraced(raw=True)`` adds to the tap dict (determinism tests)
RAW_PREFIX = "raw/"
N_STEPS = 4


def _shape(t: Any) -> Tuple[int, ...]:
    return tuple(int(d) for d in t.shape)


# --------------------------------------------------------------------------------------------------------------------
# recipe (CPU)
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class ModelRecipe:
    """The plan §2.3 flags of one (version, layout, policy): the backbone recipe plus the head-side switches."""

    version: str
    layout_name: str
    embodiment_id: int
    tower: str
    projector: str
    llm_mask: str
    rope_table: str
    final_norm: bool
    deepstack: Tuple[int, ...]
    has_vlsa: bool
    n_vlsa: int
    n_future_tokens: int
    n_blocks: int
    cross_key_sets: Tuple[Tuple[int, str], ...]  # (cross block index, key set) per cfg.dit.cross_key_subset
    euler_dtype: str
    xt_dtype: str
    n_steps: int
    mask_presence: Tuple[Tuple[str, bool], ...]  # DeviceLayout.mask_presence, fixed per layout (§10 item 4)
    trace_layout: str
    dit_backend: str
    dtype_policy: str
    mk_arena_dtype: Optional[str] = None  # the megakernel's weight arena (resolved per version); None for "ttnn"

    @classmethod
    def from_config(cls, cfg: GR00TConfig, dlayout: DeviceLayout, policy: TTPolicy) -> "ModelRecipe":
        """Derive and cross-check the flags (raises on any combination the port does not implement)."""
        bb = BackboneRecipe.from_config(cfg, dlayout, policy)  # validates cfg / dlayout / policy types and versions
        dit = cfg.dit
        cross = tuple((int(i), str(dlayout.key_set(i))) for i in dit.cross_block_indices)
        if not cross:
            raise ValueError(f"{cfg.version}: the DiT has no cross block")
        n_steps = int(cfg.sampler.num_inference_timesteps)
        if n_steps != N_STEPS:
            raise ValueError(f"{cfg.version}: {n_steps} denoising steps; the traced denoise is unrolled for {N_STEPS}")
        if int(dit.m_pad) != DIT_QUERY_ROWS_PAD:
            raise ValueError(f"{cfg.version}: DiT m_pad {dit.m_pad} != {DIT_QUERY_ROWS_PAD}")
        return cls(
            version=cfg.version,
            layout_name=dlayout.layout_name,
            embodiment_id=int(dlayout.embodiment_id),
            tower=bb.tower,
            projector=bb.projector,
            llm_mask=bb.llm_mask,
            rope_table=bb.rope_table,
            final_norm=bb.final_norm,
            deepstack=bb.deepstack,
            has_vlsa=bb.has_vlsa,
            n_vlsa=int(cfg.vl_adapter.n_blocks),
            n_future_tokens=int(dit.n_future_tokens),
            n_blocks=int(dit.n_blocks),
            cross_key_sets=cross,
            euler_dtype=str(cfg.sampler.euler_dtype),
            xt_dtype=xt_dtype(cfg),
            n_steps=n_steps,
            mask_presence=tuple((name, bool(dlayout.mask_presence[name])) for name in MASK_NAMES),
            trace_layout=policy.trace_layout,
            dit_backend=policy.dit_backend,
            dtype_policy=policy.dtype_policy,
            mk_arena_dtype=policy.mk_arena_dtype_for(cfg.version) if policy.dit_backend == "megakernel" else None,
        )

    @property
    def presence(self) -> Dict[str, bool]:
        return dict(self.mask_presence)

    @property
    def stage_names(self) -> Tuple[str, ...]:
        """Trace names in execution order for this recipe's trace layout (plan §4)."""
        return STAGE_NAMES[self.trace_layout]

    def describe(self) -> Dict[str, Any]:
        d = {k: (list(v) if isinstance(v, tuple) else v) for k, v in self.__dict__.items()}
        d["cross_key_sets"] = [list(x) for x in self.cross_key_sets]
        d["mask_presence"] = self.presence
        d["stage_names"] = list(self.stage_names)
        return d


def text_ids_spec(s_pad: int, name: str = IN_TEXT_IDS) -> InputSpec:
    """The per-call token-id input of the device embedding gather (``DeviceEmbeddingGather.spec``: uint32 ROW_MAJOR
    ``[1, 1, 1, S_pad]``)."""
    return InputSpec(name, (1, 1, 1, int(s_pad)), "uint32", True, layout="ROW_MAJOR")


def model_input_specs(
    cfg: GR00TConfig, dlayout: DeviceLayout, policy: TTPolicy, embed_gather: str = "host"
) -> Dict[str, InputSpec]:
    """Every executor buffer of the whole model in stage order: the backbone's specs (``BackboneRecipe.input_specs``)
    followed by the adapter / head per-call inputs and the constant DiT masks (plan §3.1 / §4 "per-call inputs").
    Mask entries follow ``DeviceLayout.mask_presence`` exactly (§10 item 4).  ``embed_gather="device"`` replaces the
    backbone's ``text_embeds`` by the token ids of :func:`text_ids_spec` (module docstring)."""
    if embed_gather not in EMBED_GATHERS:
        raise ValueError(f"embed_gather must be one of {EMBED_GATHERS}, got {embed_gather!r}")
    recipe = ModelRecipe.from_config(cfg, dlayout, policy)
    pres = recipe.presence
    specs: Dict[str, InputSpec] = {}
    for name, spec in BackboneRecipe.from_config(cfg, dlayout, policy).input_specs().items():
        if name == IN_TEXT and embed_gather == "device":
            specs[IN_TEXT_IDS] = text_ids_spec(dlayout.s_pad)
        else:
            specs[name] = spec
    s_vl, m_pad = int(dlayout.s_pad_vl), int(dlayout.m_pad)
    if pres["vl"]:
        specs[IN_VL_MASK] = InputSpec(IN_VL_MASK, (1, 1, s_vl, s_vl), "bf16", True)
    specs[IN_STATE_ROWS] = InputSpec(IN_STATE_ROWS, (1, 1, STATE_ROWS, int(cfg.encoders.state_dim_pad)), "bf16", True)
    specs[IN_XT] = InputSpec(IN_XT, (1, 1, m_pad, int(cfg.encoders.action_dim_pad)), recipe.xt_dtype, True)
    if pres["cross_text"]:
        if dlayout.l_txt_pad is None:
            raise ValueError(f"{dlayout.layout_name}: cross_text mask present but L_txt_pad undefined")
        specs[IN_CROSS_TEXT_MASK] = InputSpec(IN_CROSS_TEXT_MASK, (1, 1, m_pad, int(dlayout.l_txt_pad)), "bf16", True)
    if pres["cross_all"]:
        if dlayout.l_all_pad is None:
            raise ValueError(f"{dlayout.layout_name}: cross_all mask present but L_all_pad undefined")
        specs[IN_CROSS_ALL_MASK] = InputSpec(IN_CROSS_ALL_MASK, (1, 1, m_pad, int(dlayout.l_all_pad)), "bf16", True)
    specs[CONST_SELF_MASK] = InputSpec(CONST_SELF_MASK, (1, 1, m_pad, m_pad), "bf16", False)
    if pres["cross_image"]:
        if dlayout.l_img_pad is None:
            raise ValueError(f"{dlayout.layout_name}: cross_image mask present but L_img_pad undefined")
        specs[CONST_CROSS_IMAGE_MASK] = InputSpec(
            CONST_CROSS_IMAGE_MASK, (1, 1, m_pad, int(dlayout.l_img_pad)), "bf16", False
        )
    for name, present in pres.items():
        buf = MASK_BUFFER_NAMES[name]
        if (buf in specs) != present:
            raise AssertionError(f"mask {name!r}: spec presence {buf in specs} != layout presence {present}")
    return specs


def expected_tap_keys(cfg: GR00TConfig, dlayout: DeviceLayout, dit_backend: str = "ttnn") -> Tuple[str, ...]:
    """The golden-named keys ``run_untraced(return_taps=True)`` produces for this version (plan §7.2 matrix; keys in
    ``GoldenSet.threshold_key`` form).  Per-block DiT taps cover every block; the goldens hold blocks 0 / 1 / last.
    ``dit_backend`` selects the head's tap set (``tt.action_head.head_tap_names``): the megakernel exposes only the
    per-step decoder output and the final prediction (WP-K5).  The parameter default stays ``"ttnn"`` on purpose:
    it is the **full** Stage-1 set (the superset the golden tap map is checked against on CPU); model-level callers
    pass ``model.policy.dit_backend``."""
    r = ModelRecipe.from_config(cfg, dlayout, TTPolicy(dit_backend=dit_backend))
    keys: List[str] = ["vit_embeddings", "vit_block_first", "vit_block_last"]
    if r.tower == "siglip":
        keys.append("vit_post_ln")
    else:
        keys += [f"vit_deepstack_mergers[{d}]" for d in range(len(r.deepstack))]
    keys += ["projector_out", "llm_inputs_embeds", "llm_layer_first", "llm_layer_last"]
    if r.final_norm:
        keys.append("llm_final_norm")
    keys += ["backbone_features", "vlln_out"]
    keys += [f"vlsa_block{i}_out" for i in range(r.n_vlsa)]
    if r.has_vlsa:
        keys.append("vl_self_attention_out")
    keys += ["dit_encoder_hidden_states", "state_features"]
    keys += list(head_tap_names(r.dit_backend, r.n_steps, r.n_blocks))
    if len(set(keys)) != len(keys):
        raise AssertionError("duplicate tap keys")
    return tuple(keys)


# --------------------------------------------------------------------------------------------------------------------
# device opening per policy (WP-K5: the megakernel's worker_l1_size)
# --------------------------------------------------------------------------------------------------------------------


#: Bytes taken off every core's allocatable L1 (``worker_l1_size``) for the megakernel backend: the kernel-config
#: ring buffer grows by the same amount (``mk_k4_summary.md`` §7.2 fact 2).  Measured compute-core K4 program sizes
#: (``program.cpp:3224``): N1.6 91,168 B, N1.5 108,912 B (its future-token / mask paths) against the p150a default
#: ring of 70,656 B, so >= 38,256 B are needed; 64 KiB leaves 27 KB of headroom (K4's own tests use 96 KiB,
#: ``dit_program.K4_WORKER_L1_SIZE``).  Any cut >= 32 KiB breaks the N1.5 / N1.7 **adapter** stage with L1
#: intermediates: its VL-SA programs (S_pad 544 / 480) allocate static CBs up to 1,405,952 B and the stage's L1
#: activations extend down to <= 1,339,392 B (``program.cpp:2149``; WP-K5 bring-up sweeps at 96 / 48 / 40 KiB), i.e.
#: the two requirements leave no common value -- hence :data:`ADAPTER_MEM_BY_BACKEND` moves the adapter's
#: intermediates to DRAM under the megakernel backend.  ``GR00T_MK_L1_CUT_KIB`` overrides the value (sweeps).
MK_L1_CUT_BYTES_DEFAULT: int = 64 * 1024
#: Memory of the ``VLAdapterTT`` intermediates per DiT backend (``VLAdapterTT(mem=)``; the persistent K/V stay in
#: DRAM either way).  ``"DRAM"`` under the megakernel frees the ~230 KB of per-core L1 the VL-SA stage needs next to
#: its static CBs (comment above); the values are identical (same ops, same dtypes), only the placement changes.
ADAPTER_MEM_BY_BACKEND: Dict[str, str] = {"ttnn": "L1", "megakernel": "DRAM"}
MK_L1_CUT_ENV = "GR00T_MK_L1_CUT_KIB"


def mk_l1_cut_bytes() -> int:
    """:data:`MK_L1_CUT_BYTES_DEFAULT` or the :data:`MK_L1_CUT_ENV` override (KiB, positive integer)."""
    import os

    raw = os.environ.get(MK_L1_CUT_ENV)
    if not raw:
        return MK_L1_CUT_BYTES_DEFAULT
    kib = int(raw)
    if kib <= 0:
        raise ValueError(f"{MK_L1_CUT_ENV}={raw!r} must be a positive number of KiB")
    return kib * 1024


def worker_l1_size_for(policy: TTPolicy) -> Optional[int]:
    """The ``worker_l1_size`` the device must be opened with for ``policy``: ``descriptors.L1_USABLE_BYTES`` (the
    p150a default, 1,461,248 B) minus :func:`mk_l1_cut_bytes` for the megakernel backend (the K4 binaries need the
    larger kernel-config ring buffer that the cut buys, module constant above), ``None`` (the ttnn default)
    otherwise."""
    if not isinstance(policy, TTPolicy):
        raise TypeError(f"policy must be a TTPolicy, got {type(policy).__name__}")
    if policy.dit_backend != "megakernel":
        return None
    from models.experimental.gr00t.tt.megakernel.descriptors import L1_USABLE_BYTES

    return int(L1_USABLE_BYTES) - mk_l1_cut_bytes()


def device_worker_l1_size(device: Any) -> int:
    """The user-allocatable L1 per worker core the device was opened with: ``ttnn.reports.get_device_info(device)``
    ``.cb_limit`` (= the ``worker_l1_size`` of ``open_device``; the ``.worker_l1_size`` field there is the physical
    1.5 MiB).  Measured on the p150a: 1,461,248 B by default, 1,362,944 B under ``K4_WORKER_L1_SIZE``."""
    import ttnn

    return int(ttnn._ttnn.reports.get_device_info(device).cb_limit)


def open_model_device(
    policy: TTPolicy,
    *,
    trace_region_size: int = DEFAULT_TRACE_REGION_SIZE,
    l1_small_size: int = DEFAULT_L1_SMALL_SIZE,
    num_command_queues: int = 1,
    device_id: int = 0,
) -> Any:
    """``tt.device.open_gr00t_device`` (program cache, p150a grid asserts) with the ``worker_l1_size`` the policy's
    DiT backend needs: the firmware default for the ttnn backend, :func:`worker_l1_size_for` (the 64 KiB cut) for the
    megakernel -- checked against :func:`device_worker_l1_size` after the open.  ``num_command_queues=2`` is the
    production path (``Gr00tTT.from_pretrained`` default; CQ-1 input writes overlap the traces).  Callers hold the
    device lock (``bin/with-device.sh``)."""
    wl1 = worker_l1_size_for(policy)
    device = open_gr00t_device(
        trace_region_size=trace_region_size,
        l1_small_size=l1_small_size,
        num_command_queues=num_command_queues,
        device_id=device_id,
        worker_l1_size=wl1,
    )
    if wl1 is not None:
        got = device_worker_l1_size(device)
        if got != wl1:
            close_gr00t_device(device)
            raise RuntimeError(f"device opened with worker_l1_size {got}, requested {wl1} (megakernel backend)")
    return device


# --------------------------------------------------------------------------------------------------------------------
# ttnn weights the megakernel backend does not need (startup cost, WP-K5 review NB 7)
# --------------------------------------------------------------------------------------------------------------------


#: Plan categories the Stage-1 head reads from ``TTWeights`` and the megakernel arena replaces
#: (``tt.megakernel.dit_program.real_arena_plan`` packs the same host plan tensors into the DRAM arena).
MK_ARENA_CATEGORIES: Tuple[str, ...] = ("dit", "dit_precompute", "encoders")


def megakernel_needs_ttnn_tensor(name: str, category: str) -> bool:
    """Whether plan tensor ``name`` of ``category`` must stay in ``TTWeights`` under the megakernel backend: every
    tensor outside :data:`MK_ARENA_CATEGORIES`, plus the ones the *other* modules read from those categories --
    ``dit.block.{i}.Wkv`` / ``bkv`` (``VLAdapterTT`` hoists the cross K/V with them), ``enc.state.*`` (the adapter's
    state encoder) and ``enc.action.b3_pos`` (``TTWeights.constants`` builds ``b3_pos_rows`` from the resident
    tensor)."""
    if category not in MK_ARENA_CATEGORIES:
        return True
    if name.startswith("dit.block.") and (name.endswith(".Wkv") or name.endswith(".bkv")):
        return True
    if name.startswith("enc.state."):
        return True
    return name == "enc.action.b3_pos"


def megakernel_weight_names(
    version: str, embodiment_id: int, policy: TTPolicy, layout_name: Optional[str] = None
) -> Tuple[Tuple[str, ...], Dict[str, Any]]:
    """``(names to upload, facts)`` for ``TTWeights.load(names=...)`` under the megakernel backend: the full plan of
    ``(version, embodiment, policy)`` minus the tensors :func:`megakernel_needs_ttnn_tensor` rejects.  ``facts``
    records what is skipped (``n_tensors``, ``device_bytes``, per category) for ``Gr00tTT.describe``."""
    from models.experimental.gr00t.tt.policy import plan_options
    from models.experimental.gr00t.tt.weights import upload_plan

    entries = upload_plan(version, int(embodiment_id), policy, plan_options(policy, version, layout_name))
    keep: List[str] = []
    skipped: Dict[str, Dict[str, float]] = {}
    n_skip, bytes_skip = 0, 0.0
    for name, e in entries.items():
        if megakernel_needs_ttnn_tensor(name, e.category):
            keep.append(name)
            continue
        c = skipped.setdefault(e.category, {"n_tensors": 0, "device_bytes": 0.0})
        c["n_tensors"] += 1
        c["device_bytes"] += float(e.nbytes)
        n_skip += 1
        bytes_skip += float(e.nbytes)
    if not keep:
        raise AssertionError("megakernel_weight_names selected no plan tensor")
    facts = {
        "n_tensors": n_skip,
        "device_bytes": bytes_skip,
        "categories": skipped,
        "n_kept": len(keep),
        "n_full": len(entries),
        "kept_from_arena_categories": sorted(
            n for n in keep if entries[n].category in MK_ARENA_CATEGORIES and not n.startswith("dit.block.")
        ),
    }
    return tuple(keep), facts


# --------------------------------------------------------------------------------------------------------------------
# bench / serve replayer contract (benchmarks/bench_e2e.py REPLAYER_MEMBERS)
# --------------------------------------------------------------------------------------------------------------------


class TraceReplayer:
    """Per-trace view of the model's ``TraceExecutor`` for ``benchmarks/bench_e2e.py`` (its ``REPLAYER_MEMBERS``
    contract): ``stage_names`` (trace names in execution order), ``upload(mi, noise)`` (static-shape + address asserts,
    then every per-call input written on the trace CQ), ``replay(stage)`` (one ``execute_trace``, non-blocking),
    ``synchronize()``, ``read_action_pred()`` (D2H of the ``[1,1,64,A_pad]`` result in its device dtype), and the
    production path ``call(mi, noise)`` (= ``TraceExecutor.replay``: 2-CQ overlapped writes when enabled, blocking
    on the last trace).  ``upload_stats`` is the executor's per-tensor accounting of the last writes.  Captures lazily
    on the first ``upload`` / ``call``.  Anything else resolves on the underlying :class:`~tt.tracer.TraceExecutor`."""

    def __init__(self, model: "Gr00tTT") -> None:
        self.model = model
        self.executor: TraceExecutor = model.executor

    @property
    def stage_names(self) -> Tuple[str, ...]:
        return tuple(tp.name for tp in self.executor.plan().traces)

    @property
    def upload_stats(self) -> Dict[str, Dict[str, float]]:
        return {k: dict(v) for k, v in self.executor.upload_stats.items()}

    def upload(self, mi: Any, noise: Optional[torch.Tensor] = None) -> None:
        """Every per-call input written on the trace CQ (serialised; the split-pass methodology of ``bench_e2e``)."""
        ex = self.executor
        noise = self.model._resolve_noise(noise)
        if not ex.captured:
            self.model.warm_and_capture(mi, noise)
        hi = self.model.host_inputs(mi, noise)  # runs BackboneTT.check_call = the static-shape pass
        ex.check_host_inputs(hi)
        ex.check_addresses()
        ex.write_inputs(hi, ex.plan().input_names(), cq=ex.cq_id)

    def call(self, mi: Any, noise: Optional[torch.Tensor] = None) -> Any:
        """``TraceExecutor.replay`` (the ``predict_normalized`` path without the D2H); returns the result tensor."""
        noise = self.model._resolve_noise(noise)
        if not self.executor.captured:
            self.model.warm_and_capture(mi, noise)
        return self.executor.replay(mi, noise=noise)

    def replay(self, stage: str) -> None:
        ex = self.executor
        if not ex.captured:
            raise TraceStateError("replay before capture: upload() first")
        for rec in ex.records:
            if rec.plan.name == stage:
                ex.backend.execute(ex.device, rec.trace_id, ex.cq_id, False)
                ex.mark_external_call()
                return
        raise KeyError(f"unknown trace {stage!r}; have {self.stage_names}")

    def synchronize(self) -> None:
        self.executor.backend.synchronize(self.executor.device)

    def read_action_pred(self) -> torch.Tensor:
        import ttnn

        return ttnn.to_torch(self.executor.result())

    def __getattr__(self, name: str) -> Any:  # delegate the rest (plan(), describe(), timing, ...)
        return getattr(self.executor, name)


# --------------------------------------------------------------------------------------------------------------------
# Gr00tTT
# --------------------------------------------------------------------------------------------------------------------


class Gr00tTT:
    """The assembled device policy of one (version, embodiment, layout, policy) -- see the module docstring.

    Build with :meth:`from_pretrained`.  Public API (plan §2.2): :meth:`predict`, :meth:`predict_normalized`,
    :meth:`run_untraced`, :meth:`warm_and_capture`, :meth:`encode`, :meth:`decode`, :meth:`make_noise`,
    :meth:`traced_state`, :meth:`release`, :meth:`describe`; facts: ``cfg``, ``layout``, ``dlayout``, ``policy``,
    ``recipe``, ``embodiment_tag`` / ``embodiment_id``, ``W`` (``TTWeights``), ``backbone``, ``adapter``, ``head``,
    ``executor`` (``TraceExecutor``), ``tracer`` (:class:`TraceReplayer`), ``timing``.
    """

    def __init__(
        self,
        cfg: GR00TConfig,
        dlayout: DeviceLayout,
        policy: TTPolicy,
        W: Any,
        device: Any,
        *,
        embodiment_tag: str,
        ck: Any = None,
        owns_device: bool = False,
        gather_mode: str = "splice",
        backbone: Optional[BackboneTT] = None,
        embed_gather: str = "host",
        cq1_uploads: Optional[bool] = None,
        adapter_mem: Optional[str] = None,
    ) -> None:
        """Assemble the modules on an open device (``from_pretrained`` is the public entry point).

        ``backbone``: a ready :class:`~tt.backbone.BackboneTT` built for the same ``(cfg, dlayout, policy, W)`` with
        ``collect_taps=False`` (tests inject stand-in towers); default ``BackboneTT.build(...)``.
        ``embed_gather``: ``"device"`` (resident-table gather in the LLM stage; needs ``W._embed_table_tensor()``) |
        ``"host"`` (the WP-12 host gather; the constructor default, so stand-in weights keep working -- the production
        factory :meth:`from_pretrained` defaults to :data:`DEFAULT_EMBED_GATHER`; module docstring).
        ``cq1_uploads``: per-call input writes on CQ 1 (``None`` = on iff the device has two command queues).
        ``adapter_mem``: ``"L1"`` | ``"DRAM"`` for the adapter's intermediates (``None`` = :data:`ADAPTER_MEM_BY_BACKEND`
        of the policy's DiT backend; module docstring).
        """
        if not isinstance(cfg, GR00TConfig) or not isinstance(dlayout, DeviceLayout):
            raise TypeError("Gr00tTT(cfg: GR00TConfig, dlayout: DeviceLayout, policy: TTPolicy, W, device)")
        if not isinstance(policy, TTPolicy):
            raise TypeError(f"policy must be a TTPolicy, got {type(policy).__name__}")
        if embed_gather not in EMBED_GATHERS:
            raise ValueError(f"embed_gather must be one of {EMBED_GATHERS}, got {embed_gather!r}")
        if dlayout.version != cfg.version:
            raise ValueError(f"layout is for {dlayout.version}, config for {cfg.version}")
        e = cfg.embodiment_id(embodiment_tag)
        if int(dlayout.embodiment_id) != e:
            raise ValueError(
                f"{cfg.version}: layout {dlayout.layout_name!r} is baked for embodiment slot {dlayout.embodiment_id} "
                f"({dlayout.layout.embodiment_tag}), model embodiment {embodiment_tag!r} is slot {e}"
            )
        if getattr(W, "version", cfg.version) != cfg.version or int(getattr(W, "embodiment_id", e)) != e:
            raise ValueError(
                f"weights are for {getattr(W, 'version', '?')}/slot {getattr(W, 'embodiment_id', '?')}, "
                f"model is {cfg.version}/slot {e}"
            )
        wl1 = worker_l1_size_for(policy)
        if wl1 is not None:
            got = device_worker_l1_size(device)
            if got > wl1:
                raise RuntimeError(
                    f"dit_backend='megakernel' (the TTPolicy default) needs the device opened with worker_l1_size <= "
                    f"{wl1}; this device was opened with {got} (the K4 kernel binaries do not fit its kernel-config "
                    "ring buffer, mk_k4_summary.md section 7.2).  Open it with tt.model.open_model_device(policy, ...) "
                    "or let Gr00tTT.from_pretrained(device=None) open it; for the Stage-1 path on a default device "
                    "use TTPolicy(dit_backend='ttnn') / GR00T_DIT_BACKEND=ttnn"
                )
        self.cfg, self.dlayout, self.policy, self.W, self.device = cfg, dlayout, policy, W, device
        self.version: str = cfg.version
        self.layout: LayoutConfig = dlayout.layout
        self.layout_name: str = dlayout.layout_name
        self.embodiment_tag: str = embodiment_tag
        self.embodiment_id: int = e
        self.owns_device = bool(owns_device)
        self.ck = ck
        self.recipe = ModelRecipe.from_config(cfg, dlayout, policy)
        self.embed_gather = embed_gather
        self.specs: Dict[str, InputSpec] = model_input_specs(cfg, dlayout, policy, embed_gather)
        self.timing: Dict[str, float] = {}
        self.weights_skipped: Dict[str, Any] = {}  # set by from_pretrained under the megakernel (module docstring)
        self.n_predict = 0
        self._b3: Optional[torch.Tensor] = None
        self._hoist: Optional[HoistOut] = None
        self.gather: Optional[DeviceEmbeddingGather] = None

        # modules (plan §2.3: every version difference is a constructor choice)
        t0 = time.perf_counter()
        if backbone is None:
            self.backbone = BackboneTT.build(cfg, dlayout, policy, W, ck, device, collect_taps=False)
        else:
            if not isinstance(backbone, BackboneTT):
                raise TypeError(f"backbone must be a BackboneTT, got {type(backbone).__name__}")
            if backbone.recipe != BackboneRecipe.from_config(cfg, dlayout, policy):
                raise ValueError("the injected backbone was built for another (cfg, dlayout, policy)")
            if backbone.collect_taps or backbone.taps:
                raise ValueError("the injected backbone must be built with collect_taps=False and hold no taps")
            self.backbone = backbone
        t1 = time.perf_counter()
        self.adapter_mem: str = ADAPTER_MEM_BY_BACKEND[policy.dit_backend] if adapter_mem is None else str(adapter_mem)
        if self.adapter_mem not in ("L1", "DRAM"):
            raise ValueError(f"adapter_mem must be 'L1' or 'DRAM', got {adapter_mem!r}")
        self.adapter = VLAdapterTT(
            cfg,
            dlayout,
            W,
            policy,
            device,
            gather_mode=gather_mode,
            collect_taps=False,
            mem=DRAM if self.adapter_mem == "DRAM" else L1,
        )
        t2 = time.perf_counter()
        if embed_gather == "device":
            if not hasattr(W, "_embed_table_tensor"):
                raise TypeError("embed_gather='device' needs TTWeights._embed_table_tensor() (the bf16 [V, E] table)")
            self.gather = DeviceEmbeddingGather(
                W._embed_table_tensor(),
                device,
                s_pad=int(dlayout.s_pad),
                name=IN_TEXT_IDS,
                zero_slots=tuple(tuple(int(x) for x in run) for run in dlayout.image_slots),
                mem="DRAM",
            )
            if self.gather.spec != self.specs[IN_TEXT_IDS]:
                raise AssertionError(f"gather spec {self.gather.spec} != model spec {self.specs[IN_TEXT_IDS]}")
        t2b = time.perf_counter()

        # executor: the backbone's two stages, the adapter stage, the denoise stage (plan §4); ``host_inputs`` runs
        # the static-shape pass (BackboneTT.check_call), hence ``prepare_asserts``
        ex = TraceExecutor(
            device,
            dlayout,
            policy.trace_layout,
            prepare=self.host_inputs,
            prepare_constants=self.constant_host_inputs,
            loaded_embodiment_ids=(e,),
            result=OUT_ACTION_PRED,
            prepare_asserts=True,
        )
        n_cqs = int(ex.backend.num_command_queues(device))
        use_cq1 = (n_cqs >= 2) if cq1_uploads is None else bool(cq1_uploads)
        if use_cq1:
            ex.set_cq1_uploads(True)  # raises on a 1-CQ device when asked explicitly
        self._register_backbone_stages(ex)
        pres = self.recipe.presence
        adapter_in = {IN_STATE_ROWS: self.specs[IN_STATE_ROWS]}
        if pres["vl"]:
            adapter_in[IN_VL_MASK] = self.specs[IN_VL_MASK]
        ex.register_stage(
            STAGE_ADAPTER, self._adapter_stage, adapter_in, (OUT_STATE_FEATURES,), reads=(OUT_FEATURES,), trace="A"
        )
        denoise_in = {IN_XT: self.specs[IN_XT]}
        for name in (IN_CROSS_TEXT_MASK, IN_CROSS_ALL_MASK):
            if name in self.specs:
                denoise_in[name] = self.specs[name]
        denoise_const = {n: self.specs[n] for n in (CONST_SELF_MASK, CONST_CROSS_IMAGE_MASK) if n in self.specs}
        ex.register_stage(
            STAGE_DENOISE,
            self._denoise_stage,
            denoise_in,
            (OUT_ACTION_PRED,),
            constants=denoise_const,
            reads=(OUT_STATE_FEATURES,),
            trace="B",
        )
        if set(ex.specs) != set(self.specs) or any(ex.specs[n] != s for n, s in self.specs.items()):
            raise AssertionError("executor specs differ from model_input_specs")
        if tuple(tp.name for tp in ex.plan().traces) != self.recipe.stage_names:
            raise AssertionError(f"trace names {[tp.name for tp in ex.plan().traces]} != {self.recipe.stage_names}")
        ex.allocate()
        self.executor = ex

        # the head is built on the executor's mask buffers (presence pattern checked by ActionHeadTT itself)
        masks = {name: ex.buffers.get(MASK_BUFFER_NAMES[name]) for name in MASK_NAMES}
        kv = {b.block: (self.adapter.K[b.block], self.adapter.V[b.block]) for b in self.adapter.plan.blocks}
        self.head = ActionHeadTT(cfg, dlayout, W, policy, device, masks=masks, kv=kv, collect_taps=False)
        t3 = time.perf_counter()
        self.timing.update(
            backbone_build_s=t1 - t0,
            adapter_build_s=t2 - t1,
            embed_gather_build_s=t2b - t2,
            executor_and_head_build_s=t3 - t2b,
        )
        self._replayer = TraceReplayer(self)

    def _register_backbone_stages(self, ex: TraceExecutor) -> None:
        """``BackboneTT.stage_functions`` on the executor (trace "A"); with the device gather the LLM stage takes the
        token ids and gathers the text rows itself (module docstring)."""
        if self.backbone.collect_taps:
            raise RuntimeError("set_collect_taps(False) before registering the backbone stages for tracing")
        for sd in self.backbone.stage_functions():
            fn, inputs = sd.fn, dict(sd.inputs)
            if sd.name == STAGE_LLM and self.gather is not None:
                if IN_TEXT not in inputs:
                    raise AssertionError(f"llm stage inputs {sorted(inputs)} lack {IN_TEXT!r}")
                inputs.pop(IN_TEXT)
                inputs[self.gather.name] = self.gather.spec
                fn = self._llm_stage_with_gather(sd.fn)
            ex.register_stage(sd.name, fn, inputs, sd.outputs, constants=sd.constants, reads=sd.reads, trace="A")

    def _llm_stage_with_gather(self, inner: Callable[..., Dict[str, Any]]) -> Callable[..., Dict[str, Any]]:
        gather = self.gather
        if gather is None:
            raise AssertionError("no device gather")

        def llm_fn(inputs: Mapping[str, Any], state: Mapping[str, Any]) -> Dict[str, Any]:
            import ttnn

            text = gather.gather(inputs[gather.name])  # [1,1,S_pad,E] bf16 TILE, image slots zero
            ins = {k: v for k, v in inputs.items() if k != gather.name}
            ins[IN_TEXT] = text
            try:
                return inner(ins, state)
            finally:
                ttnn.deallocate(text)  # the splice copies it; nothing downstream reads the gathered rows

        return llm_fn

    # ------------------------------------------------------------------ factory
    @classmethod
    def from_pretrained(
        cls,
        version: str,
        embodiment: Optional[str] = None,
        layout_name: Optional[str] = None,
        policy: Optional[TTPolicy] = None,
        device: Any = None,
        cache_root: Any = None,
        *,
        warm_inputs: Any = None,
        noise: Optional[torch.Tensor] = None,
        embed_gather: str = DEFAULT_EMBED_GATHER,
        cq1_uploads: Optional[bool] = None,
        num_command_queues: int = 2,
        adapter_mem: Optional[str] = None,
        load_ttnn_dit_weights: Optional[bool] = None,
    ) -> "Gr00tTT":
        """Load the weights and assemble the model (plan §2.2 row ``tt/model.py``).

        Args:
            version: ``"n15" | "n16" | "n17"``.
            embodiment: embodiment tag (default: the layout's ``embodiment_tag``); must be the slot the layout is baked for.
            layout_name: one of ``cfg.layouts`` (default ``cfg.canonical_layout_name``, D2).
            policy: ``TTPolicy`` (default ``TTPolicy()``: mixed_dit, **megakernel** backend with the bfp8_b arena, meta
                RoPE, per_stage traces; ``TTPolicy().with_env_overrides()`` honours ``GR00T_DIT_BACKEND`` /
                ``GR00T_MK_ARENA_DTYPE`` -- the factory itself never reads the environment).
            device: an open mesh device; ``None`` opens one with :func:`open_model_device` (``num_command_queues``;
                the megakernel backend's ``worker_l1_size``), closed by ``release``.  A device opened another way is
                refused for the megakernel backend (constructor).
            cache_root: ``HostPlanCache`` root for ``TTWeights.load`` (default: the project cache).
            warm_inputs: a ``ModelInputs`` (or ``Observation``) to warm and capture on immediately; otherwise the
                first ``predict`` / ``predict_normalized`` / ``tracer.upload`` call warms and captures on its inputs.
            noise: the noise for that warm pass (default: the version's seeded noise).
            embed_gather, cq1_uploads, adapter_mem: see the constructor (module docstring; ``adapter_mem="DRAM"`` is
                what a ttnn-backend model needs on a device opened for the megakernel backend).
            load_ttnn_dit_weights: ``None`` (default) uploads the Stage-1 DiT / precompute / encoder tensors only for
                the ttnn backend (the megakernel packs them into its arena instead, :func:`megakernel_weight_names`;
                ``describe()["weights"]["skipped"]`` and :attr:`weights_skipped` record what was left out);
                ``True`` forces the full upload under either backend, ``False`` is an error for the ttnn backend.
        """
        from models.experimental.gr00t.common.checkpoint import LazyCheckpoint
        from models.experimental.gr00t.tt.weights import TTWeights

        cfg = get_config(version)
        pol = policy if policy is not None else TTPolicy()
        if not isinstance(pol, TTPolicy):
            raise TypeError(f"policy must be a TTPolicy, got {type(pol).__name__}")
        layout = cfg.layout(layout_name if layout_name is not None else cfg.canonical_layout_name)
        tag = embodiment if embodiment is not None else layout.embodiment_tag
        e = cfg.embodiment_id(tag)
        if int(layout.embodiment_id) != e:
            raise ValueError(
                f"{version}: layout {layout.name!r} is baked for {layout.embodiment_tag!r} (slot {layout.embodiment_id}); "
                f"embodiment {tag!r} is slot {e} -- pick a layout of that embodiment"
            )
        dlayout = DeviceLayout.from_layout(cfg, layout)
        skip_facts: Dict[str, Any] = {}
        names: Optional[Tuple[str, ...]] = None
        if load_ttnn_dit_weights is False and pol.dit_backend != "megakernel":
            raise ValueError("load_ttnn_dit_weights=False needs dit_backend='megakernel' (the ttnn head reads them)")
        if pol.dit_backend == "megakernel" and not load_ttnn_dit_weights:
            names, skip_facts = megakernel_weight_names(version, e, pol, layout.name)
        owns = device is None
        dev = open_model_device(pol, num_command_queues=int(num_command_queues)) if owns else device
        try:
            t0 = time.perf_counter()
            W = TTWeights.load(version, e, pol, dev, cache_root=cache_root, names=names, layout_name=layout.name)
            load_s = time.perf_counter() - t0
            ck = LazyCheckpoint(version)
            model = cls(
                cfg,
                dlayout,
                pol,
                W,
                dev,
                embodiment_tag=tag,
                ck=ck,
                owns_device=owns,
                embed_gather=embed_gather,
                cq1_uploads=cq1_uploads,
                adapter_mem=adapter_mem,
            )
        except Exception:
            if owns:
                close_gr00t_device(dev)
            raise
        model.timing["weights_load_s"] = load_s
        model.weights_skipped = dict(skip_facts)
        if warm_inputs is not None:
            mi = warm_inputs if not _is_observation(warm_inputs) else model.encode(warm_inputs)
            model.warm_and_capture(mi, noise)
        return model

    # ------------------------------------------------------------------ host side (plan §3.1)
    @property
    def noise_shape(self) -> Tuple[int, int, int]:
        return tuple(int(x) for x in self.dlayout.noise_shape)  # type: ignore[return-value]

    def make_noise(self, seed: Optional[int] = None) -> torch.Tensor:
        """``randn(noise_shape)`` fp32 from a CPU generator seeded with ``seed`` (default :data:`NOISE_SEED`)."""
        s = int(NOISE_SEED[self.version] if seed is None else seed)
        g = torch.Generator(device="cpu").manual_seed(s)
        return torch.randn(self.noise_shape, generator=g, dtype=torch.float32)

    def _resolve_noise(self, noise: Optional[torch.Tensor]) -> torch.Tensor:
        if noise is None:
            return self.make_noise()
        if not isinstance(noise, torch.Tensor) or not noise.is_floating_point():
            raise TypeError(f"noise must be a floating torch.Tensor, got {type(noise).__name__}")
        if _shape(noise) != self.noise_shape:
            raise ValueError(f"noise must be {self.noise_shape}, got {_shape(noise)}")
        if not bool(torch.isfinite(noise).all()):
            raise ValueError("noise contains non-finite values")
        return noise.detach().to("cpu", torch.float32)

    def encode(self, obs: Any) -> Any:
        """``common.preprocessing.encode`` for this model's embodiment and **fixed** layout (static-shape §9: the
        layout is chosen at load, never per call); an observation of another embodiment tag raises."""
        from models.experimental.gr00t.common.preprocessing import encode

        if getattr(obs, "embodiment_tag", None) != self.embodiment_tag:
            raise ValueError(
                f"model is built for embodiment {self.embodiment_tag!r} (slot {self.embodiment_id}); observation "
                f"carries {getattr(obs, 'embodiment_tag', None)!r}"
            )
        return encode(self.version, obs, embodiment=self.embodiment_tag, layout_name=self.layout_name)

    def decode(self, action_pred_normalized: torch.Tensor, obs: Any) -> Dict[str, torch.Tensor]:
        """``common.normalization.decode`` with this model's embodiment -> ``{group: (1, H, D) fp32}``."""
        from models.experimental.gr00t.common.normalization import decode

        return decode(self.version, action_pred_normalized, obs, embodiment=self.embodiment_tag)

    def vision_host_inputs(self, mi: Any) -> Dict[str, torch.Tensor]:
        """The tower's per-call input of ``mi`` (bf16 host tensor): SigLIP ``patches`` via ``tower.prepare_patches``,
        Qwen3-VL ``vis_patches = pixel_values.to(bf16)`` -- the same rounding as ``tt.qwen3vl_vision.host_inputs``
        (``patches=pv.to(bf).contiguous()``) without its per-call rebuild of the constant tables (module docstring;
        ``BackboneTT.check_call`` must have run: it asserts ``pixel_values`` shape / finiteness)."""
        if self.recipe.tower == "siglip":
            return {IN_PATCHES: self.backbone.tower.prepare_patches(mi)}
        pv = mi.pixel_values
        if tuple(pv.shape) != tuple(self.dlayout.pixel_shape) or pv.dtype != torch.float32:
            raise StaticShapeError(
                f"pixel_values {tuple(pv.shape)} {pv.dtype} != fp32 {tuple(self.dlayout.pixel_shape)} patch rows"
            )
        return {IN_VIS_PATCHES: pv.to(torch.bfloat16).contiguous()}

    def host_inputs(self, mi: Any, noise: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
        """Every **per-call** executor input of one call (the ``TraceExecutor`` ``prepare`` callable; plan §3.1):
        the backbone's (patches / pixels, text embed rows **or** token ids, N1.6 RoPE rows, N1.6 LLM mask) + VL-SA
        mask (n15 / n17), ``state_rows``, ``x_t`` (``noise_to_xt``), DiT cross masks (text | all).  Runs the
        static-shape asserts once (``BackboneTT.check_call``: ``DeviceLayout.assert_call`` + the bound constant
        tables); shape / dtype of every tensor are checked here, finiteness by the executor (``check_host_inputs``)."""
        t0 = time.perf_counter()
        noise = self._resolve_noise(noise)
        bb = self.backbone
        bb.check_call(mi)
        self._assert_call_extras(mi, noise)
        out: Dict[str, torch.Tensor] = self.vision_host_inputs(mi)
        if self.gather is not None:
            out[self.gather.name] = self.gather.host_ids(mi.input_ids_padded)
        else:
            out[IN_TEXT] = bb.text_embeds_host(mi)
        if bb.recipe.rope_per_call:
            out[IN_COS], out[IN_SIN] = bb.rope_host(mi)
        llm_mask = bb.llm_mask_host(mi)
        if llm_mask is not None:
            out[IN_LLM_MASK] = llm_mask
        masks = self.dlayout.per_call_masks(mi.slots)
        for name in ("vl", "cross_text", "cross_all"):
            buf = MASK_BUFFER_NAMES[name]
            if self.recipe.presence[name]:
                if masks[name] is None:
                    raise StaticShapeError(f"mask {name!r} is present in the layout but the call produced none")
                out[buf] = masks[name]
            elif masks[name] is not None:
                raise StaticShapeError(f"mask {name!r} produced for a layout that fixes it absent")
        out[IN_STATE_ROWS] = state_rows_host(mi.state, int(self.cfg.encoders.state_dim_pad))
        out[IN_XT] = noise_to_xt(noise, self.cfg, self.dlayout)
        for name, t in out.items():
            self.specs[name].check_host(t, finite=False)
        want = {n for n, s in self.specs.items() if s.per_call}
        if set(out) != want:
            raise AssertionError(f"host_inputs {sorted(out)} != per-call inputs {sorted(want)}")
        self.timing["last_host_inputs_s"] = time.perf_counter() - t0
        return out

    def _assert_call_extras(self, mi: Any, noise: torch.Tensor) -> None:
        """The parts of ``DeviceLayout.assert_call(mi, noise, loaded_embodiment_ids)`` that ``BackboneTT.check_call``
        (``assert_call(mi)``) does not cover: the noise shape and the embodiment slot (static-shape §10 item 6)."""
        e = int(torch.as_tensor(mi.embodiment_id).reshape(-1)[0])
        if e != self.embodiment_id:
            raise StaticShapeError(f"embodiment_id {e} is not the loaded slot {self.embodiment_id}")
        if _shape(noise) != self.noise_shape:
            raise StaticShapeError(f"noise.shape={_shape(noise)} != {self.noise_shape}")

    def constant_host_inputs(self, mi: Any) -> Dict[str, torch.Tensor]:
        """The per-layout constants (``prepare_constants``): the backbone's tables + the constant DiT masks
        (``self``; ``cross_image`` when ``L_img_pad != L_img_real``)."""
        out: Dict[str, torch.Tensor] = dict(self.backbone.constant_host_inputs(mi))
        consts = self.dlayout.constant_masks()
        out[CONST_SELF_MASK] = consts["self"]
        if self.recipe.presence["cross_image"]:
            if consts["cross_image"] is None:
                raise StaticShapeError("cross_image mask present in the layout but constant_masks gave none")
            out[CONST_CROSS_IMAGE_MASK] = consts["cross_image"]
        for name, t in out.items():
            self.specs[name].check_host(t)
        want = {n for n, s in self.specs.items() if not s.per_call}
        if set(out) != want:
            raise AssertionError(f"constants {sorted(out)} != {sorted(want)}")
        return out

    def _write_buffers(self, host: Mapping[str, torch.Tensor], names: Sequence[str]) -> None:
        """``copy_host_to_device_tensor`` of the given host tensors into the executor's persistent buffers, on the
        trace CQ (``TraceExecutor.write_inputs``; the caller validated ``host``)."""
        self.executor.write_inputs(host, names, cq=self.executor.cq_id)

    # ------------------------------------------------------------------ stage functions (device)
    def _adapter_stage(self, inputs: Mapping[str, Any], state: Mapping[str, Any]) -> Dict[str, Any]:
        out = self.adapter.forward(state[OUT_FEATURES], inputs.get(IN_VL_MASK), inputs[IN_STATE_ROWS])
        self._hoist = out
        if out.state_features is None:
            raise RuntimeError("VLAdapterTT.forward returned no state_features")
        if not self.adapter.collect_taps:
            # vl / gathered sets / the 32-row state block are consumed (K/V live in the adapter's persistent buffers)
            _deallocate_once([out.vl, *out.vl_sets.values(), out.state_rows], keep=(out.state_features,))
        return {OUT_STATE_FEATURES: out.state_features}

    def _denoise_stage(self, inputs: Mapping[str, Any], state: Mapping[str, Any]) -> Dict[str, Any]:
        import ttnn

        pred = self.head.denoise(inputs[IN_XT], state[OUT_STATE_FEATURES])
        # the head's action_pred is its own persistent buffer; the executor owns (and frees) the trace output
        return {OUT_ACTION_PRED: ttnn.clone(pred, memory_config=resolve_mem(DRAM))}

    # ------------------------------------------------------------------ traced path
    @property
    def captured(self) -> bool:
        return bool(self.executor.captured)

    @property
    def tracer(self) -> TraceReplayer:
        """The ``benchmarks/bench_e2e.py`` replayer contract over the executor."""
        return self._replayer

    def warm_and_capture(self, mi: Any, noise: Optional[torch.Tensor] = None) -> None:
        """One untraced pass with the real shapes (program cache, lazy buffers, constants) then the trace capture
        (plan §4 rules; static-shape §10 item 10).  Records ``warm_s`` / ``capture_s``."""
        import ttnn

        if self.executor.captured:
            raise TraceStateError("already captured; release() before capturing again")
        noise = self._resolve_noise(noise)
        self._set_taps(False)
        t0 = time.perf_counter()
        self.executor.warm(mi, noise)
        ttnn.synchronize_device(self.device)
        t1 = time.perf_counter()
        self.executor.capture()
        ttnn.synchronize_device(self.device)
        t2 = time.perf_counter()
        self.timing.update(warm_s=t1 - t0, capture_s=t2 - t1)

    def predict_normalized(self, mi: Any, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Traced inference on pre-encoded inputs -> fp32 ``action_pred_normalized [1, H, A_max]`` (captures on the
        first call; every call replays: static-shape asserts -> input writes -> ``execute_trace`` x n)."""
        noise = self._resolve_noise(noise)
        if not self.executor.captured:
            self.warm_and_capture(mi, noise)
        t0 = time.perf_counter()
        out = self.executor.replay(mi, noise=noise)
        pred = xt_to_action_pred(out, self.cfg, self.dlayout)  # D2H (blocking read)
        self.timing["last_replay_s"] = time.perf_counter() - t0
        self.n_predict += 1
        return pred

    def predict(self, obs: Any, noise: Optional[torch.Tensor] = None) -> Dict[str, np.ndarray]:
        """Raw observation -> ``{action key: (1, H, D) float32}`` (the policy API): ``encode`` -> replay -> ``decode``."""
        t0 = time.perf_counter()
        mi = self.encode(obs)
        t1 = time.perf_counter()
        pred = self.predict_normalized(mi, noise)
        t2 = time.perf_counter()
        acts = self.decode(pred, obs)
        t3 = time.perf_counter()
        self.timing.update(last_encode_s=t1 - t0, last_decode_s=t3 - t2, last_predict_s=t3 - t0)
        return {k: np.ascontiguousarray(v.detach().cpu().numpy(), dtype=np.float32) for k, v in acts.items()}

    def traced_state(self) -> Dict[str, torch.Tensor]:
        """Raw host copies (device dtype, device layout) of the last replay's ``action_pred`` and ``backbone_features``
        and of the hoisted ``K_0`` / ``V_0`` of the first cross block (plan §7.4 determinism / stale-state probes)."""
        import ttnn

        ex = self.executor
        if OUT_FEATURES not in ex.outputs or ex.plan().result not in ex.outputs:
            raise TraceStateError("no outputs yet: warm or replay first")
        i0 = self.recipe.cross_key_sets[0][0]
        return {
            "action_pred": ttnn.to_torch(ex.result()),
            "backbone_features": ttnn.to_torch(ex.outputs[OUT_FEATURES]),
            "K_0": ttnn.to_torch(self.adapter.K[i0]),
            "V_0": ttnn.to_torch(self.adapter.V[i0]),
        }

    # ------------------------------------------------------------------ untraced path (taps)
    def _set_taps(self, flag: bool) -> None:
        self.backbone.set_collect_taps(flag)  # raises while kept taps are alive
        self.adapter.set_collect_taps(flag)
        self.head.set_collect_taps(flag)

    def run_untraced(
        self,
        mi: Any,
        noise: Optional[torch.Tensor] = None,
        return_taps: bool = True,
        *,
        raw: bool = False,
        on_device_taps: Optional[Callable[["Gr00tTT", Mapping[str, Any]], None]] = None,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """Eager (untraced) inference on the executor's buffers with every module's taps on -> ``(pred, taps)``.

        ``pred`` is fp32 ``[1, H, A_max]``; ``taps`` maps the golden keys of :func:`expected_tap_keys` to **host**
        fp32 tensors in the goldens' shapes (``[1, S, C]`` token order for the LLM / adapter taps, raster ``[n, P, C]``
        for the SigLIP taps, ``[1, M, ...]`` / ``[1, H, ...]`` for the head taps; N1.6 vision taps are raster ordered,
        ``harness.check_taps`` maps them).  ``raw=True`` adds ``raw/action_pred``, ``raw/backbone_features``,
        ``raw/K_0``, ``raw/V_0`` (device dtype and layout).  ``on_device_taps(model, stage_outputs)`` runs while the
        device taps are still alive (``tests/tt/debug/dump_ops.py``).  Every device tap is freed before returning, so
        the captured allocation state is left as it was (an untraced run may sit between two replays).
        """
        import ttnn

        ex = self.executor
        noise = self._resolve_noise(noise)
        self.dlayout.assert_call(mi, noise, ex.loaded_embodiment_ids)
        ex.allocate()
        plan = ex.plan()
        consts = self.constant_host_inputs(mi)
        hi = self.host_inputs(mi, noise)
        ex.check_host_inputs(hi)
        if set(consts) != set(plan.constant_names()):
            raise StaticShapeError(f"constants {sorted(consts)} != registered {sorted(plan.constant_names())}")
        if ex.captured:
            ex.check_addresses()
        for name in plan.constant_names():
            ex.specs[name].check_host(consts[name], finite=True)
        ex.wait_previous_call()  # CQ-1 writes of an earlier replay must not race these CQ-0 writes (2-CQ mode)
        self._write_buffers(consts, plan.constant_names())
        self._write_buffers(hi, plan.input_names())
        ttnn.synchronize_device(self.device)
        self._set_taps(bool(return_taps))
        state: Dict[str, Any] = {}
        t0 = time.perf_counter()
        try:
            for st in ex.stages:
                outs = st.fn({n: ex.buffers[n] for n in st.inputs}, dict(state))
                if set(outs) != set(st.outputs):
                    raise RuntimeError(f"stage {st.name!r} returned {sorted(outs)}, registered {list(st.outputs)}")
                state.update(outs)
            ex.mark_external_call()
            ttnn.synchronize_device(self.device)
            self.timing["last_untraced_s"] = time.perf_counter() - t0
            pred_raw = ttnn.to_torch(state[OUT_ACTION_PRED])
            pred = xt_to_action_pred(pred_raw, self.cfg, self.dlayout)
            taps: Dict[str, torch.Tensor] = {}
            if return_taps:
                taps = self._host_taps(mi, state)
            if raw:
                i0 = self.recipe.cross_key_sets[0][0]
                taps[RAW_PREFIX + "action_pred"] = pred_raw
                taps[RAW_PREFIX + "backbone_features"] = ttnn.to_torch(state[OUT_FEATURES])
                taps[RAW_PREFIX + "K_0"] = ttnn.to_torch(self.adapter.K[i0])
                taps[RAW_PREFIX + "V_0"] = ttnn.to_torch(self.adapter.V[i0])
            if on_device_taps is not None:
                on_device_taps(self, state)
        finally:
            self._free_untraced(state)
            self._set_taps(False)
        return pred, taps

    def _host_taps(self, mi: Any, state: Mapping[str, Any]) -> Dict[str, torch.Tensor]:
        """Device taps of the three modules -> golden-keyed host tensors (module docstring, plan §7.2)."""
        cfg, dl = self.cfg, self.dlayout
        bt = self.backbone.taps
        slot = mi.slots.slot_of_token
        taps: Dict[str, torch.Tensor] = {}

        def rows(t: Any) -> torch.Tensor:
            return padded_to_rows(L.to_torch(t), slot)

        if self.recipe.tower == "siglip":
            tower = self.backbone.tower
            taps["vit_embeddings"] = tower.to_raster(bt[SIGLIP_TAP_EMBEDDINGS])
            taps["vit_block_first"] = tower.to_raster(bt[siglip_tap_block(0)])
            taps["vit_block_last"] = tower.to_raster(bt[siglip_tap_block(tower.n_layers - 1)])
            taps["vit_post_ln"] = tower.to_raster(bt[SIGLIP_TAP_POST_LN])
            vis = L.to_torch(bt[TAP_PROJECTOR_OUT])  # [n_img, 1, T(+pad), 2048]
            taps["projector_out"] = vis[:, 0, : dl.tokens_per_image].contiguous()
        else:
            for key in QV_TAP_KEYS:
                if key != "projector_out":
                    taps[key] = L.to_torch(bt[key])
            taps["projector_out"] = L.to_torch(bt[TAP_PROJECTOR_OUT]).reshape(-1, int(cfg.llm.hidden))
        n = int(cfg.llm.n_layers_executed)
        taps["llm_inputs_embeds"] = rows(bt[TAP_INPUTS_EMBEDS])
        taps["llm_layer_first"] = rows(bt["llm_layer_0"])
        taps["llm_layer_last"] = rows(bt[f"llm_layer_{n - 1}"])
        if self.recipe.final_norm:
            taps["llm_final_norm"] = rows(bt["llm_final_norm"])
        taps["backbone_features"] = rows(state[OUT_FEATURES])
        for name, t in self.adapter.debug_taps.items():
            if name == "state_features":
                taps[name] = L.to_torch(t).reshape(1, 1, -1)
            else:
                taps[name] = slots_to_rows(L.to_torch(t), mi.slots)
        taps.update(self.head.golden_taps(b3=self.b3_bias(), block_taps=True))
        missing = [k for k in expected_tap_keys(cfg, dl, self.policy.dit_backend) if k not in taps]
        if missing:
            raise RuntimeError(f"run_untraced produced no tap for {missing}")
        return taps

    def b3_bias(self) -> torch.Tensor:
        """``action_encoder.W3.b[e]`` fp32 ``[1536]`` from the checkpoint (folded into ``b3_pos`` on device; the golden
        ``action_encoder_out`` includes it, WL §4.7) -- read once, cached."""
        if self._b3 is None:
            from models.experimental.gr00t.common.checkpoint import LazyCheckpoint

            ck = self.ck if self.ck is not None else LazyCheckpoint(self.version)
            key = self.cfg.checkpoint.action_head_prefix + "action_encoder.W3.b"
            b3 = ck.slot(key, self.embodiment_id, torch.float32)
            want = (int(self.cfg.encoders.input_embedding_dim),)
            if _shape(b3) != want:
                raise ValueError(f"{key}[{self.embodiment_id}]: shape {_shape(b3)} != {want}")
            self._b3 = b3
        return self._b3

    def _free_untraced(self, state: Mapping[str, Any]) -> None:
        """Free every device tensor an untraced run left alive: backbone taps, adapter taps / outputs, head taps and
        the stage outputs (persistent buffers -- K/V, x_t, action_pred, executor inputs -- are never touched)."""
        self.backbone.free_taps()  # vis_rows / deepstack / llm taps when collected (features are skipped)
        cands: List[Any] = list(state.values())
        cands += list(self.adapter.debug_taps.values())
        h = self._hoist
        if h is not None:
            cands += [h.vl, *h.vl_sets.values(), h.state_rows, h.state_features]
        keep = tuple(self.adapter.K.values()) + tuple(self.adapter.V.values()) + tuple(self.executor.buffers.values())
        _deallocate_once(cands, keep=keep)
        self.adapter.debug_taps = {}
        self._hoist = None
        self.head.free_taps()

    # ------------------------------------------------------------------ teardown / reporting
    def release(self) -> None:
        """Release the traces and every device buffer (head, executor, adapter, backbone, weights); closes the device
        when :meth:`from_pretrained` opened it."""
        errors: List[str] = []
        steps: List[Tuple[str, Callable[[], None]]] = [
            ("head", self.head.release),
            ("executor", self.executor.release),
            ("adapter", self.adapter.release),
            ("backbone", self.backbone.release),
            ("weights", self.W.release),
        ]
        if self.gather is not None:
            steps.insert(2, ("embed_gather", self.gather.release))
        for name, fn in steps:
            try:
                fn()
            except Exception as exc:  # noqa: BLE001 - keep releasing, report at the end
                errors.append(f"{name}: {type(exc).__name__}: {exc}")
        if self.ck is not None and hasattr(self.ck, "close"):
            self.ck.close()
        if self.owns_device:
            close_gr00t_device(self.device)
            self.owns_device = False
        if errors:
            raise RuntimeError("Gr00tTT.release: " + "; ".join(errors))

    def describe(self) -> Dict[str, Any]:
        """JSON-able summary: recipe, input specs, executor plan, module facts, timings."""
        d: Dict[str, Any] = {
            "version": self.version,
            "layout_name": self.layout_name,
            "embodiment_tag": self.embodiment_tag,
            "embodiment_id": self.embodiment_id,
            "policy": self.policy.to_dict(),
            "recipe": self.recipe.describe(),
            "embed_gather": self.embed_gather,
            "adapter_mem": self.adapter_mem,
            "worker_l1_size": _worker_l1_size_or_none(self.device),
            "gather": None if self.gather is None else self.gather.describe(),
            "input_specs": {n: s.describe() for n, s in self.specs.items()},
            "executor": self.executor.describe(),
            "timing": dict(self.timing),
            "n_predict": self.n_predict,
            "weights": {
                "n_tensors": len(self.W),
                "device_bytes": float(self.W.device_bytes),
                "load_stats": dict(getattr(self.W, "load_stats", {})),
                "dtype_histogram": self.W.dtype_histogram() if hasattr(self.W, "dtype_histogram") else None,
                "skipped": dict(self.weights_skipped),
            },
        }
        for name in ("backbone", "adapter", "head"):
            comp = getattr(self, name)
            if hasattr(comp, "describe"):
                d[name] = comp.describe()
        return d


# --------------------------------------------------------------------------------------------------------------------
# helpers
# --------------------------------------------------------------------------------------------------------------------


def _worker_l1_size_or_none(device: Any) -> Optional[int]:
    """:func:`device_worker_l1_size` for a real mesh device; ``None`` for the CPU tests' stand-in devices."""
    try:
        return device_worker_l1_size(device)
    except (ImportError, TypeError, AttributeError, RuntimeError):
        return None


def _is_observation(x: Any) -> bool:
    return hasattr(x, "images") and hasattr(x, "instruction") and hasattr(x, "embodiment_tag")


def _deallocate_once(tensors: Sequence[Any], keep: Sequence[Any] = ()) -> int:
    """``ttnn.deallocate`` every distinct, still-allocated device tensor of ``tensors`` that is not in ``keep``
    (identity comparison); returns the number freed."""
    import ttnn

    keep_ids = {id(t) for t in keep if t is not None}
    seen: set = set()
    n = 0
    for t in tensors:
        if t is None or id(t) in seen or id(t) in keep_ids:
            continue
        seen.add(id(t))
        alive = t.is_allocated() if hasattr(t, "is_allocated") else True
        if not alive:
            continue
        ttnn.deallocate(t)
        n += 1
    return n


__all__ = [
    "CONST_CROSS_IMAGE_MASK",
    "CONST_SELF_MASK",
    "DECODE_TAP",
    "DEFAULT_EMBED_GATHER",
    "EMBED_GATHERS",
    "HOST_OR_UNPRODUCED_TAPS",
    "IN_CROSS_ALL_MASK",
    "IN_CROSS_TEXT_MASK",
    "IN_STATE_ROWS",
    "IN_TEXT_IDS",
    "IN_VL_MASK",
    "IN_XT",
    "MASK_BUFFER_NAMES",
    "MK_ARENA_CATEGORIES",
    "NOISE_SEED",
    "N_STEPS",
    "OUT_ACTION_PRED",
    "OUT_STATE_FEATURES",
    "RAW_PREFIX",
    "STAGE_ADAPTER",
    "STAGE_DENOISE",
    "STAGE_NAMES",
    "Gr00tTT",
    "ModelRecipe",
    "TraceReplayer",
    "ADAPTER_MEM_BY_BACKEND",
    "MK_L1_CUT_BYTES_DEFAULT",
    "MK_L1_CUT_ENV",
    "device_worker_l1_size",
    "expected_tap_keys",
    "megakernel_needs_ttnn_tensor",
    "megakernel_weight_names",
    "mk_l1_cut_bytes",
    "model_input_specs",
    "open_model_device",
    "text_ids_spec",
    "unproduced_taps",
    "worker_l1_size_for",
]