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

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