code/src/lowbit/quant/codebook.py

"""D-58 · llama.cpp's codebook formats as trainable students, started from llama.cpp's own files.

A student starts from llama.cpp's imatrix quantisation of the teacher (`scripts/ptq_baseline.py`): its expert
tensors are read from that GGUF as raw blocks and kept. Every format here stores each block as

    weight = eff(d, s) · v

where v is a small signed integer per weight (a codebook entry times its sign: the grid value of IQ2/IQ3, the
IQ4_NL table value), d is the block's fp16 scale and s its sub-blocks' integer scales:

    format   block  sub  eff                     s range   bytes  (d at bytes 0-1 in every format)
    IQ2_XXS  256    32   d·(0.5 + s)·0.25        0..15     66     s: top 4 bits of each sub-block's 2nd uint32
    IQ2_XS   256    16   d·(0.5 + s)·0.25        0..15     74     s: 16 nibbles at bytes 66-73
    IQ2_S    256    16   d·(0.5 + s)·0.25        0..15     82     s: 16 nibbles at bytes 74-81
    IQ3_XXS  256    32   d·(0.5 + s)·0.5         0..15     98     s: top 4 bits of the uint32s at bytes 66-97
    IQ3_S    256    32   d·(1 + 2s)              0..15     110    s: 8 nibbles at bytes 106-109
    IQ4_XS   256    32   d·(s − 32)              0..63     136    s: 4 low bits at bytes 4-7, 2 high at bytes 2-3
    IQ4_NL   32     32   d                       —         18     (no sub-scale)

(ggml-quants.c's dequantize_row_*; gguf-py's `dequantize_blocks` are the reference the tests compare against.)
The codes — which codebook entry and which signs — are the bytes other than d and s. They are kept unchanged
("frozen codes"): training moves d and s only, on their own grids, plus the bf16 attention and norms. The export
writes the new d and s into the original blocks, so an untrained student exports byte-identical to llama.cpp's
file, and every exported block is one llama.cpp wrote except for its scales.

Trainable parameters per tensor, both fp32: `dlog` per block (d = d0·exp(dlog), rounded to fp16 in the forward)
and `slat` per sub-block (s = round(slat) clamped to the range). The gradient reaches both through the plain
straight-through estimator, like the latents of `quant.ternary`. v·eff is exact in fp32 (d: 11 significant bits,
the s factor ≤ 7, v ≤ 7), as ggml computes it; the forward then rounds to bf16 like every other student.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import torch


@dataclass(frozen=True)
class Format:
    name: str
    block: int        # weights per block (one fp16 d)
    sub: int          # weights per sub-scale (= block when there is none)
    nbytes: int       # bytes per block
    a: float          # eff = d·a·(s + b); IQ4_NL has no s (eff = d)
    b: float
    s_max: int        # s in 0..s_max
    neutral_s: int    # the s that makes eff = d·unit when decoding v (exact)
    unit: float       # a·(neutral_s + b)


FORMATS = {
    "IQ2_XXS": Format("IQ2_XXS", 256, 32, 66, 0.25, 0.5, 15, 0, 0.125),
    "IQ2_XS": Format("IQ2_XS", 256, 16, 74, 0.25, 0.5, 15, 0, 0.125),
    "IQ2_S": Format("IQ2_S", 256, 16, 82, 0.25, 0.5, 15, 0, 0.125),
    "IQ3_XXS": Format("IQ3_XXS", 256, 32, 98, 0.5, 0.5, 15, 0, 0.25),
    "IQ3_S": Format("IQ3_S", 256, 32, 110, 2.0, 0.5, 15, 0, 1.0),
    "IQ4_XS": Format("IQ4_XS", 256, 32, 136, 1.0, -32.0, 63, 33, 1.0),
    "IQ4_NL": Format("IQ4_NL", 32, 32, 18, 1.0, 0.0, 0, 0, 1.0),
}


# ---------------------------------------------------------------- block fields (numpy, uint8 [n, nbytes])

def get_d(blocks: np.ndarray) -> np.ndarray:
    return blocks[:, :2].copy().view(np.float16)[:, 0]


def set_d(blocks: np.ndarray, d: np.ndarray) -> None:
    blocks[:, :2] = np.ascontiguousarray(d.astype(np.float16)).reshape(-1, 1).view(np.uint8)


def _nibbles(b: np.ndarray) -> np.ndarray:
    return np.stack([b & 15, b >> 4], -1).reshape(b.shape[0], -1)


def _pack_nibbles(s: np.ndarray) -> np.ndarray:
    s = s.astype(np.uint8).reshape(s.shape[0], -1, 2)
    return s[..., 0] | (s[..., 1] << 4)


def _u32(blocks: np.ndarray, lo: int, hi: int) -> np.ndarray:
    return blocks[:, lo:hi].copy().view("<u4")


def _set_u32_top4(blocks: np.ndarray, lo: int, hi: int, s: np.ndarray) -> None:
    w = _u32(blocks, lo, hi)
    w = (w & np.uint32(0x0FFFFFFF)) | (s.astype(np.uint32) << np.uint32(28))
    blocks[:, lo:hi] = w.view(np.uint8)


def get_s(fmt: Format, blocks: np.ndarray) -> np.ndarray:
    """The sub-block scales, uint8 [n, block/sub] (IQ4_NL: an empty [n, 0])."""
    n = fmt.name
    if n == "IQ2_XXS":
        return (_u32(blocks, 2, 66)[:, 1::2] >> 28).astype(np.uint8)
    if n == "IQ3_XXS":
        return (_u32(blocks, 66, 98) >> 28).astype(np.uint8)
    if n == "IQ2_XS":
        return _nibbles(blocks[:, 66:74])
    if n == "IQ2_S":
        return _nibbles(blocks[:, 74:82])
    if n == "IQ3_S":
        return _nibbles(blocks[:, 106:110])
    if n == "IQ4_XS":
        lo = _nibbles(blocks[:, 4:8])
        hi = blocks[:, 2:4].copy().view("<u2")[:, :1] >> (2 * np.arange(8, dtype=np.uint16))
        return (lo | ((hi & 3).astype(np.uint8) << 4)).astype(np.uint8)
    if n == "IQ4_NL":
        return np.zeros((blocks.shape[0], 0), np.uint8)
    raise ValueError(n)


def set_s(fmt: Format, blocks: np.ndarray, s: np.ndarray) -> None:
    n = fmt.name
    if n == "IQ2_XXS":
        w = _u32(blocks, 2, 66)
        w[:, 1::2] = (w[:, 1::2] & np.uint32(0x0FFFFFFF)) | (s.astype(np.uint32) << np.uint32(28))
        blocks[:, 2:66] = w.view(np.uint8)
    elif n == "IQ3_XXS":
        _set_u32_top4(blocks, 66, 98, s)
    elif n == "IQ2_XS":
        blocks[:, 66:74] = _pack_nibbles(s)
    elif n == "IQ2_S":
        blocks[:, 74:82] = _pack_nibbles(s)
    elif n == "IQ3_S":
        blocks[:, 106:110] = _pack_nibbles(s)
    elif n == "IQ4_XS":
        s = s.astype(np.uint16)
        blocks[:, 4:8] = _pack_nibbles(s & 15)
        hi = ((s >> 4) << (2 * np.arange(8, dtype=np.uint16))).sum(-1).astype("<u2")
        blocks[:, 2:4] = hi.reshape(-1, 1).view(np.uint8)
    elif n != "IQ4_NL":
        raise ValueError(n)


def eff_np(fmt: Format, d: np.ndarray, s: np.ndarray) -> np.ndarray:
    """fp32 effective scale per sub-block [n, block/sub], as ggml computes it."""
    d = d.astype(np.float32).reshape(-1, 1)
    if fmt.name == "IQ4_NL":
        return d
    if fmt.name == "IQ4_XS":
        return d * (s.astype(np.float32) - np.float32(32))
    if fmt.name == "IQ3_S":
        return d * (np.float32(1) + np.float32(2) * s.astype(np.float32))
    return d * (np.float32(0.5) + s.astype(np.float32)) * np.float32(fmt.a)


def values(fmt: Format, blocks: np.ndarray, gguf) -> np.ndarray:
    """int8 [n, block]: each weight's v (codebook entry × sign), decoded by gguf-py with d = 1 and a neutral s."""
    b = blocks.copy()
    set_d(b, np.ones(b.shape[0], np.float32))
    set_s(fmt, b, np.full((b.shape[0], fmt.block // fmt.sub if fmt.name != "IQ4_NL" else 0), fmt.neutral_s))
    v = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, fmt.name)).reshape(b.shape[0], fmt.block)
    v = v / np.float32(fmt.unit)
    q = np.rint(v)
    if not np.array_equal(q, v) or np.abs(q).max() > 127:
        raise ValueError(f"{fmt.name}: decoded values are not small integers")
    return q.astype(np.int8)


# ---------------------------------------------------------------- the trainable form (torch)

def round_fp16(x: torch.Tensor) -> torch.Tensor:
    """x rounded to fp16 (nearest even), as fp32, with the plain STE."""
    return x + (x.half().float() - x).detach()


def round_s(x: torch.Tensor, s_max: int) -> torch.Tensor:
    return x + (x.round().clamp(0, s_max) - x).detach()


def eff_torch(fmt: Format, d: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
    """d [..., nblk], s [..., nblk·block/sub] (both already on their grids) → eff per sub-block, fp32."""
    k = fmt.block // fmt.sub
    dd = d.unsqueeze(-1).expand(*d.shape, k).flatten(-2)
    if fmt.name == "IQ4_NL":
        return dd
    if fmt.name == "IQ4_XS":
        return dd * (s - 32.0)
    if fmt.name == "IQ3_S":
        return dd * (1.0 + 2.0 * s)
    return dd * (0.5 + s) * fmt.a


def weight(fmt: Format, v: torch.Tensor, d0: torch.Tensor, dlog: torch.Tensor, slat: torch.Tensor | None
           ) -> torch.Tensor:
    """The student's fp32 weight [..., C] from int8 v [..., C], d0 / dlog [..., C/block], slat [..., C/sub]."""
    d = round_fp16(d0 * torch.exp(dlog))
    s = round_s(slat, fmt.s_max) if slat is not None else None
    e = eff_torch(fmt, d, s)
    return (v.float().unflatten(-1, (-1, fmt.sub)) * e.unsqueeze(-1)).flatten(-2)


def trained_fields(fmt: Format, d0: torch.Tensor, dlog: torch.Tensor, slat: torch.Tensor | None
                   ) -> tuple[np.ndarray, np.ndarray]:
    """(fp16 d [..., nblk], uint8 s [..., nsub]) exactly as `weight` uses them."""
    d = (d0 * torch.exp(dlog)).half().cpu().numpy()
    s = (slat.round().clamp(0, fmt.s_max).to(torch.uint8).cpu().numpy() if slat is not None
         else np.zeros((*d.shape[:-1], 0), np.uint8))
    return d, s


def dequant(fmt: Format, blocks: np.ndarray, v: np.ndarray) -> np.ndarray:
    """fp32 weights [..., cols] from blocks [..., nblk, nbytes] and their v [..., cols] (no gguf-py needed)."""
    flat = blocks.reshape(-1, fmt.nbytes)
    e = eff_np(fmt, get_d(flat), get_s(fmt, flat)).reshape(*v.shape[:-1], -1)
    return (v.astype(np.float32).reshape(*v.shape[:-1], -1, fmt.sub) * e[..., None]).reshape(v.shape)


def write_blocks(fmt: Format, blocks: np.ndarray, d: np.ndarray, s: np.ndarray) -> np.ndarray:
    """A copy of `blocks` [..., nblk, nbytes] with d [..., nblk] and s [..., nsub] written in."""
    out = blocks.reshape(-1, fmt.nbytes).copy()
    set_d(out, d.reshape(-1))
    if fmt.name != "IQ4_NL":
        set_s(fmt, out, s.reshape(out.shape[0], -1))
    return out.reshape(blocks.shape)


# ---------------------------------------------------------------- free codes (D-58's second mode)
#
# The codes become trainable too: an fp32 latent per weight, as in `quant.ternary`, starting at llama.cpp's
# dequantised weights. The forward divides the latent by its sub-block's eff, finds the nearest codebook entry for
# each group of k magnitudes (Euclidean, exact fp32), takes the latent's signs, and multiplies back by eff; the
# backward is the plain STE to the latent and the exact gradient to eff (so d and s still train). At the start the
# latent is exactly eff·entry·sign, so every group finds its own entry at distance 0 and the student is llama.cpp's
# file. Formats whose 8 signs are stored in 7 bits (IQ2_XXS, IQ2_XS, IQ3_XXS: an even count of negatives) flip
# the sign of the smallest magnitude in an 8-weight group with an odd count, as llama.cpp's quantiser does.
# Grids come from gguf-py at import (`scripts/cb_import.py` writes grids.safetensors), so training needs no gguf-py.

GROUP = {"IQ3_S": 4, "IQ2_XXS": 8, "IQ4_NL": 1}
EVEN_SIGNS = ("IQ2_XXS",)
_CHUNK = 1 << 20  # groups per nearest-entry matmul ([chunk, entries] fp32 scores)


def _nearest(a: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
    """Index of the nearest grid row (L2) for each row of a [N, k]; grid [M, k] fp32. Ties: the lowest index."""
    half = 0.5 * (grid * grid).sum(-1)
    out = torch.empty(a.shape[0], dtype=torch.int64, device=a.device)
    for i in range(0, a.shape[0], _CHUNK):
        out[i:i + _CHUNK] = (a[i:i + _CHUNK] @ grid.T - half).argmax(-1)
    return out


_NEIGHBOURS = 16
_TABLES: dict = {}


def _lattice(grid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """(the grid's sorted coordinate values [L], for every point of values^k the indices of its 16 nearest grid
    entries [L^k, 16], nearest first), as llama.cpp's quantisers search: round each magnitude to the nearest
    coordinate value, then look only at that lattice point's neighbours. Cached per grid tensor."""
    key = (grid.data_ptr(), grid.device)
    if key not in _TABLES:
        vals = torch.unique(grid)
        k = grid.shape[1]
        pts = torch.cartesian_prod(*[vals] * k).reshape(-1, k).float()  # point p = sum r_i·L^(k-1-i)
        d2 = ((pts[:, None, :] - grid[None]) ** 2).sum(-1)
        _TABLES[key] = (vals, d2.topk(min(_NEIGHBOURS, grid.shape[0]), dim=-1, largest=False).indices)
    return _TABLES[key]


def _nearest_lattice(a: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
    """Nearest grid row for each row of a [N, k] among its rounded lattice point's neighbours (exact for a on the
    grid: the point is its own nearest neighbour). Ties: the neighbour listed first."""
    vals, table = _lattice(grid)
    L, k = vals.numel(), a.shape[1]
    mid = (vals[1:] + vals[:-1]) / 2
    place = L ** torch.arange(k - 1, -1, -1, device=a.device)
    out = torch.empty(a.shape[0], dtype=torch.int64, device=a.device)
    for i in range(0, a.shape[0], _CHUNK):
        c = a[i:i + _CHUNK]
        p = (torch.bucketize(c, mid) * place).sum(-1)
        cand = table[p]  # [n, 16]
        d2 = ((c[:, None, :] - grid[cand]) ** 2).sum(-1)
        out[i:i + _CHUNK] = cand.gather(1, d2.argmin(-1, keepdim=True))[:, 0]
    return out


@torch.compiler.disable
def free_codes(fmt: Format, x: torch.Tensor, grid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """x [..., C] (latent / eff, fp32) → (idx [..., C/k] int64, signed values [..., C] fp32) on the format's codebook.
    IQ4_NL's grid is its 16 values [16, 1], searched per weight (exactly, by the midpoints) with signs included."""
    if fmt.name == "IQ4_NL":
        g = grid[:, 0]
        idx = torch.bucketize(x, (g[1:] + g[:-1]) / 2)  # ties go to the lower value, as `_nearest`'s argmax
        return idx, g[idx]
    k = GROUP[fmt.name]
    neg = x < 0
    if fmt.name in EVEN_SIGNS:  # an odd count of negatives in an 8-group: flip the smallest magnitude
        g8 = x.abs().unflatten(-1, (-1, 8))
        odd = (neg.unflatten(-1, (-1, 8)).sum(-1) % 2) == 1
        flip = torch.nn.functional.one_hot(g8.argmin(-1), 8).bool() & odd.unsqueeze(-1)
        neg = neg ^ flip.flatten(-2)
    a = x.abs().reshape(-1, k)
    idx = _nearest_lattice(a, grid)
    vals = grid[idx].reshape(x.shape)
    return idx.reshape(*x.shape[:-1], -1), torch.where(neg, -vals, vals)


def weight_free(fmt: Format, latent: torch.Tensor, d0: torch.Tensor, dlog: torch.Tensor,
                slat: torch.Tensor | None, grid: torch.Tensor) -> torch.Tensor:
    """The free-code student's fp32 weight [..., C]: eff · (nearest entry × sign), STE to the latent."""
    d = round_fp16(d0 * torch.exp(dlog))
    s = round_s(slat, fmt.s_max) if slat is not None else None
    e = eff_torch(fmt, d, s).unsqueeze(-1)  # [..., nsub, 1]
    lat = latent.unflatten(-1, (-1, fmt.sub))
    with torch.no_grad():
        _, vals = free_codes(fmt, (lat / e.detach()).flatten(-2), grid)
    w = e * vals.unflatten(-1, (-1, fmt.sub))  # gradient to eff (d, s) through the chosen values
    return (w + (lat - lat.detach())).flatten(-2)  # and the plain STE to the latent


def encode(fmt: Format, idx: torch.Tensor, vals: torch.Tensor, d: np.ndarray, s: np.ndarray) -> np.ndarray:
    """Blocks uint8 [..., nblk, nbytes] from codebook indices, signed values (for the signs), d and s."""
    lead = vals.shape[:-1]
    n = int(np.prod(lead)) * (vals.shape[-1] // fmt.block)
    neg = (vals < 0).reshape(n, fmt.block).cpu().numpy()
    out = np.zeros((n, fmt.nbytes), np.uint8)
    set_d(out, d.reshape(-1))
    if fmt.name == "IQ4_NL":
        q = idx.reshape(n, 32).cpu().numpy().astype(np.uint8)
        out[:, 2:18] = q[:, :16] | (q[:, 16:] << 4)
        return out.reshape(*lead, -1, fmt.nbytes)
    gi = idx.reshape(n, -1).cpu().numpy()
    if fmt.name == "IQ3_S":
        out[:, 2:66] = (gi & 0xFF).astype(np.uint8)
        out[:, 66:74] = np.packbits(((gi >> 8) & 1).astype(np.uint8), axis=-1, bitorder="little")
        out[:, 74:106] = np.packbits(neg.astype(np.uint8), axis=-1, bitorder="little")
    elif fmt.name == "IQ2_XXS":
        pat = np.packbits(neg.astype(np.uint8).reshape(n, 32, 8), axis=-1, bitorder="little")[..., 0]  # [n, 32]
        if (np.unpackbits(pat[..., None], axis=-1).sum(-1) % 2).any():
            raise ValueError("IQ2_XXS sign groups must have an even count of negatives")
        sidx = (pat & 0x7F).astype(np.uint32).reshape(n, 8, 4)
        aux1 = (sidx << np.array([0, 7, 14, 21], np.uint32)).sum(-1).astype(np.uint32)  # [n, 8]
        aux0 = gi.astype(np.uint8).reshape(n, 8, 4).copy().view("<u4")[..., 0]
        out[:, 2:66] = np.stack([aux0, aux1], -1).astype("<u4").view(np.uint8).reshape(n, 64)
    else:
        raise NotImplementedError(f"free codes for {fmt.name}")
    set_s(fmt, out, s.reshape(n, -1))
    return out.reshape(*lead, -1, fmt.nbytes)


def grid_from_gguf(fmt: Format, gguf) -> torch.Tensor:
    """The format's codebook as fp32 [entries, k] (IQ4_NL and IQ4_XS: their 16 values as [16, 1])."""
    if fmt.name in ("IQ4_NL", "IQ4_XS"):
        from lowbit.quant.ternary import IQ4NL_VALUES
        return torch.tensor(IQ4NL_VALUES, dtype=torch.float32)[:, None]
    cls = getattr(gguf.quants, fmt.name)
    cls.init_grid()
    return torch.from_numpy(np.ascontiguousarray(cls.grid.reshape(cls.grid_shape), dtype=np.float32))