dotcache/modes/m4_key_project.py

from __future__ import annotations

from functools import lru_cache
from math import ceil

import numpy as np

from .m0_affine import pad_last_dim
from .turbo3 import fwht_last_dim

_VALID_M4_BASIS_FAMILIES = ("hadamard", "dct", "svd", "svd_shared")


def valid_m4_basis_families() -> tuple[str, ...]:
    return _VALID_M4_BASIS_FAMILIES


def _dct_basis(group_size: int) -> np.ndarray:
    positions = np.arange(group_size, dtype=np.float32)[None, :]
    frequencies = np.arange(group_size, dtype=np.float32)[:, None]
    basis = np.cos((np.pi / np.float32(group_size)) * (positions + np.float32(0.5)) * frequencies).astype(np.float32)
    basis[0] *= np.float32(np.sqrt(1.0 / group_size))
    if group_size > 1:
        basis[1:] *= np.float32(np.sqrt(2.0 / group_size))
    return basis


@lru_cache(maxsize=None)
def fixed_project_basis(group_size: int, rank: int, basis_family: str = "hadamard") -> np.ndarray:
    if group_size <= 0 or (group_size & (group_size - 1)):
        raise ValueError("M4 fixed-project requires a power-of-two group_size")
    if basis_family not in _VALID_M4_BASIS_FAMILIES:
        raise ValueError(f"M4 fixed-project basis_family must be one of {', '.join(_VALID_M4_BASIS_FAMILIES)}")
    usable_rank = max(1, min(int(rank), group_size - 1))
    if basis_family == "hadamard":
        basis = fwht_last_dim(np.eye(group_size, dtype=np.float32))
    else:
        basis = _dct_basis(group_size)
    # Skip the DC row because the page mean already captures the constant offset.
    return np.asarray(basis[1 : 1 + usable_rank], dtype=np.float32)


def quantize_tensor_m4(
    values: np.ndarray,
    *,
    group_size: int,
    project_dim: int,
    basis_family: str = "hadamard",
    basis_override: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray | None, np.ndarray, int]:
    array = np.asarray(values, dtype=np.float32)
    if array.ndim != 2:
        raise ValueError("values must have shape [token_count, head_dim]")

    token_count, head_dim = array.shape
    num_groups = ceil(head_dim / group_size)
    padded_head_dim = num_groups * group_size
    padded = pad_last_dim(array, padded_head_dim)
    grouped = padded.reshape(token_count, num_groups, group_size)

    if basis_override is not None:
        stored_basis = np.asarray(basis_override)
        learned_basis = stored_basis.astype(np.float32, copy=False)
        if learned_basis.ndim != 3 or learned_basis.shape[0] != num_groups or learned_basis.shape[2] != group_size:
            raise ValueError("basis_override must have shape [num_groups, rank, group_size]")
        coeffs = np.zeros((token_count, num_groups, learned_basis.shape[1]), dtype=np.float32)
        mean = np.zeros((num_groups, group_size), dtype=np.float32)
        for group_index in range(num_groups):
            group_values = grouped[:, group_index, :]
            group_mean = group_values.mean(axis=0, dtype=np.float32)
            centered = group_values - group_mean[None, :]
            coeffs[:, group_index, :] = centered @ learned_basis[group_index].T
            mean[group_index, :] = group_mean
        return (
            coeffs.astype(np.float16, copy=False),
            stored_basis.astype(np.float16, copy=False),
            mean.astype(np.float16, copy=False),
            padded_head_dim,
        )

    if basis_family in {"svd", "svd_shared"}:
        usable_rank = max(1, min(int(project_dim), group_size, token_count))
        coeffs = np.zeros((token_count, num_groups, usable_rank), dtype=np.float32)
        learned_basis = np.zeros((num_groups, usable_rank, group_size), dtype=np.float32)
        mean = np.zeros((num_groups, group_size), dtype=np.float32)
        for group_index in range(num_groups):
            group_values = grouped[:, group_index, :]
            group_mean = group_values.mean(axis=0, dtype=np.float32)
            centered = group_values - group_mean[None, :]
            u, s, vt = np.linalg.svd(centered, full_matrices=False)
            group_rank = max(1, min(usable_rank, int(vt.shape[0]), int(u.shape[1])))
            coeffs[:, group_index, :group_rank] = (u[:, :group_rank] * s[:group_rank]).astype(np.float32, copy=False)
            learned_basis[group_index, :group_rank, :] = vt[:group_rank, :].astype(np.float32, copy=False)
            mean[group_index, :] = group_mean
        return (
            coeffs.astype(np.float16, copy=False),
            learned_basis.astype(np.float16, copy=False),
            mean.astype(np.float16, copy=False),
            padded_head_dim,
        )

    basis = fixed_project_basis(group_size, project_dim, basis_family)
    coeffs = np.zeros((token_count, num_groups, basis.shape[0]), dtype=np.float32)
    mean = np.zeros((num_groups, group_size), dtype=np.float32)

    for group_index in range(num_groups):
        group_values = grouped[:, group_index, :]
        group_mean = group_values.mean(axis=0, dtype=np.float32)
        centered = group_values - group_mean[None, :]
        coeffs[:, group_index, :] = centered @ basis.T
        mean[group_index, :] = group_mean

    return (
        coeffs.astype(np.float16, copy=False),
        None,
        mean.astype(np.float16, copy=False),
        padded_head_dim,
    )


def fit_shared_project_basis(
    values: np.ndarray,
    *,
    group_size: int,
    project_dim: int,
    page_size: int,
) -> np.ndarray:
    array = np.asarray(values, dtype=np.float32)
    if array.ndim != 2:
        raise ValueError("values must have shape [token_count, head_dim]")
    if page_size <= 0:
        raise ValueError("page_size must be positive")

    token_count, head_dim = array.shape
    num_groups = ceil(head_dim / group_size)
    padded_head_dim = num_groups * group_size
    padded = pad_last_dim(array, padded_head_dim)
    grouped = padded.reshape(token_count, num_groups, group_size)
    usable_rank = max(1, min(int(project_dim), group_size, token_count))
    basis = np.zeros((num_groups, usable_rank, group_size), dtype=np.float32)

    for group_index in range(num_groups):
        group_values = grouped[:, group_index, :]
        residual_chunks: list[np.ndarray] = []
        for page_start in range(0, token_count, page_size):
            page_values = group_values[page_start : page_start + page_size]
            if page_values.shape[0] == 0:
                continue
            page_mean = page_values.mean(axis=0, dtype=np.float32)
            residual_chunks.append(page_values - page_mean[None, :])
        residual = np.concatenate(residual_chunks, axis=0) if residual_chunks else group_values
        u, s, vt = np.linalg.svd(residual, full_matrices=False)
        group_rank = max(1, min(usable_rank, int(vt.shape[0]), int(u.shape[1])))
        basis[group_index, :group_rank, :] = vt[:group_rank, :].astype(np.float32, copy=False)
    return basis


def reconstruct_group_m4(
    coefficients: np.ndarray,
    *,
    mean: np.ndarray,
    group_size: int,
    basis_family: str = "hadamard",
    basis: np.ndarray | None = None,
) -> np.ndarray:
    coeff_array = np.asarray(coefficients, dtype=np.float32)
    if basis is None:
        rank = int(coeff_array.shape[-1])
        basis_array = fixed_project_basis(int(group_size), rank, basis_family)
    else:
        basis_array = np.asarray(basis, dtype=np.float32)
    reconstructed = coeff_array @ basis_array
    return reconstructed + np.asarray(mean, dtype=np.float32)[None, :]