# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 import json import math from dataclasses import asdict, dataclass, replace from pathlib import Path import torch import ttnn from models.common.lightweightmodule import LightweightModule TILE = 32 SEQ_ALIGN = 128 SEQ_ALIGN_LARGE = 256 LARGE_PATCHES = 2048 CU_WINDOW_ENTRIES = TILE * TILE MASK_NEG = -1e30 VISION_PREFIX = "vision_tower." EMBEDDER_PREFIX = "embed_vision." @dataclass(frozen=True) class VisionConfig: hidden_size: int intermediate_size: int num_hidden_layers: int num_attention_heads: int head_dim: int patch_size: int pooling_kernel_size: int position_embedding_size: int rope_theta: float rms_norm_eps: float text_hidden_size: int @classmethod def from_hf(cls, hf_config): if isinstance(hf_config, VisionConfig): return hf_config if isinstance(hf_config, (str, Path)): hf_config = json.loads(Path(hf_config).read_text()) if not isinstance(hf_config, dict): hf_config = hf_config.to_dict() v = hf_config["vision_config"] if v["rope_parameters"]["rope_type"] != "axial": raise ValueError(f"unsupported vision rope {v['rope_parameters']}") if v.get("use_clipped_linears") or v.get("standardize"): raise ValueError("clipped linears and standardize are not part of this checkpoint") if v["num_key_value_heads"] != v["num_attention_heads"]: raise ValueError("vision attention is MHA") return cls( hidden_size=v["hidden_size"], intermediate_size=v["intermediate_size"], num_hidden_layers=v["num_hidden_layers"], num_attention_heads=v["num_attention_heads"], head_dim=v["head_dim"], patch_size=v["patch_size"], pooling_kernel_size=v["pooling_kernel_size"], position_embedding_size=v["position_embedding_size"], rope_theta=float(v["rope_parameters"]["rope_theta"]), rms_norm_eps=v["rms_norm_eps"], text_hidden_size=hf_config["text_config"]["hidden_size"], ) WEIGHT_GROUPS = ("qkv", "o", "gate", "up", "down") DTYPES = {"bf16": ttnn.bfloat16, "bfp8": ttnn.bfloat8_b, "bfp4": ttnn.bfloat4_b, "fp32": ttnn.float32} FIDELITIES = {"lofi": ttnn.MathFidelity.LoFi, "hifi2": ttnn.MathFidelity.HiFi2, "hifi4": ttnn.MathFidelity.HiFi4} @dataclass(frozen=True) class VisionPolicy: weight_dtype: tuple = (("qkv", "bf16"), ("o", "bf16"), ("gate", "bf16"), ("up", "bf16"), ("down", "bf16")) weight_fidelity: tuple = (("qkv", "hifi2"), ("o", "hifi2"), ("gate", "hifi2"), ("up", "hifi2"), ("down", "hifi2")) residual: str = "fp32" sdpa_fidelity: str = "hifi4" sdpa_fp32_acc: bool = True sdpa_exp_approx: bool = True single_image_mask: bool = True mask_patches: tuple = (640, 1280, 5120) mask_dtype: str = "bfp4" windowed_chunks: tuple = (128, 256) windowed_chunks_short: tuple = (64, 128) short_patches: int = 2048 mask_chunks: tuple = (128, 256) minimal_matmul: bool = True gelu_in_matmul: bool = False packed_gate_up: bool = True l1_max_rows: int = 1280 l1_wide_max_rows: int = 640 attention_impl: str = "sdpa" def dtype(self, group): return DTYPES[dict(self.weight_dtype)[group]] def fidelity(self, group): return FIDELITIES[dict(self.weight_fidelity)[group]] def with_weights(self, dtype=None, fidelity=None, groups=WEIGHT_GROUPS): wd, wf = dict(self.weight_dtype), dict(self.weight_fidelity) for g in groups: if dtype is not None: wd[g] = dtype if fidelity is not None: wf[g] = fidelity return replace( self, weight_dtype=tuple((g, wd[g]) for g in WEIGHT_GROUPS), weight_fidelity=tuple((g, wf[g]) for g in WEIGHT_GROUPS), ) def describe(self): d = asdict(self) d["weight_dtype"] = dict(self.weight_dtype) d["weight_fidelity"] = dict(self.weight_fidelity) return d def head_channel_order(head_dim): q = head_dim // 4 return torch.cat( [torch.arange(0, q), torch.arange(2 * q, 3 * q), torch.arange(q, 2 * q), torch.arange(3 * q, 4 * q)] ) def axial_inv_freq(config): spatial = config.head_dim // 2 return 1.0 / (config.rope_theta ** (torch.arange(0, spatial, 2, dtype=torch.float) / spatial)) def rope_tables(config): inv = axial_inv_freq(config) n = config.position_embedding_size ang = torch.arange(n, dtype=torch.float32)[:, None] * inv[None, :] q = inv.numel() zero = torch.zeros(n, q) tables = [] for axis in range(2): parts = [] for fn in (torch.cos, torch.sin): v = fn(ang) parts.append(torch.cat([v, zero, v, zero] if axis == 0 else [zero, v, zero, v], dim=-1)) t = torch.cat(parts, dim=-1) tables.append(torch.cat([t, torch.zeros(1, t.shape[-1])], dim=0)) return tables def padded_patches(num_patches): align = SEQ_ALIGN_LARGE if num_patches >= LARGE_PATCHES else SEQ_ALIGN return align * math.ceil(num_patches / align) def padded_slots(num_slots): return TILE * math.ceil(num_slots / TILE) @dataclass class VisionHostInputs: pixels: torch.Tensor x_idx: torch.Tensor y_idx: torch.Tensor cu_window: torch.Tensor pool: torch.Tensor gather_idx: torch.Tensor key_bias: torch.Tensor num_images: int num_patches: int patches_padded: int slots: int slots_padded: int valid_patches: tuple tokens_per_image: tuple tokens_padded: int @dataclass class VisionDeviceInputs: pixels: ttnn.Tensor x_idx: ttnn.Tensor y_idx: ttnn.Tensor cu_window: ttnn.Tensor pool: ttnn.Tensor gather_idx: ttnn.Tensor key_bias: ttnn.Tensor host: VisionHostInputs def prepare_host_inputs(config, pixel_values, position_ids, tokens_padded=None): n, p, d = pixel_values.shape k = config.pooling_kernel_size slots = p // (k * k) if slots * k * k != p: raise ValueError(f"patch count {p} is not a multiple of {k * k}") pp = padded_patches(p) sp = padded_slots(slots) pad_row = config.position_embedding_size pixels = torch.zeros(1, 1, n * pp, d) x_idx = torch.full((n, pp), pad_row, dtype=torch.int32) y_idx = torch.full((n, pp), pad_row, dtype=torch.int32) pool = torch.zeros(n, sp, pp) key_bias = torch.full((n, 1, TILE, pp), MASK_NEG) scale = config.hidden_size**0.5 / (k * k) cu = [0] valid_counts, token_counts, gather = [], [], [] for i in range(n): pos = position_ids[i].long() valid = (pos != -1).all(dim=-1) nv = int(valid.sum()) if not bool(valid[:nv].all()): raise ValueError("padding patches must follow valid patches") pixels[0, 0, i * pp : i * pp + p] = pixel_values[i].float() key_bias[i, 0, :, :nv] = 0.0 x_idx[i, :nv] = pos[:nv, 0].int() y_idx[i, :nv] = pos[:nv, 1].int() xs, ys = pos[:nv, 0], pos[:nv, 1] max_x = int(xs.max()) + 1 slot = torch.div(xs, k, rounding_mode="floor") + (max_x // k) * torch.div(ys, k, rounding_mode="floor") pool[i, slot, torch.arange(nv)] = scale used = torch.unique(slot).tolist() gather.extend(i * sp + s for s in used) valid_counts.append(nv) token_counts.append(len(used)) cu.extend([i * pp + nv, (i + 1) * pp]) total = len(gather) tp = tokens_padded if tokens_padded is not None else padded_slots(n * slots) if total > tp: raise ValueError(f"{total} soft tokens exceed tokens_padded {tp}") gather_idx = torch.full((1, tp), n * sp, dtype=torch.int32) gather_idx[0, :total] = torch.tensor(gather, dtype=torch.int32) if len(cu) > CU_WINDOW_ENTRIES: raise ValueError(f"{n} images need {len(cu)} window boundaries, limit {CU_WINDOW_ENTRIES}") cu = cu + [cu[-1]] * (CU_WINDOW_ENTRIES - len(cu)) return VisionHostInputs( pixels=pixels, x_idx=x_idx.reshape(1, n * pp), y_idx=y_idx.reshape(1, n * pp), cu_window=torch.tensor(cu, dtype=torch.int32), pool=pool, gather_idx=gather_idx, key_bias=key_bias, num_images=n, num_patches=p, patches_padded=pp, slots=slots, slots_padded=sp, valid_patches=tuple(valid_counts), tokens_per_image=tuple(token_counts), tokens_padded=tp, ) INPUT_SPEC = { "pixels": (ttnn.bfloat16, ttnn.TILE_LAYOUT), "x_idx": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), "y_idx": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), "cu_window": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), "pool": (ttnn.bfloat16, ttnn.TILE_LAYOUT), "gather_idx": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), "key_bias": (ttnn.bfloat16, ttnn.TILE_LAYOUT), } def host_inputs_to_device(host, mesh_device): tensors = { name: ttnn.from_torch( getattr(host, name).contiguous(), dtype=dt, layout=layout, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) for name, (dt, layout) in INPUT_SPEC.items() } return VisionDeviceInputs(host=host, **tensors) def copy_host_inputs(host, device_inputs): old = device_inputs.host for f in ("num_images", "patches_padded", "slots_padded", "tokens_padded"): if getattr(old, f) != getattr(host, f): raise ValueError(f"{f} differs from the captured bucket: {getattr(host, f)} vs {getattr(old, f)}") for name, (dt, layout) in INPUT_SPEC.items(): src = ttnn.from_torch(getattr(host, name).contiguous(), dtype=dt, layout=layout) ttnn.copy_host_to_device_tensor(src, getattr(device_inputs, name)) device_inputs.host = host def _compute_config(mesh_device, fidelity=ttnn.MathFidelity.HiFi4, fp32_acc=True): return ttnn.init_device_compute_kernel_config( mesh_device.arch(), math_fidelity=fidelity, math_approx_mode=False, fp32_dest_acc_en=fp32_acc, packer_l1_acc=True, ) def _put_weight(t, mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT): return ttnn.from_torch( t.contiguous(), dtype=dtype, layout=layout, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG ) def _norm_weight(w, mesh_device): return _put_weight(w.float().reshape(1, 1, w.numel() // TILE, TILE), mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT) MATMUL_BLOCKS = { "qkv": ((8, 4, 8, 2, 2), (16, 4, 8, 1, 4)), "o": ((8, 4, 8, 1, 4), (16, 8, 4, 1, 4)), "gate": ((8, 8, 12, 2, 2), (16, 8, 4, 1, 4)), "up": ((8, 8, 12, 2, 2), (16, 8, 4, 1, 4)), "down": ((8, 4, 8, 1, 4), (8, 12, 8, 1, 4)), "gate_up": ((4, 4, 24, 2, 2), (16, 8, 4, 1, 4)), } LARGE_M = 8192 def matmul_config(group, m, grid): mb, kb, nb, sh, sw = MATMUL_BLOCKS[group][1 if m >= LARGE_M else 0] return ttnn.MinimalMatmulConfig( M_block_size=mb, K_block_size=kb, N_block_size=nb, subblock_h=sh, subblock_w=sw, compute_with_storage_grid_size=grid, ) @dataclass class AttentionPlan: mode: str l1_activations: bool = False cu_window: ttnn.Tensor = None mask: ttnn.Tensor = None program_config: object = None key_bias: ttnn.Tensor = None num_images: int = 1 patches_padded: int = 0 class VisionEncoderLayer(LightweightModule): def __init__(self, config, layer_idx, mesh_device, weights, policy): super().__init__() self.policy = policy self.residual_dtype = DTYPES[policy.residual] self.config = config self.layer_idx = layer_idx self.mesh_device = mesh_device self.w = weights self.eps = config.rms_norm_eps self.num_heads = config.num_attention_heads self.norm_cfg = _compute_config(mesh_device) self.mm_cfg = {g: _compute_config(mesh_device, fidelity=policy.fidelity(g)) for g in WEIGHT_GROUPS} self.sdpa_compute_cfg = _compute_config( mesh_device, fidelity=FIDELITIES[policy.sdpa_fidelity], fp32_acc=policy.sdpa_fp32_acc ) self.grid = mesh_device.compute_with_storage_grid_size() self._l1_activations = False @classmethod def from_state_dict(cls, state_dict, *, hf_config, layer_idx, mesh_device, prefix=None, policy=None): config = VisionConfig.from_hf(hf_config) policy = policy or VisionPolicy() if prefix is None: prefix = f"{VISION_PREFIX}encoder.layers.{layer_idx}." if prefix + "input_layernorm.weight" not in state_dict: prefix = "" order = head_channel_order(config.head_dim) h = config.num_attention_heads def g(name): return state_dict[prefix + name].float() def per_head_rows(w): return w.reshape(h, config.head_dim, -1)[:, order].reshape(h * config.head_dim, -1) def linear_w(w, group): return _put_weight(w.T, mesh_device, dtype=policy.dtype(group)) wq = per_head_rows(g("self_attn.q_proj.linear.weight")) wk = per_head_rows(g("self_attn.k_proj.linear.weight")) wv = g("self_attn.v_proj.linear.weight") weights = { "qkv": linear_w(torch.cat([wq, wk, wv], dim=0), "qkv"), "o": linear_w(g("self_attn.o_proj.linear.weight"), "o"), "q_norm": _norm_weight(g("self_attn.q_norm.weight")[order], mesh_device), "k_norm": _norm_weight(g("self_attn.k_norm.weight")[order], mesh_device), "input_norm": _norm_weight(g("input_layernorm.weight"), mesh_device), "post_attn_norm": _norm_weight(g("post_attention_layernorm.weight"), mesh_device), "pre_ff_norm": _norm_weight(g("pre_feedforward_layernorm.weight"), mesh_device), "post_ff_norm": _norm_weight(g("post_feedforward_layernorm.weight"), mesh_device), "gate": linear_w(g("mlp.gate_proj.linear.weight"), "gate"), "up": linear_w(g("mlp.up_proj.linear.weight"), "up"), "down": linear_w(g("mlp.down_proj.linear.weight"), "down"), } if policy.packed_gate_up: weights["gate_up"] = _put_weight( torch.cat([g("mlp.gate_proj.linear.weight"), g("mlp.up_proj.linear.weight")], dim=0).T, mesh_device, dtype=policy.dtype("gate"), ) return cls(config, layer_idx, mesh_device, weights, policy) def _norm(self, x, weight=None, dtype=None): mem = self._mem(x) y = ttnn.rms_norm(x, epsilon=self.eps, weight=weight, compute_kernel_config=self.norm_cfg, memory_config=mem) if dtype is not None and y.dtype != dtype: z = ttnn.typecast(y, dtype, memory_config=mem) ttnn.deallocate(y) return z return y def _mem(self, x, wide=False): limit = self.policy.l1_wide_max_rows if wide else self.policy.l1_max_rows use_l1 = self._l1_activations and x.padded_shape[-2] <= limit return ttnn.L1_MEMORY_CONFIG if use_l1 else ttnn.DRAM_MEMORY_CONFIG def _linear(self, x, group, activation=None, wide=False): if self.policy.minimal_matmul: return ttnn.experimental.minimal_matmul( input_tensor=x, weight_tensor=self.w[group], fused_activation=activation, config=matmul_config(group, x.padded_shape[-2], self.grid), compute_kernel_config=self.mm_cfg[group], dtype=ttnn.bfloat16, memory_config=self._mem(x, wide), ) return ttnn.linear( x, self.w[group], activation=activation, dtype=ttnn.bfloat16, compute_kernel_config=self.mm_cfg[group], memory_config=self._mem(x, wide), ) def attention(self, x, cos, sin, plan): qkv = self._linear(x, "qkv") q, k, v = ttnn.experimental.nlp_create_qkv_heads( qkv, num_heads=self.num_heads, num_kv_heads=self.num_heads, transpose_k_heads=False, memory_config=self._mem(x), ) ttnn.deallocate(qkv) qn = self._norm(q, self.w["q_norm"]) kn = self._norm(k, self.w["k_norm"]) vn = self._norm(v) ttnn.deallocate(q) ttnn.deallocate(k) ttnn.deallocate(v) qr = ttnn.experimental.rotary_embedding_hf(qn, cos, sin, is_decode_mode=False, memory_config=self._mem(x)) kr = ttnn.experimental.rotary_embedding_hf(kn, cos, sin, is_decode_mode=False, memory_config=self._mem(x)) ttnn.deallocate(qn) ttnn.deallocate(kn) if self.policy.attention_impl == "fp32_reference": ctx = self._fp32_attention(qr, kr, vn, plan) ttnn.deallocate(qr) ttnn.deallocate(kr) ttnn.deallocate(vn) out = self._linear(ctx, "o") ttnn.deallocate(ctx) return out extra = {"cu_window_seqlens": plan.cu_window} if plan.mode == "windowed" else {"attn_mask": plan.mask} ctx = ttnn.transformer.scaled_dot_product_attention( qr, kr, vn, is_causal=False, scale=1.0, program_config=plan.program_config, compute_kernel_config=self.sdpa_compute_cfg, output_concat_heads=True, memory_config=self._mem(x), **extra, ) ttnn.deallocate(qr) ttnn.deallocate(kr) ttnn.deallocate(vn) out = self._linear(ctx, "o") ttnn.deallocate(ctx) return out def _fp32_attention(self, q, k, v, plan): n, pp, d = plan.num_images, plan.patches_padded, self.config.head_dim shape = (self.num_heads, n, pp, d) q4, k4, v4 = [ttnn.reshape(ttnn.typecast(t, ttnn.float32), shape) for t in (q, k, v)] scores = ttnn.matmul(q4, k4, transpose_b=True, dtype=ttnn.float32, compute_kernel_config=self.norm_cfg) bias = ttnn.reshape(ttnn.slice(plan.key_bias, (0, 0, 0, 0), (n, 1, 1, pp)), (1, n, 1, pp)) scores = ttnn.add(scores, ttnn.typecast(bias, ttnn.float32)) probs = ttnn.softmax(scores, dim=-1, numeric_stable=True, compute_kernel_config=self.norm_cfg) ctx = ttnn.matmul(probs, v4, dtype=ttnn.float32, compute_kernel_config=self.norm_cfg) ctx = ttnn.reshape(ctx, (1, self.num_heads, n * pp, d)) merged = ttnn.experimental.nlp_concat_heads( ttnn.typecast(ctx, ttnn.bfloat16), memory_config=ttnn.DRAM_MEMORY_CONFIG ) for t in (q4, k4, v4, scores, probs, ctx): ttnn.deallocate(t) return merged def mlp(self, x): gelu = ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU_TANH) if self.policy.packed_gate_up: gate, up = ttnn.experimental.minimal_matmul_split( x, self.w["gate_up"], chunks=2, dim=-1, config=matmul_config("gate_up", x.padded_shape[-2], self.grid), compute_kernel_config=self.mm_cfg["gate"], dtype=ttnn.bfloat16, memory_config=self._mem(x, wide=True), ) h = ttnn.multiply(gate, up, input_tensor_a_activations=[gelu], memory_config=self._mem(x, wide=True)) elif self.policy.gelu_in_matmul: gate = self._linear(x, "gate", activation=gelu, wide=True) up = self._linear(x, "up", wide=True) h = ttnn.multiply(gate, up, memory_config=self._mem(x, wide=True)) else: gate = self._linear(x, "gate", wide=True) up = self._linear(x, "up", wide=True) h = ttnn.multiply(gate, up, input_tensor_a_activations=[gelu], memory_config=self._mem(x, wide=True)) ttnn.deallocate(gate) ttnn.deallocate(up) out = self._linear(h, "down") ttnn.deallocate(h) return out def forward(self, hidden, cos, sin, plan): self._l1_activations = plan.l1_activations rd = self.residual_dtype h = self._norm(hidden, self.w["input_norm"], dtype=ttnn.bfloat16) a = self.attention(h, cos, sin, plan) ttnn.deallocate(h) h = self._norm(a, self.w["post_attn_norm"]) ttnn.deallocate(a) x = ttnn.add(hidden, h, dtype=rd, memory_config=self._mem(hidden)) ttnn.deallocate(h) h = self._norm(x, self.w["pre_ff_norm"], dtype=ttnn.bfloat16) m = self.mlp(h) ttnn.deallocate(h) h = self._norm(m, self.w["post_ff_norm"]) ttnn.deallocate(m) out = ttnn.add(x, h, dtype=rd, memory_config=self._mem(x)) ttnn.deallocate(h) ttnn.deallocate(x) return out class VisionEncoder(LightweightModule): def __init__(self, config, mesh_device, weights, layers, policy): super().__init__() self.policy = policy self.residual_dtype = DTYPES[policy.residual] self.config = config self.mesh_device = mesh_device self.w = weights self.layers = layers self.eps = config.rms_norm_eps self.compute_cfg = _compute_config(mesh_device) self.grid = mesh_device.compute_with_storage_grid_size() self.zero_rows = ttnn.zeros( (1, 1, TILE, config.text_hidden_size), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) @classmethod def from_state_dict(cls, state_dict, *, hf_config, mesh_device, policy=None, num_layers=None): config = VisionConfig.from_hf(hf_config) policy = policy or VisionPolicy() def g(name): return state_dict[name].float() w_in = g(VISION_PREFIX + "patch_embedder.input_proj.weight") table = g(VISION_PREFIX + "patch_embedder.position_embedding_table") zero = torch.zeros(1, config.hidden_size) rope_x, rope_y = rope_tables(config) weights = { "w_patch": _put_weight((2.0 * w_in).T, mesh_device), "b_patch": _put_weight(-w_in.sum(dim=1).reshape(1, -1), mesh_device), "pos_x": _put_weight(torch.cat([table[0], zero]), mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), "pos_y": _put_weight(torch.cat([table[1], zero]), mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), "rope_x": _put_weight(rope_x, mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), "rope_y": _put_weight(rope_y, mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), "w_embed": _put_weight(g(EMBEDDER_PREFIX + "embedding_projection.weight").T, mesh_device), } n = config.num_hidden_layers if num_layers is None else num_layers layers = [ VisionEncoderLayer.from_state_dict( state_dict, hf_config=config, layer_idx=i, mesh_device=mesh_device, prefix=f"{VISION_PREFIX}encoder.layers.{i}.", policy=policy, ) for i in range(n) ] return cls(config, mesh_device, weights, layers, policy) def attention_mode(self, num_images, patches_padded): p = self.policy return "mask" if p.single_image_mask and num_images == 1 and patches_padded in p.mask_patches else "windowed" def _sdpa_program_config(self, chunks, seq_len): q, k = chunks while seq_len % q: q //= 2 while seq_len % k: k //= 2 return ttnn.SDPAProgramConfig( compute_with_storage_grid_size=self.grid, q_chunk_size=q, k_chunk_size=k, exp_approx_mode=self.policy.sdpa_exp_approx, ) def attention_plan(self, inputs): host = inputs.host if self.attention_mode(host.num_images, host.patches_padded) == "mask": pp = host.patches_padded full = ttnn.repeat(inputs.key_bias, (1, 1, pp // TILE, 1)) mask = ttnn.typecast(full, DTYPES[self.policy.mask_dtype]) ttnn.deallocate(full) return AttentionPlan( "mask", l1_activations=True, mask=mask, program_config=self._sdpa_program_config(self.policy.mask_chunks, pp), key_bias=inputs.key_bias, num_images=host.num_images, patches_padded=pp, ) seq = host.num_images * host.patches_padded short = host.patches_padded < self.policy.short_patches chunks = self.policy.windowed_chunks_short if short else self.policy.windowed_chunks return AttentionPlan( "windowed", cu_window=inputs.cu_window, program_config=self._sdpa_program_config(chunks, seq), key_bias=inputs.key_bias, num_images=host.num_images, patches_padded=host.patches_padded, ) def _linear(self, x, w, bias=None): return ttnn.linear( x, w, bias=bias, compute_kernel_config=self.compute_cfg, memory_config=ttnn.DRAM_MEMORY_CONFIG ) def _lookup(self, idx, table): out = ttnn.embedding(idx, table, layout=ttnn.TILE_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG) return ttnn.reshape(out, (1, 1, out.shape[-2], out.shape[-1])) def rope(self, inputs): cs = ttnn.add(self._lookup(inputs.x_idx, self.w["rope_x"]), self._lookup(inputs.y_idx, self.w["rope_y"])) d = self.config.head_dim cos = ttnn.slice(cs, (0, 0, 0, 0), (1, 1, cs.shape[-2], d)) sin = ttnn.slice(cs, (0, 0, 0, d), (1, 1, cs.shape[-2], 2 * d)) ttnn.deallocate(cs) return cos, sin def patch_embed(self, inputs): if self.policy.minimal_matmul: h = ttnn.experimental.minimal_matmul( input_tensor=inputs.pixels, weight_tensor=self.w["w_patch"], bias_tensor=self.w["b_patch"], config=matmul_config("o", inputs.pixels.padded_shape[-2], self.grid), compute_kernel_config=self.compute_cfg, dtype=ttnn.bfloat16, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) else: h = self._linear(inputs.pixels, self.w["w_patch"], bias=self.w["b_patch"]) px = self._lookup(inputs.x_idx, self.w["pos_x"]) py = self._lookup(inputs.y_idx, self.w["pos_y"]) pos = ttnn.add(px, py) ttnn.deallocate(px) ttnn.deallocate(py) out = ttnn.add(h, pos, dtype=self.residual_dtype) ttnn.deallocate(h) ttnn.deallocate(pos) return out def pool_and_embed(self, hidden, inputs): host = inputs.host h = ttnn.reshape(hidden, (host.num_images, host.patches_padded, self.config.hidden_size)) pooled = ttnn.matmul(inputs.pool, h, compute_kernel_config=self.compute_cfg, dtype=ttnn.bfloat16) flat = ttnn.reshape(pooled, (1, 1, host.num_images * host.slots_padded, self.config.hidden_size)) normed = ttnn.rms_norm(flat, epsilon=self.eps, compute_kernel_config=self.compute_cfg) proj = self._linear(normed, self.w["w_embed"]) ttnn.deallocate(normed) table = ttnn.concat([proj, self.zero_rows], dim=2) ttnn.deallocate(proj) table_rm = ttnn.to_layout(ttnn.reshape(table, (table.shape[-2], table.shape[-1])), ttnn.ROW_MAJOR_LAYOUT) ttnn.deallocate(table) tokens = self._lookup(inputs.gather_idx, table_rm) ttnn.deallocate(table_rm) return tokens, flat def forward(self, inputs, return_intermediates=False): cos, sin = self.rope(inputs) plan = self.attention_plan(inputs) h = self.patch_embed(inputs) inter = {"patch_embed": ttnn.clone(h, memory_config=ttnn.DRAM_MEMORY_CONFIG)} if return_intermediates else None for i, layer in enumerate(self.layers): nxt = layer(h, cos, sin, plan) if return_intermediates: inter[f"layer{i:02d}"] = ttnn.clone(nxt, memory_config=ttnn.DRAM_MEMORY_CONFIG) ttnn.deallocate(h) h = nxt ttnn.deallocate(cos) ttnn.deallocate(sin) if plan.mask is not None: ttnn.deallocate(plan.mask) tokens, pooled = self.pool_and_embed(h, inputs) if not return_intermediates: ttnn.deallocate(h) ttnn.deallocate(pooled) return {"tokens": tokens} inter["pooled"] = pooled inter["tokens"] = tokens return inter