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

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

"""Vision -> LLM projectors and the token splicer on device (IMPLEMENTATION_PLAN.md §2.2 row ``tt/projectors.py``,
§3.2 "N1.6 projector" / "Splice"; WP-07).  Class names mirror ``reference/projectors.py``.

* :class:`LinearProjectorTT` (N1.5): one biased ``linear(1152 -> 2048)`` on the post-LN rows (plan `projector.fc1.W/.b`,
  bf16 under every policy, f1 §2.1); output ``[n_img, 1, P_pad, 2048]``.
* :class:`UnshuffleMLPProjectorTT` (N1.6): ``ttnn.slice`` rows ``0:324`` -> ``ttnn.reshape([n, 1, 81, 4608])`` -> ``layer_norm
  (projector.ln, eps 1e-5)`` -> ``linear(projector.fc1, GELU-erf fused = UnaryWithParam(GELU, 0.0))`` -> ``linear(projector.fc2)``;
  output logical 81 rows in a 96-row tile tensor.  The reshape **is** the 2x2 pixel-unshuffle because the tower rows are
  in block-major order (``tt/siglip.py`` :func:`~models.experimental.gr00t.tt.siglip.block_major_perm`, static-shape-plan
  §3) and the plan was built with ``TransformOptions(n16_unshuffle_order="(i,j,c)")`` -- the ``(i, j, c)`` column order is
  folded into ``projector.ln.w/.b`` and the input columns of ``projector.fc1.W`` (f1 T14, HL §3.3); the constructor asserts
  that option on the weights (``TTWeights.options`` and, when present, the tensorbin manifest).
* :class:`IdentityProjectorTT` (N1.7): the Qwen3-VL merger lives in the tower; the rows pass through.
* :class:`TokenSplicer`: ``inputs_embeds [1, 1, S_pad, 2048]`` from the token-embedding rows (image slots zero, uploaded
  per call) and the vision rows by **constant** slices and one ``concat(dim=2)`` of <= ``2 n_img + 1`` pieces
  (static-shape-plan §4.5, plan §3.2 "Splice"); :meth:`TokenSplicer.place_into_zeros` places the vision rows into zeros
  for the N1.7 deepstack adds (same pieces, constant zero runs).  Slot boundaries are not tile aligned (N1.6 ``[44, 125)``,
  N1.5 ``[20, 276)``), so the splice runs in ROW_MAJOR: ``to_layout(RM)`` of both inputs, row slices (RM pages = rows,
  exact copies), ``concat``, ``to_layout(TILE)`` -- the same detour ``ttnn.slice`` takes internally for unaligned TILE
  begins (``slice.cpp`` ``rm_only``), done once instead of per piece.  Every value is copied, never recomputed: the
  result equals the host splice bit for bit.

``ttnn`` is imported lazily; the piece plan (:func:`splice_plan`) and the shape checks are CPU-testable.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any, Dict, List, Sequence, Tuple

import torch

from models.experimental.gr00t.common.configs import TILE, GR00TConfig
from models.experimental.gr00t.tt import layers as L
from models.experimental.gr00t.tt.device import CKC_HIFI2_FP32ACC, DRAM, L1, resolve_mem
from models.experimental.gr00t.tt.layout import DeviceLayout

PROJECTOR_KINDS: Tuple[str, ...] = ("linear", "unshuffle_ln_mlp", "qwen3vl_merger")
#: the unshuffle column order the device reshape produces (HL §3.3; ``common.weights.UNSHUFFLE_ORDERS``)
UNSHUFFLE_ORDER_DEVICE = "(i,j,c)"
#: ``tt.layers`` activation name of the exact erf GELU (``UnaryWithParam(GELU, 0.0)``; synthesis §3 N9: never tanh here)
GELU_ERF = "gelu"
#: plan §3.3: LLM residual stream in L1 at S_pad <= 256, DRAM above (the splice output is the LLM input)
SPLICE_L1_MAX_ROWS = 256


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


def _expect(t: Any, shape: Sequence[int], name: str) -> None:
    s = _shape(t)
    if len(s) == 4 and s[0] == 1 and s[1] == 1:
        s = s[2:]
    if s != tuple(int(d) for d in shape):
        raise ValueError(f"{name}: shape {s} != expected {tuple(shape)}")


def _weights_version(W: Any, cfg: GR00TConfig) -> None:
    v = getattr(W, "version", cfg.version)
    if v != cfg.version:
        raise ValueError(f"weights are for {v!r}, config is {cfg.version!r}")


# --------------------------------------------------------------------------------------------------------------------
# projectors
# --------------------------------------------------------------------------------------------------------------------


class LinearProjectorTT:
    """N1.5 ``mlp1``: ``linear(post_ln [n, 1, P_pad, 1152], projector.fc1.W [1152, 2048], .b)`` (plan §3.2)."""

    def __init__(
        self, cfg: GR00TConfig, dlayout: DeviceLayout, W: Any, *, mem: Any = L1, ckc: Any = CKC_HIFI2_FP32ACC
    ) -> None:
        pj = cfg.projector
        if pj.kind != "linear":
            raise ValueError(f"{cfg.version}: LinearProjectorTT needs projector kind 'linear', got {pj.kind!r}")
        if dlayout.version != cfg.version:
            raise ValueError(f"dlayout is for {dlayout.version!r}, config is {cfg.version!r}")
        _weights_version(W, cfg)
        self.cfg, self.dlayout = cfg, dlayout
        self.in_dim, self.out_dim = int(pj.in_dim), int(pj.out_dim)
        self.rows = int(dlayout.layout.patch_rows_pad)
        self.n_images = int(dlayout.n_images)
        self.mem, self.ckc = mem, ckc
        self.w = W.get("projector.fc1.W")
        self.b = W.get("projector.fc1.b")
        _expect(self.w, (self.in_dim, self.out_dim), "projector.fc1.W")
        _expect(self.b, (1, self.out_dim), "projector.fc1.b")
        self.weight_dtype = L.dtype_name(self.w.dtype)
        L.matmul_spec(self.rows, self.in_dim, self.out_dim, self.weight_dtype)  # raises for a missing table row
        self.batched = self.n_images == 1 or _table_ok(
            self.n_images * self.rows, self.in_dim, self.out_dim, self.weight_dtype
        )
        self.last_ops: List[str] = []
        self.debug_taps: Dict[str, Any] = {}

    def forward(self, post_ln: Any) -> Any:
        """``[n_img, 1, P_pad, 1152] -> [n_img, 1, P_pad, 2048]`` (batched, or per image when the table lacks the row)."""
        xs = _shape(post_ln)
        if xs != (self.n_images, 1, self.rows, self.in_dim):
            raise ValueError(f"post_ln must be [{self.n_images}, 1, {self.rows}, {self.in_dim}], got {xs}")
        self.last_ops = []
        if self.batched:
            self.last_ops.append("linear")
            return L.linear(post_ln, self.w, self.b, mem=self.mem, ckc=self.ckc)
        return _per_image(post_ln, self.n_images, self.mem, self.last_ops, self._one)

    def _one(self, x: Any) -> Any:
        self.last_ops.append("linear")
        return L.linear(x, self.w, self.b, mem=self.mem, ckc=self.ckc)

    __call__ = forward

    def describe(self) -> Dict[str, Any]:
        return {
            "kind": "linear",
            "in_dim": self.in_dim,
            "out_dim": self.out_dim,
            "rows": self.rows,
            "n_images": self.n_images,
            "weight_dtype": self.weight_dtype,
            "batched": self.batched,
            "primitive_ops": 1 if self.batched else 2 * self.n_images + 1,
        }


class UnshuffleMLPProjectorTT:
    """N1.6 ``mlp1``: slice -> reshape (= 2x2 unshuffle, ``(i,j,c)`` columns) -> LN(4608) -> fc1 + GELU(erf) -> fc2.

    Input ``post_ln [n_img, 1, P_pad=352, 1152]`` in the tower's block-major row order; output ``[n_img, 1, 81, 2048]``
    (logical 81 rows, tile padded to 96).  ``debug_taps`` (when ``collect_taps``): ``merged`` (the reshape output, the
    device-order counterpart of the golden ``vision_pixel_shuffled``), ``ln``, ``fc1`` (pre-GELU is not observable:
    the GELU is fused, so ``fc1`` is post-GELU).
    """

    def __init__(
        self,
        cfg: GR00TConfig,
        dlayout: DeviceLayout,
        W: Any,
        *,
        mem: Any = L1,
        ckc: Any = CKC_HIFI2_FP32ACC,
        collect_taps: bool = False,
    ) -> None:
        pj, vis = cfg.projector, cfg.vision
        if pj.kind != "unshuffle_ln_mlp":
            raise ValueError(f"{cfg.version}: UnshuffleMLPProjectorTT needs kind 'unshuffle_ln_mlp', got {pj.kind!r}")
        if dlayout.version != cfg.version:
            raise ValueError(f"dlayout is for {dlayout.version!r}, config is {cfg.version!r}")
        _weights_version(W, cfg)
        opts = getattr(W, "options", None)
        order = getattr(opts, "n16_unshuffle_order", None)
        if order != UNSHUFFLE_ORDER_DEVICE:
            raise ValueError(
                f"the weights were built with n16_unshuffle_order={order!r}; the device reshape produces "
                f"{UNSHUFFLE_ORDER_DEVICE!r} (fold it into mlp1 with TransformOptions, f1 T14 / tt.policy.plan_options)"
            )
        manifest_fn = getattr(W, "manifest", None)
        is_complete = getattr(W, "is_complete", None)
        if callable(manifest_fn) and callable(is_complete) and is_complete():
            m_order = manifest_fn().get("options", {}).get("n16_unshuffle_order")
            if m_order != UNSHUFFLE_ORDER_DEVICE:
                raise ValueError(
                    f"tensorbin manifest says n16_unshuffle_order={m_order!r}, need {UNSHUFFLE_ORDER_DEVICE!r}"
                )
        lay = dlayout.layout
        self.cfg, self.dlayout = cfg, dlayout
        self.n_images = int(dlayout.n_images)
        self.hidden = int(vis.hidden)
        self.merge = int(vis.merge_size)
        self.p, self.p_pad = int(lay.patches_per_image), int(lay.patch_rows_pad)
        grid = math.isqrt(self.p)
        if grid * grid != self.p or grid % self.merge:
            raise ValueError(f"{self.p} patches are not a {self.merge}x{self.merge}-mergeable square grid")
        self.tokens = (grid // self.merge) ** 2
        if self.tokens != int(lay.tokens_per_image) or self.tokens != int(dlayout.tokens_per_image):
            raise ValueError(f"merged tokens {self.tokens} != layout tokens_per_image {lay.tokens_per_image}")
        self.tokens_pad = -(-self.tokens // TILE) * TILE
        self.in_dim, self.hidden_dim, self.out_dim = int(pj.in_dim), int(pj.hidden_dim), int(pj.out_dim)
        if self.in_dim != self.merge * self.merge * self.hidden or pj.ln_dim != self.in_dim:
            raise ValueError(f"projector in_dim {self.in_dim} / ln_dim {pj.ln_dim} != merge^2 * hidden")
        if pj.ln_position != "post_concat" or pj.activation != "gelu_erf" or pj.ln_eps is None:
            raise ValueError(f"unexpected N1.6 projector config {pj}")
        self.ln_eps = float(pj.ln_eps)
        self.mem, self.ckc = mem, ckc
        self.collect_taps = bool(collect_taps)
        self.ln_w, self.ln_b = W.get("projector.ln.w"), W.get("projector.ln.b")
        self.w1, self.b1 = W.get("projector.fc1.W"), W.get("projector.fc1.b")
        self.w2, self.b2 = W.get("projector.fc2.W"), W.get("projector.fc2.b")
        _expect(self.ln_w, (1, self.in_dim), "projector.ln.w")
        _expect(self.ln_b, (1, self.in_dim), "projector.ln.b")
        _expect(self.w1, (self.in_dim, self.hidden_dim), "projector.fc1.W")
        _expect(self.b1, (1, self.hidden_dim), "projector.fc1.b")
        _expect(self.w2, (self.hidden_dim, self.out_dim), "projector.fc2.W")
        _expect(self.b2, (1, self.out_dim), "projector.fc2.b")
        self.dtype1, self.dtype2 = L.dtype_name(self.w1.dtype), L.dtype_name(self.w2.dtype)
        L.matmul_spec(self.tokens, self.in_dim, self.hidden_dim, self.dtype1)
        L.matmul_spec(self.tokens, self.hidden_dim, self.out_dim, self.dtype2)
        self.batched = self.n_images == 1 or (
            _table_ok(self.n_images * self.tokens, self.in_dim, self.hidden_dim, self.dtype1)
            and _table_ok(self.n_images * self.tokens, self.hidden_dim, self.out_dim, self.dtype2)
        )
        self.last_ops: List[str] = []
        self.debug_taps: Dict[str, Any] = {}

    def merge_rows(self, post_ln: Any) -> Any:
        """``slice`` rows ``0:P`` -> ``reshape [n, 1, P/4, 4*hidden]``: the ``(i, j, c)`` unshuffle of block-major rows."""
        import ttnn

        xs = _shape(post_ln)
        if xs != (self.n_images, 1, self.p_pad, self.hidden):
            raise ValueError(f"post_ln must be [{self.n_images}, 1, {self.p_pad}, {self.hidden}], got {xs}")
        n = xs[0]
        cut = ttnn.slice(post_ln, [0, 0, 0, 0], [n, 1, self.p, self.hidden])
        self.last_ops.append("slice")
        merged = ttnn.reshape(cut, (n, 1, self.tokens, self.in_dim), memory_config=resolve_mem(self.mem), pad_value=0.0)
        self.last_ops.append("reshape")
        ttnn.deallocate(cut)
        return merged

    def forward(self, post_ln: Any) -> Any:
        """``[n_img, 1, P_pad, 1152] -> [n_img, 1, 81, 2048]`` (5 ops for one image: slice, reshape, LN, fc1, fc2)."""
        import ttnn

        self.last_ops = []
        self.debug_taps = {}
        merged = self.merge_rows(post_ln)
        h = L.layer_norm(merged, self.ln_w, self.ln_b, eps=self.ln_eps, mem=self.mem)
        self.last_ops.append("layer_norm")
        self._tap("merged", merged)
        if self.batched:
            out = self._mlp(h)
        else:
            out = _per_image(h, self.n_images, self.mem, self.last_ops, self._mlp)
        self._tap("ln", h)
        return out

    def _mlp(self, h: Any) -> Any:
        f = L.linear(h, self.w1, self.b1, act=GELU_ERF, mem=self.mem, ckc=self.ckc)
        self.last_ops.append("linear")
        out = L.linear(f, self.w2, self.b2, mem=self.mem, ckc=self.ckc)
        self.last_ops.append("linear")
        self._tap("fc1", f)  # kept or deallocated
        return out

    def _tap(self, name: str, t: Any) -> bool:
        """Keep ``t`` when taps are collected; otherwise deallocate it.  Returns whether it was kept."""
        import ttnn

        if self.collect_taps:
            self.debug_taps[name] = t
            return True
        ttnn.deallocate(t)
        return False

    __call__ = forward

    def free_taps(self) -> None:
        import ttnn

        for t in self.debug_taps.values():
            ttnn.deallocate(t)
        self.debug_taps = {}

    def describe(self) -> Dict[str, Any]:
        return {
            "kind": "unshuffle_ln_mlp",
            "unshuffle_order": UNSHUFFLE_ORDER_DEVICE,
            "patches": self.p,
            "rows_pad": self.p_pad,
            "tokens": self.tokens,
            "tokens_pad": self.tokens_pad,
            "in_dim": self.in_dim,
            "hidden_dim": self.hidden_dim,
            "out_dim": self.out_dim,
            "ln_eps": self.ln_eps,
            "activation": GELU_ERF,
            "weight_dtypes": [self.dtype1, self.dtype2],
            "n_images": self.n_images,
            "batched": self.batched,
            "primitive_ops": 5 if self.batched else 3 + 3 * self.n_images + 1,
        }


class IdentityProjectorTT:
    """N1.7: the merger is part of ``Qwen3VLVisionTT``; ``forward`` validates and returns the ``[1, 1, R, 2048]`` rows."""

    def __init__(self, cfg: GR00TConfig, dlayout: DeviceLayout) -> None:
        if cfg.projector.kind != "qwen3vl_merger":
            raise ValueError(
                f"{cfg.version}: IdentityProjectorTT is for kind 'qwen3vl_merger', got {cfg.projector.kind!r}"
            )
        if dlayout.version != cfg.version:
            raise ValueError(f"dlayout is for {dlayout.version!r}, config is {cfg.version!r}")
        self.cfg, self.dlayout = cfg, dlayout
        self.out_dim = int(cfg.projector.out_dim)
        self.rows = int(dlayout.n_images) * int(dlayout.tokens_per_image)
        self.last_ops: List[str] = []

    def forward(self, merged: Any) -> Any:
        """``[1, 1, R, 2048]`` (or ``[R, 2048]``, viewed) with ``R >= n_img * tokens_per_image``; returned as is."""
        import ttnn

        self.last_ops = []
        s = _shape(merged)
        if len(s) == 2:
            merged = ttnn.reshape(merged, (1, 1, s[0], s[1]))  # 0-cost view (last dim unchanged)
            self.last_ops.append("reshape")
            s = _shape(merged)
        if len(s) != 4 or s[0] != 1 or s[1] != 1 or s[2] < self.rows or s[3] != self.out_dim:
            raise ValueError(f"merged rows must be [1, 1, >= {self.rows}, {self.out_dim}], got {s}")
        return merged

    __call__ = forward

    def describe(self) -> Dict[str, Any]:
        return {"kind": "identity", "rows": self.rows, "out_dim": self.out_dim, "primitive_ops": 0}


def build_projector(cfg: GR00TConfig, dlayout: DeviceLayout, W: Any, **kw: Any) -> Any:
    """The version's projector by ``cfg.projector.kind`` (the only version switch, for ``tt/backbone.py`` / ``tt/model.py``)."""
    kind = cfg.projector.kind
    if kind == "linear":
        return LinearProjectorTT(cfg, dlayout, W, **kw)
    if kind == "unshuffle_ln_mlp":
        return UnshuffleMLPProjectorTT(cfg, dlayout, W, **kw)
    if kind == "qwen3vl_merger":
        return IdentityProjectorTT(cfg, dlayout)
    raise ValueError(f"unknown projector kind {kind!r}; expected one of {PROJECTOR_KINDS}")


def _table_ok(m: int, k: int, n: int, dtype: str) -> bool:
    try:
        L.matmul_spec(m, k, n, dtype)
    except KeyError:
        return False
    return True


def _per_image(x: Any, n_images: int, mem: Any, ops: List[str], fn: Any) -> Any:
    """Apply ``fn`` to every batch element of ``x [n, 1, R, C]`` and concatenate along dim 0 (multi-image layouts whose
    batched row count has no ``MATMUL_CONFIGS`` row; decided at construction, recorded in ``describe()``)."""
    import ttnn

    s = _shape(x)
    outs: List[Any] = []
    for i in range(n_images):
        xi = ttnn.slice(x, [i, 0, 0, 0], [i + 1, s[1], s[2], s[3]])
        ops.append("slice")
        outs.append(fn(xi))
        ttnn.deallocate(xi)
    out = ttnn.concat(outs, dim=0, memory_config=resolve_mem(mem))
    ops.append("concat")
    for o in outs:
        ttnn.deallocate(o)
    return out


# --------------------------------------------------------------------------------------------------------------------
# token splicer
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class Piece:
    """Rows ``[lo, hi)`` of ``inputs_embeds`` come from the token embeddings (``image < 0``) or from image ``image``."""

    lo: int
    hi: int
    image: int = -1

    @property
    def rows(self) -> int:
        return self.hi - self.lo

    @property
    def is_image(self) -> bool:
        return self.image >= 0


def splice_plan(dlayout: DeviceLayout) -> Tuple[Piece, ...]:
    """The ordered pieces of the splice (static-shape-plan §4.5): image runs ``dlayout.image_slots`` (image index in slot
    order) and their complement ``dlayout.nonimage_runs``; they tile ``[0, S_pad)`` exactly, ``<= 2 n_img + 1`` pieces.
    """
    pieces = [Piece(lo, hi, k) for k, (lo, hi) in enumerate(dlayout.image_slots)]
    pieces += [Piece(lo, hi) for lo, hi in dlayout.nonimage_runs]
    pieces.sort(key=lambda p: p.lo)
    cursor = 0
    for p in pieces:
        if p.lo != cursor or p.hi <= p.lo:
            raise ValueError(f"splice pieces do not tile [0, {dlayout.s_pad}): gap/overlap at {p}")
        if p.is_image and p.rows != dlayout.tokens_per_image:
            raise ValueError(f"image piece {p} has {p.rows} rows != tokens_per_image {dlayout.tokens_per_image}")
        cursor = p.hi
    if cursor != dlayout.s_pad:
        raise ValueError(f"splice pieces end at {cursor} != S_pad {dlayout.s_pad}")
    if len(pieces) > 2 * dlayout.n_images + 1:
        raise AssertionError(f"{len(pieces)} pieces > 2 n_img + 1 = {2 * dlayout.n_images + 1}")
    return tuple(pieces)


def vis_mode(shape: Sequence[int], n_images: int, tokens: int, hidden: int) -> str:
    """Layout of a vision-rows tensor: ``"batched"`` = ``[n_img, 1, >= T, D]`` (batch = image, SigLIP projectors) or
    ``"flat"`` = ``[1, 1, >= n_img * T, D]`` (images consecutive, N1.7 merger).  For one image both coincide (rows ``0:T``
    of batch 0) and ``"batched"`` is returned.  Anything else raises."""
    s = tuple(int(d) for d in shape)
    if len(s) != 4 or s[1] != 1 or s[3] != hidden:
        raise ValueError(f"vis_rows must be [n_img|1, 1, rows, {hidden}], got {s}")
    if s[0] == n_images and s[2] >= tokens:
        return "batched"
    if s[0] == 1 and s[2] >= n_images * tokens:
        return "flat"
    raise ValueError(
        f"vis_rows {s}: expected [{n_images}, 1, >= {tokens}, {hidden}] or [1, 1, >= {n_images * tokens}, {hidden}]"
    )


def host_splice(text_rows: torch.Tensor, vis_rows: torch.Tensor, dlayout: DeviceLayout) -> torch.Tensor:
    """CPU oracle of :meth:`TokenSplicer.forward`: ``text_rows [1, 1, S_pad, D]`` with image slots replaced by the
    vision rows (``[n_img, 1, T(+pad), D]`` batched or ``[1, 1, n_img*T(+pad), D]`` flat)."""
    plan = splice_plan(dlayout)
    t = int(dlayout.tokens_per_image)
    if tuple(text_rows.shape) != (1, 1, dlayout.s_pad, text_rows.shape[-1]):
        raise ValueError(f"text_rows must be [1, 1, {dlayout.s_pad}, D], got {tuple(text_rows.shape)}")
    mode = vis_mode(vis_rows.shape, dlayout.n_images, t, int(text_rows.shape[-1]))
    out = text_rows.clone()
    for p in plan:
        if p.is_image:
            src = vis_rows[p.image, 0, :t] if mode == "batched" else vis_rows[0, 0, p.image * t : (p.image + 1) * t]
            out[0, 0, p.lo : p.hi] = src.to(out.dtype)
    return out


class TokenSplicer:
    """Constant slot splice of one static layout (module docstring; plan §2.2 ``TokenSplicer``).

    ``forward(text_embeds [1, 1, S_pad, 2048], vis_rows) -> inputs_embeds [1, 1, S_pad, 2048]`` and
    ``place_into_zeros(vis_rows) -> [1, 1, S_pad, 2048]``.  ``vis_rows`` is ``[n_img, 1, T_pad >= T, 2048]`` (batch = image:
    SigLIP projectors) or ``[1, 1, R >= n_img * T, 2048]`` (flat, images consecutive: N1.7 merger); for one image both
    coincide.  Output memory: ``mem`` (default L1 when ``S_pad <= 256`` else DRAM, plan §3.3).  ``last_ops`` lists the
    ttnn ops of the last call.
    """

    def __init__(self, dlayout: DeviceLayout, device: Any = None, *, mem: Any = None) -> None:
        if not isinstance(dlayout, DeviceLayout):
            raise TypeError(f"TokenSplicer needs a DeviceLayout, got {type(dlayout).__name__}")
        self.dlayout = dlayout
        self.s_pad = int(dlayout.s_pad)
        self.n_images = int(dlayout.n_images)
        self.tokens = int(dlayout.tokens_per_image)
        self.hidden = int(dlayout.cfg.llm.hidden)
        self.plan = splice_plan(dlayout)
        self.n_pieces = len(self.plan)
        self.mem = mem if mem is not None else (L1 if self.s_pad <= SPLICE_L1_MAX_ROWS else DRAM)
        self.device = device
        self._zeros: Dict[Tuple[int, int], Any] = {}
        self.last_ops: List[str] = []

    # ------------------------------------------------------------------ checks
    def _check_text(self, t: Any) -> None:
        s = _shape(t)
        if s != (1, 1, self.s_pad, self.hidden):
            raise ValueError(f"text_embeds must be [1, 1, {self.s_pad}, {self.hidden}], got {s}")

    def vis_mode(self, vis_rows: Any) -> str:
        """``"batched"`` (``[n_img, 1, >= T, D]``) or ``"flat"`` (``[1, 1, >= n_img*T, D]``); raises otherwise."""
        return vis_mode(_shape(vis_rows), self.n_images, self.tokens, self.hidden)

    # ------------------------------------------------------------------ device pieces
    def _to_rm(self, t: Any, created: List[Any]) -> Any:
        import ttnn

        if t.layout == ttnn.ROW_MAJOR_LAYOUT:
            return t
        rm = ttnn.to_layout(t, ttnn.ROW_MAJOR_LAYOUT, memory_config=resolve_mem(self.mem))
        self.last_ops.append("to_layout(RM)")
        created.append(rm)
        return rm

    def _rows(self, t: Any, b: int, lo: int, hi: int, created: List[Any]) -> Any:
        """Rows ``[lo, hi)`` of batch element ``b`` of a ROW_MAJOR ``[B, 1, R, D]`` tensor (identity when it is all of it)."""
        import ttnn

        s = _shape(t)
        if s[0] == 1 and lo == 0 and hi == s[2]:
            return t
        piece = ttnn.slice(t, [b, 0, lo, 0], [b + 1, 1, hi, s[3]], memory_config=resolve_mem(self.mem))
        self.last_ops.append("slice")
        created.append(piece)
        return piece

    def _image_piece(self, vis_rm: Any, mode: str, k: int, created: List[Any]) -> Any:
        t = self.tokens
        if mode == "batched":
            return self._rows(vis_rm, k, 0, t, created)
        return self._rows(vis_rm, 0, k * t, (k + 1) * t, created)

    def _zero_piece(self, p: Piece, device: Any) -> Any:
        key = (p.lo, p.hi)
        if key not in self._zeros:
            z = torch.zeros(1, 1, p.rows, self.hidden, dtype=torch.bfloat16)
            self._zeros[key] = L.to_device(z, device, dtype="bf16", mem=self.mem, layout="ROW_MAJOR")
        return self._zeros[key]

    def _assemble(self, pieces: List[Any], created: List[Any]) -> Any:
        import ttnn

        mem = resolve_mem(self.mem)
        rm = ttnn.concat(pieces, dim=2, memory_config=mem)
        self.last_ops.append("concat")
        out = ttnn.to_layout(rm, ttnn.TILE_LAYOUT, memory_config=mem)
        self.last_ops.append("to_layout(TILE)")
        ttnn.deallocate(rm)
        for t in created:
            ttnn.deallocate(t)
        if _shape(out) != (1, 1, self.s_pad, self.hidden):
            raise AssertionError(f"splice produced {_shape(out)}, expected {(1, 1, self.s_pad, self.hidden)}")
        return out

    def forward(self, text_embeds: Any, vis_rows: Any) -> Any:
        """``inputs_embeds``: token-embedding rows on the non-image slots, vision rows (in slot order) on the image slots."""
        self._check_text(text_embeds)
        mode = self.vis_mode(vis_rows)
        self.last_ops = []
        created: List[Any] = []
        text_rm = self._to_rm(text_embeds, created)
        vis_rm = self._to_rm(vis_rows, created)
        pieces = [
            (
                self._image_piece(vis_rm, mode, p.image, created)
                if p.is_image
                else self._rows(text_rm, 0, p.lo, p.hi, created)
            )
            for p in self.plan
        ]
        return self._assemble(pieces, created)

    __call__ = forward

    def place_into_zeros(self, vis_rows: Any) -> Any:
        """Vision rows on the image slots, exact zeros elsewhere (the N1.7 deepstack residual adds, plan §3.2)."""
        mode = self.vis_mode(vis_rows)
        device = self.device if self.device is not None else vis_rows.device()
        self.last_ops = []
        created: List[Any] = []
        vis_rm = self._to_rm(vis_rows, created)
        pieces = [
            self._image_piece(vis_rm, mode, p.image, created) if p.is_image else self._zero_piece(p, device)
            for p in self.plan
        ]
        return self._assemble(pieces, created)

    def release(self) -> None:
        import ttnn

        for t in self._zeros.values():
            ttnn.deallocate(t)
        self._zeros = {}

    def describe(self) -> Dict[str, Any]:
        return {
            "s_pad": self.s_pad,
            "n_images": self.n_images,
            "tokens_per_image": self.tokens,
            "hidden": self.hidden,
            "pieces": [(p.lo, p.hi, p.image) for p in self.plan],
            "n_pieces": self.n_pieces,
            "max_pieces": 2 * self.n_images + 1,
            "output_mem": getattr(self.mem, "kind", str(self.mem)),
        }


__all__ = [
    "GELU_ERF",
    "PROJECTOR_KINDS",
    "SPLICE_L1_MAX_ROWS",
    "UNSHUFFLE_ORDER_DEVICE",
    "IdentityProjectorTT",
    "LinearProjectorTT",
    "Piece",
    "TokenSplicer",
    "UnshuffleMLPProjectorTT",
    "build_projector",
    "host_splice",
    "splice_plan",
    "vis_mode",
]