code/alpamayo_tt/vision_ref.py

"""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())