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

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

"""N1.7 Qwen3-VL vision tower on one Blackhole: a thin wrapper over the **verbatim** tt_dit tower (WP-08).

Plan basis: IMPLEMENTATION_PLAN.md §1.1 D9 (Stage 1 = ``models.tt_dit.encoders.qwen3vl.vision_qwen3vl.
Qwen3VlVisionModel`` unchanged, measured GO in device-stack-smoke.md §1 with PCC 0.9998 single-device), §2.2 row
``tt/qwen3vl_vision.py`` and §3.2 (N1.7 paragraph: patch ``Linear(1536->1024)`` -> ``+pos`` -> 24 blocks -> merger +
3 deepstack mergers; per-image SDPA over ``cu_seqlens``).  Research: ttnn-blocks-vlm.md §5.2 (constructor, host
``prepare_pos_embeds`` / ``prepare_rope`` helpers, strict load of the ``visual.*`` sub-tree) and §7.3 option A.

Two towers live here (plan D9):

* :class:`Qwen3VLVisionTT` -- Stage 1: the **verbatim** tt_dit op sequence (the oracle);
* :class:`TunedQwen3VLVisionTT` (WP-P, Stage 3) -- the same weights and the same math with explicit program configs:
  the 4 images run as the batch dimension of ``[n_img, 1, 448, 1024]`` (a free ``ttnn.reshape`` view of the
  ``[1792, 1024]`` residual), ``nlp_create_qkv_heads`` / ``nlp_concat_heads`` replace the 9 slice / reshape / permute
  head ops, one HiFi2 SDPA over ``[4, 16, 448, 64]`` with ``SDPAProgramConfig(448, 448)`` replaces the 4 per-image calls
  (12 slices + concat), the FF GELU-tanh is fused into the fc1 program config (the verbatim tower runs a separate
  **erf** ``ttnn.gelu`` -- the fork therefore matches ``gelu_pytorch_tanh`` of the reference more closely), every
  matmul takes the ``tt.layers.MATMUL_CONFIGS`` row of its ``(1792 | 448, K, N)`` shape, and the mergers fuse their
  erf GELU.  The 2-D RoPE is either the tower's composite (``rope_mode="composite"``: the identical 7 eltwise ops per
  tensor) or ``rotary_embedding_hf`` (``"hf_op"``, one op per tensor; the vision tables are already in the HF
  rotate-half layout ``[freqs, freqs]``).  Selected by ``TTPolicy.vit_fork`` (default **True** since 2026-09-17)
  through :func:`build_vision_tower` / :func:`resolve_vit_fork` -- the call ``tt/backbone.py`` ``BackboneTT.build``
  makes; precedence: an explicit ``fork=`` argument > the ``GR00T_VIT_FORK=0|1`` environment override > the policy
  (``TTPolicy.vit_fork_for(family)``); a bare ``Qwen3VLVisionTT(...)`` without a policy is the verbatim tower unless
  the environment says otherwise.  Measured in ``tests/tt/results/p_summary.md`` (WP-P) and
  ``docs/plan/n17-fidelity.md`` (per-block fidelity: :class:`TunedVisionConfig` ``hifi_blocks`` -- HiFi3 + fp32 acc
  on the blocks feeding deepstack merger 0, the tap the all-HiFi2 fork missed by 1.2e-4).  ``GR00T_VIT_TUNED`` (a JSON
  object of :class:`TunedVisionConfig` fields) overrides the fork defaults for experiments; anything unparsable
  raises.

What the wrapper adds on top of the tower:

* **weights**: the raw HF ``visual.*`` sub-tree is read through ``common.checkpoint.LazyCheckpoint`` (one tensor in
  flight), the container prefix is stripped with ``common.checkpoint.canonical_key`` and the 315 tensors are handed to
  ``Qwen3VlVisionModel.load_torch_state_dict(strict=True)`` -- the tower's own ``_prepare_torch_state`` hooks flatten
  the Conv3d kernel, transpose the linears and pop ``pos_embed.weight`` to the host.
* **host tables per call** (``prepare_inputs``): ``ModelInputs.pixel_values`` (fp32 ``[P, 1536]`` patch rows) ->
  bf16 patches; the position embedding ``common.prompts.vision_pos_embed(table, mi.vision_pos_embed_indices,
  mi.vision_pos_embed_weights)`` from the fp32 plan table ``vision.pos_embed_table`` (``memory == "HOST"``, f1 §2);
  ``mi.vision_rope_cos / sin`` ``[P, 64]``; ``mi.vision_cu_seqlens`` (one attention block per image: ``(0, 448, 896,
  1344, 1792)`` for DROID).  Every shape is checked against the baked :class:`~tt.layout.DeviceLayout`.
* **block-level taps**: :meth:`Qwen3VLVisionTT.forward` runs exactly the tower's op sequence (``ttnn.add(patch_embed,
  pos)`` -> blocks one by one -> deepstack mergers at ``(5, 11, 17)`` -> merger) and can record the golden taps of
  golden-tap-map §2 (``vit_embeddings`` = the PRE-pos-add patch embedding, ``vit_block_first``, ``vit_block_last``,
  ``vit_deepstack_mergers[d]``, ``projector_out``); :meth:`forward_blocks` exposes the block loop on its own.
* **warm-up** (:meth:`warm`): one untraced pass so every program (the merger ``ttnn.reshape`` included) is compiled
  and its allocations exist before a trace is captured (plan §4 / §7.5 item 3).

Static facts of the DROID layout (configs.N17_CONFIG ``droid``): 4 images x 448 patches = 1792 rows, 112 tokens per
image -> ``merged [448, 2048]`` and three deepstack features ``[448, 2048]``.

``ttnn`` and the tt_dit module are imported lazily (inside ``__init__`` / functions): CPU tests import this module
without a device.  No silent fallbacks: unexpected shapes, keys, dtypes or a non-replicated tower raise.
"""

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass, fields
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

import torch

from models.experimental.gr00t.common import weights as W
from models.experimental.gr00t.common.checkpoint import LazyCheckpoint, canonical_key
from models.experimental.gr00t.common.configs import GR00TConfig
from models.experimental.gr00t.common.prompts import vision_cu_seqlens, vision_pos_embed
from models.experimental.gr00t.tt.device import (
    CKC_HIFI2_FP32ACC,
    CKC_HIFI4_FP32ACC,
    CKC_LOFI_FP32ACC,
    CKC_SDPA,
    CKC_SDPA_HIFI4,
    DRAM,
    CKCPreset,
    MemPreset,
)
from models.experimental.gr00t.tt.layout import DeviceLayout, StaticShapeError

# --------------------------------------------------------------------------------------------------------------------
# Constants (IMPLEMENTATION_PLAN §1.1 D9 / §2.2 row tt/qwen3vl_vision.py; arch-n17 §2.1-2.2)
# --------------------------------------------------------------------------------------------------------------------

POS_TABLE_NAME = "vision.pos_embed_table"  # plan tensor: fp32 [2304, 1024], memory HOST (common/weights.py f1 §2)
HIDDEN_ACT = "gelu_pytorch_tanh"  # cfg.vision.activation == "gelu_tanh" (block MLP); mergers use erf GELU inside tt_dit
SPATIAL_MERGE_SIZE = 2
TOWER_KWARGS_D9: Mapping[str, Any] = {
    "hidden_size": 1024,
    "num_heads": 16,
    "depth": 24,
    "intermediate_size": 4096,
    "in_channels": 3,
    "patch_size": 16,
    "temporal_patch_size": 2,
    "spatial_merge_size": SPATIAL_MERGE_SIZE,
    "num_position_embeddings": 2304,
    "out_hidden_size": 2048,
    "hidden_act": HIDDEN_ACT,
    "norm_eps": 1e-6,
    "deepstack_visual_indexes": (5, 11, 17),
}
# 12 tensors per block, 6 per merger (4 mergers), 2 patch-embed, 1 position table (checkpoint header, WP-08 §prompt)
N_VISUAL_TENSORS = 12 * 24 + 6 * 4 + 2 + 1
TAP_EMBEDDINGS = "vit_embeddings"  # PRE pos-add (golden-tap-map: `vis_patch_embed_out` "BEFORE pos-emb add")
TAP_BLOCK_FIRST = "vit_block_first"
TAP_BLOCK_LAST = "vit_block_last"
TAP_PROJECTOR = "projector_out"
TAP_DEEPSTACK = "vit_deepstack_mergers"  # threshold key `vit_deepstack_mergers[d]` (GoldenSet.threshold_key)
TAP_KEYS: Tuple[str, ...] = (
    TAP_EMBEDDINGS,
    TAP_BLOCK_FIRST,
    TAP_BLOCK_LAST,
    f"{TAP_DEEPSTACK}[0]",
    f"{TAP_DEEPSTACK}[1]",
    f"{TAP_DEEPSTACK}[2]",
    TAP_PROJECTOR,
)


def deepstack_key(d: int) -> str:
    """Threshold key of deepstack merger ``d`` (``GoldenSet.threshold_key('vit_deepstack_mergers', deepstack=d)``)."""
    return f"{TAP_DEEPSTACK}[{int(d)}]"


# --------------------------------------------------------------------------------------------------------------------
# Config -> tower hyper-parameters and the raw HF state dict (CPU only)
# --------------------------------------------------------------------------------------------------------------------


def tower_kwargs(cfg: GR00TConfig) -> Dict[str, Any]:
    """``Qwen3VlVisionModel`` constructor arguments derived from ``cfg`` (ttnn-blocks-vlm.md §5.2 last bullet).

    Every value is cross-checked against the D9 constants of :data:`TOWER_KWARGS_D9`; a mismatch raises because the
    verbatim tower is only validated for that configuration (device-stack-smoke.md §1).
    """
    v = cfg.vision
    if v.family != "qwen3vl":
        raise ValueError(f"{cfg.version}: vision family {v.family!r} is not the Qwen3-VL tower")
    if v.activation != "gelu_tanh":
        raise ValueError(f"{cfg.version}: block activation {v.activation!r}, the tt_dit tower expects gelu_tanh")
    kw: Dict[str, Any] = {
        "hidden_size": int(v.hidden),
        "num_heads": int(v.n_heads),
        "depth": int(v.n_layers),
        "intermediate_size": int(v.mlp_dim),
        "in_channels": 3,
        "patch_size": int(v.patch_size),
        "temporal_patch_size": int(v.temporal_patch),
        "spatial_merge_size": int(v.merge_size),
        "num_position_embeddings": int(v.pos_table_rows),
        "out_hidden_size": int(cfg.projector.out_dim),
        "hidden_act": HIDDEN_ACT,
        "norm_eps": float(v.ln_eps),
        "deepstack_visual_indexes": tuple(int(i) for i in v.deepstack_layers),
    }
    if v.hidden // v.n_heads != v.head_dim or v.head_dim != 64:
        raise ValueError(f"{cfg.version}: head_dim {v.head_dim} (hidden {v.hidden} / heads {v.n_heads}) must be 64")
    if v.patch_in_features != 3 * v.temporal_patch * v.patch_size * v.patch_size:
        raise ValueError(f"{cfg.version}: patch_in_features {v.patch_in_features} != 3*t*p*p")
    if cfg.projector.ln_eps != v.ln_eps or cfg.projector.in_dim != v.hidden * v.merge_size**2:
        raise ValueError(
            f"{cfg.version}: projector (ln_eps {cfg.projector.ln_eps}, in {cfg.projector.in_dim}) mismatch"
        )
    if cfg.projector.n_deepstack != len(kw["deepstack_visual_indexes"]):
        raise ValueError(
            f"{cfg.version}: {cfg.projector.n_deepstack} deepstack mergers for {kw['deepstack_visual_indexes']}"
        )
    bad = {k: (kw[k], TOWER_KWARGS_D9[k]) for k in TOWER_KWARGS_D9 if kw[k] != TOWER_KWARGS_D9[k]}
    if bad:
        raise ValueError(f"{cfg.version}: tower hyper-parameters differ from plan D9: {bad}")
    return kw


def visual_state_dict(
    cfg: GR00TConfig, ck: LazyCheckpoint, dtype: Optional[torch.dtype] = None
) -> Dict[str, torch.Tensor]:
    """The raw HF ``visual.*`` sub-tree with the container prefix stripped (``canonical_key``), as stored (bf16).

    Iterates ``ck.iter_by_prefix(cfg.checkpoint.vision_prefix)`` (shard order, one tensor in flight) and checks the
    tensor count against :data:`N_VISUAL_TENSORS`; the tower's ``load_torch_state_dict(strict=True)`` then reports any
    name that does not match its module tree.
    """
    if ck.version is not None and ck.version != cfg.version:
        raise ValueError(f"checkpoint is {ck.version}, config is {cfg.version}")
    prefix = cfg.checkpoint.vision_prefix
    state: Dict[str, torch.Tensor] = {}
    for key, tensor in ck.iter_by_prefix(prefix, dtype):
        name = canonical_key(cfg.version, key)
        if name in state:
            raise KeyError(f"duplicate canonical key {name!r} from {key!r}")
        state[name] = tensor
    if len(state) != N_VISUAL_TENSORS:
        raise KeyError(f"{cfg.version}: {len(state)} tensors under {prefix!r}, expected {N_VISUAL_TENSORS}")
    return state


def pos_table_from_plan(cfg: GR00TConfig, ck: LazyCheckpoint) -> torch.Tensor:
    """The fp32 ``[2304, 1024]`` position table ``vision.pos_embed_table`` built by the host plan (``memory == "HOST"``).

    Same bytes as ``ck.get(<vision_prefix>pos_embed.weight).float()`` -- read through ``common.weights.build_device_plan``
    so the tower consumes the plan's tensor (f1 §2 / weight-loader §4.2 N1.7), not a private copy of the rule.
    """
    plan = W.build_device_plan(cfg.version, None, "bf16", ck=ck, names=[POS_TABLE_NAME])
    pt = plan[POS_TABLE_NAME]
    if pt.memory != "HOST" or pt.dtype != "fp32" or pt.tensor.dtype != torch.float32:
        raise ValueError(f"{POS_TABLE_NAME}: expected a HOST fp32 plan tensor, got memory={pt.memory} dtype={pt.dtype}")
    expected = (cfg.vision.pos_table_rows, cfg.vision.hidden)
    if tuple(pt.tensor.shape) != expected:
        raise ValueError(f"{POS_TABLE_NAME}: shape {tuple(pt.tensor.shape)} != {expected}")
    return pt.tensor


def expected_cu_seqlens(dlayout: DeviceLayout) -> Tuple[int, ...]:
    """Static ``cu_seqlens`` of a layout: one block of ``h * w`` patches per image (``common.prompts.vision_cu_seqlens``)."""
    if dlayout.grid_thw is None:
        raise ValueError(f"{dlayout.version}/{dlayout.layout_name}: no grid_thw (not a Qwen3-VL layout)")
    grid = torch.tensor([dlayout.grid_thw] * dlayout.n_images, dtype=torch.int64)
    return vision_cu_seqlens(grid)


# --------------------------------------------------------------------------------------------------------------------
# Per-call inputs
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class VisionHostInputs:
    """Host bf16 tensors of one call, in the tower's call convention (plan §3.2 N1.7 / ttnn-blocks-vlm §5.2)."""

    patches: torch.Tensor  # [P, 1536] bf16 (ModelInputs.pixel_values rounded once)
    pos: torch.Tensor  # [P, 1024] bf16 (vision_pos_embed of the fp32 plan table)
    cos: torch.Tensor  # [P, 64] bf16
    sin: torch.Tensor  # [P, 64] bf16
    cu_seqlens: Tuple[int, ...]

    @property
    def n_patches(self) -> int:
        return int(self.patches.shape[0])


@dataclass
class VisionDeviceInputs:
    """The persistent device buffers (bf16 TILE, DRAM interleaved) the tower reads; addresses are baked into traces."""

    patches: Any
    pos: Any
    cos: Any
    sin: Any
    cu_seqlens: Tuple[int, ...]

    @property
    def rope(self) -> Tuple[Any, Any]:
        return self.cos, self.sin

    def tensors(self) -> Dict[str, Any]:
        return {"patches": self.patches, "pos": self.pos, "cos": self.cos, "sin": self.sin}


def host_inputs(cfg: GR00TConfig, dlayout: DeviceLayout, pos_table: torch.Tensor, mi: Any) -> VisionHostInputs:
    """``ModelInputs`` -> :class:`VisionHostInputs` (pure torch; the host half of :meth:`Qwen3VLVisionTT.prepare_inputs`).

    Checks (static-shape-plan §10 item 1 for the vision tables): ``pixel_values`` is fp32 ``dlayout.pixel_shape``; the
    four N1.7 host tables are present with ``(4, P)`` / ``(P, head_dim)`` shapes; ``cu_seqlens`` equals the layout's
    static boundaries.  Rounds every table to bf16 exactly once.
    """
    if mi.version != cfg.version:
        raise StaticShapeError(f"inputs are for {mi.version}, tower for {cfg.version}")
    dlayout.assert_call(mi)
    pv = mi.pixel_values
    if tuple(pv.shape) != tuple(dlayout.pixel_shape):
        raise StaticShapeError(f"pixel_values {tuple(pv.shape)} != layout {tuple(dlayout.pixel_shape)}")
    if pv.dtype != torch.float32:
        raise StaticShapeError(f"pixel_values dtype {pv.dtype}, expected fp32 patch rows")
    n = int(pv.shape[0])
    idx, wgt = mi.vision_pos_embed_indices, mi.vision_pos_embed_weights
    cos, sin = mi.vision_rope_cos, mi.vision_rope_sin
    if idx is None or wgt is None or cos is None or sin is None or mi.vision_cu_seqlens is None:
        raise StaticShapeError("ModelInputs lacks the N1.7 vision host tables (vision_pos_embed_*/vision_rope_*/cu)")
    if tuple(idx.shape) != (4, n) or tuple(wgt.shape) != (4, n) or idx.dtype != torch.int64:
        raise StaticShapeError(f"vision_pos_embed tables {tuple(idx.shape)} / {tuple(wgt.shape)} != (4, {n}) int64")
    hd = cfg.vision.head_dim
    if tuple(cos.shape) != (n, hd) or tuple(sin.shape) != (n, hd):
        raise StaticShapeError(f"vision_rope_cos/sin {tuple(cos.shape)} / {tuple(sin.shape)} != ({n}, {hd})")
    cu = tuple(int(c) for c in mi.vision_cu_seqlens)
    if cu != expected_cu_seqlens(dlayout):
        raise StaticShapeError(f"vision_cu_seqlens {cu} != layout {expected_cu_seqlens(dlayout)}")
    pos = vision_pos_embed(pos_table, idx, wgt)  # fp32 [P, hidden]
    if tuple(pos.shape) != (n, cfg.vision.hidden):
        raise StaticShapeError(f"pos embed {tuple(pos.shape)} != ({n}, {cfg.vision.hidden})")
    for name, t in (("pixel_values", pv), ("pos", pos), ("cos", cos), ("sin", sin)):
        if not torch.isfinite(t).all():
            raise StaticShapeError(f"{name} is not finite")
    bf = torch.bfloat16
    return VisionHostInputs(
        patches=pv.to(bf).contiguous(),
        pos=pos.to(bf).contiguous(),
        cos=cos.to(bf).contiguous(),
        sin=sin.to(bf).contiguous(),
        cu_seqlens=cu,
    )


# --------------------------------------------------------------------------------------------------------------------
# Device helpers (lazy ttnn)
# --------------------------------------------------------------------------------------------------------------------


def to_torch(t: Any) -> torch.Tensor:
    """Device tensor -> fp32 host tensor (``ttnn.to_torch``; bf16 promoted once, no other change)."""
    import ttnn

    return ttnn.to_torch(t).float()


def traced_seconds(device: Any, fn: Callable[[], Any], reps: int = 5, cq_id: int = 0) -> Dict[str, Any]:
    """Capture one call of ``fn`` in a Metal trace (programs must be warm) and replay it ``reps`` times.

    Returns ``{"capture_s", "min_s", "mean_s", "reps_s", "output"}`` where ``output`` is what ``fn`` returned during
    the capture -- after every replay its buffers hold the replayed result (the ttnn trace pattern of
    ``benchmarks/mb1_common.time_traced``; plan D12 "traced == untraced").  The trace is released before returning.
    """
    import ttnn

    if reps < 1:
        raise ValueError("reps must be >= 1")
    ttnn.synchronize_device(device)
    t0 = time.perf_counter()
    tid = ttnn.begin_trace_capture(device, cq_id=cq_id)
    try:
        out = fn()
    finally:
        ttnn.end_trace_capture(device, tid, cq_id=cq_id)
    ttnn.synchronize_device(device)
    capture_s = time.perf_counter() - t0
    times: List[float] = []
    try:
        ttnn.execute_trace(device, tid, cq_id=cq_id, blocking=False)  # untimed replay
        ttnn.synchronize_device(device)
        for _ in range(reps):
            t1 = time.perf_counter()
            ttnn.execute_trace(device, tid, cq_id=cq_id, blocking=False)
            ttnn.synchronize_device(device)
            times.append(time.perf_counter() - t1)
    finally:
        ttnn.release_trace(device, tid)
    return {
        "capture_s": capture_s,
        "min_s": min(times),
        "mean_s": sum(times) / len(times),
        "reps_s": times,
        "output": out,
    }


# --------------------------------------------------------------------------------------------------------------------
# The wrapper
# --------------------------------------------------------------------------------------------------------------------


class Qwen3VLVisionTT:
    """``Qwen3VLVisionTT(cfg, dlayout, ck, device)``: the verbatim tt_dit tower with GR00T weights, host tables and taps.

    ::

        tower = Qwen3VLVisionTT(cfg, DeviceLayout.for_inputs(cfg, mi), LazyCheckpoint("n17"), device)
        host = tower.prepare_inputs(mi)              # bf16 patches / pos / cos / sin + cu_seqlens (torch)
        dev = tower.upload(host)                     # persistent DRAM buffers
        tower.warm(dev)                              # one untraced pass (compiles every program, merger reshape incl.)
        merged, deepstack = tower.forward(dev, taps=taps)   # merged [448, 2048], deepstack 3 x [448, 2048]
        tower.refresh(dev, tower.prepare_inputs(mi2))       # next call: copy_host_to_device_tensor into the same buffers

    ``pos_table`` may be passed by a caller that already holds the plan's host tensor (``TTWeights.host(
    "vision.pos_embed_table")``); it must equal the table the tower popped from the state dict.  ``policy`` /
    ``fork`` select the WP-P fork (:func:`resolve_vit_fork`; the constructor then returns a
    :class:`TunedQwen3VLVisionTT`).
    """

    def __init__(
        self,
        cfg: GR00TConfig,
        dlayout: DeviceLayout,
        ck: LazyCheckpoint,
        device: Any,
        *,
        pos_table: Optional[torch.Tensor] = None,
        merger_ckc: Any = None,
        policy: Any = None,
        fork: Optional[bool] = None,
    ) -> None:
        # ``policy`` / ``fork`` are the WP-P selector arguments consumed by ``__new__`` (resolve_vit_fork); the
        # verbatim tower itself ignores them.
        import ttnn  # lazy (plan §2.1 rules)
        from models.tt_dit.encoders.qwen3vl.vision_qwen3vl import Qwen3VlVisionModel

        if dlayout.version != cfg.version:
            raise ValueError(f"layout is for {dlayout.version}, config for {cfg.version}")
        if dlayout.grid_thw is None:
            raise ValueError(f"{cfg.version}/{dlayout.layout_name}: layout has no grid_thw (not a Qwen3-VL layout)")
        self.cfg = cfg
        self.dlayout = dlayout
        self.device = device
        self.ttnn = ttnn
        self.kwargs: Dict[str, Any] = tower_kwargs(cfg)
        self.depth: int = int(self.kwargs["depth"])
        self.deepstack_indexes: Tuple[int, ...] = tuple(self.kwargs["deepstack_visual_indexes"])
        self.head_dim: int = cfg.vision.head_dim
        self.n_patches: int = int(dlayout.n_images * dlayout.layout.patches_per_image)
        self.n_tokens: int = self.n_patches // (int(self.kwargs["spatial_merge_size"]) ** 2)
        self.cu_seqlens: Tuple[int, ...] = expected_cu_seqlens(dlayout)
        if dlayout.n_images * dlayout.tokens_per_image != self.n_tokens:
            raise ValueError(f"{dlayout.n_images} x {dlayout.tokens_per_image} tokens != {self.n_tokens} merged rows")
        self.merged_shape: Tuple[int, int] = (self.n_tokens, int(self.kwargs["out_hidden_size"]))
        self.hidden_shape: Tuple[int, int] = (self.n_patches, int(self.kwargs["hidden_size"]))

        t0 = time.perf_counter()
        self.model = Qwen3VlVisionModel(**self.kwargs, mesh_device=device, parallel_config=None)
        p = self.model._p  # noqa: SLF001 - the replicated (single device) path is a hard requirement here
        if p.tp or p.sp:
            raise RuntimeError(f"tower resolved a parallel placement {p}; WP-08 wraps the replicated tower only")
        state = visual_state_dict(cfg, ck)
        t1 = time.perf_counter()
        self.model.load_torch_state_dict(state, strict=True)
        del state
        t2 = time.perf_counter()

        popped = self.model._pos_embed_weight  # noqa: SLF001 - the tower keeps the table on the host (§5.2)
        if popped is None:
            raise RuntimeError("Qwen3VlVisionModel did not pop pos_embed.weight; the strict load is incomplete")
        table = pos_table_from_plan(cfg, ck) if pos_table is None else pos_table
        if table.dtype != torch.float32 or not torch.equal(table, popped.float()):
            raise ValueError(f"{POS_TABLE_NAME} differs from the tower's popped pos_embed.weight (dtype {table.dtype})")
        self.pos_table: torch.Tensor = table
        scale = float(self.model.blocks[0].attn.scale)
        if abs(scale - cfg.vision.attention_scale) > 1e-12:
            raise RuntimeError(f"tower SDPA scale {scale} != cfg {cfg.vision.attention_scale}")
        if int(self.model.blocks[0].attn.padded_head_dim) != self.head_dim:
            raise RuntimeError(f"tower padded head_dim {self.model.blocks[0].attn.padded_head_dim} != {self.head_dim}")
        self.load_timing: Dict[str, float] = {"read_state_s": t1 - t0, "convert_upload_s": t2 - t1, "total_s": t2 - t0}
        self.warmed = False
        self.merger_ckc = merger_ckc
        if merger_ckc is not None:
            self._set_merger_compute_config(merger_ckc)

    def _set_merger_compute_config(self, ckc: Any) -> None:
        """D0 precision lever (``tests/tt/results/d0_summary.md`` §5): override ``Linear.compute_config`` of the four
        patch mergers (``merger`` + ``deepstack_merger_list[0..2]``, ``linear_fc1`` / ``linear_fc2``) with a
        ``tt.device`` preset / preset name / built ``ttnn`` compute-kernel config.  The tt_dit ``Linear`` default is
        HiFi2 + fp32 acc (``models/tt_dit/layers/linear.py``); the blocks and the SDPA are untouched (plan D9)."""
        from models.experimental.gr00t.tt.device import resolve_ckc

        cfg = resolve_ckc(ckc)
        mergers = [self.model.merger, *self.model.deepstack_merger_list]
        if len(mergers) != 1 + len(self.deepstack_indexes):
            raise RuntimeError(f"expected {1 + len(self.deepstack_indexes)} mergers, found {len(mergers)}")
        for m in mergers:
            for name in ("linear_fc1", "linear_fc2"):
                lin = getattr(m, name)
                if not hasattr(lin, "compute_config"):
                    raise AttributeError(f"tt_dit Linear {name} has no compute_config attribute (API drift)")
                lin.compute_config = cfg

    # ------------------------------------------------------------------ host side
    def prepare_inputs(self, mi: Any) -> VisionHostInputs:
        """``ModelInputs`` -> bf16 host tensors ``(patches [P,1536], pos [P,1024], cos, sin [P,64])`` + ``cu_seqlens``."""
        host = host_inputs(self.cfg, self.dlayout, self.pos_table, mi)
        if host.n_patches != self.n_patches or host.cu_seqlens != self.cu_seqlens:
            raise StaticShapeError(
                f"{host.n_patches} patches / {host.cu_seqlens} != baked {self.n_patches} / {self.cu_seqlens}"
            )
        return host

    def upload(self, host: VisionHostInputs, memory_config: Any = None) -> VisionDeviceInputs:
        """Create the persistent device buffers (bf16 TILE, DRAM interleaved by default) from ``host``."""
        ttnn = self.ttnn
        mem = ttnn.DRAM_MEMORY_CONFIG if memory_config is None else memory_config

        def up(t: torch.Tensor) -> Any:
            return ttnn.from_torch(
                t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.device, memory_config=mem
            )

        self._check_host(host)
        return VisionDeviceInputs(
            patches=up(host.patches), pos=up(host.pos), cos=up(host.cos), sin=up(host.sin), cu_seqlens=host.cu_seqlens
        )

    def refresh(
        self, dev: VisionDeviceInputs, host: VisionHostInputs, cq_id: Optional[int] = None
    ) -> VisionDeviceInputs:
        """Write the next call's tables into the existing buffers (``ttnn.copy_host_to_device_tensor``; same shapes)."""
        ttnn = self.ttnn
        self._check_host(host)
        if host.cu_seqlens != dev.cu_seqlens:
            raise StaticShapeError(f"cu_seqlens {host.cu_seqlens} != buffers {dev.cu_seqlens}")
        for name, t in (("patches", host.patches), ("pos", host.pos), ("cos", host.cos), ("sin", host.sin)):
            buf = getattr(dev, name)
            if tuple(buf.shape) != tuple(t.shape):
                raise StaticShapeError(f"{name}: device buffer {tuple(buf.shape)} != host {tuple(t.shape)}")
            h = ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
            if cq_id is None:
                ttnn.copy_host_to_device_tensor(h, buf)
            else:
                ttnn.copy_host_to_device_tensor(h, buf, cq_id)
        return dev

    def _check_host(self, host: VisionHostInputs) -> None:
        shapes = {
            "patches": (self.n_patches, self.cfg.vision.patch_in_features),
            "pos": self.hidden_shape,
            "cos": (self.n_patches, self.head_dim),
            "sin": (self.n_patches, self.head_dim),
        }
        for name, shape in shapes.items():
            t = getattr(host, name)
            if tuple(t.shape) != shape or t.dtype != torch.bfloat16:
                raise StaticShapeError(f"{name}: {tuple(t.shape)} {t.dtype} != {shape} bf16")
        if host.cu_seqlens != self.cu_seqlens:
            raise StaticShapeError(f"cu_seqlens {host.cu_seqlens} != baked {self.cu_seqlens}")

    # ------------------------------------------------------------------ device side (the tower's own op sequence)
    def embed(self, patches: Any) -> Any:
        """Patch embedding ``Linear(1536 -> 1024)`` **before** the position add (golden ``vit_embeddings`` is pre-add)."""
        return self.model.patch_embed.forward(patches)

    def add_pos(self, lin: Any, pos: Any) -> Any:
        """``ttnn.add(patch_embed, pos)`` -- the residual stream entering block 0 (tower ``forward`` line 1)."""
        return self.ttnn.add(lin, pos)

    def block(self, i: int, x: Any, rope: Tuple[Any, Any], cu_seqlens: Optional[Sequence[int]] = None) -> Any:
        """Block ``i`` exactly as the tower calls it: ``blocks[i].forward(x, pos_embeds=rope, cu_seqlens=cu_seqlens)``."""
        if not 0 <= i < self.depth:
            raise IndexError(f"block {i} out of range 0..{self.depth - 1}")
        cu = self.cu_seqlens if cu_seqlens is None else tuple(cu_seqlens)
        return self.model.blocks[i].forward(x, pos_embeds=rope, cu_seqlens=cu)

    def forward_blocks(
        self,
        x: Any,
        n: int,
        rope: Tuple[Any, Any],
        cu_seqlens: Optional[Sequence[int]] = None,
        *,
        start: int = 0,
        collect: Iterable[int] = (),
    ) -> Tuple[Any, Dict[int, Any]]:
        """Run blocks ``start .. start + n - 1`` one by one from the residual ``x``; return ``(x, {i: block_i_out})``.

        ``collect`` names the block indices whose outputs are kept (``0`` -> ``vit_block_first``, ``depth - 1`` ->
        ``vit_block_last``); every block output is a fresh tensor so the kept handles stay valid.
        """
        if n < 0 or start < 0 or start + n > self.depth:
            raise IndexError(f"blocks {start}..{start + n - 1} out of range 0..{self.depth - 1}")
        keep = set(int(i) for i in collect)
        bad = keep - set(range(start, start + n))
        if bad:
            raise IndexError(f"collect indices {sorted(bad)} are not run by blocks {start}..{start + n - 1}")
        out: Dict[int, Any] = {}
        for i in range(start, start + n):
            x = self.block(i, x, rope, cu_seqlens)
            if i in keep:
                out[i] = x
        return x, out

    def deepstack_merge(self, d: int, x: Any) -> Any:
        """Deepstack merger ``d`` (post-shuffle LN(4096) -> fc1 -> GELU -> fc2) on the block output it taps."""
        if not 0 <= d < len(self.deepstack_indexes):
            raise IndexError(f"deepstack merger {d} out of range 0..{len(self.deepstack_indexes) - 1}")
        return self.model.deepstack_merger_list[d].forward(x)

    def merge(self, x: Any) -> Any:
        """Final merger (pre-shuffle LN(1024) -> reshape [P/4, 4096] -> fc1 -> GELU -> fc2) -> ``[P/4, 2048]``."""
        return self.model.merger.forward(x)

    def forward(self, inputs: VisionDeviceInputs, taps: Optional[Dict[str, Any]] = None) -> Tuple[Any, List[Any]]:
        """``(merged [448, 2048], deepstack [3 x [448, 2048]])`` -- the tower's ``forward`` op for op, with optional taps.

        The sequence is ``Qwen3VlVisionModel.forward`` verbatim (``ttnn.add(patch_embed(patches), pos)``, the block
        loop with a deepstack merger after blocks ``(5, 11, 17)``, the final merger; ``_gather_tokens`` is the identity
        on one device).  When ``taps`` is a dict the golden tap keys of :data:`TAP_KEYS` are stored (device tensors).
        """
        if inputs.cu_seqlens != self.cu_seqlens:
            raise StaticShapeError(f"cu_seqlens {inputs.cu_seqlens} != baked {self.cu_seqlens}")
        rec = taps is not None
        lin = self.embed(inputs.patches)
        if rec:
            taps[TAP_EMBEDDINGS] = lin
        x = self.add_pos(lin, inputs.pos)
        deepstack: List[Any] = []
        for i in range(self.depth):
            x = self.block(i, x, inputs.rope, inputs.cu_seqlens)
            if rec and i == 0:
                taps[TAP_BLOCK_FIRST] = x
            if i in self.deepstack_indexes:
                d = self.deepstack_indexes.index(i)
                feat = self.deepstack_merge(d, x)
                deepstack.append(feat)
                if rec:
                    taps[deepstack_key(d)] = feat
        if rec:
            taps[TAP_BLOCK_LAST] = x
        merged = self.merge(x)
        if rec:
            taps[TAP_PROJECTOR] = merged
        if len(deepstack) != len(self.deepstack_indexes):
            raise RuntimeError(f"{len(deepstack)} deepstack features for indexes {self.deepstack_indexes}")
        return merged, deepstack

    def warm(self, inputs: VisionDeviceInputs) -> Dict[str, float]:
        """One untraced, synchronised pass so every program is compiled and allocated before a trace capture.

        Also the natural place to pay the first-call JIT (device-stack-smoke §1: ~14 s for one block's op set on a cold
        ``TT_METAL_CACHE``).  Outputs are deallocated; returns ``{"warm_forward_s": seconds}``.
        """
        ttnn = self.ttnn
        ttnn.synchronize_device(self.device)
        t0 = time.perf_counter()
        merged, deepstack = self.forward(inputs)
        ttnn.synchronize_device(self.device)
        dt = time.perf_counter() - t0
        for t in [merged, *deepstack]:
            ttnn.deallocate(t)
        self.warmed = True
        return {"warm_forward_s": dt}

    def forward_host(
        self, inputs: VisionDeviceInputs, taps: Optional[Dict[str, Any]] = None
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """``forward`` + readback to fp32 host tensors (taps, when given, are read back in place)."""
        merged, deepstack = self.forward(inputs, taps)
        self.ttnn.synchronize_device(self.device)
        if taps is not None:
            for k in list(taps):
                taps[k] = to_torch(taps[k])
        return to_torch(merged), [to_torch(t) for t in deepstack]

    def release(self) -> None:
        """Free the tower's device weights."""
        self.model.deallocate_weights()

    def describe(self) -> Dict[str, Any]:
        """JSON-able summary for results files (the fork overrides it with its configuration)."""
        return {
            "tower": "verbatim",
            "vit_fork": False,
            "class": type(self).__name__,
            "depth": self.depth,
            "n_patches": self.n_patches,
            "n_tokens": self.n_tokens,
            "cu_seqlens": list(self.cu_seqlens),
        }

    # ------------------------------------------------------------------ WP-P selector (adapter, see resolve_vit_fork)
    def __new__(cls, *args: Any, **kwargs: Any) -> "Qwen3VLVisionTT":
        """``Qwen3VLVisionTT(...)`` builds the **fork** when :func:`resolve_vit_fork` says so.

        Adapter for callers that construct the base class directly: ``policy=`` / ``fork=`` keyword arguments and the
        ``GR00T_VIT_FORK`` environment override are honoured (``BackboneTT.build`` goes through
        :func:`build_vision_tower` instead).  Without a policy and without the environment variable the class you
        named is what you get (the verbatim tower).
        """
        if cls is Qwen3VLVisionTT:
            cfg = args[0] if args else kwargs.get("cfg")
            family = getattr(getattr(cfg, "vision", None), "family", None)
            if resolve_vit_fork(kwargs.get("policy"), kwargs.get("fork"), family):
                return object.__new__(TunedQwen3VLVisionTT)
        return object.__new__(cls)


# --------------------------------------------------------------------------------------------------------------------
# WP-P fork: TunedQwen3VLVisionTT (plan D9 / §3.2 "WP-P fork"; Stage 3 row of §6)
# --------------------------------------------------------------------------------------------------------------------

VIT_FORK_ENV = "GR00T_VIT_FORK"  # "1" / "0": explicit override of TTPolicy.vit_fork for policy-less constructors
VIT_ROPE_MODES: Tuple[str, ...] = ("composite", "hf_op")
VIT_SDPA_CLASS = "vit_qwen3vl"  # tt.layers.SDPA_CHUNKS class of the (448, 448, 64) B=4 row (mb1 Table 3)
VIT_FF_ACT = "gelu_tanh"  # cfg.vision.activation; fused into fc1 (plan D9)
MERGER_ACT = "gelu"  # the HF merger uses nn.GELU() (erf); fused into fc1 of the mergers (exact erf, mb2 §B.4)
#: The verbatim tower's SDPA precision (``vision_qwen3vl.py`` ``_sdpa_compute_kernel_config``: HiFi4, approx off,
#: fp32 acc on) as a ``tt.device`` preset so the fork can reproduce it (``tt.device.CKC_SDPA_HIFI4`` has fp32 acc off).
CKC_VIT_SDPA_HIFI4_FP32ACC = CKCPreset("vit_sdpa_hifi4_fp32acc", "HiFi4", math_approx_mode=False, fp32_dest_acc_en=True)
#: Fidelity candidates of the N1.7 merger-0 investigation (docs/plan/n17-fidelity.md).  HiFi3 consumes the full
#: mantissa of both bf16 operands except the (LSB x LSB) cross term (tech_reports/matrix_engine: HiFi2 truncates the
#: SrcB operand to 1 + 6 mantissa bits, HiFi4 = exact bf16 products); the fp32-acc SDPA presets keep the running
#: softmax / P.V accumulation in fp32 (``fp32_dest_acc_en``; the plan-§1.2 ``CKC_SDPA`` accumulates in 16 bit).
CKC_HIFI3_FP32ACC = CKCPreset("hifi3_fp32acc", "HiFi3", math_approx_mode=False, fp32_dest_acc_en=True)
CKC_SDPA_HIFI2_FP32ACC = CKCPreset("sdpa_hifi2_fp32acc", "HiFi2", math_approx_mode=False, fp32_dest_acc_en=True)
CKC_SDPA_HIFI3_FP32ACC = CKCPreset("sdpa_hifi3_fp32acc", "HiFi3", math_approx_mode=False, fp32_dest_acc_en=True)
CKC_SDPA_HIFI3 = CKCPreset("sdpa_hifi3", "HiFi3", math_approx_mode=False, fp32_dest_acc_en=False)
_VIT_CKC_NAMES: Dict[str, Any] = {
    "sdpa_hifi2": CKC_SDPA,
    "sdpa_hifi3": CKC_SDPA_HIFI3,
    "sdpa_hifi4": CKC_SDPA_HIFI4,
    "sdpa_hifi2_fp32acc": CKC_SDPA_HIFI2_FP32ACC,
    "sdpa_hifi3_fp32acc": CKC_SDPA_HIFI3_FP32ACC,
    "vit_sdpa_hifi4_fp32acc": CKC_VIT_SDPA_HIFI4_FP32ACC,
    "hifi2_fp32acc": CKC_HIFI2_FP32ACC,
    "hifi3_fp32acc": CKC_HIFI3_FP32ACC,
    "hifi4_fp32acc": CKC_HIFI4_FP32ACC,
    "lofi_fp32acc": CKC_LOFI_FP32ACC,
}
VIT_TUNED_ENV = "GR00T_VIT_TUNED"  # JSON object of TunedVisionConfig fields: experiment override of the fork defaults


def _env_flag(name: str) -> Optional[bool]:
    raw = os.environ.get(name)
    if raw is None or raw.strip() == "":
        return None
    if raw.strip() in ("1", "true", "True", "yes"):
        return True
    if raw.strip() in ("0", "false", "False", "no"):
        return False
    raise ValueError(f"{name}={raw!r} must be 1/0 (true/false)")


def resolve_vit_fork(policy: Any = None, fork: Optional[bool] = None, family: Optional[str] = None) -> bool:
    """Whether the N1.7 tower is the WP-P fork.

    Precedence: an explicit ``fork`` argument > the ``GR00T_VIT_FORK=0|1`` environment **override** > the policy
    (``TTPolicy.vit_fork_for(family)`` when ``family`` -- ``GR00TConfig.vision.family`` -- is given, else
    ``policy.vit_fork``) > ``False`` when there is no policy at all (a bare constructor builds the class it names).
    ``TTPolicy()`` defaults to ``vit_fork=True``, so a policy-driven build of the Qwen3-VL tower is the fork unless the
    caller or the environment says otherwise; SigLIP families never fork.  An unparsable environment value raises.
    """
    if fork is not None:
        return bool(fork)
    env = _env_flag(VIT_FORK_ENV)
    if env is not None:
        return env
    if policy is None:
        return False
    if not hasattr(policy, "vit_fork"):
        raise TypeError(f"policy must be a TTPolicy (has vit_fork), got {type(policy).__name__}")
    if family is not None:
        if not hasattr(policy, "vit_fork_for"):
            raise TypeError(f"policy must be a TTPolicy (has vit_fork_for), got {type(policy).__name__}")
        return bool(policy.vit_fork_for(family))
    return bool(policy.vit_fork)


def resolve_vit_lofi(policy: Any = None, lofi: Optional[bool] = None) -> bool:
    """LoFi on the fork's fc1 / fc2 only (plan §1.2 row 3 gate): explicit ``lofi`` > ``policy.vit_lofi`` > False."""
    if lofi is not None:
        return bool(lofi)
    if policy is not None and bool(getattr(policy, "vit_lofi", False)):
        return True
    return False


def _ckc_of(name: Any) -> Any:
    """Preset name (this module's table or ``tt.device.CKC_PRESETS``) / ``CKCPreset`` / built config -> preset."""
    if isinstance(name, str):
        if name in _VIT_CKC_NAMES:
            return _VIT_CKC_NAMES[name]
        from models.experimental.gr00t.tt.device import CKC_PRESETS

        if name not in CKC_PRESETS:
            raise KeyError(
                f"unknown compute-kernel preset {name!r}; known: {sorted(set(_VIT_CKC_NAMES) | set(CKC_PRESETS))}"
            )
        return CKC_PRESETS[name]
    return name


@dataclass(frozen=True)
class TunedVisionConfig:
    """Program-config knobs of :class:`TunedQwen3VLVisionTT` (every value is measured by ``benchmarks/sweep_backbone.py
    --what vit``; defaults = the winners recorded in ``tests/tt/results/p_summary.md``).

    * ``sdpa_chunks``: ``(q_chunk, k_chunk)`` of the batched ``[4, 16, 448, 64]`` SDPA -- must be in the
      ``tt.layers.SDPA_CHUNKS`` row ``("vit_qwen3vl", 448, 448, 64)`` (mb1 Table 3: (256, 256) 84.7 µs for 4 images).
    * ``sdpa_ckc``: preset name; ``"vit_sdpa_hifi4_fp32acc"`` reproduces the verbatim tower, ``"sdpa_hifi2"`` is the
      plan §1.2 default of the other ViT.
    * ``rope_mode``: ``"composite"`` (the tower's 7-op rotate-half per tensor, bit-identical math) | ``"hf_op"``
      (``rotary_embedding_hf``, one op per tensor, plan D8 "hf" alternative).
    * ``mem`` / ``ff_mem``: memory preset name of the residual + attention intermediates / of the ``[4, 1, 448, 4096]``
      FF hidden (D0 fact: the fc1 ``ibw=8`` CBs clash with L1-resident activations at M = 1792, so the FF hidden
      defaults to DRAM).
    * ``ckc`` / ``ff_ckc``: projections / fc1+fc2 presets (``"lofi_fp32acc"`` for fc1 / fc2 is the ``TTPolicy.vit_lofi``
      lever, applied by the constructor).
    * ``tuned_embed`` / ``tuned_mergers``: patch embed / mergers through ``tt.layers.linear`` (fused erf GELU) instead of
      the tt_dit modules; ``merger_ckc`` their preset.
    * ``hifi_blocks`` + ``hifi_ckc`` / ``hifi_ff_ckc`` / ``hifi_sdpa_ckc`` / ``hifi_sdpa_chunks``: per-block fidelity
      overrides (docs/plan/n17-fidelity.md) -- the listed blocks take these presets for their projections / fc1+fc2 /
      SDPA (``None`` keeps the base value); ``hifi_mergers`` (``0`` = final merger, ``1 + d`` = deepstack merger
      ``d``) take ``hifi_merger_ckc``.  Default: HiFi3 + fp32 acc projections and FF on blocks 0-5 (the blocks feeding
      deepstack merger 0, whose gate 0.9998 the all-HiFi2 fork missed at 0.99968; measured in the fidelity table).
    """

    # p_summary.md §1 / §4 (session 3a on the golden pixels, min of 5 traced replays): (448,448) HiFi2 16.97 ms,
    # (448,448) HiFi4 17.23, (256,256) HiFi2 17.56, (256,256) HiFi4 17.77, (256,256) HiFi4 + fp32 acc (the verbatim
    # tower's config) 19.02 -- gates with HiFi2 (448,448): block_last 0.99195 / mergers 0.999681, 0.99742, 0.99739 /
    # projector 0.99346 vs the verbatim tower's 0.99044 / 0.999652, 0.99665, 0.99678 / 0.99366; plan §1.2 sets the
    # ViT SDPA to HiFi2 (exp approx off).
    sdpa_chunks: Tuple[int, int] = (448, 448)
    sdpa_ckc: str = "sdpa_hifi2"
    rope_mode: str = "hf_op"  # p_summary §2: 29.6 vs 59.2 µs per tensor, tower -1.4 ms, PCC vs golden >= composite
    mem: str = "L1"
    ff_mem: str = "DRAM"
    ckc: str = "hifi2_fp32acc"
    ff_ckc: str = "hifi2_fp32acc"
    tuned_embed: bool = True
    tuned_mergers: bool = True
    merger_ckc: str = "hifi2_fp32acc"
    merger_gelu_fused: bool = False  # p_summary §5: fused erf GELU 182 µs vs matmul 164 + separate gelu at [448, 4096]
    # --- per-block fidelity (docs/plan/n17-fidelity.md §2): HiFi3 + fp32 acc on the projections and the FF of the six
    # blocks feeding deepstack merger 0 (deepstack index 5).  HiFi2 consumes only 1 + 6 mantissa bits of one bf16
    # operand per product (tech_reports/matrix_engine); the CPU emulation of the block chain shows that truncation
    # alone takes the canonical merger-0 PCC from 0.99996 to 0.99987 while bf16 RoPE tables / 16-bit SDPA accumulation
    # cost < 1e-6 each.  Device (test_tt_vit_fidelity, 7 goldens): HiFi2 min 0.999644 (miss), HiFi3 blocks 0-5 min
    # 0.999845 for +0.48 ms, HiFi4 blocks 0-5 min 0.999843 for +0.94 ms -- the cheapest passing candidate ships.
    hifi_blocks: Tuple[int, ...] = (0, 1, 2, 3, 4, 5)
    hifi_ckc: Optional[str] = "hifi3_fp32acc"
    hifi_ff_ckc: Optional[str] = "hifi3_fp32acc"
    hifi_sdpa_ckc: Optional[str] = None
    hifi_sdpa_chunks: Optional[Tuple[int, int]] = None
    hifi_mergers: Tuple[int, ...] = ()
    hifi_merger_ckc: Optional[str] = None

    def __post_init__(self) -> None:
        object.__setattr__(self, "sdpa_chunks", _chunk_pair("sdpa_chunks", self.sdpa_chunks))
        if self.hifi_sdpa_chunks is not None:
            object.__setattr__(self, "hifi_sdpa_chunks", _chunk_pair("hifi_sdpa_chunks", self.hifi_sdpa_chunks))
        if self.rope_mode not in VIT_ROPE_MODES:
            raise ValueError(f"rope_mode must be one of {VIT_ROPE_MODES}, got {self.rope_mode!r}")
        for f in ("mem", "ff_mem"):
            if getattr(self, f) not in ("L1", "DRAM"):
                raise ValueError(f"{f} must be 'L1' or 'DRAM', got {getattr(self, f)!r}")
        for f in ("sdpa_ckc", "ckc", "ff_ckc", "merger_ckc"):
            _ckc_of(getattr(self, f))
        for f in ("hifi_ckc", "hifi_ff_ckc", "hifi_sdpa_ckc", "hifi_merger_ckc"):
            if getattr(self, f) is not None:
                _ckc_of(getattr(self, f))
        for f in ("hifi_blocks", "hifi_mergers"):
            object.__setattr__(self, f, _index_tuple(f, getattr(self, f)))
        if self.hifi_blocks and all(
            getattr(self, f) is None for f in ("hifi_ckc", "hifi_ff_ckc", "hifi_sdpa_ckc", "hifi_sdpa_chunks")
        ):
            raise ValueError("hifi_blocks given without any hifi_* override")
        if self.hifi_mergers and self.hifi_merger_ckc is None:
            raise ValueError("hifi_mergers given without hifi_merger_ckc")
        if (self.hifi_sdpa_chunks is not None) and not self.hifi_blocks:
            raise ValueError("hifi_sdpa_chunks without hifi_blocks")

    def to_dict(self) -> Dict[str, Any]:
        d = dict(self.__dict__)
        for k in ("sdpa_chunks", "hifi_sdpa_chunks", "hifi_blocks", "hifi_mergers"):
            if d[k] is not None:
                d[k] = list(d[k])
        return d

    @classmethod
    def from_mapping(cls, d: Mapping[str, Any]) -> "TunedVisionConfig":
        """Build from a JSON-like mapping (lists become tuples); unknown keys raise ``ValueError``."""
        known = {f.name for f in fields(cls)}
        unknown = sorted(set(d) - known)
        if unknown:
            raise ValueError(f"unknown TunedVisionConfig fields {unknown}; known: {sorted(known)}")
        kw: Dict[str, Any] = {}
        for k, v in d.items():
            kw[k] = tuple(v) if isinstance(v, list) else v
        return cls(**kw)

    @classmethod
    def from_env(cls, name: str = VIT_TUNED_ENV) -> Optional["TunedVisionConfig"]:
        """The ``GR00T_VIT_TUNED`` experiment override (a JSON object of fields) or ``None`` when unset."""
        raw = os.environ.get(name)
        if raw is None or raw.strip() == "":
            return None
        try:
            d = json.loads(raw)
        except json.JSONDecodeError as e:
            raise ValueError(f"{name} is not a JSON object: {e}") from e
        if not isinstance(d, dict):
            raise ValueError(f"{name} must be a JSON object of TunedVisionConfig fields, got {type(d).__name__}")
        return cls.from_mapping(d)


def _chunk_pair(name: str, v: Any) -> Tuple[int, int]:
    if len(v) != 2:
        raise ValueError(f"{name} must be (q_chunk, k_chunk), got {v}")
    return (int(v[0]), int(v[1]))


def _index_tuple(name: str, v: Any) -> Tuple[int, ...]:
    out = tuple(int(i) for i in v)
    if any(i < 0 for i in out) or len(set(out)) != len(out) or list(out) != sorted(out):
        raise ValueError(f"{name} must be a sorted tuple of distinct non-negative block indices, got {v}")
    return out


def build_vision_tower(
    cfg: GR00TConfig, dlayout: DeviceLayout, ck: LazyCheckpoint, device: Any, policy: Any = None, **kw: Any
) -> "Qwen3VLVisionTT":
    """The N1.7 tower for ``policy`` (plan D9): :class:`TunedQwen3VLVisionTT` iff :func:`resolve_vit_fork`, else the
    verbatim :class:`Qwen3VLVisionTT`.  ``policy=None`` means the default ``tt.policy.TTPolicy()`` (fork).  ``kw`` is
    forwarded (``pos_table``, ``tuned``, ``merger_ckc``; ``fork=`` is the explicit selector).  This is the call
    ``tt/backbone.py`` ``BackboneTT.build`` makes."""
    if policy is None:
        from models.experimental.gr00t.tt.policy import TTPolicy

        policy = TTPolicy()
    if resolve_vit_fork(policy, kw.pop("fork", None), cfg.vision.family):
        return TunedQwen3VLVisionTT(cfg, dlayout, ck, device, policy=policy, fork=True, **kw)
    kw.pop("tuned", None)
    return Qwen3VLVisionTT(cfg, dlayout, ck, device, policy=policy, fork=False, **kw)


class TunedQwen3VLVisionTT(Qwen3VLVisionTT):
    """The WP-P fork of the N1.7 tower (module docstring; plan D9, §3.2 "WP-P fork", §6 Stage 3).

    Same constructor as :class:`Qwen3VLVisionTT` plus ``tuned: TunedVisionConfig`` and ``policy`` (``vit_lofi``);
    same public API (``prepare_inputs`` / ``upload`` / ``refresh`` / ``embed`` / ``add_pos`` / ``block`` /
    ``forward_blocks`` / ``deepstack_merge`` / ``merge`` / ``forward`` / ``warm`` / ``forward_host`` / ``release``)
    and the same tap keys / shapes, so it is a drop-in for the wrapper everywhere.  The verbatim op sequence stays
    available on the same weights as :meth:`forward_verbatim` (the oracle of ``sweep_backbone --what vit``).

    Layouts: the residual stream is ``[n_img, 1, s_img, hidden]`` inside the blocks (``s_img`` = patches per image,
    448 for DROID; a free view of the tower's ``[P, hidden]``); ``block`` / ``forward_blocks`` accept either layout and
    return the layout they were given, ``forward`` returns ``[P/4, 2048]`` tensors like the tower.  The batched SDPA
    is exact block-diagonal attention because every image is one ``cu_seqlens`` block of the same length (checked at
    construction) and the 2-D RoPE tables of the images are identical (checked per call on the host in
    :meth:`prepare_inputs` and once on device in :meth:`upload`).
    """

    def __init__(
        self,
        cfg: GR00TConfig,
        dlayout: DeviceLayout,
        ck: LazyCheckpoint,
        device: Any,
        *,
        pos_table: Optional[torch.Tensor] = None,
        merger_ckc: Any = None,
        tuned: Optional[TunedVisionConfig] = None,
        policy: Any = None,
        fork: Optional[bool] = None,
        lofi: Optional[bool] = None,
    ) -> None:
        super().__init__(cfg, dlayout, ck, device, pos_table=pos_table, merger_ckc=merger_ckc)
        from models.experimental.gr00t.tt.layers import dtype_name, matmul_spec

        if tuned is None:
            tuned = TunedVisionConfig.from_env()  # GR00T_VIT_TUNED experiment override, else the measured defaults
            self.tuned_source = VIT_TUNED_ENV if tuned is not None else "default"
            if tuned is None:
                tuned = TunedVisionConfig()
        else:
            self.tuned_source = "argument"
        self.tuned: TunedVisionConfig = tuned
        if not isinstance(self.tuned, TunedVisionConfig):
            raise TypeError(f"tuned must be a TunedVisionConfig, got {type(self.tuned).__name__}")
        if cfg.vision.activation != "gelu_tanh":
            raise ValueError(
                f"{cfg.version}: block activation {cfg.vision.activation!r} != gelu_pytorch_tanh (D9 fork)"
            )
        self.n_img: int = int(dlayout.n_images)
        self.s_img: int = int(dlayout.layout.patches_per_image)
        self.hidden: int = int(self.kwargs["hidden_size"])
        self.n_heads: int = int(self.kwargs["num_heads"])
        self.inter: int = int(self.kwargs["intermediate_size"])
        self.eps: float = float(self.kwargs["norm_eps"])
        self.merge_group: int = int(self.kwargs["spatial_merge_size"]) ** 2
        self.out_hidden: int = int(self.kwargs["out_hidden_size"])
        self.scale: float = float(self.head_dim) ** -0.5
        if abs(self.scale - cfg.vision.attention_scale) > 1e-12:
            raise RuntimeError(f"fork scale {self.scale} != cfg {cfg.vision.attention_scale}")
        if self.n_img * self.s_img != self.n_patches:
            raise ValueError(f"{self.n_img} images x {self.s_img} patches != {self.n_patches}")
        if self.s_img % 32:
            raise ValueError(f"patches per image {self.s_img} is not a tile multiple (nlp_create_qkv_heads needs it)")
        if self.cu_seqlens != tuple(range(0, self.n_patches + 1, self.s_img)):
            raise ValueError(f"cu_seqlens {self.cu_seqlens} are not {self.n_img} equal blocks of {self.s_img}")
        if self.head_dim != 64:
            raise ValueError(f"head_dim {self.head_dim}: the fork assumes 64 (= padded, no RoPE tail)")
        if self.n_heads * self.head_dim != self.hidden:
            raise ValueError("hidden != n_heads * head_dim")
        self.reconfigure(self.tuned, lofi=resolve_vit_lofi(policy, lofi))
        # --- per-block weight handles of the tt_dit tower (same bytes as the verbatim path)
        self.bw: List[Dict[str, Any]] = []
        for i in range(self.depth):
            b = self.model.blocks[i]
            w = {
                "ln1.w": b.norm1.weight.data,
                "ln1.b": b.norm1.bias.data,
                "Wqkv": b.attn.qkv.weight.data,
                "bqkv": b.attn.qkv.bias.data,
                "Wo": b.attn.proj.weight.data,
                "bo": b.attn.proj.bias.data,
                "ln2.w": b.norm2.weight.data,
                "ln2.b": b.norm2.bias.data,
                "fc1.W": b.mlp.linear_fc1.weight.data,
                "fc1.b": b.mlp.linear_fc1.bias.data,
                "fc2.W": b.mlp.linear_fc2.weight.data,
                "fc2.b": b.mlp.linear_fc2.bias.data,
            }
            self._check_block_weights(i, w)
            self.bw.append(w)
        self.pe_w = self.model.patch_embed.proj.weight.data
        self.pe_b = self.model.patch_embed.proj.bias.data
        pes = tuple(int(d) for d in self.pe_w.shape)
        if pes != (cfg.vision.patch_in_features, self.hidden):
            raise ValueError(f"patch embed weight {pes} != ({cfg.vision.patch_in_features}, {self.hidden})")
        self.mw: List[Dict[str, Any]] = []
        for m in [self.model.merger, *self.model.deepstack_merger_list]:
            mw = {
                "post": bool(m.use_postshuffle_norm),
                "norm.w": m.norm.weight.data,
                "norm.b": m.norm.bias.data,
                "fc1.W": m.linear_fc1.weight.data,
                "fc1.b": m.linear_fc1.bias.data,
                "fc2.W": m.linear_fc2.weight.data,
                "fc2.b": m.linear_fc2.bias.data,
            }
            merged = self.hidden * self.merge_group
            if tuple(int(d) for d in mw["fc1.W"].shape) != (merged, merged):
                raise ValueError(f"merger fc1 {tuple(mw['fc1.W'].shape)} != ({merged}, {merged})")
            if tuple(int(d) for d in mw["fc2.W"].shape) != (merged, self.out_hidden):
                raise ValueError(f"merger fc2 {tuple(mw['fc2.W'].shape)} != ({merged}, {self.out_hidden})")
            self.mw.append(mw)
        # --- table rows must exist (KeyError here, never an auto config at run time; plan D3)
        dt = dtype_name(self.bw[0]["Wqkv"].dtype)
        self.weight_dtype = dt
        p = self.n_patches
        matmul_spec(p, self.hidden, 3 * self.hidden, dt)
        matmul_spec(p, self.hidden, self.hidden, dt)
        matmul_spec(p, self.hidden, self.inter, dt)
        matmul_spec(p, self.inter, self.hidden, dt)
        if self.tuned.tuned_embed:
            matmul_spec(p, cfg.vision.patch_in_features, self.hidden, dtype_name(self.pe_w.dtype))
        if self.tuned.tuned_mergers:
            merged = self.hidden * self.merge_group
            matmul_spec(self.n_tokens, merged, merged, dtype_name(self.mw[0]["fc1.W"].dtype))
            matmul_spec(self.n_tokens, merged, self.out_hidden, dtype_name(self.mw[0]["fc2.W"].dtype))
        self._rope_checked_addr: Optional[Tuple[int, int]] = None
        self.n_blocks_run = 0

    def reconfigure(self, tuned: TunedVisionConfig, *, lofi: Optional[bool] = None) -> None:
        """Re-derive the program-config knobs from ``tuned`` (no weight reload; the sweep's variant loop).  ``lofi``
        keeps the current value when ``None``.  Program-cache entries of earlier configurations stay valid."""
        from models.experimental.gr00t.tt.layers import sdpa_chunk_ok, sdpa_chunks_allowed

        if not isinstance(tuned, TunedVisionConfig):
            raise TypeError(f"tuned must be a TunedVisionConfig, got {type(tuned).__name__}")
        chunks = tuned.sdpa_chunks
        if not sdpa_chunk_ok(self.s_img, self.s_img, self.head_dim, chunks):
            raise ValueError(
                f"sdpa_chunks {chunks} not allowed for (Sq=Sk={self.s_img}, d={self.head_dim}); allowed: "
                f"{sdpa_chunks_allowed(self.s_img, self.s_img, self.head_dim)}"
            )
        if tuned.hifi_sdpa_chunks is not None and not sdpa_chunk_ok(
            self.s_img, self.s_img, self.head_dim, tuned.hifi_sdpa_chunks
        ):
            raise ValueError(
                f"hifi_sdpa_chunks {tuned.hifi_sdpa_chunks} not allowed for (Sq=Sk={self.s_img}, d={self.head_dim}); "
                f"allowed: {sdpa_chunks_allowed(self.s_img, self.s_img, self.head_dim)}"
            )
        if tuned.hifi_blocks and tuned.hifi_blocks[-1] >= self.depth:
            raise ValueError(f"hifi_blocks {tuned.hifi_blocks} exceed depth {self.depth}")
        n_mergers = 1 + len(self.deepstack_indexes)
        if tuned.hifi_mergers and tuned.hifi_mergers[-1] >= n_mergers:
            raise ValueError(f"hifi_mergers {tuned.hifi_mergers} exceed the {n_mergers} mergers (0 = final, 1+d)")
        self.tuned = tuned
        self.sdpa_chunks = chunks
        self.rope_mode = tuned.rope_mode
        if lofi is not None:
            self.lofi = bool(lofi)
        self.mem = MemPreset(tuned.mem)
        self.ff_mem = MemPreset(tuned.ff_mem)
        self.ckc = _ckc_of(tuned.ckc)
        self.ff_ckc = CKC_LOFI_FP32ACC if self.lofi else _ckc_of(tuned.ff_ckc)
        self.sdpa_ckc = _ckc_of(tuned.sdpa_ckc)
        self.merger_ckc_preset = _ckc_of(tuned.merger_ckc)
        # per-block / per-merger resolution (docs/plan/n17-fidelity.md): hifi overrides on the listed indices only
        hi = set(tuned.hifi_blocks)
        hifi_ckc = _ckc_of(tuned.hifi_ckc) if tuned.hifi_ckc is not None else self.ckc
        hifi_ff = _ckc_of(tuned.hifi_ff_ckc) if tuned.hifi_ff_ckc is not None else self.ff_ckc
        hifi_sdpa = _ckc_of(tuned.hifi_sdpa_ckc) if tuned.hifi_sdpa_ckc is not None else self.sdpa_ckc
        hifi_chunks = tuned.hifi_sdpa_chunks if tuned.hifi_sdpa_chunks is not None else chunks
        self.block_ckc: List[Any] = [hifi_ckc if i in hi else self.ckc for i in range(self.depth)]
        self.block_ff_ckc: List[Any] = [hifi_ff if i in hi else self.ff_ckc for i in range(self.depth)]
        self.block_sdpa: List[Tuple[Any, Tuple[int, int]]] = [
            (hifi_sdpa, hifi_chunks) if i in hi else (self.sdpa_ckc, chunks) for i in range(self.depth)
        ]
        hm = set(tuned.hifi_mergers)
        hifi_merger = _ckc_of(tuned.hifi_merger_ckc) if tuned.hifi_merger_ckc is not None else self.merger_ckc_preset
        self.merger_ckcs: List[Any] = [hifi_merger if m in hm else self.merger_ckc_preset for m in range(n_mergers)]

    def _check_block_weights(self, i: int, w: Mapping[str, Any]) -> None:
        shapes = {
            "ln1.w": (1, self.hidden),
            "ln1.b": (1, self.hidden),
            "Wqkv": (self.hidden, 3 * self.hidden),
            "bqkv": (1, 3 * self.hidden),
            "Wo": (self.hidden, self.hidden),
            "bo": (1, self.hidden),
            "ln2.w": (1, self.hidden),
            "ln2.b": (1, self.hidden),
            "fc1.W": (self.hidden, self.inter),
            "fc1.b": (1, self.inter),
            "fc2.W": (self.inter, self.hidden),
            "fc2.b": (1, self.hidden),
        }
        for k, want in shapes.items():
            got = tuple(int(d) for d in w[k].shape)
            if got != want:
                raise ValueError(f"block {i} weight {k}: shape {got} != {want}")

    # ------------------------------------------------------------------ host side additions
    def prepare_inputs(self, mi: Any) -> VisionHostInputs:
        host = super().prepare_inputs(mi)
        s = self.s_img
        for name, t in (("cos", host.cos), ("sin", host.sin)):
            for k in range(1, self.n_img):
                if not torch.equal(t[:s], t[k * s : (k + 1) * s]):
                    raise StaticShapeError(
                        f"{name}: image {k} RoPE rows differ from image 0 (batched fork needs equal grids)"
                    )
        return host

    def upload(self, host: VisionHostInputs, memory_config: Any = None) -> VisionDeviceInputs:
        dev = super().upload(host, memory_config)
        self._check_rope_blocks(dev)
        return dev

    def _check_rope_blocks(self, dev: VisionDeviceInputs) -> None:
        """Device-side check (once per buffer pair) that the per-image RoPE rows are identical."""
        addr = (int(dev.cos.buffer_address()), int(dev.sin.buffer_address()))
        if self._rope_checked_addr == addr:
            return
        s = self.s_img
        for name in ("cos", "sin"):
            t = to_torch(getattr(dev, name))
            for k in range(1, self.n_img):
                if not torch.equal(t[:s], t[k * s : (k + 1) * s]):
                    raise StaticShapeError(f"device {name}: image {k} RoPE rows differ from image 0")
        self._rope_checked_addr = addr

    def _tap_copy(self, t: Any) -> Any:
        """DRAM copy of an L1 tap (``ttnn.clone``): the fork's activations live in L1, and the untraced tap-collecting
        run of ``Gr00tTT.run_untraced`` keeps every tap alive through the LLM / adapter stages -- 3 x 3.6 MB of L1 taps
        pushed the VL-SA ff1 CBs into the L1 buffers (``test_e2e_untraced_taps_vs_golden[n17]``, session 3b).  The traced
        path (``taps=None``) never calls this; values are unchanged (a copy)."""
        return self.ttnn.clone(t, memory_config=self.ttnn.DRAM_MEMORY_CONFIG)

    # ------------------------------------------------------------------ views
    def _view4(self, x: Any) -> Any:
        """``[P, hidden]`` -> ``[n_img, 1, s_img, hidden]`` (free view: same buffer, tile-aligned row split)."""
        s = tuple(int(d) for d in x.shape)
        if s == (self.n_img, 1, self.s_img, self.hidden):
            return x
        if s != (self.n_patches, self.hidden):
            raise ValueError(f"expected [{self.n_patches}, {self.hidden}] or the batched view, got {s}")
        return self.ttnn.reshape(x, (self.n_img, 1, self.s_img, self.hidden))

    def _view2(self, x: Any) -> Any:
        """``[n_img, 1, s_img, hidden]`` -> ``[P, hidden]`` (free view)."""
        s = tuple(int(d) for d in x.shape)
        if s == (self.n_patches, self.hidden):
            return x
        if s != (self.n_img, 1, self.s_img, self.hidden):
            raise ValueError(f"expected the batched view or [{self.n_patches}, {self.hidden}], got {s}")
        return self.ttnn.reshape(x, (self.n_patches, self.hidden))

    def rope_views(self, rope: Tuple[Any, Any]) -> Tuple[Any, Any]:
        """Per-image ``cos / sin [1, 1, s_img, head_dim]`` from the tower tables ``[P, head_dim]`` (2 slices + 2 views;
        tables already in the per-image shape pass through)."""
        ttnn = self.ttnn
        out = []
        for t in rope:
            s = tuple(int(d) for d in t.shape)
            if s == (1, 1, self.s_img, self.head_dim):
                out.append(t)
                continue
            if s != (self.n_patches, self.head_dim):
                raise ValueError(
                    f"rope table must be [{self.n_patches}, {self.head_dim}] or [1,1,{self.s_img},{self.head_dim}], got {s}"
                )
            sl = ttnn.slice(t, [0, 0], [self.s_img, self.head_dim], memory_config=self.mem.build())
            out.append(ttnn.reshape(sl, (1, 1, self.s_img, self.head_dim)))
        return out[0], out[1]

    # ------------------------------------------------------------------ device side
    def embed(self, patches: Any) -> Any:
        if not self.tuned.tuned_embed:
            return super().embed(patches)
        from models.experimental.gr00t.tt.layers import linear

        return linear(patches, self.pe_w, self.pe_b, mem=self.mem, ckc=self.ckc)

    def add_pos(self, lin: Any, pos: Any) -> Any:
        return self.ttnn.add(lin, pos, memory_config=self.mem.build())

    def _rope(self, x: Any, cos: Any, sin: Any) -> Any:
        """2-D RoPE on ``x [B, H, S, 64]`` with ``cos / sin [1, 1, S, 64]``; deallocates ``x``."""
        ttnn = self.ttnn
        mem = self.mem.build()
        if self.rope_mode == "hf_op":
            from models.experimental.gr00t.tt.layers import rope

            out = rope(x, cos, sin, layout="hf", mem=self.mem)
            ttnn.deallocate(x)
            return out
        # composite == vision_qwen3vl._apply_vision_rope with head_dim == padded head_dim (no tail)
        b, h, s, d = (int(v) for v in x.shape)
        half = d // 2
        hi = ttnn.slice(x, [0, 0, 0, half], [b, h, s, d], memory_config=mem)
        lo = ttnn.slice(x, [0, 0, 0, 0], [b, h, s, half], memory_config=mem)
        neg = ttnn.neg(hi, memory_config=mem)
        ttnn.deallocate(hi)
        rotated = ttnn.concat([neg, lo], dim=-1, memory_config=mem)
        ttnn.deallocate(neg)
        ttnn.deallocate(lo)
        xc = ttnn.mul(x, cos, memory_config=mem)
        rs = ttnn.mul(rotated, sin, memory_config=mem)
        ttnn.deallocate(rotated)
        ttnn.deallocate(x)
        out = ttnn.add(xc, rs, memory_config=mem)
        ttnn.deallocate(xc)
        ttnn.deallocate(rs)
        return out

    def block4(self, i: int, x: Any, cos4: Any, sin4: Any) -> Any:
        """Block ``i`` on the batched layout ``x [n_img, 1, s_img, hidden]`` with per-image ``cos4 / sin4``."""
        from models.experimental.gr00t.tt.layers import (
            concat_heads,
            layer_norm,
            linear,
            residual_add,
            sdpa,
            split_heads,
        )

        ttnn = self.ttnn
        if not 0 <= i < self.depth:
            raise IndexError(f"block {i} out of range 0..{self.depth - 1}")
        xs = tuple(int(d) for d in x.shape)
        if xs != (self.n_img, 1, self.s_img, self.hidden):
            raise ValueError(f"block4 expects [{self.n_img}, 1, {self.s_img}, {self.hidden}], got {xs}")
        w = self.bw[i]
        mem = self.mem
        ckc, ff_ckc = self.block_ckc[i], self.block_ff_ckc[i]
        sdpa_ckc, sdpa_chunks = self.block_sdpa[i]
        h = layer_norm(x, w["ln1.w"], w["ln1.b"], eps=self.eps, mem=mem)
        qkv = linear(h, w["Wqkv"], w["bqkv"], mem=mem, ckc=ckc)
        ttnn.deallocate(h)
        q, k, v = split_heads(qkv, self.n_heads, self.n_heads, mem=mem)
        ttnn.deallocate(qkv)
        q = self._rope(q, cos4, sin4)
        k = self._rope(k, cos4, sin4)
        o = sdpa(q, k, v, scale=self.scale, mask=None, is_causal=False, chunks=sdpa_chunks, ckc=sdpa_ckc, mem=mem)
        ttnn.deallocate(q)
        ttnn.deallocate(k)
        ttnn.deallocate(v)
        a = concat_heads(o, mem=mem)
        ttnn.deallocate(o)
        p = linear(a, w["Wo"], w["bo"], mem=mem, ckc=ckc)
        ttnn.deallocate(a)
        x1 = residual_add(x, p, mem=mem)
        ttnn.deallocate(p)
        h2 = layer_norm(x1, w["ln2.w"], w["ln2.b"], eps=self.eps, mem=mem)
        f = linear(h2, w["fc1.W"], w["fc1.b"], act=VIT_FF_ACT, mem=self.ff_mem, ckc=ff_ckc)
        ttnn.deallocate(h2)
        g = linear(f, w["fc2.W"], w["fc2.b"], mem=mem, ckc=ff_ckc)
        ttnn.deallocate(f)
        out = residual_add(x1, g, mem=mem)
        ttnn.deallocate(g)
        ttnn.deallocate(x1)
        self.n_blocks_run += 1
        return out

    def block(self, i: int, x: Any, rope: Tuple[Any, Any], cu_seqlens: Optional[Sequence[int]] = None) -> Any:
        """Drop-in for the wrapper's ``block``: accepts ``[P, hidden]`` or the batched view (returns the same layout)
        and the tower tables or per-image ``cos / sin``.  Per-call table slicing happens here; :meth:`forward_blocks` /
        :meth:`forward` slice once per call instead."""
        cu = self.cu_seqlens if cu_seqlens is None else tuple(int(c) for c in cu_seqlens)
        if cu != self.cu_seqlens:
            raise StaticShapeError(f"cu_seqlens {cu} != baked {self.cu_seqlens}")
        flat = len(x.shape) == 2
        cos4, sin4 = self.rope_views(rope)
        out = self.block4(i, self._view4(x), cos4, sin4)
        return self._view2(out) if flat else out

    def forward_blocks(
        self,
        x: Any,
        n: int,
        rope: Tuple[Any, Any],
        cu_seqlens: Optional[Sequence[int]] = None,
        *,
        start: int = 0,
        collect: Iterable[int] = (),
    ) -> Tuple[Any, Dict[int, Any]]:
        """As the wrapper's: blocks ``start .. start+n-1``; kept outputs come back in the layout of ``x``."""
        if n < 0 or start < 0 or start + n > self.depth:
            raise IndexError(f"blocks {start}..{start + n - 1} out of range 0..{self.depth - 1}")
        keep = set(int(i) for i in collect)
        bad = keep - set(range(start, start + n))
        if bad:
            raise IndexError(f"collect indices {sorted(bad)} are not run by blocks {start}..{start + n - 1}")
        cu = self.cu_seqlens if cu_seqlens is None else tuple(int(c) for c in cu_seqlens)
        if cu != self.cu_seqlens:
            raise StaticShapeError(f"cu_seqlens {cu} != baked {self.cu_seqlens}")
        flat = len(x.shape) == 2
        cos4, sin4 = self.rope_views(rope)
        cur = self._view4(x)
        out: Dict[int, Any] = {}
        for i in range(start, start + n):
            cur = self.block4(i, cur, cos4, sin4)
            if i in keep:
                out[i] = self._view2(cur) if flat else cur
        return (self._view2(cur) if flat else cur), out

    def _merge(self, m: int, x: Any) -> Any:
        """Merger ``m`` (0 = final, ``1 + d`` = deepstack ``d``) on ``x`` in either layout -> ``[P/4, out_hidden]``."""
        if not self.tuned.tuned_mergers:
            x2 = self._view2(x)
            return self.model.merger.forward(x2) if m == 0 else self.model.deepstack_merger_list[m - 1].forward(x2)
        from models.experimental.gr00t.tt.layers import layer_norm, linear

        ttnn = self.ttnn
        w = self.mw[m]
        ckc = self.merger_ckcs[m]
        merged = self.hidden * self.merge_group
        x2 = self._view2(x)
        if w["post"]:
            r = ttnn.reshape(x2, (self.n_tokens, merged), memory_config=self.mem.build())  # real op (row fold)
            h = layer_norm(r, w["norm.w"], w["norm.b"], eps=self.eps, mem=self.mem)
            ttnn.deallocate(r)
        else:
            hn = layer_norm(x2, w["norm.w"], w["norm.b"], eps=self.eps, mem=self.mem)
            h = ttnn.reshape(hn, (self.n_tokens, merged), memory_config=self.mem.build())
            ttnn.deallocate(hn)
        if self.tuned.merger_gelu_fused:
            f = linear(h, w["fc1.W"], w["fc1.b"], act=MERGER_ACT, mem=self.mem, ckc=ckc)
        else:
            from models.experimental.gr00t.tt.layers import gelu

            f0 = linear(h, w["fc1.W"], w["fc1.b"], mem=self.mem, ckc=ckc)
            f = gelu(f0, "erf", mem=self.mem)  # exact erf GELU (nn.GELU()), mb2 §B.4
            ttnn.deallocate(f0)
        ttnn.deallocate(h)
        out = linear(f, w["fc2.W"], w["fc2.b"], mem=DRAM, ckc=ckc)
        ttnn.deallocate(f)
        return out

    def deepstack_merge(self, d: int, x: Any) -> Any:
        if not 0 <= d < len(self.deepstack_indexes):
            raise IndexError(f"deepstack merger {d} out of range 0..{len(self.deepstack_indexes) - 1}")
        return self._merge(1 + d, x)

    def merge(self, x: Any) -> Any:
        return self._merge(0, x)

    def forward(self, inputs: VisionDeviceInputs, taps: Optional[Dict[str, Any]] = None) -> Tuple[Any, List[Any]]:
        """``(merged [P/4, 2048], deepstack 3 x [P/4, 2048])`` with the fork's op sequence; taps as the wrapper's
        (``[P, hidden]`` DRAM copies of the L1 block taps, see :meth:`_tap_copy`).  Every intermediate is freed."""
        if inputs.cu_seqlens != self.cu_seqlens:
            raise StaticShapeError(f"cu_seqlens {inputs.cu_seqlens} != baked {self.cu_seqlens}")
        ttnn = self.ttnn
        rec = taps is not None
        lin = self.embed(inputs.patches)
        if rec:
            taps[TAP_EMBEDDINGS] = self._tap_copy(lin)
        x2 = self.add_pos(lin, inputs.pos)
        ttnn.deallocate(lin)
        cos4, sin4 = self.rope_views(inputs.rope)
        x = self._view4(x2)
        deepstack: List[Any] = []
        for i in range(self.depth):
            nxt = self.block4(i, x, cos4, sin4)
            if rec and i == 0:
                taps[TAP_BLOCK_FIRST] = self._tap_copy(self._view2(nxt))
            if i in self.deepstack_indexes:
                d = self.deepstack_indexes.index(i)
                feat = self.deepstack_merge(d, nxt)
                deepstack.append(feat)
                if rec:
                    taps[deepstack_key(d)] = feat
            ttnn.deallocate(x)
            x = nxt
        if rec:
            taps[TAP_BLOCK_LAST] = self._tap_copy(self._view2(x))
        merged = self.merge(x)
        if rec:
            taps[TAP_PROJECTOR] = merged
        ttnn.deallocate(x)
        for t in (cos4, sin4):
            if t is not inputs.cos and t is not inputs.sin:
                ttnn.deallocate(t)
        if len(deepstack) != len(self.deepstack_indexes):
            raise RuntimeError(f"{len(deepstack)} deepstack features for indexes {self.deepstack_indexes}")
        return merged, deepstack

    def forward_verbatim(
        self, inputs: VisionDeviceInputs, taps: Optional[Dict[str, Any]] = None
    ) -> Tuple[Any, List[Any]]:
        """The **verbatim** tt_dit op sequence on the same weights (the oracle of ``sweep_backbone --what vit``).

        Explicitly the base-class methods (``Qwen3VLVisionTT.embed / add_pos / block / deepstack_merge / merge`` --
        the tt_dit modules); ``Qwen3VLVisionTT.forward(self, ...)`` would dispatch through this class' overrides.
        """
        if inputs.cu_seqlens != self.cu_seqlens:
            raise StaticShapeError(f"cu_seqlens {inputs.cu_seqlens} != baked {self.cu_seqlens}")
        base = Qwen3VLVisionTT
        rec = taps is not None
        lin = base.embed(self, inputs.patches)
        if rec:
            taps[TAP_EMBEDDINGS] = lin
        x = base.add_pos(self, lin, inputs.pos)
        deepstack: List[Any] = []
        for i in range(self.depth):
            x = base.block(self, i, x, inputs.rope, inputs.cu_seqlens)
            if rec and i == 0:
                taps[TAP_BLOCK_FIRST] = x
            if i in self.deepstack_indexes:
                d = self.deepstack_indexes.index(i)
                feat = base.deepstack_merge(self, d, x)
                deepstack.append(feat)
                if rec:
                    taps[deepstack_key(d)] = feat
        if rec:
            taps[TAP_BLOCK_LAST] = x
        merged = base.merge(self, x)
        if rec:
            taps[TAP_PROJECTOR] = merged
        return merged, deepstack

    def describe(self) -> Dict[str, Any]:
        d = super().describe()

        def nm(c: Any) -> str:
            return str(getattr(c, "name", c))

        d.update(
            tower="fork",
            vit_fork=True,
            n_img=self.n_img,
            s_img=self.s_img,
            tuned=self.tuned.to_dict(),
            tuned_source=self.tuned_source,
            lofi=self.lofi,
            weight_dtype=self.weight_dtype,
            sdpa_chunks=list(self.sdpa_chunks),
            rope_mode=self.rope_mode,
            ff_ckc=nm(self.ff_ckc),
            sdpa_ckc=nm(self.sdpa_ckc),
            per_block={
                "ckc": [nm(c) for c in self.block_ckc],
                "ff_ckc": [nm(c) for c in self.block_ff_ckc],
                "sdpa_ckc": [nm(c) for c, _ in self.block_sdpa],
                "sdpa_chunks": [list(ch) for _, ch in self.block_sdpa],
            },
            merger_ckcs=[nm(c) for c in self.merger_ckcs],
        )
        return d


__all__ = [
    "CKC_HIFI3_FP32ACC",
    "CKC_SDPA_HIFI2_FP32ACC",
    "CKC_SDPA_HIFI3",
    "CKC_SDPA_HIFI3_FP32ACC",
    "CKC_VIT_SDPA_HIFI4_FP32ACC",
    "VIT_TUNED_ENV",
    "MERGER_ACT",
    "TunedQwen3VLVisionTT",
    "TunedVisionConfig",
    "VIT_FF_ACT",
    "VIT_FORK_ENV",
    "VIT_ROPE_MODES",
    "VIT_SDPA_CLASS",
    "build_vision_tower",
    "resolve_vit_fork",
    "resolve_vit_lofi",
    "HIDDEN_ACT",
    "N_VISUAL_TENSORS",
    "POS_TABLE_NAME",
    "SPATIAL_MERGE_SIZE",
    "TAP_BLOCK_FIRST",
    "TAP_BLOCK_LAST",
    "TAP_DEEPSTACK",
    "TAP_EMBEDDINGS",
    "TAP_KEYS",
    "TAP_PROJECTOR",
    "TOWER_KWARGS_D9",
    "Qwen3VLVisionTT",
    "VisionDeviceInputs",
    "VisionHostInputs",
    "deepstack_key",
    "expected_cu_seqlens",
    "host_inputs",
    "pos_table_from_plan",
    "to_torch",
    "tower_kwargs",
    "traced_seconds",
    "visual_state_dict",
]