# 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", ]