# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """``EncoderBlockTT`` -- the unified pre-LN encoder block of SigLIP (N1.5 / N1.6 tower) and the VL self-attention adapter (N1.5 / N1.7), built only from :mod:`tt.layers` primitives (IMPLEMENTATION_PLAN.md §2.2 ``tt/blocks.py``, §3.2 / §3.3 op tables, graft "C: unified EncoderBlockTT"). x -> LN(ln1) -> linear(Wqkv, bqkv) -> split_heads -> sdpa(scale, [mask]) -> concat_heads -> linear(Wo, bo) -> +x -> LN(ln2) -> linear(ff1, act fused) -> linear(ff2) -> +x Weights are the plan tensors of one block (``common.weights`` names without the ``vision.layer.{i}.`` / ``vlsa.block.{i}.`` prefix, or with it when ``prefix`` is given): * SigLIP family: ``ln1.w, ln1.b, Wqkv, bqkv, Wo, bo, ln2.w, ln2.b, fc1.W, fc1.b, fc2.W, fc2.b`` (heads padded 72 -> 96: ``Wqkv [1152, 4608]``, ``Wo [1536, 1152]``; ``eps 1e-6``; ``scale = 72 ** -0.5``; S = 256 / 352). * VL-SA family: ``norm1.w, norm1.b, Wqkv, bqkv, Wo, bo, norm3.w, norm3.b, ff1.W, ff1.b, ff2.W, ff2.b`` (``[2048, 6144]``, 32 x 64, ``ff 8192``, ``eps 1e-5``, ``scale = 64 ** -0.5``, key-valid mask over ``S_pad``). ``ttnn`` is imported lazily inside ``forward``; the constructor only validates shapes (duck-typed ``.shape``), so the shape contract is CPU-testable. Head-op variants (WP-P, ``benchmarks/sweep_backbone.py --what heads``; ``tests/tt/results/p_summary.md``): the fused-QKV split and the heads concat are selectable per block (``split_impl`` / ``concat_impl``) so the measured winner can be applied without touching the op vocabulary of ``tt/layers.py``: ``"nlp"`` = ``nlp_create_qkv_heads`` / ``nlp_concat_heads`` (plan D7, bit-exact), ``"slices"`` = slice + reshape + permute per tensor; ``"permute"`` = permute + reshape, ``"concatenate_heads"`` = ``ttnn.transformer.concatenate_heads``. (``ttnn.experimental. split_query_key_value_and_split_heads`` is the BERT-large TM op -- batch 7..9 only, MHA only -- and does not apply.) :data:`DEFAULT_SPLIT_IMPL` / :data:`DEFAULT_CONCAT_IMPL` hold the winners. """ from __future__ import annotations from typing import Any, Dict, Mapping, Optional, Sequence, Tuple from models.experimental.gr00t.tt.device import CKC_HIFI2_FP32ACC, CKC_SDPA, L1 from models.experimental.gr00t.tt.layers import ( ACTIVATIONS, concat_heads, layer_norm, linear, residual_add, sdpa, sdpa_chunk_ok, sdpa_chunks_allowed, split_heads, ) MASK_MODES: Tuple[Optional[str], ...] = (None, "key_valid") #: How a **constant** key-valid mask is applied (WP-P, p_summary.md §4 / §9): ``"additive"`` = the ``[1,1,S,S]`` additive #: mask handed to the SDPA (D5; the mask itself costs ~40 µs on the N1.6 SigLIP block, D0 §3.4 note); ``"key_bias"`` = #: no SDPA mask -- the zero pad channel of the padded heads (SigLIP 72 -> 96) carries ``+1`` in every Q row and ``-1e4`` #: in the pad-key rows of K (two constant adds), so pad keys score ``-1e4 * scale`` and get softmax weight exactly 0 #: while real keys are untouched (their pad channel is 0). Only for constructor-constant masks on padded heads. #: ``"auto"`` = ``key_bias`` when it applies (constructor-constant pure key mask, padded heads), else ``additive``. #: Measured (p_summary.md §4b, N1.6 SigLIP block 352 rows / 324 valid): 389.0 vs 414.2 µs, outputs **bit-exact**. MASK_IMPLS: Tuple[str, ...] = ("auto", "additive", "key_bias") DEFAULT_MASK_IMPL = "auto" KEY_BIAS_VALUE = -1.0e4 SPLIT_IMPLS: Tuple[str, ...] = ("nlp", "slices") CONCAT_IMPLS: Tuple[str, ...] = ("nlp", "permute", "concatenate_heads") #: WP-P winners (sweep_backbone --what heads, p_summary.md §heads): ``nlp_create_qkv_heads`` from L1 and #: ``nlp_concat_heads`` stay the fastest exact variants at every backbone shape (D0 §B agreed at the DiT shape). DEFAULT_SPLIT_IMPL = "nlp" DEFAULT_CONCAT_IMPL = "nlp" # family -> (ln1, ln2, ff1, ff2) tensor-name stems (plan §2.2 blocks row; common.weights plan names) WEIGHT_FAMILIES: Dict[str, Tuple[str, str, str, str]] = { "siglip": ("ln1", "ln2", "fc1", "fc2"), "vlsa": ("norm1", "norm3", "ff1", "ff2"), } def _shape(t: Any) -> Tuple[int, ...]: return tuple(int(d) for d in t.shape) def _vec_ok(shape: Tuple[int, ...], n: int) -> bool: if shape[-1] != n: return False rows = 1 for d in shape[:-1]: rows *= d return rows == 1 def _mat_ok(shape: Tuple[int, ...], k: int, n: int) -> bool: if len(shape) == 4 and shape[0] == 1 and shape[1] == 1: shape = shape[2:] return shape == (k, n) def split_heads_impl(qkv: Any, n_heads: int, n_kv_heads: int, impl: str, mem: Any) -> Tuple[Any, Any, Any]: """Fused ``[B, 1, S, (H + 2 Hkv) d]`` -> ``Q, K, V`` heads by ``impl`` (:data:`SPLIT_IMPLS`; all bit-exact).""" if impl == "nlp": return split_heads(qkv, n_heads, n_kv_heads, mem=mem) import ttnn from models.experimental.gr00t.tt.device import resolve_mem mc = resolve_mem(mem) b, _, s, width = _shape(qkv) d = width // (n_heads + 2 * n_kv_heads) if impl == "slices": outs = [] off = 0 for nh in (n_heads, n_kv_heads, n_kv_heads): part = ttnn.slice(qkv, [0, 0, 0, off], [b, 1, s, off + nh * d], memory_config=mc) off += nh * d outs.append(ttnn.permute(ttnn.reshape(part, (b, s, nh, d)), (0, 2, 1, 3), memory_config=mc)) ttnn.deallocate(part) return outs[0], outs[1], outs[2] raise ValueError(f"split_impl must be one of {SPLIT_IMPLS}, got {impl!r}") def concat_heads_impl(o: Any, impl: str, mem: Any) -> Any: """``[B, H, S, d]`` -> ``[B, 1, S, H d]`` by ``impl`` (:data:`CONCAT_IMPLS`; all bit-exact).""" if impl == "nlp": return concat_heads(o, mem=mem) import ttnn from models.experimental.gr00t.tt.device import resolve_mem mc = resolve_mem(mem) b, h, s, d = _shape(o) if impl == "permute": p = ttnn.permute(o, (0, 2, 1, 3), memory_config=mc) return ttnn.reshape(p, (b, 1, s, h * d)) if impl == "concatenate_heads": return ttnn.transformer.concatenate_heads(o, memory_config=mc) raise ValueError(f"concat_impl must be one of {CONCAT_IMPLS}, got {impl!r}") class EncoderBlockTT: """One encoder block; ``forward(x, mask=None)`` maps ``[B, 1, m_rows, dim]`` -> the same shape. Args: dim: residual width (1152 SigLIP / 2048 VL-SA). n_heads: attention heads (16 / 32). head_dim_pad: padded per-head width on device (96 / 64); ``Wqkv`` has ``3 * n_heads * head_dim_pad`` columns and ``Wo`` has ``n_heads * head_dim_pad`` rows. ff_dim: FF hidden width (4304 / 8192). ln_eps: LayerNorm epsilon of both norms (1e-6 / 1e-5). scale: explicit attention scale (``real_head_dim ** -0.5``: the kernel default would use the padded dim). activation: fused FF activation (``"gelu_tanh"`` for SigLIP, Qwen3-VL ViT and VL-SA). mask_mode: ``None`` (no mask; N1.5 SigLIP with 256 real patches) or ``"key_valid"`` (an additive ``[1, 1, m_rows, m_rows]`` mask is mandatory -- N1.6 SigLIP 324 logical rows in 352, VL-SA key-valid). sdpa_chunks: ``(q_chunk, k_chunk)``; validated against ``tt.layers.SDPA_CHUNKS`` for ``(m_rows, m_rows, head_dim_pad)`` at construction. m_rows: rows per batch element (padded sequence). weights: the block's device tensors (see module docstring); ``prefix`` is stripped from the names. family: ``"siglip"`` / ``"vlsa"`` / ``None`` (auto-detected from the ``ln1`` vs ``norm1`` keys). mask: an optional constant mask kept by the block (N1.6 SigLIP); ``forward(mask=)`` overrides per call. collect_taps: keep ``debug_taps`` (attention output before ``Wo``, residual after attention, FF hidden, ...) instead of deallocating the intermediates. split_impl / concat_impl: head-op variants (module docstring; defaults :data:`DEFAULT_SPLIT_IMPL` / :data:`DEFAULT_CONCAT_IMPL`). mask_impl: ``"auto"`` | ``"additive"`` | ``"key_bias"`` (:data:`MASK_IMPLS`) for a constructor-constant ``mask``; ``"key_bias"`` needs ``head_dim_pad > real head dim`` (``scale ** -2``), a pure key mask and forbids a per-call ``forward(mask=)``; ``"auto"`` picks it when those hold, else ``"additive"`` (resolved value in ``mask_impl``). device: the device for the constant bias tensors (``"key_bias"`` / ``"auto"``; default: the mask's device). """ TAP_NAMES: Tuple[str, ...] = ("ln1", "qkv", "attn", "attn_proj", "res1", "ln2", "ff_hidden", "ff_out", "out") def __init__( self, dim: int, n_heads: int, head_dim_pad: int, ff_dim: int, ln_eps: float, scale: float, activation: str = "gelu_tanh", mask_mode: Optional[str] = None, sdpa_chunks: Sequence[int] = (64, 128), m_rows: int = 256, weights: Optional[Mapping[str, Any]] = None, *, prefix: str = "", family: Optional[str] = None, mask: Any = None, mem: Any = L1, ckc: Any = CKC_HIFI2_FP32ACC, sdpa_ckc: Any = CKC_SDPA, collect_taps: bool = False, split_impl: str = DEFAULT_SPLIT_IMPL, concat_impl: str = DEFAULT_CONCAT_IMPL, mask_impl: str = DEFAULT_MASK_IMPL, device: Any = None, ) -> None: if weights is None: raise ValueError("EncoderBlockTT needs the block's weight dict") if split_impl not in SPLIT_IMPLS: raise ValueError(f"split_impl must be one of {SPLIT_IMPLS}, got {split_impl!r}") if concat_impl not in CONCAT_IMPLS: raise ValueError(f"concat_impl must be one of {CONCAT_IMPLS}, got {concat_impl!r}") if mask_impl not in MASK_IMPLS: raise ValueError(f"mask_impl must be one of {MASK_IMPLS}, got {mask_impl!r}") self.split_impl, self.concat_impl, self.mask_impl = split_impl, concat_impl, mask_impl self.q_bias: Any = None self.k_bias: Any = None if activation not in ACTIVATIONS: raise ValueError(f"activation must be one of {ACTIVATIONS}, got {activation!r}") if mask_mode not in MASK_MODES: raise ValueError(f"mask_mode must be one of {MASK_MODES}, got {mask_mode!r}") if scale is None or scale <= 0: raise ValueError("scale must be an explicit positive float (plan D5)") if dim <= 0 or n_heads <= 0 or head_dim_pad <= 0 or ff_dim <= 0 or m_rows <= 0: raise ValueError("dim, n_heads, head_dim_pad, ff_dim and m_rows must be positive") if head_dim_pad % 32: raise ValueError(f"head_dim_pad must be a tile multiple, got {head_dim_pad}") if m_rows % 32: raise ValueError(f"m_rows must be a tile multiple, got {m_rows}") chunks = (int(sdpa_chunks[0]), int(sdpa_chunks[1])) if len(sdpa_chunks) == 2 else None if chunks is None or not sdpa_chunk_ok(m_rows, m_rows, head_dim_pad, chunks): try: allowed: Any = sdpa_chunks_allowed(m_rows, m_rows, head_dim_pad) except KeyError: allowed = "none (shape not in SDPA_CHUNKS)" raise ValueError( f"sdpa_chunks {tuple(sdpa_chunks)} not allowed for (Sq=Sk={m_rows}, d={head_dim_pad}); allowed: {allowed}" ) self.dim, self.n_heads, self.head_dim_pad, self.ff_dim = int(dim), int(n_heads), int(head_dim_pad), int(ff_dim) self.ln_eps, self.scale, self.activation = float(ln_eps), float(scale), activation self.mask_mode, self.sdpa_chunks, self.m_rows = mask_mode, chunks, int(m_rows) self.mem, self.ckc, self.sdpa_ckc = mem, ckc, sdpa_ckc self.collect_taps = collect_taps self.debug_taps: Dict[str, Any] = {} names = {k[len(prefix) :]: v for k, v in weights.items() if k.startswith(prefix)} if prefix else dict(weights) if family is None: has_siglip = "ln1.w" in names has_vlsa = "norm1.w" in names if has_siglip == has_vlsa: raise KeyError( "cannot detect the weight family: expected exactly one of 'ln1.w' (siglip) / 'norm1.w' (vlsa) " f"in {sorted(names)[:12]}..." ) family = "siglip" if has_siglip else "vlsa" if family not in WEIGHT_FAMILIES: raise KeyError(f"family must be one of {sorted(WEIGHT_FAMILIES)}, got {family!r}") self.family = family ln1, ln2, ff1, ff2 = WEIGHT_FAMILIES[family] inner = self.n_heads * self.head_dim_pad required: Dict[str, Tuple[str, Tuple[int, ...]]] = { f"{ln1}.w": ("vec", (self.dim,)), f"{ln1}.b": ("vec", (self.dim,)), "Wqkv": ("mat", (self.dim, 3 * inner)), "bqkv": ("vec", (3 * inner,)), "Wo": ("mat", (inner, self.dim)), "bo": ("vec", (self.dim,)), f"{ln2}.w": ("vec", (self.dim,)), f"{ln2}.b": ("vec", (self.dim,)), f"{ff1}.W": ("mat", (self.dim, self.ff_dim)), f"{ff1}.b": ("vec", (self.ff_dim,)), f"{ff2}.W": ("mat", (self.ff_dim, self.dim)), f"{ff2}.b": ("vec", (self.dim,)), } missing = [k for k in required if k not in names] if missing: raise KeyError(f"EncoderBlockTT({family}) is missing weights {missing}") for key, (kind, dims) in required.items(): s = _shape(names[key]) ok = _vec_ok(s, dims[0]) if kind == "vec" else _mat_ok(s, dims[0], dims[1]) if not ok: raise ValueError(f"weight {key!r} has shape {s}, expected {kind} {dims}") self.w: Dict[str, Any] = { "ln1.w": names[f"{ln1}.w"], "ln1.b": names[f"{ln1}.b"], "Wqkv": names["Wqkv"], "bqkv": names["bqkv"], "Wo": names["Wo"], "bo": names["bo"], "ln2.w": names[f"{ln2}.w"], "ln2.b": names[f"{ln2}.b"], "ff1.W": names[f"{ff1}.W"], "ff1.b": names[f"{ff1}.b"], "ff2.W": names[f"{ff2}.W"], "ff2.b": names[f"{ff2}.b"], } self.mask = None if mask is not None: self._check_mask(mask) self.mask = mask if self.mask_impl == "auto": self.mask_impl = "key_bias" if self._key_bias_applies(mask) else "additive" if self.mask_impl == "key_bias": self._build_key_bias(mask, device) elif self.mask_impl == "key_bias": raise ValueError("mask_impl='key_bias' needs the constant mask at construction") else: self.mask_impl = "additive" # per-call masks (VL-SA) are always additive # --------------------------------------------------------------------------------------------------------------- def _key_bias_applies(self, mask: Any) -> bool: """``"auto"`` rule: a zero pad channel exists (``scale ** -2 < head_dim_pad``) and the mask is a pure key mask.""" import torch from models.experimental.gr00t.tt.layers import to_torch real_d = int(round(self.scale**-2)) if real_d >= self.head_dim_pad: return False m = to_torch(mask) if not isinstance(mask, torch.Tensor) else mask.float() m = m.reshape(-1, self.m_rows, self.m_rows)[0] key_pad = m[0] < 0 return bool(torch.equal((m < 0), key_pad.unsqueeze(0).expand(self.m_rows, -1))) def _build_key_bias(self, mask: Any, device: Any) -> None: """Constant ``q_bias [1,1,1,d_pad]`` (+1 on the pad channel) and ``k_bias [1,1,S,d_pad]`` (``KEY_BIAS_VALUE`` on the pad channel of the masked key rows) from the additive key mask (host readback once; module docstring).""" import torch from models.experimental.gr00t.tt.layers import to_device, to_torch real_d = int(round(self.scale**-2)) if real_d >= self.head_dim_pad: raise ValueError( f"mask_impl='key_bias' needs a zero pad channel (real head dim {real_d} vs padded {self.head_dim_pad})" ) m = to_torch(mask) if not isinstance(mask, torch.Tensor) else mask.float() m = m.reshape(-1, self.m_rows, self.m_rows)[0] key_pad = m[0] < 0 if not torch.equal((m < 0), key_pad.unsqueeze(0).expand(self.m_rows, -1)): raise ValueError("mask_impl='key_bias' needs a pure key mask (every query row masks the same keys)") q_bias = torch.zeros(1, 1, 1, self.head_dim_pad) q_bias[..., real_d] = 1.0 k_bias = torch.zeros(1, 1, self.m_rows, self.head_dim_pad) k_bias[0, 0, key_pad, real_d] = KEY_BIAS_VALUE dev = device if device is not None else getattr(mask, "device", None) if dev is None: raise ValueError("mask_impl='key_bias' needs the device (pass device=)") if callable(dev): dev = dev() self.q_bias = to_device(q_bias, dev, dtype="bf16", mem="DRAM") self.k_bias = to_device(k_bias, dev, dtype="bf16", mem="DRAM") self.n_pad_keys = int(key_pad.sum()) def _apply_key_bias(self, q: Any, k: Any) -> Tuple[Any, Any]: import ttnn from models.experimental.gr00t.tt.device import resolve_mem mc = resolve_mem(self.mem) qb = ttnn.add(q, self.q_bias, memory_config=mc) kb = ttnn.add(k, self.k_bias, memory_config=mc) self._free(q, k) return qb, kb # --------------------------------------------------------------------------------------------------------------- def _check_mask(self, mask: Any) -> None: if self.mask_mode is None: raise ValueError("a mask was given but mask_mode is None (mask presence is fixed per block, plan D5)") ms = _shape(mask) if len(ms) != 4 or ms[2] != self.m_rows or ms[3] != self.m_rows: raise ValueError(f"mask must be [1|B, 1|H, {self.m_rows}, {self.m_rows}], got {ms}") 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 forward(self, x: Any, mask: Any = None) -> Any: """Run the block on ``x [B, 1, m_rows, dim]`` (bf16 TILE); returns a new tensor of the same shape.""" xs = _shape(x) if len(xs) != 4 or xs[1] != 1 or xs[2] != self.m_rows or xs[3] != self.dim: raise ValueError(f"x must be [B, 1, {self.m_rows}, {self.dim}], got {xs}") if self.mask_impl == "key_bias": if mask is not None and mask is not self.mask: raise ValueError("mask_impl='key_bias' applies the constructor mask; no per-call mask override") mask = None elif mask is None: mask = self.mask if self.mask_mode == "key_valid" and mask is None and self.mask_impl != "key_bias": raise ValueError("mask_mode='key_valid' requires a mask (constructor or forward)") if mask is not None: self._check_mask(mask) self.debug_taps = {} w = self.w h = layer_norm(x, w["ln1.w"], w["ln1.b"], eps=self.ln_eps, mem=self.mem) self._tap("ln1", h) qkv = linear(h, w["Wqkv"], w["bqkv"], mem=self.mem, ckc=self.ckc) self._free(h) self._tap("qkv", qkv) q, k, v = split_heads_impl(qkv, self.n_heads, self.n_heads, self.split_impl, self.mem) self._free(qkv) if self.mask_impl == "key_bias": q, k = self._apply_key_bias(q, k) o = sdpa(q, k, v, scale=self.scale, mask=mask, is_causal=False, chunks=self.sdpa_chunks, ckc=self.sdpa_ckc) self._free(q, k, v) a = concat_heads_impl(o, self.concat_impl, self.mem) self._free(o) self._tap("attn", a) p = linear(a, w["Wo"], w["bo"], mem=self.mem, ckc=self.ckc) self._free(a) self._tap("attn_proj", p) x1 = residual_add(x, p, mem=self.mem) self._free(p) self._tap("res1", x1) h2 = layer_norm(x1, w["ln2.w"], w["ln2.b"], eps=self.ln_eps, mem=self.mem) self._tap("ln2", h2) f = linear(h2, w["ff1.W"], w["ff1.b"], act=self.activation, mem=self.mem, ckc=self.ckc) self._free(h2) self._tap("ff_hidden", f) g = linear(f, w["ff2.W"], w["ff2.b"], mem=self.mem, ckc=self.ckc) self._free(f) self._tap("ff_out", g) out = residual_add(x1, g, mem=self.mem) self._free(g, x1) self._tap("out", out) return out __call__ = forward def describe(self) -> Dict[str, Any]: """JSON-able summary (used by the benchmarks / results files).""" return { "family": self.family, "dim": self.dim, "n_heads": self.n_heads, "head_dim_pad": self.head_dim_pad, "ff_dim": self.ff_dim, "ln_eps": self.ln_eps, "scale": self.scale, "activation": self.activation, "mask_mode": self.mask_mode, "sdpa_chunks": list(self.sdpa_chunks), "m_rows": self.m_rows, "split_impl": self.split_impl, "concat_impl": self.concat_impl, "mask_impl": self.mask_impl, } __all__ = [ "CONCAT_IMPLS", "DEFAULT_CONCAT_IMPL", "DEFAULT_MASK_IMPL", "DEFAULT_SPLIT_IMPL", "EncoderBlockTT", "KEY_BIAS_VALUE", "MASK_IMPLS", "MASK_MODES", "SPLIT_IMPLS", "WEIGHT_FAMILIES", "concat_heads_impl", "split_heads_impl", ]