code/alpamayo_tt/expert_ref.py

"""Torch reference of one expert (Qwen3-VL text) decoder layer with HF semantics, used as the oracle for the
device implementation when the golden dump does not contain the full VLM KV cache.

Layer i of the expert attends to [VLM cache rows of layer i (already q/k-normed and roped by the VLM) ; its own
64 tokens], non-causally. Own k gets k_norm + rope, cached keys are used as stored.
"""
from __future__ import annotations

import torch

from .action_proj import rms_norm, rope_cos_sin
from .config import TextConfig


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


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """x [B, H, L, D], cos/sin [L, D] (broadcast over heads)."""
    return x * cos[None, None] + rotate_half(x) * sin[None, None]


def layer_forward(x: torch.Tensor, w: dict[str, torch.Tensor], k_cache: torch.Tensor, v_cache: torch.Tensor,
                  cos: torch.Tensor, sin: torch.Tensor, tc: TextConfig,
                  dtype: torch.dtype = torch.float32, return_taps: bool = False):
    """x [1, L, hidden]; w = Checkpoint.text_layer(i, prefix=EXPERT) (HF [out, in] weights);
    k_cache/v_cache [1, n_kv, S, head_dim] (HF head order); cos/sin [L, head_dim].
    Returns [1, L, hidden] (and intermediate taps if requested)."""
    H, KV, D = tc.n_heads, tc.n_kv_heads, tc.head_dim
    eps = tc.rms_eps
    x = x.to(dtype)
    w = {k: v.to(dtype) for k, v in w.items()}
    B, L, _ = x.shape

    h = rms_norm(x, w["input_layernorm.weight"], eps)
    q = (h @ w["self_attn.q_proj.weight"].t()).view(B, L, H, D)
    k = (h @ w["self_attn.k_proj.weight"].t()).view(B, L, KV, D)
    v = (h @ w["self_attn.v_proj.weight"].t()).view(B, L, KV, D)
    q = rms_norm(q, w["self_attn.q_norm.weight"], eps).transpose(1, 2)  # [B, H, L, D]
    k = rms_norm(k, w["self_attn.k_norm.weight"], eps).transpose(1, 2)
    v = v.transpose(1, 2)
    q = apply_rope(q, cos.to(dtype), sin.to(dtype))
    k = apply_rope(k, cos.to(dtype), sin.to(dtype))

    K = torch.cat([k_cache.to(dtype), k], dim=2)  # [B, KV, S+L, D]
    V = torch.cat([v_cache.to(dtype), v], dim=2)
    rep = H // KV
    K = K[:, :, None].expand(B, KV, rep, K.shape[2], D).reshape(B, H, -1, D)
    V = V[:, :, None].expand(B, KV, rep, V.shape[2], D).reshape(B, H, -1, D)
    scores = (q.float() @ K.float().transpose(-1, -2)) * (D ** -0.5)
    p = torch.softmax(scores, dim=-1).to(dtype)
    attn = (p @ V).transpose(1, 2).reshape(B, L, H * D)
    o = attn @ w["self_attn.o_proj.weight"].t()
    x1 = x + o

    h2 = rms_norm(x1, w["post_attention_layernorm.weight"], eps)
    g = torch.nn.functional.silu(h2 @ w["mlp.gate_proj.weight"].t())
    u = h2 @ w["mlp.up_proj.weight"].t()
    m = (g * u) @ w["mlp.down_proj.weight"].t()
    out = x1 + m
    if return_taps:
        return out, {"h": h, "q": q, "k": k, "v": v, "attn": attn, "o": o, "x1": x1, "h2": h2, "g": g, "u": u, "m": m}
    return out


def expert_cos_sin(S: int, rope_delta: int, tc: TextConfig, n_tokens: int = 64):
    pos = torch.arange(n_tokens) + int(rope_delta) + int(S)
    return rope_cos_sin(pos, tc.head_dim, tc.rope_theta)


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