runtime/orbitquant_wan_a2/reference.py

from __future__ import annotations

import math
import torch
import torch.nn.functional as F

from .nibbles import pack_uint4, unpack_uint4

EPS = 1e-10


def fwht_last_dim(x: torch.Tensor, block_size: int) -> torch.Tensor:
    if block_size <= 0 or block_size & (block_size - 1):
        raise ValueError("block_size must be a power of two")
    if x.shape[-1] % block_size:
        raise ValueError("last dimension must divide block_size")
    y = x.reshape(*x.shape[:-1], x.shape[-1] // block_size, block_size)
    step = 1
    while step < block_size:
        shape = y.shape
        z = y.reshape(*shape[:-1], -1, 2, step)
        a, b = z[..., 0, :], z[..., 1, :]
        y = torch.stack((a + b, a - b), dim=-2).reshape(*shape)
        step *= 2
    return (y / math.sqrt(block_size)).reshape_as(x)


def rpbh(x: torch.Tensor, perm: torch.Tensor, signs: torch.Tensor, block_size: int) -> torch.Tensor:
    p = perm.to(device=x.device, dtype=torch.long)
    s = signs.to(device=x.device, dtype=torch.float32)
    y = x.float().index_select(-1, p) * s
    return fwht_last_dim(y, block_size)


def nearest_codes(unit: torch.Tensor, codebook: torch.Tensor) -> torch.Tensor:
    cb = codebook.to(device=unit.device, dtype=torch.float32)
    thresholds = 0.5 * (cb[:-1] + cb[1:])
    return torch.bucketize(unit.contiguous(), thresholds).to(torch.uint8)


def a4_pack_reference(
    x: torch.Tensor,
    perm: torch.Tensor,
    signs: torch.Tensor,
    block_size: int,
    codebook: torch.Tensor,
    eps: float = EPS,
) -> tuple[torch.Tensor, torch.Tensor]:
    rot = rpbh(x, perm, signs, block_size)
    scale = torch.linalg.vector_norm(rot, ord=2, dim=-1).float()
    unit = rot / (scale.unsqueeze(-1) + eps)
    codes = nearest_codes(unit, codebook)
    return pack_uint4(codes), scale


def dequant_packed(
    packed: torch.Tensor,
    scale: torch.Tensor,
    codebook: torch.Tensor,
    logical_k: int,
    dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    codes = unpack_uint4(packed, logical_k).long()
    cb = codebook.to(device=packed.device, dtype=torch.float32)
    return (cb[codes] * scale.float().unsqueeze(-1)).to(dtype)


def w4a4_linear_reference(
    a_packed: torch.Tensor,
    a_scale: torch.Tensor,
    w_packed: torch.Tensor,
    w_scale: torch.Tensor,
    codebook: torch.Tensor,
    logical_k: int,
    bias: torch.Tensor | None = None,
    out_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    a = dequant_packed(a_packed, a_scale, codebook, logical_k, torch.bfloat16)
    w = dequant_packed(w_packed, w_scale, codebook, logical_k, torch.bfloat16)
    b = bias.to(torch.bfloat16) if bias is not None else None
    # Match the Project-A fake-quant runtime: the quantized operands enter
    # F.linear as BF16.  CUDA tensor cores accumulate internally; the public
    # numerical contract is BF16 operands and BF16 output.
    out = F.linear(a, w, b)
    return out.to(out_dtype)