"""Pure-torch re-implementation of the Qwen3-VL vision tower (HF transformers 4.57.1 semantics). Used for two things: * host-side pre-processing shared with the ttnn path (`vision.py`): position-embedding interpolation (`pos_embed_interpolate`) and 2-D rotary tables (`rot_pos_freqs`); * a CPU reference (`VisionTorchRef`) for tests with random weights / small inputs, and for validating this file itself against the HF golden dump. Nothing here imports transformers; weights come from any object with `.get(key) -> Tensor` (e.g. `weights.Checkpoint`) using the HF key names under `vlm.model.visual.`. """ from __future__ import annotations import math import torch import torch.nn.functional as F from .config import VisionConfig from .weights import VLM_VIS # ---------------------------------------------------------------------------------------------- # host-side helpers shared with the device implementation # ---------------------------------------------------------------------------------------------- def pos_embed_interpolate(table: torch.Tensor, grid_thw: torch.Tensor, cfg: VisionConfig) -> torch.Tensor: """HF `Qwen3VLVisionModel.fast_pos_embed_interpolate`. table: [pos_grid**2, hidden] embedding weight (its dtype is used for the arithmetic, like HF). Returns [sum(t*h*w), hidden] in patch order (rows already permuted into 2x2 merge-block order). """ side = cfg.pos_grid out = [] for t, h, w in grid_thw.tolist(): h_idx = torch.linspace(0, side - 1, h) w_idx = torch.linspace(0, side - 1, w) h_floor = h_idx.int() w_floor = w_idx.int() h_ceil = (h_floor + 1).clip(max=side - 1) w_ceil = (w_floor + 1).clip(max=side - 1) dh = h_idx - h_floor dw = w_idx - w_floor base_h = h_floor * side base_h_ceil = h_ceil * side indices = [ (base_h[None].T + w_floor[None]).flatten(), (base_h[None].T + w_ceil[None]).flatten(), (base_h_ceil[None].T + w_floor[None]).flatten(), (base_h_ceil[None].T + w_ceil[None]).flatten(), ] weights = [ ((1 - dh)[None].T * (1 - dw)[None]).flatten(), ((1 - dh)[None].T * dw[None]).flatten(), (dh[None].T * (1 - dw)[None]).flatten(), (dh[None].T * dw[None]).flatten(), ] pe = None for idx, wt in zip(indices, weights): term = table[idx.long()] * wt.to(table.dtype)[:, None] pe = term if pe is None else pe + term m = cfg.merge pe = pe.repeat(t, 1) pe = pe.view(t, h // m, m, w // m, m, -1).permute(0, 1, 3, 2, 4, 5).reshape(-1, pe.shape[-1]) out.append(pe) return torch.cat(out, 0) def rot_pos_freqs(grid_thw: torch.Tensor, cfg: VisionConfig, theta: float = 10000.0) -> torch.Tensor: """HF `Qwen3VLVisionModel.rot_pos_emb`: per-patch (row, col) angles, [N, head_dim // 2] float32. Columns [0, head_dim//4) are row angles, [head_dim//4, head_dim//2) are column angles. """ dim = cfg.head_dim // 2 # rotary dim = 36; inv_freq has dim//2 = 18 entries inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) max_hw = int(grid_thw[:, 1:].max().item()) freq_table = torch.outer(torch.arange(max_hw, dtype=torch.float32), inv_freq) # [max_hw, 18] m = cfg.merge pos = [] for t, h, w in grid_thw.tolist(): mh, mw = h // m, w // m row = (torch.arange(mh)[:, None, None, None] * m + torch.arange(m)[None, None, :, None]).expand(mh, mw, m, m) col = (torch.arange(mw)[None, :, None, None] * m + torch.arange(m)[None, None, None, :]).expand(mh, mw, m, m) coords = torch.stack((row.reshape(-1), col.reshape(-1)), -1) if t > 1: coords = coords.repeat(t, 1) pos.append(coords) pos = torch.cat(pos, 0) return freq_table[pos].flatten(1) # [N, 36] def rope_cos_sin(freqs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """HF: emb = cat(freqs, freqs) -> cos, sin of shape [N, head_dim] (rotate_half convention).""" emb = torch.cat((freqs, freqs), -1) return emb.cos(), emb.sin() def rotate_half(x: torch.Tensor) -> torch.Tensor: x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), -1) def cu_seqlens(grid_thw: torch.Tensor) -> list[int]: """Per-image (per-frame) token boundaries [0, n0, n0+n1, ...] like HF's cu_seqlens.""" lens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).tolist() out = [0] for n in lens: out.append(out[-1] + n) return out def patch_embed_weight(conv_w: torch.Tensor) -> torch.Tensor: """Conv3d weight [hidden, C, T, P, P] -> [hidden, C*T*P*P] (matches HF's view of pixel_values).""" return conv_w.reshape(conv_w.shape[0], -1) # ---------------------------------------------------------------------------------------------- # torch reference model # ---------------------------------------------------------------------------------------------- class VisionTorchRef: """Float32 (by default) CPU reference of the whole tower. `getter(key)` returns HF tensors.""" def __init__(self, cfg: VisionConfig, getter, dtype: torch.dtype = torch.float32, prefix: str = VLM_VIS): self.cfg = cfg self.dtype = dtype self.prefix = prefix self._get = getter def w(self, name: str) -> torch.Tensor: return self._get(self.prefix + name).to(self.dtype) # -- pieces ---------------------------------------------------------------------------------- def patch_embed(self, pixel_values: torch.Tensor) -> torch.Tensor: w = patch_embed_weight(self.w("patch_embed.proj.weight")) return F.linear(pixel_values.to(self.dtype), w, self.w("patch_embed.proj.bias")) def pos_embed(self, grid_thw: torch.Tensor) -> torch.Tensor: # HF does this arithmetic in the table's dtype (bf16 in the checkpoint); keep that behaviour return pos_embed_interpolate(self._get(self.prefix + "pos_embed.weight"), grid_thw, self.cfg).to(self.dtype) def attention(self, i: int, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, cu: list[int]) -> torch.Tensor: c = self.cfg n = x.shape[0] qkv = F.linear(x, self.w(f"blocks.{i}.attn.qkv.weight"), self.w(f"blocks.{i}.attn.qkv.bias")) q, k, v = qkv.reshape(n, 3, c.n_heads, c.head_dim).permute(1, 0, 2, 3).unbind(0) # [n, heads, hd] cos_ = cos.unsqueeze(-2).float() sin_ = sin.unsqueeze(-2).float() qf, kf = q.float(), k.float() q = (qf * cos_ + rotate_half(qf) * sin_).to(self.dtype) k = (kf * cos_ + rotate_half(kf) * sin_).to(self.dtype) q, k, v = (t.transpose(0, 1).unsqueeze(0) for t in (q, k, v)) # [1, heads, n, hd] outs = [] for a, b in zip(cu[:-1], cu[1:]): o = F.scaled_dot_product_attention(q[:, :, a:b], k[:, :, a:b], v[:, :, a:b], is_causal=False, scale=c.head_dim ** -0.5) outs.append(o) o = torch.cat(outs, 2)[0].transpose(0, 1).reshape(n, -1) return F.linear(o, self.w(f"blocks.{i}.attn.proj.weight"), self.w(f"blocks.{i}.attn.proj.bias")) def mlp(self, i: int, x: torch.Tensor) -> torch.Tensor: h = F.linear(x, self.w(f"blocks.{i}.mlp.linear_fc1.weight"), self.w(f"blocks.{i}.mlp.linear_fc1.bias")) h = F.gelu(h, approximate="tanh") return F.linear(h, self.w(f"blocks.{i}.mlp.linear_fc2.weight"), self.w(f"blocks.{i}.mlp.linear_fc2.bias")) def block(self, i: int, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, cu: list[int]) -> torch.Tensor: c = self.cfg h = F.layer_norm(x, (c.hidden,), self.w(f"blocks.{i}.norm1.weight"), self.w(f"blocks.{i}.norm1.bias"), c.ln_eps) x = x + self.attention(i, h, cos, sin, cu) h = F.layer_norm(x, (c.hidden,), self.w(f"blocks.{i}.norm2.weight"), self.w(f"blocks.{i}.norm2.bias"), c.ln_eps) return x + self.mlp(i, h) def merger(self, which: str, x: torch.Tensor, postshuffle: bool) -> torch.Tensor: c = self.cfg d = c.merged_dim if postshuffle: x = x.reshape(-1, d) x = F.layer_norm(x, (d,), self.w(f"{which}.norm.weight"), self.w(f"{which}.norm.bias"), 1e-6) else: x = F.layer_norm(x, (c.hidden,), self.w(f"{which}.norm.weight"), self.w(f"{which}.norm.bias"), 1e-6) x = x.reshape(-1, d) h = F.gelu(F.linear(x, self.w(f"{which}.linear_fc1.weight"), self.w(f"{which}.linear_fc1.bias"))) return F.linear(h, self.w(f"{which}.linear_fc2.weight"), self.w(f"{which}.linear_fc2.bias")) # -- whole tower -------------------------------------------------------------------------------- def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, taps: dict | None = None, ) -> tuple[torch.Tensor, list[torch.Tensor]]: c = self.cfg x = self.patch_embed(pixel_values) if taps is not None: taps["patch_embed.out"] = x.clone() x = x + self.pos_embed(grid_thw) freqs = rot_pos_freqs(grid_thw, c) cos, sin = rope_cos_sin(freqs) cu = cu_seqlens(grid_thw) if taps is not None: taps["embed"] = x.clone() deep = [] for i in range(c.depth): x = self.block(i, x, cos, sin, cu) if taps is not None: taps[f"block{i}.out"] = x.clone() if i in c.deepstack_indexes: j = c.deepstack_indexes.index(i) deep.append(self.merger(f"deepstack_merger_list.{j}", x, postshuffle=True)) return self.merger("merger", x, postshuffle=False), deep def pcc(a: torch.Tensor, b: torch.Tensor) -> float: """Pearson correlation of the flattened tensors.""" a = a.detach().double().flatten() b = b.detach().double().flatten() a = a - a.mean() b = b - b.mean() den = a.norm() * b.norm() if den == 0: return float(torch.equal(a, b)) return float((a @ b) / den) def max_abs_err(a: torch.Tensor, b: torch.Tensor) -> float: return float((a.detach().float() - b.detach().float()).abs().max())