code/alpamayo_tt/torch_ref.py

"""Pure-torch reference of one Qwen3-VL text decoder layer (fp32 math, real checkpoint weights).

Used as the layer-level oracle before/besides the CPU golden run. Follows modeling_qwen3_vl.py exactly
(RMSNorm in fp32, per-head q/k norm before rope, HF rotate_half rope, GQA causal attention, SwiGLU).
"""
from __future__ import annotations

import math

import torch

from .config import TextConfig


def rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
    xf = x.float()
    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
    return (w.float() * xf).to(x.dtype)


def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)


def ref_layer(x: torch.Tensor, w: dict, cos: torch.Tensor, sin: torch.Tensor, tc: TextConfig,
              dtype: torch.dtype = torch.float32, return_kv: bool = False):
    """x: [1, L, hidden]; w: Checkpoint.text_layer(i) dict (HF names); cos/sin: [L, head_dim].
    Computes in `dtype` (float32 default = tight oracle; bfloat16 mimics the HF CPU run)."""
    L = x.shape[1]
    H, KV, D = tc.n_heads, tc.n_kv_heads, tc.head_dim
    x = x.to(dtype)
    W = {k: v.to(dtype) for k, v in w.items()}
    h = rms_norm(x, W["input_layernorm.weight"], tc.rms_eps)
    q = (h @ W["self_attn.q_proj.weight"].t()).view(1, L, H, D)
    k = (h @ W["self_attn.k_proj.weight"].t()).view(1, L, KV, D)
    v = (h @ W["self_attn.v_proj.weight"].t()).view(1, L, KV, D)
    q = rms_norm(q, W["self_attn.q_norm.weight"], tc.rms_eps).transpose(1, 2)  # [1, H, L, D]
    k = rms_norm(k, W["self_attn.k_norm.weight"], tc.rms_eps).transpose(1, 2)
    v = v.transpose(1, 2)
    c, s = cos.to(dtype)[None, None], sin.to(dtype)[None, None]
    q = q * c + rotate_half(q) * s
    k = k * c + rotate_half(k) * s
    rep = H // KV
    kk = k[:, :, None].expand(1, KV, rep, L, D).reshape(1, H, L, D)
    vv = v[:, :, None].expand(1, KV, rep, L, D).reshape(1, H, L, D)
    att = (q.float() @ kk.float().transpose(-1, -2)) / math.sqrt(D)
    mask = torch.full((L, L), float("-inf")).triu(1)
    att = torch.softmax(att + mask, dim=-1).to(dtype)
    o = (att @ vv).transpose(1, 2).reshape(1, L, H * D)
    o = o @ W["self_attn.o_proj.weight"].t()
    x = x + o
    h = rms_norm(x, W["post_attention_layernorm.weight"], tc.rms_eps)
    g = torch.nn.functional.silu(h @ W["mlp.gate_proj.weight"].t())
    u = h @ W["mlp.up_proj.weight"].t()
    x = x + (g * u) @ W["mlp.down_proj.weight"].t()
    if return_kv:
        return x, k, v
    return x


def pcc(a: torch.Tensor, b: torch.Tensor) -> float:
    a = a.double().flatten()
    b = b.double().flatten()
    a = a - a.mean()
    b = b - b.mean()
    return float((a @ b) / (a.norm() * b.norm() + 1e-12))