# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """The trace-A tail: ``vlln`` -> VL self-attention x4 -> ``[image | non-image]`` gather -> hoisted cross K/V -> state encoder (IMPLEMENTATION_PLAN.md §2.2 row ``tt/adapter.py``, §3.3 "vlln / VL-SA / hoist" paragraph, decisions D5 / D7; static-shape-plan.md §5, §6.2-6.3; foundation-f3b §6 item 1; WP-10). Data flow of :meth:`VLAdapterTT.forward` (bf16 TILE device tensors; ``S_pad`` = ``DeviceLayout.s_pad_vl``):: features [1,1,S_pad,2048] --layer_norm(vlln.w, vlln.b, eps 1e-5; HiFi4 fp32acc)--> vlln_out --4 x EncoderBlockTT(2048, 32 x 64, ff 8192, eps 1e-5, scale 64**-0.5, key-valid mask [1,1,S_pad,S_pad])--> vl = dit_encoder_hidden_states rows (N1.6: vl = vlln_out, no VL-SA) --gather (static slot permutation, static-shape §6.2)--> vl_sets: {"image": [1,1,L_img_pad,2048], "text": [1,1,L_txt_pad,2048]} (N1.5: {"all": vl}) per cross block i (cfg.dit.cross_block_indices, key set dlayout.key_set(i)): kv = linear(vl_sets[ks], dit.block.i.Wkv [2048,4096], dit.block.i.bkv) (M = L_pad, table row) K, V = split_kv_heads(kv, 32, q_dummy_) -- nlp_create_qkv_heads(q_zero, input_kv=kv), plan D7 ttnn.copy(K, K_i); ttnn.copy(V, V_i) -- persistent [1,32,L_pad,64] DRAM buffers (addresses captured by trace B never change; static-shape §9) state_rows [1,1,32,D_pad] (row 0 = state) --linear(enc.state.W1, b1, RELU fused)--> linear(W2, b2) --> state_features [1,1,1,1536] (+ the full 32-row tile block) Gather (static-shape-plan §6.2 "gather, not mask"; both modes give logical ``[1,1,L_pad,2048]`` tensors whose pad rows ``>= L_real`` are exactly zero, so the hoisted K/V of the pad rows equal the bias -- finite, masked in the DiT): * ``gather_mode="splice"`` (plan §3.3 default): one ``ttnn.to_layout(ROW_MAJOR)`` of ``vl``, then per key set ``tt.layers.splice_rows`` (``ttnn.slice`` of every run + a constant zero-row piece + one ``ttnn.concat``) and one ``ttnn.to_layout(TILE)``. The runs of ``DeviceLayout.gather_pieces()`` are not tile aligned (N1.6 image slots ``[44,125)``, N1.7 ``[4,116)`` ...), and a TILE-layout ``ttnn.slice`` with an unaligned begin untilizes the whole input for every piece (``slice.cpp`` ``rm_only`` path) -- the explicit row-major round trip does that once. * ``gather_mode="onehot"``: one ``ttnn.matmul(P_ks [1,1,L_pad,S_pad], vl)`` per key set with a constant 0/1 matrix (row ``r < L_real`` selects slot ``index[r]``, rows ``>= L_real`` are zero) under an explicit program config of the plan-D3 family (1-D ``mcast_in0`` at ``L_pad <= 64``, full-grid 2-D above) and HiFi4 / fp32 acc, which reproduces every bf16 input exactly (one non-zero product per output). ``(m_bucket(L_pad), S_pad, 2048)`` has no ``MATMUL_CONFIGS`` row, so the configs are built by :func:`gather_matmul_plan` here (WP-01 may adopt them as table rows). Fidelity: ``vlln`` and the VL-SA norms HiFi4 fp32acc approx-off (plan §1.2); VL-SA projections HiFi2 fp32acc; the hoist projection defaults to **HiFi4** fp32acc (``hoist_ckc``) -- it runs once per inference (16 matmuls at most) and the WP-10 acceptance gates the hoisted K/V at PCC >= 0.9999 against ``vl_set @ Wkv`` on the device-rounded weights, which HiFi2's reduced-mantissa passes on the activation operand cannot guarantee; the per-step DiT projections (WP-13) keep the plan's HiFi2. ``policy.kv_dtype`` must be ``"bf16"`` (a bfp8 K/V buffer needs a typecast after the head split -- not implemented, raised). Host helpers (CPU, torch only): :func:`rows_to_slots` / :func:`slots_to_rows` move golden ``[1, S, D]`` tensors into / out of the padded slot layout through ``SlotLayout.slot_of_token``; :func:`state_rows_host` builds the ``[1,1,32,D_pad]`` state input; :func:`gather_index` / :func:`onehot_gather_matrix` expose the permutation. ``ttnn`` is imported lazily (inside ``__init__`` / ``forward``); :class:`AdapterPlan` is CPU-testable. """ from __future__ import annotations from dataclasses import dataclass, field from typing import Any, Dict, List, Mapping, Optional, Tuple import torch from models.experimental.gr00t.common.configs import TILE, GR00TConfig from models.experimental.gr00t.common.prompts import SlotLayout from models.experimental.gr00t.tt.blocks import EncoderBlockTT from models.experimental.gr00t.tt.device import ( CKC_HIFI2_FP32ACC, CKC_HIFI4_FP32ACC, DRAM, EXPECTED_COMPUTE_GRID, L1, resolve_ckc, resolve_mem, ) from models.experimental.gr00t.tt.layers import ( WEIGHT_DTYPES, MatmulPlan, MatmulSpec, RowPiece, default_sdpa_chunks, divisors, dtype_name, layer_norm, linear, m_bucket, matmul_plan, matmul_spec, splice_rows, split_kv_heads, subblock_2d, tiles, to_device, ) from models.experimental.gr00t.tt.layout import DeviceLayout, Run, runs_index, runs_length from models.experimental.gr00t.tt.policy import TTPolicy GATHER_MODES: Tuple[str, ...] = ("splice", "onehot") KEY_SETS: Tuple[str, ...] = ("all", "text", "image") STATE_ROWS: int = TILE # the state encoder runs on one tile row block: row 0 = state, rows 1..31 = zeros VLSA_WEIGHT_NAMES: Tuple[str, ...] = ( "norm1.w", "norm1.b", "Wqkv", "bqkv", "Wo", "bo", "norm3.w", "norm3.b", "ff1.W", "ff1.b", "ff2.W", "ff2.b", ) TAP_NAMES: Tuple[str, ...] = ( "vlln_out", "vlsa_block0_out", "vl_self_attention_out", "dit_encoder_hidden_states", "state_features", ) GATHER_ONEHOT_IBW_MAX = 8 # in0_block_w cap of the one-hot gather matmul (mb1 Table 2 uses ibw 8 for K = 2048) def _shape(t: Any) -> Tuple[int, ...]: return tuple(int(d) for d in t.shape) # -------------------------------------------------------------------------------------------------------------------- # host helpers (torch only) # -------------------------------------------------------------------------------------------------------------------- def _check_slots_arg(slots: Any) -> SlotLayout: if not isinstance(slots, SlotLayout): raise TypeError(f"expected a common.prompts.SlotLayout, got {type(slots).__name__}") sot = slots.slot_of_token if sot.ndim != 1 or int(sot.numel()) != int(slots.seq_len): raise ValueError(f"slot_of_token must be (S={slots.seq_len},), got {tuple(sot.shape)}") return slots def rows_to_slots(x: torch.Tensor, slots: SlotLayout, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """``[1, S, D]`` (or ``[S, D]``) real-token rows -> ``[1, 1, S_pad, D]`` in the static slot layout: row ``t`` goes to slot ``slots.slot_of_token[t]`` (static-shape-plan §4.3 / §4.5); pad slots are zero. This is how a golden ``backbone_features`` tap becomes the device input of the adapter (plan §7.2 module-vs-golden tier).""" slots = _check_slots_arg(slots) if x.ndim == 2: x = x[None] if x.ndim != 3 or x.shape[0] != 1: raise ValueError(f"rows_to_slots expects [1, S, D] or [S, D], got {tuple(x.shape)}") s, d = int(x.shape[1]), int(x.shape[2]) if s != slots.seq_len: raise ValueError(f"rows_to_slots: {s} rows but the slot layout has S={slots.seq_len}") s_pad = int(slots.shape_key.s_pad) out = torch.zeros(1, 1, s_pad, d, dtype=dtype if dtype is not None else x.dtype) out[0, 0, slots.slot_of_token] = x[0].to(out.dtype) return out def slots_to_rows(x: torch.Tensor, slots: SlotLayout) -> torch.Tensor: """Inverse of :func:`rows_to_slots`: ``[1, 1, S_pad, D]`` (or ``[1, S_pad, D]`` / ``[S_pad, D]``) -> ``[1, S, D]`` real-token rows in token order (the golden's layout).""" slots = _check_slots_arg(slots) while x.ndim > 2 and x.shape[0] == 1: x = x[0] if x.ndim != 2: raise ValueError(f"slots_to_rows expects [1, 1, S_pad, D], got a tensor that squeezes to {tuple(x.shape)}") s_pad = int(slots.shape_key.s_pad) if int(x.shape[0]) != s_pad: raise ValueError(f"slots_to_rows: {int(x.shape[0])} rows != S_pad={s_pad}") return x[slots.slot_of_token][None] def state_rows_host(state: torch.Tensor, state_dim_pad: int, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: """``state [1, 1, D_max]`` (``ModelInputs.state``, fp32) -> ``[1, 1, 32, D_pad]`` bf16: row 0 holds the state zero padded to ``EncoderConfig.state_dim_pad`` (the padded rows of ``enc.state.W1``, f1 §2 / weight-loader §4.7); rows 1..31 are zero. The device input of :class:`StateEncoderTT` (a persistent per-call buffer under trace).""" st = torch.as_tensor(state) if st.ndim != 3 or st.shape[0] != 1 or st.shape[1] != 1: raise ValueError(f"state must be [1, 1, D_max], got {tuple(st.shape)}") d_max = int(st.shape[2]) if d_max > state_dim_pad or state_dim_pad % TILE: raise ValueError(f"state width {d_max} does not fit state_dim_pad={state_dim_pad} (a multiple of {TILE})") if not bool(torch.isfinite(st).all()): raise ValueError("state contains non-finite values") out = torch.zeros(1, 1, STATE_ROWS, state_dim_pad, dtype=dtype) out[0, 0, 0, :d_max] = st[0, 0].to(dtype) return out def gather_index(dlayout: DeviceLayout, key_set: str) -> torch.Tensor: """``[L_real]`` int64 slot index of key set ``key_set`` (``"image"`` / ``"text"`` from ``gather_pieces()``, ``"all"`` = ``arange(S_pad_vl)``): row ``r`` of ``vl_sets[key_set]`` is slot ``index[r]`` of ``vl`` (static-shape §6.2).""" if key_set == "all": return torch.arange(dlayout.s_pad_vl, dtype=torch.int64) image_runs, nonimage_runs = dlayout.gather_pieces() if key_set == "image": return runs_index(image_runs) if key_set == "text": return runs_index(nonimage_runs) raise KeyError(f"unknown key set {key_set!r}; expected one of {KEY_SETS}") def onehot_gather_matrix(dlayout: DeviceLayout, key_set: str, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: """Constant ``[1, 1, L_pad, S_pad_vl]`` 0/1 matrix ``P`` with ``P[r, index[r]] = 1`` for ``r < L_real`` and zero rows ``>= L_real``, so that ``P @ vl == vl_sets[key_set]`` exactly (``gather_mode="onehot"``).""" ks = _key_set_lengths(dlayout, key_set) idx = gather_index(dlayout, key_set) if int(idx.numel()) != ks[0]: raise AssertionError(f"{key_set}: index has {int(idx.numel())} entries, expected L_real={ks[0]}") p = torch.zeros(1, 1, ks[1], dlayout.s_pad_vl, dtype=dtype) p[0, 0, torch.arange(ks[0]), idx] = 1.0 return p def _key_set_lengths(dlayout: DeviceLayout, key_set: str) -> Tuple[int, int]: """``(L_real, L_pad)`` of a key set (``"all"``: both ``S_pad_vl`` -- every row is a key, the per-call mask hides the pads, static-shape §6.2 N1.5 row).""" if key_set == "all": if dlayout.l_all_pad is None: raise KeyError( f"{dlayout.layout_name}: key set 'all' is not defined (cross_key_set={dlayout.cross_key_set})" ) return dlayout.s_pad_vl, int(dlayout.l_all_pad) if key_set == "image": if dlayout.l_img_pad is None: raise KeyError(f"{dlayout.layout_name}: key set 'image' is not defined") return int(dlayout.l_img_real), int(dlayout.l_img_pad) if key_set == "text": if dlayout.l_txt_pad is None: raise KeyError(f"{dlayout.layout_name}: key set 'text' is not defined") return int(dlayout.l_txt_real), int(dlayout.l_txt_pad) raise KeyError(f"unknown key set {key_set!r}; expected one of {KEY_SETS}") # -------------------------------------------------------------------------------------------------------------------- # one-hot gather matmul configs (adapter-local, plan D3 family) # -------------------------------------------------------------------------------------------------------------------- def _gather_in0_block_w(kt: int) -> int: """Largest divisor of ``Kt`` not above :data:`GATHER_ONEHOT_IBW_MAX` (``in0_block_w`` must divide ``Kt``).""" return max(d for d in divisors(kt) if d <= GATHER_ONEHOT_IBW_MAX) def gather_matmul_plan(l_pad: int, s_pad: int, width: int) -> MatmulPlan: """Explicit program config of ``P [L_pad, S_pad] @ vl [S_pad, width]`` (the one-hot gather; plan D3 recipe): ``L_pad <= 64`` -> 1-D ``mcast_in0`` (``per_core_M = Mt``, ``per_core_N = 1``, ``Nt`` cores), else the full-grid 2-D multicast (``per_core_M = ceil(Mt / 10)``, ``per_core_N = ceil(Nt / 11)``, :func:`tt.layers.subblock_2d`). Never the ``ttnn`` auto config.""" if l_pad <= 0 or l_pad % TILE or s_pad <= 0 or s_pad % TILE or width <= 0 or width % TILE: raise ValueError(f"gather matmul needs tile-multiple dims, got L_pad={l_pad}, S_pad={s_pad}, width={width}") mt, kt, nt = tiles(l_pad), tiles(s_pad), tiles(width) ibw = _gather_in0_block_w(kt) cols, rows = EXPECTED_COMPUTE_GRID key = (m_bucket(l_pad), int(s_pad), int(width), "bf16") src = "tt/adapter.py one-hot gather (WP-10; no MATMUL_CONFIGS row)" if mt <= 2: if nt > cols * rows: raise ValueError(f"one-hot gather: Nt={nt} exceeds the {cols}x{rows} grid at per_core_N=1") spec = MatmulSpec("mcast1d", nt, mt, 1, ibw, (1, 1), EXPECTED_COMPUTE_GRID, True, None, None, src) return MatmulPlan(key, spec, l_pad, s_pad, width, "mcast1d", spec.grid, mt, 1, ibw, (1, 1), None) per_core_m = -(-mt // rows) per_core_n = -(-nt // cols) sub = subblock_2d(per_core_m, per_core_n) spec = MatmulSpec("mcast2d", cols * rows, per_core_m, per_core_n, ibw, sub, (cols, rows), True, None, None, src) return MatmulPlan(key, spec, l_pad, s_pad, width, "mcast2d", spec.grid, per_core_m, per_core_n, ibw, sub, None) # -------------------------------------------------------------------------------------------------------------------- # CPU-side plan # -------------------------------------------------------------------------------------------------------------------- @dataclass(frozen=True) class KeySetPlan: """One cross-attention key set of the layout (static-shape §6.2): its gather runs, lengths and mask presence.""" name: str runs: Tuple[Run, ...] l_real: int l_pad: int identity: bool # True for "all": vl_sets["all"] is vl itself (no gather op) mask_present: bool # DeviceLayout.mask_presence["cross_"] @dataclass(frozen=True) class HoistBlockPlan: """One cross block's hoist: which key set feeds ``dit.block.{i}.Wkv`` and the persistent buffer length.""" block: int key_set: str l_pad: int @dataclass(frozen=True) class AdapterPlan: """Everything :class:`VLAdapterTT` decides before touching the device (validated in :meth:`build`).""" version: str layout_name: str gather_mode: str s_pad: int dim: int vlln_eps: float n_vlsa: int vlsa_heads: int vlsa_head_dim: int vlsa_ff: int vlsa_eps: float vlsa_scale: float vlsa_activation: str vlsa_chunks: Optional[Tuple[int, int]] n_heads: int # DiT heads (32) head_dim_pad: int # 64 kv_width: int # 2 * n_heads * head_dim_pad = 4096 key_sets: Mapping[str, KeySetPlan] blocks: Tuple[HoistBlockPlan, ...] hoist_lengths: Tuple[int, ...] state_dim_max: int state_dim_pad: int state_hidden: int state_out: int vl_mask_present: bool gather_plans: Mapping[str, MatmulPlan] = field(default_factory=dict) @classmethod def build( cls, cfg: GR00TConfig, dlayout: DeviceLayout, policy: TTPolicy, gather_mode: str = "splice" ) -> "AdapterPlan": """Cross-check config, layout and policy; resolve every table row the forward will need (raise early).""" if not isinstance(cfg, GR00TConfig) or not isinstance(dlayout, DeviceLayout): raise TypeError("AdapterPlan.build(cfg: GR00TConfig, dlayout: DeviceLayout, policy: TTPolicy)") if not isinstance(policy, TTPolicy): raise TypeError(f"policy must be a TTPolicy, got {type(policy).__name__}") if dlayout.version != cfg.version: raise ValueError(f"layout is for {dlayout.version}, config for {cfg.version}") if gather_mode not in GATHER_MODES: raise ValueError(f"gather_mode must be one of {GATHER_MODES}, got {gather_mode!r}") if policy.kv_dtype != "bf16": raise NotImplementedError( f"kv_dtype={policy.kv_dtype!r}: the hoisted K/V buffers are bf16 in Stage 1 (plan §1.2); a bfp8 " "buffer needs a typecast after the head split (WP-Q lever), not implemented" ) vla, dit, enc = cfg.vl_adapter, cfg.dit, cfg.encoders dim = int(vla.dim) if dim != dit.cross_attention_dim or dim != cfg.backbone_feature_dim: raise ValueError( f"{cfg.version}: vl_adapter.dim={dim}, dit.cross_attention_dim={dit.cross_attention_dim}, " f"backbone_feature_dim={cfg.backbone_feature_dim} must agree" ) s_pad = int(dlayout.s_pad_vl) n_vlsa = int(vla.n_blocks) if (n_vlsa > 0) != dlayout.has_vlsa or dlayout.mask_presence["vl"] != (n_vlsa > 0): raise ValueError(f"{cfg.version}: vl_adapter.n_blocks={n_vlsa} disagrees with the layout's has_vlsa") chunks: Optional[Tuple[int, int]] = None if n_vlsa > 0: if vla.head_dim % TILE: raise ValueError(f"VL-SA head_dim {vla.head_dim} is not a tile multiple (the block pads nothing)") if vla.activation not in ("gelu_tanh", "gelu", "relu", "silu"): raise ValueError(f"unsupported VL-SA activation {vla.activation!r}") try: chunks = default_sdpa_chunks("vlsa", s_pad, s_pad, int(vla.head_dim)) except KeyError: raise KeyError( f"{dlayout.layout_name}: no SDPA_CHUNKS row ('vlsa', {s_pad}, {s_pad}, {vla.head_dim}); measure " "it in test_tt_layers before using this layout (plan D6)" ) from None n_heads, hdp = int(dit.n_heads), int(dit.head_dim_pad) kv_width = 2 * n_heads * hdp # key sets (static-shape §6.2) and the per-block hoist image_runs, nonimage_runs = dlayout.gather_pieces() key_sets: Dict[str, KeySetPlan] = {} if dlayout.cross_key_set == "all": l_real, l_pad = _key_set_lengths(dlayout, "all") key_sets["all"] = KeySetPlan("all", ((0, s_pad),), l_real, l_pad, True, dlayout.mask_presence["cross_all"]) else: for name, runs in (("text", nonimage_runs), ("image", image_runs)): l_real, l_pad = _key_set_lengths(dlayout, name) if runs_length(runs) != l_real: raise AssertionError(f"{name}: runs cover {runs_length(runs)} rows, layout says {l_real}") if l_pad < l_real or l_pad % TILE: raise ValueError(f"{name}: L_pad={l_pad} must be a tile multiple >= L_real={l_real}") key_sets[name] = KeySetPlan(name, runs, l_real, l_pad, False, dlayout.mask_presence[f"cross_{name}"]) blocks: List[HoistBlockPlan] = [] for i in dit.cross_block_indices: ks = dlayout.key_set(i) if ks not in key_sets: raise KeyError(f"block {i}: key set {ks!r} not among {sorted(key_sets)}") l_pad = key_sets[ks].l_pad for dt in WEIGHT_DTYPES: # both dtypes: the policy decides at load time (plan §1.2) matmul_spec(l_pad, dim, kv_width, dt) blocks.append(HoistBlockPlan(int(i), ks, l_pad)) if not blocks: raise ValueError(f"{cfg.version}: no cross blocks") lengths = tuple(sorted({b.l_pad for b in blocks})) from models.experimental.gr00t.tt.weights import hoist_lengths # local: tt.weights is torch-only too if lengths != tuple(hoist_lengths(dlayout.layout)): raise AssertionError(f"hoist lengths {lengths} != tt.weights.hoist_lengths {hoist_lengths(dlayout.layout)}") # state encoder rows (f1 §2 enc.state.*; plan §3.3 last line) d_max, d_pad = int(enc.max_state_dim), int(enc.state_dim_pad) hid, emb = int(enc.hidden_size), int(enc.input_embedding_dim) if tuple(enc.state_encoder_dims) != (d_max, hid, emb): raise ValueError(f"{cfg.version}: state_encoder_dims {enc.state_encoder_dims} != ({d_max}, {hid}, {emb})") if d_pad < d_max or d_pad % TILE: raise ValueError(f"{cfg.version}: state_dim_pad={d_pad} must be a tile multiple >= max_state_dim={d_max}") matmul_plan(STATE_ROWS, d_pad, hid, "bf16", act="relu") matmul_plan(STATE_ROWS, hid, emb, "bf16") gather_plans: Dict[str, MatmulPlan] = {} if gather_mode == "onehot": for name, ksp in key_sets.items(): if not ksp.identity: gather_plans[name] = gather_matmul_plan(ksp.l_pad, s_pad, dim) return cls( version=cfg.version, layout_name=dlayout.layout_name, gather_mode=gather_mode, s_pad=s_pad, dim=dim, vlln_eps=float(vla.vlln_eps), n_vlsa=n_vlsa, vlsa_heads=int(vla.n_heads), vlsa_head_dim=int(vla.head_dim), vlsa_ff=int(vla.ff_dim), vlsa_eps=float(vla.ln_eps), vlsa_scale=float(vla.head_dim) ** -0.5, vlsa_activation=str(vla.activation), vlsa_chunks=chunks, n_heads=n_heads, head_dim_pad=hdp, kv_width=kv_width, key_sets=key_sets, blocks=tuple(blocks), hoist_lengths=lengths, state_dim_max=d_max, state_dim_pad=d_pad, state_hidden=hid, state_out=emb, vl_mask_present=n_vlsa > 0, gather_plans=gather_plans, ) def cross_mask_present(self, block: int) -> bool: """Whether cross block ``block`` reads a mask (fixed per layout, plan D5 / static-shape §10 item 4).""" for b in self.blocks: if b.block == block: return self.key_sets[b.key_set].mask_present raise KeyError(f"block {block} is not a cross block of {self.version}") def gathered_sets(self) -> Tuple[str, ...]: """Key sets that need a gather op (the ``"all"`` set is the identity).""" return tuple(n for n, k in self.key_sets.items() if not k.identity) def describe(self) -> Dict[str, Any]: """JSON-able summary for the results files.""" return { "version": self.version, "layout": self.layout_name, "gather_mode": self.gather_mode, "s_pad": self.s_pad, "n_vlsa": self.n_vlsa, "vlsa_chunks": list(self.vlsa_chunks) if self.vlsa_chunks else None, "key_sets": { n: {"l_real": k.l_real, "l_pad": k.l_pad, "runs": [list(r) for r in k.runs], "mask": k.mask_present} for n, k in self.key_sets.items() }, "blocks": [{"block": b.block, "key_set": b.key_set, "l_pad": b.l_pad} for b in self.blocks], "hoist_lengths": list(self.hoist_lengths), "state": {"d_max": self.state_dim_max, "d_pad": self.state_dim_pad, "rows": STATE_ROWS}, "gather_matmul": {n: {"kind": p.kind, "key": list(p.key)} for n, p in self.gather_plans.items()}, } # -------------------------------------------------------------------------------------------------------------------- # state encoder # -------------------------------------------------------------------------------------------------------------------- class StateEncoderTT: """``state_features = linear(relu(linear(state)))`` on one tile-row block (plan §2.2 ``StateEncoderTT``; synthesis §1.6 ``CategorySpecificMLP`` slot; f1 §2 ``enc.state.W1 [D_pad,1024] / b1 / W2 [1024,1536] / b2``). ``forward(state_rows [1,1,32,D_pad])`` -> ``(state_features [1,1,1,1536], rows [1,1,32,1536])``: the RELU is fused into the first matmul's program config (D3), both matmuls run HiFi4 / fp32 acc (bf16-forever class, plan §1.2), ``state_features`` is the tile-aligned ``ttnn.slice`` of row 0; ``rows`` is the whole block (rows 1..31 = ``relu(b1) @ W2 + b2``, never zero -- consumers that need zero pad rows use ``state_features``). """ def __init__(self, cfg: GR00TConfig, W: Any, *, mem: Any = L1, ckc: Any = CKC_HIFI4_FP32ACC) -> None: enc = cfg.encoders self.d_pad, self.hidden, self.out = int(enc.state_dim_pad), int(enc.hidden_size), int(enc.input_embedding_dim) self.d_max = int(enc.max_state_dim) self.mem, self.ckc = mem, ckc self.w1, self.b1 = W.get("enc.state.W1"), W.get("enc.state.b1") self.w2, self.b2 = W.get("enc.state.W2"), W.get("enc.state.b2") for name, t, want in ( ("enc.state.W1", self.w1, (self.d_pad, self.hidden)), ("enc.state.b1", self.b1, (1, self.hidden)), ("enc.state.W2", self.w2, (self.hidden, self.out)), ("enc.state.b2", self.b2, (1, self.out)), ): if _shape(t) != want: raise ValueError(f"{name}: device shape {_shape(t)} != {want}") matmul_plan(STATE_ROWS, self.d_pad, self.hidden, dtype_name(self.w1.dtype), act="relu") matmul_plan(STATE_ROWS, self.hidden, self.out, dtype_name(self.w2.dtype)) def host_rows(self, state: torch.Tensor) -> torch.Tensor: """:func:`state_rows_host` for this version's ``D_pad``.""" return state_rows_host(state, self.d_pad) def forward(self, state_rows: Any) -> Tuple[Any, Any]: import ttnn s = _shape(state_rows) if s != (1, 1, STATE_ROWS, self.d_pad): raise ValueError(f"state_rows must be [1, 1, {STATE_ROWS}, {self.d_pad}] (state_rows_host), got {s}") h = linear(state_rows, self.w1, self.b1, act="relu", mem=self.mem, ckc=self.ckc) rows = linear(h, self.w2, self.b2, mem=self.mem, ckc=self.ckc) ttnn.deallocate(h) feat = ttnn.slice(rows, [0, 0, 0, 0], [1, 1, 1, self.out]) return feat, rows __call__ = forward # -------------------------------------------------------------------------------------------------------------------- # VL adapter # -------------------------------------------------------------------------------------------------------------------- @dataclass class HoistOut: """Outputs of :meth:`VLAdapterTT.forward` (device tensors). * ``vl`` ``[1,1,S_pad_vl,2048]``: the DiT ``encoder_hidden_states`` rows in the padded slot layout (tap ``dit_encoder_hidden_states`` after :func:`slots_to_rows`); * ``vl_sets``: key set -> ``[1,1,L_pad,2048]`` gathered rows (``{"all": vl}`` for N1.5; zero pad rows); * ``kv``: cross block index -> ``(K_i, V_i)`` **persistent** ``[1,32,L_pad,64]`` buffers (the same objects on every call -- never deallocate them, :meth:`VLAdapterTT.release` does); * ``state_features`` ``[1,1,1,1536]`` / ``state_rows`` ``[1,1,32,1536]`` when ``forward`` got ``state_rows``. """ vl: Any vl_sets: Dict[str, Any] kv: Dict[int, Tuple[Any, Any]] state_features: Optional[Any] = None state_rows: Optional[Any] = None class VLAdapterTT: """``vlln`` -> VL-SA x n -> gather -> hoisted cross K/V -> state encoder (module docstring). Args: cfg, dlayout: the version config and the baked static layout (``DeviceLayout.for_inputs`` / ``from_layout``). W: ``tt.weights.TTWeights`` (or any object with ``get(name)`` and ``constants(layout)``) holding ``vlln.*``, ``vlsa.block.{i}.*``, ``dit.block.{i}.Wkv/bkv`` (even ``i``) and ``enc.state.*``. policy: ``TTPolicy`` (``kv_dtype`` must be ``"bf16"``). device: the open mesh device (persistent K/V buffers and gather constants are allocated here). gather_mode: ``"splice"`` (default, plan §3.3) | ``"onehot"``. mem: memory of the intermediates (L1); ``kv_mem``: memory of the persistent K/V (DRAM, plan §3.3). hoist_ckc: compute config of the hoist projection (HiFi4 fp32acc, see the module docstring). collect_taps: keep ``debug_taps`` (:data:`TAP_NAMES` + every ``vlsa_block{i}_out``) instead of freeing them. """ def __init__( self, cfg: GR00TConfig, dlayout: DeviceLayout, W: Any, policy: TTPolicy, device: Any, *, gather_mode: str = "splice", mem: Any = L1, kv_mem: Any = DRAM, hoist_ckc: Any = CKC_HIFI4_FP32ACC, vlsa_ckc: Any = CKC_HIFI2_FP32ACC, collect_taps: bool = False, ) -> None: if not (hasattr(W, "get") and hasattr(W, "constants")): raise TypeError("W must expose get(name) and constants(layout) (tt.weights.TTWeights)") self.plan = AdapterPlan.build(cfg, dlayout, policy, gather_mode) self.cfg, self.dlayout, self.W, self.policy, self.device = cfg, dlayout, W, policy, device self.mem, self.kv_mem, self.hoist_ckc = mem, kv_mem, hoist_ckc self.collect_taps = bool(collect_taps) self.debug_taps: Dict[str, Any] = {} p = self.plan # vlln self.w_vlln, self.b_vlln = W.get("vlln.w"), W.get("vlln.b") for name, t in (("vlln.w", self.w_vlln), ("vlln.b", self.b_vlln)): if _shape(t)[-1] != p.dim or _shape(t)[:-1] not in ((1,), (1, 1, 1)): raise ValueError(f"{name}: device shape {_shape(t)} != [1, {p.dim}]") # VL-SA blocks (plan §3.3: EncoderBlockTT(2048, 32, 64, 8192, eps 1e-5, scale 64**-0.5, key_valid mask)) self.blocks: List[EncoderBlockTT] = [] for i in range(p.n_vlsa): prefix = f"vlsa.block.{i}." weights = {n: W.get(prefix + n) for n in VLSA_WEIGHT_NAMES} self.blocks.append( EncoderBlockTT( p.dim, p.vlsa_heads, p.vlsa_head_dim, p.vlsa_ff, p.vlsa_eps, p.vlsa_scale, p.vlsa_activation, "key_valid", p.vlsa_chunks, # type: ignore[arg-type] p.s_pad, weights, family="vlsa", mem=mem, ckc=vlsa_ckc, collect_taps=False, ) ) # hoist weights: the actual dtype must have a table row at M = L_pad (plan §3.4) self.w_kv: Dict[int, Any] = {} self.b_kv: Dict[int, Any] = {} for b in p.blocks: w, bias = W.get(f"dit.block.{b.block}.Wkv"), W.get(f"dit.block.{b.block}.bkv") if _shape(w) != (p.dim, p.kv_width): raise ValueError(f"dit.block.{b.block}.Wkv: device shape {_shape(w)} != {(p.dim, p.kv_width)}") if _shape(bias)[-1] != p.kv_width: raise ValueError(f"dit.block.{b.block}.bkv: device shape {_shape(bias)} != [1, {p.kv_width}]") matmul_plan(b.l_pad, p.dim, p.kv_width, dtype_name(w.dtype)) self.w_kv[b.block], self.b_kv[b.block] = w, bias # constants: zero Q dummies (TTWeights.constants, D7), persistent K/V, gather constants consts = W.constants(dlayout) self.q_dummy: Dict[int, Any] = {} for l_pad in p.hoist_lengths: key = f"q_dummy_{l_pad}" if key not in consts: raise KeyError(f"TTWeights.constants({dlayout.layout_name!r}) lacks {key!r}") qd = consts[key] if _shape(qd) != (1, 1, l_pad, p.n_heads * p.head_dim_pad): raise ValueError(f"{key}: shape {_shape(qd)} != {(1, 1, l_pad, p.n_heads * p.head_dim_pad)}") self.q_dummy[l_pad] = qd self.K: Dict[int, Any] = {} self.V: Dict[int, Any] = {} for b in p.blocks: zeros = torch.zeros(1, p.n_heads, b.l_pad, p.head_dim_pad, dtype=torch.bfloat16) self.K[b.block] = to_device(zeros, device, dtype="bf16", mem=kv_mem) self.V[b.block] = to_device(zeros, device, dtype="bf16", mem=kv_mem) self._zero_rows_rm: Dict[str, Any] = {} self._onehot: Dict[str, Any] = {} for name in p.gathered_sets(): ksp = p.key_sets[name] if p.gather_mode == "splice": if ksp.l_pad > ksp.l_real: z = torch.zeros(1, 1, ksp.l_pad - ksp.l_real, p.dim, dtype=torch.bfloat16) self._zero_rows_rm[name] = to_device(z, device, dtype="bf16", mem=DRAM, layout="ROW_MAJOR") else: self._onehot[name] = to_device(onehot_gather_matrix(dlayout, name), device, dtype="bf16", mem=DRAM) self.state_encoder = StateEncoderTT(cfg, W, mem=mem) # --------------------------------------------------------------------------------------------------------------- def set_collect_taps(self, flag: bool) -> None: """Toggle tap collection (e.g. off before trace timing so intermediates are freed inside the capture).""" self.collect_taps = bool(flag) def _tap(self, name: str, t: Any) -> None: if self.collect_taps: self.debug_taps[name] = t def _free(self, *ts: Any) -> None: if self.collect_taps: return import ttnn for t in ts: ttnn.deallocate(t) def _gather(self, vl: Any) -> Dict[str, Any]: import ttnn p = self.plan sets: Dict[str, Any] = {} if not p.gathered_sets(): sets["all"] = vl return sets mem_cfg = resolve_mem(self.mem) if p.gather_mode == "splice": vl_rm = ttnn.to_layout(vl, ttnn.ROW_MAJOR_LAYOUT, memory_config=mem_cfg) for name in p.gathered_sets(): ksp = p.key_sets[name] pieces: List[RowPiece] = [] cursor = 0 for lo, hi in ksp.runs: pieces.append(RowPiece(cursor, cursor + (hi - lo), vl_rm, lo)) cursor += hi - lo if cursor != ksp.l_real: raise AssertionError(f"{name}: runs give {cursor} rows, expected {ksp.l_real}") if ksp.l_pad > ksp.l_real: pieces.append(RowPiece(ksp.l_real, ksp.l_pad, self._zero_rows_rm[name], 0)) g_rm = splice_rows(None, pieces, mem=self.mem, allow_unaligned=True) g = ttnn.to_layout(g_rm, ttnn.TILE_LAYOUT, memory_config=mem_cfg) if g_rm is not vl_rm and g_rm is not self._zero_rows_rm.get(name): ttnn.deallocate(g_rm) sets[name] = g ttnn.deallocate(vl_rm) else: ck = resolve_ckc(CKC_HIFI4_FP32ACC) for name in p.gathered_sets(): sets[name] = ttnn.matmul( self._onehot[name], vl, program_config=p.gather_plans[name].to_ttnn(), memory_config=mem_cfg, dtype=ttnn.bfloat16, compute_kernel_config=ck, ) for name, g in sets.items(): want = (1, 1, p.key_sets[name].l_pad, p.dim) if _shape(g) != want: raise AssertionError(f"gathered {name}: shape {_shape(g)} != {want}") return sets def forward(self, features: Any, vl_mask: Any = None, state_rows: Any = None) -> HoistOut: """``features [1,1,S_pad_vl,2048]`` (backbone output in the padded slot layout) + the per-call VL-SA mask (``layout.materialise(...)["vl"]``; required for N1.5 / N1.7, must be ``None`` for N1.6 -- presence is fixed per layout, D5) [+ ``state_rows`` from :func:`state_rows_host`] -> :class:`HoistOut`.""" import ttnn p = self.plan fs = _shape(features) if fs != (1, 1, p.s_pad, p.dim): raise ValueError(f"features must be [1, 1, {p.s_pad}, {p.dim}], got {fs}") if (vl_mask is not None) != p.vl_mask_present: raise ValueError( f"{p.version}/{p.layout_name}: vl_mask presence {vl_mask is not None} != fixed {p.vl_mask_present}" ) self.debug_taps = {} x = layer_norm(features, self.w_vlln, self.b_vlln, eps=p.vlln_eps, mem=self.mem) self._tap("vlln_out", x) for i, blk in enumerate(self.blocks): y = blk.forward(x, vl_mask) self._free(x) x = y self._tap(f"vlsa_block{i}_out", x) if self.blocks: self._tap("vl_self_attention_out", x) self._tap("dit_encoder_hidden_states", x) vl = x vl_sets = self._gather(vl) kv: Dict[int, Tuple[Any, Any]] = {} for b in p.blocks: proj = linear(vl_sets[b.key_set], self.w_kv[b.block], self.b_kv[b.block], mem=self.mem, ckc=self.hoist_ckc) k, v = split_kv_heads(proj, p.n_heads, self.q_dummy[b.l_pad], mem=self.kv_mem) ttnn.copy(k, self.K[b.block]) ttnn.copy(v, self.V[b.block]) ttnn.deallocate(k) ttnn.deallocate(v) ttnn.deallocate(proj) kv[b.block] = (self.K[b.block], self.V[b.block]) state_features = state_block = None if state_rows is not None: state_features, state_block = self.state_encoder.forward(state_rows) self._tap("state_features", state_features) return HoistOut(vl, vl_sets, kv, state_features, state_block) __call__ = forward # --------------------------------------------------------------------------------------------------------------- def free_outputs(self, out: HoistOut) -> None: """Deallocate the per-call outputs of ``out`` (``vl``, gathered sets, state tensors) and the intermediates kept by ``collect_taps`` -- never the persistent K/V; every object is freed once (``vl`` is also a tap).""" import ttnn seen: List[int] = [] for t in [out.vl, *out.vl_sets.values(), out.state_features, out.state_rows, *self.debug_taps.values()]: if t is None or id(t) in seen: continue seen.append(id(t)) ttnn.deallocate(t) self.debug_taps = {} def release(self) -> None: """Deallocate the persistent K/V buffers and the gather constants (the weights belong to ``W``).""" import ttnn for d in (self.K, self.V, self._zero_rows_rm, self._onehot): for t in d.values(): ttnn.deallocate(t) d.clear() def describe(self) -> Dict[str, Any]: d = self.plan.describe() d["hoist_ckc"] = getattr(self.hoist_ckc, "name", str(self.hoist_ckc)) d["blocks_vlsa"] = [b.describe() for b in self.blocks] return d __all__ = [ "GATHER_MODES", "GATHER_ONEHOT_IBW_MAX", "KEY_SETS", "STATE_ROWS", "TAP_NAMES", "VLSA_WEIGHT_NAMES", "AdapterPlan", "HoistBlockPlan", "HoistOut", "KeySetPlan", "StateEncoderTT", "VLAdapterTT", "gather_index", "gather_matmul_plan", "onehot_gather_matrix", "rows_to_slots", "slots_to_rows", "state_rows_host", ]