# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """Static device layout of one GR00T shape key: additive masks, gather pieces, static-shape asserts, materialisation. This module is the *only* layout code in ``tt/`` (IMPLEMENTATION_PLAN.md §1.1 D2 / D5, §2.2 row ``tt/layout.py``). It never re-derives slots, position ids or RoPE tables -- those come from ``common.configs.LayoutConfig``, ``common.prompts.ShapeKey`` / ``SlotLayout`` and ``common.preprocessing.ModelInputs``. It adds: * the four additive mask builders of static-shape-plan.md §6.4 (bf16, ``-1e9`` on masked keys, never ``-inf`` -- §4.3 item 4), a constant DiT self mask (§6.1) and the constant image mask (§6.2); * the gather pieces of the cross K/V hoist (§6.2: image runs and their complement, the non-image runs); * the DiT row bookkeeping (velocity / noise rows, §6.1) and the per-layout pixel shape; * :meth:`DeviceLayout.assert_call`: items 1-8 of static-shape-plan.md §10 (the pi05 "baked Python ints" lesson); items 9-11 (buffer identity, program-cache warmth, device grid) belong to ``tt/tracer.py`` / ``tt/device.py``; * :func:`materialise`: upload of the masks as bf16 TILE DRAM tensors, constants once and per-call masks into persistent buffers via ``ttnn.copy_host_to_device_tensor`` (static-shape-plan.md §9). ``ttnn`` is imported lazily inside :func:`materialise` only; the module (and its tests) import without a device. Mask names (fixed per layout -- ``mask_presence`` never toggles between ``None`` and a tensor, §10 item 4): ============= =========================== ======================================= ========== name shape masked (``-1e9``) columns changes ============= =========================== ======================================= ========== ``self`` ``[1,1,64,64]`` DiT pad keys ``>= M`` (§6.1) constant ``cross_text`` ``[1,1,64,L_txt_pad]`` gathered non-image slot not key-valid, per call or zero pad row ``>= L_txt_real`` (§6.2) ``cross_image`` ``[1,1,64,L_img_pad]`` zero pad rows ``>= L_img_real``; ``None`` constant when ``L_img_real == L_img_pad`` (§6.2) ``cross_all`` ``[1,1,64,S_pad_vl]`` N1.5 all-token keys not key-valid (§6.2) per call ``llm`` ``[1,1,S_pad,S_pad]`` N1.6 ``text_slot`` padding: not causal per call or key not valid (§4.3 item 3) ``vl`` ``[1,1,S_pad_vl,S_pad_vl]`` VL-SA (N1.5 / N1.7) keys not valid (§5) per call ============= =========================== ======================================= ========== """ from __future__ import annotations from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Collection, Dict, List, Mapping, Optional, Sequence, Tuple import torch from models.experimental.gr00t.common.configs import ( DIT_HEAD_DIM_PAD, DIT_QUERY_ROWS_PAD, LLM_SEQ_PAD_MULTIPLE, TILE, GR00TConfig, LayoutConfig, ) from models.experimental.gr00t.common.prompts import ShapeKey, SlotLayout, token_runs if TYPE_CHECKING: # pragma: no cover - typing only; keeps the module import free of cv2 / transformers from models.experimental.gr00t.common.preprocessing import ModelInputs Run = Tuple[int, int] # static-shape-plan.md §4.3 item 4: a large finite negative in bf16 (-1e9 -> -998244352), never -inf / finfo.min. MASK_NEG: float = -1e9 MASK_DTYPE = torch.bfloat16 # mask names and their update class (static-shape-plan.md §6.4 / §9) CONSTANT_MASKS: Tuple[str, ...] = ("self", "cross_image") PER_CALL_MASKS: Tuple[str, ...] = ("cross_text", "cross_all", "llm", "vl") MASK_NAMES: Tuple[str, ...] = CONSTANT_MASKS + PER_CALL_MASKS # DiT sampler constants the traced path is built for (static-shape-plan.md §10 item 7) _DT = 0.25 _N_STEPS = 4 _BUCKETS = (0, 250, 500, 750) class StaticShapeError(AssertionError): """A static-shape assert of static-shape-plan.md §10 failed (raised, never silently handled).""" # -------------------------------------------------------------------------------------------------------------------- # run helpers # -------------------------------------------------------------------------------------------------------------------- def pad_to_tile(n: int) -> int: """Smallest multiple of the 32-row tile that is ``>= n`` (static-shape-plan.md §1).""" if n < 0: raise ValueError(f"negative length {n}") return ((n + TILE - 1) // TILE) * TILE def check_runs(runs: Sequence[Run], length: int, what: str) -> Tuple[Run, ...]: """Validate half-open ``[lo, hi)`` runs: sorted, non-empty, non-overlapping, inside ``[0, length)``.""" out: List[Run] = [] prev_hi = 0 for r in runs: if len(r) != 2: raise ValueError(f"{what}: run {r!r} is not (lo, hi)") lo, hi = int(r[0]), int(r[1]) if not (0 <= lo < hi <= length): raise ValueError(f"{what}: run [{lo},{hi}) outside [0,{length}) or empty") if lo < prev_hi: raise ValueError(f"{what}: run [{lo},{hi}) overlaps / is not sorted after {prev_hi}") out.append((lo, hi)) prev_hi = hi return tuple(out) def complement_runs(runs: Sequence[Run], length: int) -> Tuple[Run, ...]: """The maximal runs of ``[0, length)`` not covered by ``runs`` (static-shape-plan.md §6.2 'non-image runs').""" runs = check_runs(runs, length, "complement_runs") out: List[Run] = [] cursor = 0 for lo, hi in runs: if lo > cursor: out.append((cursor, lo)) cursor = hi if cursor < length: out.append((cursor, length)) return tuple(out) def runs_length(runs: Sequence[Run]) -> int: return sum(hi - lo for lo, hi in runs) def runs_index(runs: Sequence[Run]) -> torch.Tensor: """Concatenated slot indices of ``runs`` in run order: the gather index of the hoist (§6.2).""" if not runs: return torch.zeros(0, dtype=torch.int64) return torch.cat([torch.arange(lo, hi, dtype=torch.int64) for lo, hi in runs]) def additive_key_mask(valid: torch.Tensor, rows: int) -> torch.Tensor: """``[1, 1, rows, L]`` bf16 additive mask: ``0`` where ``valid[L]`` is True, :data:`MASK_NEG` elsewhere.""" if valid.dtype != torch.bool or valid.ndim != 1: raise TypeError(f"valid must be a 1-D bool tensor, got {valid.dtype} {tuple(valid.shape)}") row = torch.where(valid, torch.zeros((), dtype=torch.float32), torch.full((), MASK_NEG, dtype=torch.float32)) return row.to(MASK_DTYPE)[None, None, None, :].expand(1, 1, rows, valid.numel()).clone() def _check_slots(dlayout: "DeviceLayout", slots: SlotLayout) -> torch.Tensor: """Validate ``slots`` against the baked shape key; return ``key_valid[S_pad]`` (bool).""" if not isinstance(slots, SlotLayout): raise TypeError(f"expected a SlotLayout, got {type(slots).__name__}") if slots.shape_key != dlayout.shape_key: raise StaticShapeError(f"slots built for {slots.shape_key} but layout is {dlayout.shape_key}") kv = slots.key_valid if kv.dtype != torch.bool or tuple(kv.shape) != (1, dlayout.s_pad): raise StaticShapeError(f"key_valid must be bool (1, {dlayout.s_pad}), got {kv.dtype} {tuple(kv.shape)}") if slots.text_padding != dlayout.text_padding: raise StaticShapeError(f"slots text_padding {slots.text_padding!r} != layout {dlayout.text_padding!r}") return kv[0] # -------------------------------------------------------------------------------------------------------------------- # DeviceLayout # -------------------------------------------------------------------------------------------------------------------- @dataclass(frozen=True) class DeviceLayout: """Every static integer of one shape key plus the mask / gather builders (IMPLEMENTATION_PLAN.md §2.2). Built by :meth:`from_layout`; all cross-checks between ``GR00TConfig`` and ``LayoutConfig`` run there and raise. """ version: str layout_name: str shape_key: ShapeKey embodiment_id: int # sequence s_pad: int # S_pad_llm (static-shape-plan §4.2) s_pad_vl: int # rows of vlln / VL-SA / cross K/V (== s_pad in the first implementation) l_max: int text_padding: str # "right" | "text_slot" n_images: int grid_thw: Optional[Tuple[int, int, int]] tokens_per_image: int image_slots: Tuple[Run, ...] # == LayoutConfig.image_slots nonimage_runs: Tuple[Run, ...] # complement of image_slots in [0, s_pad) # cross K/V key sets (§6.2) cross_key_set: str # DiTConfig.cross_key_set: "all" | "alternate_text_image" l_img_real: int l_img_pad: Optional[int] # None for the "all" key set (N1.5) l_txt_real: int l_txt_pad: Optional[int] # None for the "all" key set (N1.5) l_all_pad: Optional[int] # s_pad_vl for the "all" key set, None otherwise # DiT rows (§6.1) m_logical: int m_pad: int action_horizon: int velocity_rows: slice # decoder rows that form the velocity: [m_logical - H, m_logical) noise_rows: slice # rows of x_t that receive the noise / actions: identical range (tt/encoders.py noise_to_xt) attention_scale: float # 48 ** -0.5, passed explicitly (D5) # shapes pixel_shape: Tuple[int, ...] state_shape: Tuple[int, int, int] noise_shape: Tuple[int, int, int] # tokens / tables image_token_id: int pad_id: int rope_table: str # "constant_arange" | "per_call_rows" | "constant_mrope_per_layout" llm_head_dim: int has_vlsa: bool mask_presence: Mapping[str, bool] # name -> tensor (True) or None (False); fixed per layout (§10 item 4) cfg: GR00TConfig = field(repr=False, compare=False) layout: LayoutConfig = field(repr=False, compare=False) # ------------------------------------------------------------------ construction @classmethod def from_layout(cls, cfg: GR00TConfig, layout: LayoutConfig) -> "DeviceLayout": """Bake the static layout of ``layout`` (one of ``cfg.layouts``); raises on any inconsistency.""" if not isinstance(cfg, GR00TConfig) or not isinstance(layout, LayoutConfig): raise TypeError("from_layout(cfg: GR00TConfig, layout: LayoutConfig)") if not any(lay is layout or lay == layout for lay in cfg.layouts): raise ValueError( f"layout {layout.name!r} is not one of {cfg.version}'s layouts {[l.name for l in cfg.layouts]}" ) s_pad = int(layout.s_pad_llm) s_pad_vl = int(layout.s_pad_vl) if s_pad <= 0 or s_pad % LLM_SEQ_PAD_MULTIPLE != 0: raise ValueError(f"{layout.name}: s_pad_llm={s_pad} is not a positive multiple of {LLM_SEQ_PAD_MULTIPLE}") if s_pad_vl % TILE != 0: raise ValueError(f"{layout.name}: s_pad_vl={s_pad_vl} is not a multiple of {TILE}") if s_pad_vl != s_pad: # static-shape-plan §4.2: S_pad_vl == S_pad_llm in the first implementation (no slice); the per-call masks # below are built over all S_pad slots, so a tighter S_pad_vl needs a slot selection this module lacks. raise NotImplementedError(f"{layout.name}: s_pad_vl={s_pad_vl} != s_pad_llm={s_pad} is not supported") if layout.text_padding not in ("right", "text_slot"): raise ValueError(f"{layout.name}: unknown text_padding {layout.text_padding!r}") if layout.l_max <= 0 or layout.l_max > s_pad: raise ValueError(f"{layout.name}: l_max={layout.l_max} out of range for s_pad={s_pad}") # image slots and their complement (§4.5 / §6.2) image_slots = check_runs(layout.image_slots, s_pad, f"{layout.name}.image_slots") if len(image_slots) != layout.n_images: raise ValueError(f"{layout.name}: {len(image_slots)} image runs for n_images={layout.n_images}") for lo, hi in image_slots: if hi - lo != layout.tokens_per_image: raise ValueError( f"{layout.name}: image run [{lo},{hi}) has {hi - lo} != {layout.tokens_per_image} tokens" ) nonimage_runs = complement_runs(image_slots, s_pad) l_img_real = runs_length(image_slots) l_txt_real = runs_length(nonimage_runs) if l_img_real + l_txt_real != s_pad: raise AssertionError("image + non-image slots must partition [0, s_pad)") if layout.non_image_slots is not None and layout.non_image_slots != l_txt_real: raise ValueError(f"{layout.name}: non_image_slots={layout.non_image_slots} != {l_txt_real}") # cross-attention key sets (§6.2); presence of the image mask is fixed here (§10 item 4) dit = cfg.dit if dit.cross_key_set == "all": if ( layout.cross_all_keys is None or layout.cross_image_keys is not None or layout.cross_text_keys is not None ): raise ValueError(f"{layout.name}: key set 'all' needs cross_all_keys only") if layout.cross_all_keys[1] != s_pad_vl or not (0 < layout.cross_all_keys[0] <= s_pad_vl): raise ValueError( f"{layout.name}: cross_all_keys={layout.cross_all_keys} inconsistent with s_pad_vl={s_pad_vl}" ) l_img_pad: Optional[int] = None l_txt_pad: Optional[int] = None l_all_pad: Optional[int] = s_pad_vl image_mask_present = False text_mask_present = False all_mask_present = True elif dit.cross_key_set == "alternate_text_image": if layout.cross_all_keys is not None or layout.cross_image_keys is None or layout.cross_text_keys is None: raise ValueError( f"{layout.name}: key set 'alternate_text_image' needs cross_image_keys and cross_text_keys" ) l_img_pad = pad_to_tile(l_img_real) l_txt_pad = pad_to_tile(l_txt_real) if tuple(layout.cross_image_keys) != (l_img_real, l_img_pad): raise ValueError( f"{layout.name}: cross_image_keys={layout.cross_image_keys} != {(l_img_real, l_img_pad)}" ) if tuple(layout.cross_text_keys) != (l_txt_real, l_txt_pad): raise ValueError( f"{layout.name}: cross_text_keys={layout.cross_text_keys} != {(l_txt_real, l_txt_pad)}" ) l_all_pad = None image_mask_present = l_img_pad != l_img_real text_mask_present = True all_mask_present = False else: raise ValueError(f"{cfg.version}: unknown DiTConfig.cross_key_set {dit.cross_key_set!r}") for i in dit.cross_block_indices: if dit.cross_key_subset(i) not in ("all", "text", "image"): raise ValueError(f"block {i}: unknown key subset {dit.cross_key_subset(i)!r}") # DiT rows (§6.1) and sampler constants (§10 item 7) horizon = cfg.encoders.action_horizon m_logical = int(layout.dit_m_logical) m_pad = int(layout.dit_m_pad) if m_pad != DIT_QUERY_ROWS_PAD or dit.m_pad != m_pad: raise ValueError(f"{layout.name}: dit_m_pad={m_pad} / DiTConfig.m_pad={dit.m_pad} != {DIT_QUERY_ROWS_PAD}") if ( m_logical != dit.n_query_tokens or m_logical != 1 + dit.n_future_tokens + horizon or not 0 < m_logical <= m_pad ): raise ValueError( f"{layout.name}: dit_m_logical={m_logical} != 1 + {dit.n_future_tokens} future + {horizon} actions " f"(DiTConfig.n_query_tokens={dit.n_query_tokens}, m_pad={m_pad})" ) velocity_rows = slice(m_logical - horizon, m_logical) if list(range(m_logical)[dit.velocity_rows]) != list(range(m_logical - horizon, m_logical)): raise ValueError( f"{cfg.version}: DiTConfig.velocity_rows={dit.velocity_rows} != rows [{m_logical - horizon},{m_logical})" ) if dit.head_dim_pad != DIT_HEAD_DIM_PAD: raise ValueError(f"{cfg.version}: DiTConfig.head_dim_pad={dit.head_dim_pad} != {DIT_HEAD_DIM_PAD}") if ( cfg.sampler.dt != _DT or cfg.sampler.num_inference_timesteps != _N_STEPS or tuple(cfg.sampler.timestep_buckets) != _BUCKETS ): raise ValueError( f"{cfg.version}: sampler (dt={cfg.sampler.dt}, steps={cfg.sampler.num_inference_timesteps}, " f"buckets={cfg.sampler.timestep_buckets}) != ({_DT}, {_N_STEPS}, {_BUCKETS})" ) # shapes (§10 items 1 / 6) vis = cfg.vision if vis.image_size is not None: if layout.grid_thw is not None: raise ValueError(f"{layout.name}: grid_thw must be None for a fixed-size SigLIP tower") pixel_shape: Tuple[int, ...] = (layout.n_images, 3, vis.image_size, vis.image_size) grid_thw = None else: if layout.grid_thw is None: raise ValueError(f"{layout.name}: a dynamic-resolution tower needs grid_thw") grid_thw = tuple(int(g) for g in layout.grid_thw) if len(grid_thw) != 3: raise ValueError(f"{layout.name}: grid_thw={grid_thw} is not (t, h, w)") t, h, w = grid_thw if ( t * h * w != layout.patches_per_image or (h // vis.merge_size) * (w // vis.merge_size) * t != layout.tokens_per_image ): raise ValueError( f"{layout.name}: grid {grid_thw} gives {t * h * w} patches / " f"{t * (h // vis.merge_size) * (w // vis.merge_size)} tokens, layout says " f"{layout.patches_per_image} / {layout.tokens_per_image}" ) pixel_shape = (layout.n_images * layout.patches_per_image, vis.patch_in_features) state_shape = tuple(int(x) for x in layout.state_shape) noise_shape = tuple(int(x) for x in layout.noise_shape) if state_shape != (1, 1, cfg.encoders.max_state_dim): raise ValueError(f"{layout.name}: state_shape={state_shape} != (1, 1, {cfg.encoders.max_state_dim})") if noise_shape != (1, horizon, cfg.encoders.max_action_dim) or tuple(cfg.sampler.noise_shape) != noise_shape: raise ValueError( f"{layout.name}: noise_shape={noise_shape} != (1, {horizon}, {cfg.encoders.max_action_dim})" ) has_vlsa = cfg.vl_adapter.n_blocks > 0 mask_presence = { "self": True, "cross_text": text_mask_present, "cross_image": image_mask_present, "cross_all": all_mask_present, "llm": layout.text_padding == "text_slot", "vl": has_vlsa, } key = ShapeKey( version=cfg.version, n_images=layout.n_images, grid_thw=grid_thw, tokens_per_image=layout.tokens_per_image, s_pad=s_pad, l_max=layout.l_max, layout_name=layout.name, ) return cls( version=cfg.version, layout_name=layout.name, shape_key=key, embodiment_id=int(layout.embodiment_id), s_pad=s_pad, s_pad_vl=s_pad_vl, l_max=int(layout.l_max), text_padding=layout.text_padding, n_images=int(layout.n_images), grid_thw=grid_thw, tokens_per_image=int(layout.tokens_per_image), image_slots=image_slots, nonimage_runs=nonimage_runs, cross_key_set=dit.cross_key_set, l_img_real=l_img_real, l_img_pad=l_img_pad, l_txt_real=l_txt_real, l_txt_pad=l_txt_pad, l_all_pad=l_all_pad, m_logical=m_logical, m_pad=m_pad, action_horizon=int(horizon), velocity_rows=velocity_rows, noise_rows=slice(m_logical - horizon, m_logical), attention_scale=float(dit.attention_scale), pixel_shape=pixel_shape, state_shape=state_shape, # type: ignore[arg-type] noise_shape=noise_shape, # type: ignore[arg-type] image_token_id=int(cfg.prompt.image_token_id), pad_id=int(cfg.prompt.pad_id), rope_table=layout.rope_table, llm_head_dim=int(cfg.llm.head_dim), has_vlsa=has_vlsa, mask_presence=mask_presence, cfg=cfg, layout=layout, ) @classmethod def for_inputs(cls, cfg: GR00TConfig, mi: "ModelInputs") -> "DeviceLayout": """The layout ``encode`` selected for ``mi`` (``mi.shape_key.layout_name``).""" if mi.version != cfg.version: raise ValueError(f"inputs are for {mi.version}, config is {cfg.version}") return cls.from_layout(cfg, cfg.layout(mi.shape_key.layout_name)) # ------------------------------------------------------------------ key sets and gather pieces def key_set(self, block_idx: int) -> str: """Key set of cross block ``block_idx``: ``"all"`` | ``"text"`` | ``"image"`` (``DiTConfig.cross_key_subset``).""" return self.cfg.dit.cross_key_subset(block_idx) def cross_mask_name(self, block_idx: int) -> str: """Mask name (:data:`MASK_NAMES`) a cross block reads; the entry may be ``None`` (see ``mask_presence``).""" return f"cross_{self.key_set(block_idx)}" def gather_pieces(self) -> Tuple[Tuple[Run, ...], Tuple[Run, ...]]: """``(image_runs, nonimage_runs)`` of the once-per-inference slot permutation (static-shape-plan.md §6.2): ``vl_img = concat(rows of image_runs)`` (``L_img_real`` rows, zero-padded to ``L_img_pad``) and ``vl_txt = concat(rows of nonimage_runs)`` (``L_txt_real`` -> ``L_txt_pad``). N1.5 keeps the identity.""" return self.image_slots, self.nonimage_runs def image_slot_index(self) -> torch.Tensor: """``[L_img_real]`` int64: slot of every gathered image key column (§6.2).""" return runs_index(self.image_slots) def nonimage_slot_index(self) -> torch.Tensor: """``[L_txt_real]`` int64: slot of every gathered non-image (text) key column (§6.2).""" return runs_index(self.nonimage_runs) # ------------------------------------------------------------------ masks def self_mask(self) -> torch.Tensor: """Constant ``[1, 1, 64, 64]`` bf16 DiT self-attention mask: ``-1e9`` on pad key columns ``>= M`` (§6.1).""" valid = torch.arange(self.m_pad) < self.m_logical return additive_key_mask(valid, self.m_pad) def cross_image_mask(self) -> Optional[torch.Tensor]: """Constant ``[1, 1, 64, L_img_pad]`` mask on the zero pad rows of ``vl_img``; ``None`` when ``L_img_real == L_img_pad`` (N1.7 DROID 448) or for the ``"all"`` key set (§6.2).""" if not self.mask_presence["cross_image"]: return None assert self.l_img_pad is not None valid = torch.arange(self.l_img_pad) < self.l_img_real return additive_key_mask(valid, self.m_pad) def cross_masks(self, slots: SlotLayout) -> Dict[str, Optional[torch.Tensor]]: """Per-call cross-attention masks ``{"text", "image", "all"}`` (static-shape-plan.md §6.2 / §6.4). * ``text`` ``[1,1,64,L_txt_pad]``: ``-1e9`` where the gathered non-image slot is not ``key_valid`` or the column is a zero pad row ``>= L_txt_real`` -- the reference's ``~image_mask & attention_mask`` key set; * ``image``: :meth:`cross_image_mask` (constant, possibly ``None``); * ``all`` ``[1,1,64,S_pad_vl]``: ``-1e9`` where ``key_valid`` is False (N1.5; the reference has no mask because it has no pads). Keys the version's DiT never reads are ``None`` (fixed per layout, ``mask_presence``). """ kv = _check_slots(self, slots) out: Dict[str, Optional[torch.Tensor]] = {"text": None, "image": self.cross_image_mask(), "all": None} if self.mask_presence["cross_text"]: assert self.l_txt_pad is not None gathered = kv[self.nonimage_slot_index()] valid = torch.zeros(self.l_txt_pad, dtype=torch.bool) valid[: self.l_txt_real] = gathered if not bool(valid.any()): raise StaticShapeError("text key set has no valid key") out["text"] = additive_key_mask(valid, self.m_pad) if self.mask_presence["cross_all"]: if not bool(kv.any()): raise StaticShapeError("all-token key set has no valid key") out["all"] = additive_key_mask(kv, self.m_pad) return out def llm_mask(self, slots: SlotLayout) -> Optional[torch.Tensor]: """``[1, 1, S_pad, S_pad]`` additive causal AND key-valid mask for ``text_padding == "text_slot"`` (N1.6, static-shape-plan.md §4.3 item 3), else ``None`` (right padding runs SDPA ``is_causal=True`` without a mask).""" kv = _check_slots(self, slots) if not self.mask_presence["llm"]: return None causal = torch.tril(torch.ones(self.s_pad, self.s_pad, dtype=torch.bool)) allowed = causal & kv[None, :] if not bool(allowed.any(dim=-1).all()): raise StaticShapeError("LLM mask leaves a query row without any valid key (softmax would be NaN)") neg = torch.full((), MASK_NEG, dtype=torch.float32) return torch.where(allowed, torch.zeros((), dtype=torch.float32), neg).to(MASK_DTYPE)[None, None] def vl_mask(self, slots: SlotLayout) -> Optional[torch.Tensor]: """``[1, 1, S_pad_vl, S_pad_vl]`` key-valid mask of the VL self-attention blocks (N1.5 / N1.7, static-shape-plan.md §5 / §4.3 item 5), ``None`` when the version has no VL-SA (N1.6).""" kv = _check_slots(self, slots) if not self.mask_presence["vl"]: return None return additive_key_mask(kv, self.s_pad_vl) def constant_masks(self) -> Dict[str, Optional[torch.Tensor]]: """``{"self", "cross_image"}`` -- built once at load (static-shape-plan.md §9).""" return {"self": self.self_mask(), "cross_image": self.cross_image_mask()} def per_call_masks(self, slots: SlotLayout) -> Dict[str, Optional[torch.Tensor]]: """``{"cross_text", "cross_all", "llm", "vl"}`` for one call (contents change with ``n_text``, §9).""" cross = self.cross_masks(slots) return { "cross_text": cross["text"], "cross_all": cross["all"], "llm": self.llm_mask(slots), "vl": self.vl_mask(slots), } def all_masks(self, slots: SlotLayout) -> Dict[str, Optional[torch.Tensor]]: """Every mask of :data:`MASK_NAMES`; ``None`` entries follow ``mask_presence`` exactly.""" out = dict(self.constant_masks()) out.update(self.per_call_masks(slots)) for name in MASK_NAMES: if (out[name] is not None) != self.mask_presence[name]: raise StaticShapeError( f"mask {name!r} presence {out[name] is not None} != baked {self.mask_presence[name]}" ) return out # ------------------------------------------------------------------ static-shape asserts (§10 items 1-8) def assert_call( self, mi: "ModelInputs", noise: Optional[torch.Tensor] = None, loaded_embodiment_ids: Optional[Collection[int]] = None, ) -> None: """Host-side checks before every trace replay (static-shape-plan.md §10 items 1-8); raises :class:`StaticShapeError` on the first violation. 1. ``n_images``, every ``image_grid_thw[i]`` (N1.7), ``tokens_per_image`` and ``pixel_values.shape``; 2. ``S <= S_pad``, ``n_text <= L_max``, image slot runs recomputed from ``input_ids_padded == image_token_id`` (and from the unpadded ids through ``slot_of_token``) equal the baked ones; 3. pad ids exactly on the non-key-valid slots, which are where the padding policy puts them (``right``: ``[S, S_pad)``; ``text_slot``: ``[t0 + n_text, t0 + L_max)`` and the trailing run); 4. mask presence per name equals the baked pattern, ``key_valid`` is bool ``(1, S_pad)``; 5. positions: ``per_call_rows`` (N1.6) real tokens carry ``arange(S)`` in slot order (max ``S - 1``), pads are finite and ``>= S``; ``constant_arange`` (N1.5) ``position_ids == arange(S_pad)``; ``constant_mrope_per_layout`` (N1.7) ``(3, 1, S_pad)`` with the grid of item 1; RoPE tables ``(1, S_pad, head_dim)`` and finite; 6. ``state`` / ``state_mask`` ``(1, 1, D_max)`` finite, ``embodiment_id`` ``(1,)`` int64 (and in ``loaded_embodiment_ids`` when given), ``noise.shape == noise_shape`` when given; 7. the DiT / sampler constants were verified at construction (:meth:`from_layout`); 8. hoist lengths: image / non-image key counts of this call equal the baked ``L_img_real`` / ``L_txt_real``, the valid text keys number ``S - L_img_real`` and the mask is built for the same ``n_text`` as the ids. """ fail = StaticShapeError if mi.version != self.version: raise fail(f"inputs are for {mi.version}, layout for {self.version}") if mi.shape_key != self.shape_key or mi.slots.shape_key != self.shape_key: raise fail(f"shape key {mi.shape_key} / {mi.slots.shape_key} != baked {self.shape_key}") slots = mi.slots s_pad = self.s_pad # 1. camera layout if mi.n_images != self.n_images: raise fail(f"n_images={mi.n_images} != {self.n_images}") if tuple(int(t) for t in mi.tokens_per_image) != (self.tokens_per_image,) * self.n_images: raise fail(f"tokens_per_image={tuple(mi.tokens_per_image)} != {(self.tokens_per_image,) * self.n_images}") if self.grid_thw is not None: g = mi.image_grid_thw if g is None or tuple(g.shape) != (self.n_images, 3): raise fail(f"image_grid_thw must be ({self.n_images}, 3), got {None if g is None else tuple(g.shape)}") for i in range(self.n_images): if tuple(int(x) for x in g[i].tolist()) != self.grid_thw: raise fail(f"image {i} grid {g[i].tolist()} != baked {self.grid_thw}") elif mi.image_grid_thw is not None: raise fail("image_grid_thw given for a fixed-size tower") if tuple(mi.pixel_values.shape) != self.pixel_shape: raise fail(f"pixel_values.shape={tuple(mi.pixel_values.shape)} != {self.pixel_shape}") if not bool(torch.isfinite(mi.pixel_values).all()): raise fail("pixel_values contain non-finite values") # 2. sequence lengths and image slot runs s = int(mi.seq_len) if tuple(mi.input_ids.shape) != (1, s) or s <= 0 or s > s_pad or slots.seq_len != s: raise fail( f"S={s} (input_ids {tuple(mi.input_ids.shape)}, slots.seq_len={slots.seq_len}) not in (0, {s_pad}]" ) if mi.n_text != slots.n_text: raise fail(f"ModelInputs.n_text={mi.n_text} != slots.n_text={slots.n_text} (mask built for another n_text)") if not 0 <= mi.n_text <= self.l_max: raise fail(f"n_text={mi.n_text} exceeds L_max={self.l_max}") ids_padded = slots.input_ids_padded if tuple(ids_padded.shape) != (1, s_pad) or ids_padded.dtype != torch.int64: raise fail(f"input_ids_padded must be int64 (1, {s_pad}), got {ids_padded.dtype} {tuple(ids_padded.shape)}") runs = token_runs(ids_padded[0], self.image_token_id) if runs != self.image_slots: raise fail(f"image slot runs {runs} != baked {self.image_slots} (template / tokenizer drift?)") if tuple(tuple(r) for r in slots.image_slots) != self.image_slots: raise fail(f"slots.image_slots={slots.image_slots} != baked {self.image_slots}") sot = slots.slot_of_token if ( tuple(sot.shape) != (s,) or int(sot.min()) < 0 or int(sot.max()) >= s_pad or bool((sot[1:] <= sot[:-1]).any()) ): raise fail("slot_of_token must be a strictly increasing (S,) map into [0, S_pad)") raw_runs = token_runs(mi.input_ids[0], self.image_token_id) if raw_runs != tuple(tuple(r) for r in mi.image_slots): raise fail(f"unpadded image runs {raw_runs} != ModelInputs.image_slots {mi.image_slots}") if tuple((int(sot[a]), int(sot[b - 1]) + 1) for a, b in raw_runs) != self.image_slots: raise fail("unpadded image runs do not map onto the baked slots through slot_of_token") # 3. pad ids and pad positions kv = slots.key_valid if kv.dtype != torch.bool or tuple(kv.shape) != (1, s_pad): raise fail(f"key_valid must be bool (1, {s_pad}), got {kv.dtype} {tuple(kv.shape)}") kv0 = kv[0] if int(kv0.sum()) != s: raise fail(f"key_valid has {int(kv0.sum())} valid slots, S={s}") if not bool(kv0[sot].all()): raise fail("slot_of_token points at a non-valid slot") pad_idx = torch.nonzero(~kv0).reshape(-1) if not bool((ids_padded[0, pad_idx] == self.pad_id).all()): raise fail(f"pad slots do not all hold pad id {self.pad_id}") expected_pads = torch.zeros(s_pad, dtype=torch.bool) if self.text_padding == "right": expected_pads[s:] = True else: t0 = int(slots.text_budget_slots[0]) if tuple(slots.text_budget_slots) != (t0, t0 + self.l_max): raise fail(f"text_budget_slots={slots.text_budget_slots} != ({t0}, {t0 + self.l_max})") expected_pads[t0 + mi.n_text : t0 + self.l_max] = True expected_pads[s + self.l_max - mi.n_text :] = True if not torch.equal(~kv0, expected_pads): raise fail( f"pad positions {token_runs(~kv0, True)} != policy {self.text_padding!r} {token_runs(expected_pads, True)}" ) if tuple(tuple(r) for r in slots.pad_slots) != token_runs(expected_pads, True): raise fail(f"slots.pad_slots={slots.pad_slots} != {token_runs(expected_pads, True)}") if not bool(kv0[self.image_slot_index()].all()): raise fail("an image slot is not key-valid") # 4. mask presence (never toggles) and shapes masks = self.all_masks(slots) # raises on a presence mismatch shapes = { "self": (1, 1, self.m_pad, self.m_pad), "cross_text": (1, 1, self.m_pad, self.l_txt_pad), "cross_image": (1, 1, self.m_pad, self.l_img_pad), "cross_all": (1, 1, self.m_pad, self.l_all_pad), "llm": (1, 1, s_pad, s_pad), "vl": (1, 1, self.s_pad_vl, self.s_pad_vl), } for name, t in masks.items(): if t is not None and ( tuple(t.shape) != shapes[name] or t.dtype != MASK_DTYPE or not bool(torch.isfinite(t).all()) ): raise fail( f"mask {name!r}: shape {tuple(t.shape)} dtype {t.dtype} (expected {shapes[name]} {MASK_DTYPE}, finite)" ) # 5. positions and RoPE tables pos = mi.position_ids if self.rope_table == "per_call_rows": if tuple(pos.shape) != (1, s_pad): raise fail(f"position_ids.shape={tuple(pos.shape)} != (1, {s_pad})") real = pos[0][kv0] if not torch.equal(real, torch.arange(s, dtype=real.dtype)) or int(real.max()) != s - 1: raise fail("real-token positions are not arange(S) in slot order (max must be S - 1)") pads = pos[0][~kv0] if pads.numel() and (int(pads.min()) < s or pads.unique().numel() != pads.numel()): raise fail("pad positions must be unique and >= S") if not torch.equal(slots.position_ids, pos): raise fail("ModelInputs.position_ids != slots.position_ids") elif self.rope_table == "constant_arange": if tuple(pos.shape) != (1, s_pad) or not torch.equal(pos[0], torch.arange(s_pad, dtype=pos.dtype)): raise fail("constant_arange layout needs position_ids == arange(S_pad) (right padding)") elif self.rope_table == "constant_mrope_per_layout": if tuple(pos.shape) != (3, 1, s_pad): raise fail(f"mRoPE position_ids.shape={tuple(pos.shape)} != (3, 1, {s_pad})") else: raise fail(f"unknown rope_table {self.rope_table!r}") for name, t in (("rope_cos", mi.rope_cos), ("rope_sin", mi.rope_sin)): if tuple(t.shape) != (1, s_pad, self.llm_head_dim) or not bool(torch.isfinite(t).all()): raise fail(f"{name}.shape={tuple(t.shape)} != (1, {s_pad}, {self.llm_head_dim}) or non-finite") # 6. state / embodiment / noise if tuple(mi.state.shape) != self.state_shape or not bool(torch.isfinite(mi.state).all()): raise fail(f"state.shape={tuple(mi.state.shape)} != {self.state_shape} or non-finite") if tuple(mi.state_mask.shape) != self.state_shape: raise fail(f"state_mask.shape={tuple(mi.state_mask.shape)} != {self.state_shape}") if tuple(mi.embodiment_id.shape) != (1,) or mi.embodiment_id.dtype != torch.int64: raise fail( f"embodiment_id must be int64 (1,), got {mi.embodiment_id.dtype} {tuple(mi.embodiment_id.shape)}" ) if loaded_embodiment_ids is not None and int(mi.embodiment_id[0]) not in set( int(e) for e in loaded_embodiment_ids ): raise fail( f"embodiment_id {int(mi.embodiment_id[0])} has no loaded slice buffers ({sorted(loaded_embodiment_ids)})" ) if noise is not None: if tuple(noise.shape) != self.noise_shape or not bool(torch.isfinite(noise).all()): raise fail(f"noise.shape={tuple(noise.shape)} != {self.noise_shape} or non-finite") # 8. hoist lengths (7. is construction-time) n_img_keys = runs_length(runs) if n_img_keys != self.l_img_real or s_pad - n_img_keys != self.l_txt_real: raise fail( f"hoist lengths ({n_img_keys}, {s_pad - n_img_keys}) != baked ({self.l_img_real}, {self.l_txt_real})" ) n_valid_text = int(kv0[self.nonimage_slot_index()].sum()) if n_valid_text != s - self.l_img_real: raise fail(f"{n_valid_text} valid text keys != S - L_img_real = {s - self.l_img_real}") # -------------------------------------------------------------------------------------------------------------------- # device materialisation (ttnn imported lazily) # -------------------------------------------------------------------------------------------------------------------- def materialise( device: Any, dlayout: DeviceLayout, slots: SlotLayout, buffers: Optional[Dict[str, Any]] = None, *, memory_config: Any = None, cq_id: Any = None, ) -> Dict[str, Any]: """Upload the masks of ``dlayout`` for ``slots`` as bf16 TILE DRAM tensors (static-shape-plan.md §9). * ``buffers is None`` (load time): every mask of :data:`MASK_NAMES` is created with ``ttnn.from_torch`` on ``device`` (``None`` entries stay ``None`` -- the presence pattern is fixed per layout) and the dict is returned; keep it, its addresses are baked into the traces. * ``buffers`` given (per call): the constant masks are left untouched and each per-call mask is written into its persistent device buffer with ``ttnn.copy_host_to_device_tensor`` (same shape, dtype, layout); the same dict is returned. A presence mismatch or a foreign key set raises. """ import ttnn # lazy: CPU tests import this module without a device mem = ttnn.DRAM_MEMORY_CONFIG if memory_config is None else memory_config masks = dlayout.all_masks(slots) def upload(t: torch.Tensor) -> Any: return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=mem) if buffers is None: return {name: (None if t is None else upload(t)) for name, t in masks.items()} if set(buffers) != set(MASK_NAMES): raise StaticShapeError(f"buffers keys {sorted(buffers)} != {sorted(MASK_NAMES)}") for name in MASK_NAMES: if (buffers[name] is None) != (masks[name] is None): raise StaticShapeError( f"mask {name!r}: buffer presence {buffers[name] is not None} != call {masks[name] is not None}" ) for name in PER_CALL_MASKS: t = masks[name] if t is None: continue dev = buffers[name] if tuple(dev.shape) != tuple(t.shape): raise StaticShapeError(f"mask {name!r}: device buffer {tuple(dev.shape)} != host {tuple(t.shape)}") host = ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT) if cq_id is None: ttnn.copy_host_to_device_tensor(host, dev) else: ttnn.copy_host_to_device_tensor(host, dev, cq_id) return buffers __all__ = [ "CONSTANT_MASKS", "MASK_DTYPE", "MASK_NAMES", "MASK_NEG", "PER_CALL_MASKS", "DeviceLayout", "StaticShapeError", "additive_key_mask", "check_runs", "complement_runs", "materialise", "pad_to_tile", "runs_index", "runs_length", ]