vqkv/quantizers.py

"""
vqkv.quantizers — KV-cache quantizers for AttnVQ.

ScalarKV, KIVIScalarKV, ProductVQKV (attention-weighted batched LBG), RoPESplitVQKV,
SignScalarKV, TernaryScalarKV. Each exposes fit() + roundtrip_k/v() and
bits_per_element(). ProductVQ / RoPESplit Lloyd updates are weighted by key
attention mass (see vqkv.metrics.calibration_sample_weights); quality is
reported via attention-output error, not cache MSE alone.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import torch


# ----------------------------------------------------------------------------
# Utilities
# ----------------------------------------------------------------------------
def _affine_quantize(x: torch.Tensor, nbits: int, dim: int):
    """Symmetric-range affine quantization along `dim`. Returns (q, scale, zero)."""
    qmax = (1 << nbits) - 1
    xmin = x.amin(dim=dim, keepdim=True)
    xmax = x.amax(dim=dim, keepdim=True)
    scale = (xmax - xmin).clamp_min(1e-8) / qmax
    zero = torch.round(-xmin / scale)
    q = torch.clamp(torch.round(x / scale) + zero, 0, qmax)
    return q, scale, zero


def _affine_dequantize(q, scale, zero):
    return (q - zero) * scale


# ----------------------------------------------------------------------------
# Scalar baselines
# ----------------------------------------------------------------------------
@dataclass
class ScalarKV:
    """Per-token affine quantization of both K and V (the transformers default)."""
    nbits: int = 4

    def fit(self, k_calib, v_calib):  # scalar quant is tuning-free
        return self

    def roundtrip_k(self, k):
        # k: (..., head_dim) -> quantize along head_dim (per-token, per-head)
        q, s, z = _affine_quantize(k, self.nbits, dim=-1)
        return _affine_dequantize(q, s, z)

    def roundtrip_v(self, v):
        q, s, z = _affine_quantize(v, self.nbits, dim=-1)
        return _affine_dequantize(q, s, z)

    def bits_per_element(self, head_dim):
        # nbits per element + scale/zero (fp16 each) amortized over head_dim
        return self.nbits + (2 * 16) / head_dim


@dataclass
class KIVIScalarKV:
    """KIVI-style: keys quantized per-channel, values per-token.

    Keys are quantized along the TOKEN axis (per channel); values along the
    head_dim axis (per token). Requires a token-axis view, so fit/roundtrip
    operate on a full (N_tokens, head_dim) block per head.
    """
    nbits: int = 4

    def fit(self, k_calib, v_calib):
        return self

    def roundtrip_k(self, k):
        # k: (N_tokens, head_dim). per-channel == quantize along token axis (dim=0)
        q, s, z = _affine_quantize(k, self.nbits, dim=0)
        return _affine_dequantize(q, s, z)

    def roundtrip_v(self, v):
        q, s, z = _affine_quantize(v, self.nbits, dim=-1)
        return _affine_dequantize(q, s, z)

    def bits_per_element(self, head_dim):
        return self.nbits + (2 * 16) / head_dim


# ----------------------------------------------------------------------------
# LBG / k-means codebook (the 1980 algorithm, the engine of the project)
# ----------------------------------------------------------------------------
def lbg_codebook(data: torch.Tensor, n_codes: int, iters: int = 25,
                 seed: int = 0, sample_weights: torch.Tensor | None = None) -> torch.Tensor:
    """Linde-Buzo-Gray / Lloyd design on sub-vectors.

    Assignment: nearest centroid (squared L2). Update: weighted mean when
    ``sample_weights`` (N,) is given — attention-weighted LBG for AttnVQ.
    """
    g = torch.Generator().manual_seed(seed)
    N, d = data.shape
    dev = data.device
    w = sample_weights
    if w is not None:
        w = w.to(device=dev, dtype=data.dtype).clamp_min(1e-8)
    idx = torch.randperm(N, generator=g)[:n_codes].to(dev)
    cb = data[idx].clone()
    ones = torch.ones(N, dtype=data.dtype, device=dev)
    for _ in range(iters):
        # Assignment: BLAS GEMM path avoids the O(N*K*d) intermediate that
        # torch.cdist allocates for small d (e.g. d=8 for n_sub=16).
        assign = _sq_l2(data, cb).argmin(dim=1)

        # Update: vectorized scatter_add replaces the Python loop over n_codes.
        # Was: 256 Python iterations; now: 3 torch ops.
        new_cb = torch.zeros_like(cb)
        counts = torch.zeros(n_codes, dtype=data.dtype, device=dev)
        pts = data if w is None else data * w.unsqueeze(1)
        new_cb.scatter_add_(0, assign.unsqueeze(1).expand(-1, d), pts)
        counts.scatter_add_(0, assign, ones if w is None else w)
        live = counts > 0
        new_cb[live] /= counts[live].unsqueeze(1)
        # Re-seed dead centroids (LBG splitting heuristic)
        dead = (~live).nonzero(as_tuple=True)[0]
        if dead.numel() > 0:
            ri = torch.randint(0, N, (dead.numel(),), generator=g).to(dev)
            new_cb[dead] = data[ri]

        shift = (new_cb - cb).norm()
        cb = new_cb
        if shift < 1e-5:
            break
    return cb


def lbg_codebook_batched(xb: torch.Tensor, n_codes: int, iters: int = 25,
                          seed: int = 0,
                          sample_weights: torch.Tensor | None = None) -> torch.Tensor:
    """Fit n_sub independent attention-weighted LBG codebooks in one batched pass.

    Mirrors the structure of ProductVQKV._roundtrip: replaces the Python loop
    over sub-blocks with a single bmm-based assignment and a scatter_add-based
    update, so n_sub sub-block fits become one operation at each Lloyd step.

    xb:      (n_sub, N, sub_dim)  training vectors, pre-normalized if needed
    sample_weights: optional (N,) masses from key_attention_mass (AttnVQ fit)
    returns: (n_sub, n_codes, sub_dim) codebooks, one per sub-block
    """
    g = torch.Generator().manual_seed(seed)
    n_sub, N, sub_dim = xb.shape
    dev = xb.device
    w = sample_weights
    if w is not None:
        w = w.to(device=dev, dtype=xb.dtype).clamp_min(1e-8)

    # Initialization: random subset per sub-block (same RNG sequence as serial)
    idx = torch.stack([torch.randperm(N, generator=g)[:n_codes]
                       for _ in range(n_sub)]).to(dev)          # (n_sub, K)
    cb = xb[torch.arange(n_sub, device=dev).unsqueeze(1), idx].clone()  # (n_sub, K, sub_dim)

    ones = torch.ones(N, dtype=xb.dtype, device=dev)

    for _ in range(iters):
        # Assignment — one bmm replaces n_sub independent GEMMs
        x_sq  = (xb * xb).sum(-1, keepdim=True)            # (n_sub, N, 1)
        c_sq  = (cb * cb).sum(-1).unsqueeze(1)              # (n_sub, 1, K)
        cross = torch.bmm(xb, cb.transpose(1, 2))           # (n_sub, N, K)
        assign = (x_sq - 2 * cross + c_sq).argmin(dim=-1)  # (n_sub, N)

        # Update — batched scatter_add (optionally attention-weighted)
        new_cb = torch.zeros_like(cb)
        counts = torch.zeros(n_sub, n_codes, dtype=xb.dtype, device=dev)
        assign_exp = assign.unsqueeze(-1).expand(-1, -1, sub_dim)
        pts = xb if w is None else xb * w.view(1, N, 1)
        new_cb.scatter_add_(1, assign_exp, pts)
        cnt_src = ones if w is None else w
        counts.scatter_add_(1, assign, cnt_src.unsqueeze(0).expand(n_sub, -1))

        live = counts > 0
        new_cb[live] /= counts[live].unsqueeze(-1)

        # Re-seed dead centroids per sub-block (rare; loop is fine)
        dead_any = ~live
        if dead_any.any():
            for s in range(n_sub):
                dead = dead_any[s].nonzero(as_tuple=True)[0]
                if dead.numel():
                    ri = torch.randint(0, N, (dead.numel(),), generator=g).to(dev)
                    new_cb[s, dead] = xb[s, ri]

        shift = (new_cb - cb).norm()
        cb = new_cb
        if shift < 1e-5:
            break

    return cb


def _sq_l2(x: torch.Tensor, cb: torch.Tensor) -> torch.Tensor:
    """Squared L2 distance matrix (N, K) via BLAS GEMM.

    torch.cdist for small d (e.g. d=8 for n_sub=16) falls back to a naive
    expand-and-subtract path that allocates an (N, K, d) intermediate.
    For N=131072, K=256, d=8 that is 1.07 GB per call, causing massive
    CPU allocation pressure and 100x slowdowns versus the BLAS path.

    ||x-c||^2 = ||x||^2 - 2*(x @ c.T) + ||c||^2 uses only (N,K) memory.
    """
    return ((x * x).sum(1, keepdim=True)
            + (cb * cb).sum(1)
            - 2 * (x @ cb.T))


def vq_encode(x: torch.Tensor, cb: torch.Tensor) -> torch.Tensor:
    """Nearest-codeword indices. x:(N,d) cb:(K,d) -> (N,) long."""
    return _sq_l2(x, cb).argmin(dim=1)


# ----------------------------------------------------------------------------
# Product VQ
# ----------------------------------------------------------------------------
@dataclass
class ProductVQKV:
    """Product vector quantization of head-dim sub-blocks.

    Each head's `head_dim` vector is split into `n_sub` contiguous sub-vectors
    of length `sub_dim = head_dim / n_sub`; each sub-vector is quantized
    against its own codebook of size `n_codes`.

    If `normalize=True`, each sub-vector is standardized by its per-dimension
    calibration mean/std before codebook design and restored after decode.
    This decorrelates the scale variation RoPE introduces in the rotated half
    of the key, and is the mechanism that lets RoPE-split specialize.

    Rate (bits/element) = n_sub * log2(n_codes) / head_dim.
    e.g. head_dim=128, n_sub=8, n_codes=256  -> 8*8/128 = 0.5 bits/element.
    """
    n_sub: int = 8
    n_codes: int = 256
    iters: int = 25
    normalize: bool = False
    k_codebooks: list = field(default_factory=list)
    v_codebooks: list = field(default_factory=list)
    _k_stats: list = field(default_factory=list)
    _v_stats: list = field(default_factory=list)
    _k_stacked: tuple = None   # lazily-built (cb, mu, sd) batched tensors
    _v_stacked: tuple = None

    def _split(self, x):
        # x: (N, head_dim) -> list of (N, sub_dim)
        return list(torch.chunk(x, self.n_sub, dim=-1))

    def _fit_one(self, x, sample_weights=None):
        N, head_dim = x.shape
        sub_dim = head_dim // self.n_sub
        # (n_sub, N, sub_dim) -- same layout _roundtrip uses, so batching mirrors inference
        xb = x.reshape(N, self.n_sub, sub_dim).permute(1, 0, 2).contiguous()

        if self.normalize:
            mu = xb.mean(dim=1, keepdim=True)          # (n_sub, 1, sub_dim)
            sd = xb.std(dim=1, keepdim=True).clamp_min(1e-6)
            xb = (xb - mu) / sd
            # unstack into (1, sub_dim) tuples so _stack / to() are unchanged
            stats = [(mu[s], sd[s]) for s in range(self.n_sub)]
        else:
            stats = [None] * self.n_sub

        cb_batched = lbg_codebook_batched(
            xb, self.n_codes, self.iters, sample_weights=sample_weights)
        cbs = list(cb_batched.unbind(dim=0))            # n_sub x (K, sub_dim)
        return cbs, stats

    def fit(self, k_calib, v_calib, sample_weights=None, n_q_heads=None,
            k_struct=None):
        if sample_weights is None and k_struct is not None and k_struct.dim() == 3:
            from vqkv.metrics import calibration_sample_weights
            sample_weights = calibration_sample_weights(k_struct, n_q_heads)
        self.k_codebooks, self._k_stats = self._fit_one(k_calib, sample_weights)
        self.v_codebooks, self._v_stats = self._fit_one(v_calib, sample_weights)
        return self

    def _stack(self, codebooks, stats):
        """Lazily stack per-sub-block codebooks/stats into batched tensors so the
        whole product-VQ encode is a few large ops instead of n_sub small ones.

        Returns:
          cb_stacked: (n_sub, K, sub_dim)
          mu_stacked: (1, n_sub, sub_dim) or None   (None => no normalization)
          sd_stacked: (1, n_sub, sub_dim) or None
        Requires uniform sub_dim and n_codes across sub-blocks, which ProductVQ
        guarantees (torch.chunk into equal pieces, single n_codes).
        """
        cb_stacked = torch.stack(codebooks, dim=0)  # (n_sub, K, sub_dim)
        if any(st is not None for st in stats):
            mu = torch.stack([st[0].reshape(-1) for st in stats], dim=0)  # (n_sub, sub_dim)
            sd = torch.stack([st[1].reshape(-1) for st in stats], dim=0)
            mu_stacked = mu.unsqueeze(0)  # (1, n_sub, sub_dim)
            sd_stacked = sd.unsqueeze(0)
        else:
            mu_stacked = sd_stacked = None
        return cb_stacked, mu_stacked, sd_stacked

    def _ensure_stacked(self):
        if getattr(self, "_k_stacked", None) is None:
            self._k_stacked = self._stack(self.k_codebooks, self._k_stats)
        if getattr(self, "_v_stacked", None) is None:
            self._v_stacked = self._stack(self.v_codebooks, self._v_stats)

    def _roundtrip(self, x, stacked):
        """Batched product-VQ round-trip.

        x: (N, head_dim). Splits into (N, n_sub, sub_dim), then does ONE batched
        squared-L2 (n_sub, N, K), one argmin, one gather -- replacing the Python
        loop over sub-blocks and its 3*n_sub small kernels.

        Chunks over N so peak memory (the (n_sub, chunk, K) distance tensor)
        stays bounded: at N=131072, n_sub=16, K=256 the un-chunked tensor is
        ~2 GB in fp32. The batched path is a GPU optimization -- on CPU it is
        roughly on par with the per-sub-block loop (no launch overhead to hide).
        """
        cb, mu, sd = stacked                      # cb: (n_sub, K, sub_dim)
        n_sub, K, sub_dim = cb.shape
        N = x.shape[0]
        c_sq = (cb * cb).sum(-1).unsqueeze(1)           # (n_sub, 1, K)
        if mu is not None:
            mu_b = mu.permute(1, 0, 2)                  # (n_sub, 1, sub_dim)
            sd_b = sd.permute(1, 0, 2)

        # cap the distance tensor at ~256 MB fp32: n_sub * chunk * K * 4 bytes
        chunk = max(1, (256 * 1024 * 1024) // (n_sub * K * 4))
        out_chunks = []
        for start in range(0, N, chunk):
            xc = x[start:start + chunk]                 # (c, head_dim)
            c = xc.shape[0]
            xb = xc.reshape(c, n_sub, sub_dim).permute(1, 0, 2).contiguous()
            if mu is not None:
                xb = (xb - mu_b) / sd_b
            x_sq = (xb * xb).sum(-1, keepdim=True)      # (n_sub, c, 1)
            cross = torch.bmm(xb, cb.transpose(1, 2))   # (n_sub, c, K)
            d2 = x_sq - 2 * cross + c_sq
            idx = d2.argmin(dim=-1)                     # (n_sub, c)
            idx_exp = idx.unsqueeze(-1).expand(-1, -1, sub_dim)
            rec = torch.gather(cb, 1, idx_exp)          # (n_sub, c, sub_dim)
            if mu is not None:
                rec = rec * sd_b + mu_b
            out_chunks.append(rec.permute(1, 0, 2).reshape(c, n_sub * sub_dim))
        return torch.cat(out_chunks, dim=0)

    def roundtrip_k(self, k):
        self._ensure_stacked()
        return self._roundtrip(k, self._k_stacked)

    def roundtrip_v(self, v):
        self._ensure_stacked()
        return self._roundtrip(v, self._v_stacked)

    def to(self, device):
        self.k_codebooks = [cb.to(device) for cb in self.k_codebooks]
        self.v_codebooks = [cb.to(device) for cb in self.v_codebooks]
        self._k_stats = [(st[0].to(device), st[1].to(device)) if st is not None else None
                         for st in self._k_stats]
        self._v_stats = [(st[0].to(device), st[1].to(device)) if st is not None else None
                         for st in self._v_stats]
        # invalidate cached stacks; they rebuild on next roundtrip on the new device
        self._k_stacked = None
        self._v_stacked = None
        return self

    def bits_per_element(self, head_dim):
        sub_dim = head_dim / self.n_sub
        return math.log2(self.n_codes) / sub_dim


@dataclass
class TurboQuantKV:
    """Data-oblivious rotation + per-coordinate scalar quantization, in the
    spirit of TurboQuant (Zandieh et al., ICLR 2026).

    Pipeline: random rotation Pi (so coordinates become near-iid in high dim),
    then a per-coordinate scalar codebook. We use a uniform Lloyd-Max-style
    codebook on the rotated coordinates as a faithful stand-in for their
    precomputed Beta-optimal codebook (the exact codebook is an implementation
    detail; the rotation is the load-bearing idea). We store the per-vector L2
    norm in fp16 and rescale on dequant, as the paper specifies.

    This is included so we can compare a rotation-based SCALAR method against
    product VQ on the attention-output COSINE metric -- the comparison the
    TurboQuant paper does not make (it optimizes cache-vector MSE / inner prod).
    """
    nbits: int = 4
    seed: int = 0
    _rot: torch.Tensor = None
    _levels: torch.Tensor = None

    def _make_rotation(self, d):
        g = torch.Generator().manual_seed(self.seed)
        a = torch.randn(d, d, generator=g)
        q, _ = torch.linalg.qr(a)
        return q

    def fit(self, k_calib, v_calib):
        d = k_calib.shape[-1]
        self._rot = self._make_rotation(d)
        # Fit levels on UNIT-NORMALIZED, rotated coordinates -- the same domain
        # the round-trip quantizes in. (Earlier bug: levels were fit on
        # un-normalized rotated data, so the per-vector-normalized values fell
        # outside the level range and collapsed to one bin.)
        kn = k_calib / k_calib.norm(dim=-1, keepdim=True).clamp_min(1e-8)
        rk = kn @ self._rot
        lo = torch.quantile(rk.flatten()[:200000], 0.001)
        hi = torch.quantile(rk.flatten()[:200000], 0.999)
        self._levels = torch.linspace(lo.item(), hi.item(), (1 << self.nbits))
        return self

    def _roundtrip(self, x):
        norms = x.norm(dim=-1, keepdim=True).clamp_min(1e-8)
        xn = x / norms
        y = xn @ self._rot                              # rotate
        idx = torch.bucketize(y, self._levels)
        idx = idx.clamp(0, self._levels.numel() - 1)
        y_hat = self._levels[idx]                       # per-coord scalar quant
        x_hat = y_hat @ self._rot.T                     # un-rotate
        return x_hat * norms                            # rescale by stored norm

    def roundtrip_k(self, k):
        return self._roundtrip(k)

    def roundtrip_v(self, v):
        return self._roundtrip(v)

    def to(self, device):
        self._rot    = self._rot.to(device)
        self._levels = self._levels.to(device)
        return self

    def bits_per_element(self, head_dim):
        # nbits/coord + fp16 norm amortized over head_dim
        return self.nbits + 16 / head_dim


@dataclass
class RoPESplitVQKV:
    """ProductVQ with separate codebooks for the RoPE'd vs pass-through halves
    of each KEY head.

    Laguna full-attention layers use partial_rotary_factor=0.5: the first half
    of head_dim is rotated (position-dependent, broad distribution), the second
    half is identity (position-independent). A single codebook must straddle
    two regimes; splitting lets each codebook specialize.

    Values receive no RoPE, so V uses a plain ProductVQ.
    """
    n_sub_half: int = 4          # sub-vectors PER HALF for keys
    n_codes: int = 256
    iters: int = 25
    rotary_fraction: float = 0.5
    _rope_vq: ProductVQKV = None
    _pass_vq: ProductVQKV = None
    _v_vq: ProductVQKV = None

    def fit(self, k_calib, v_calib, sample_weights=None, n_q_heads=None,
            k_struct=None):
        if sample_weights is None and k_struct is not None and k_struct.dim() == 3:
            from vqkv.metrics import calibration_sample_weights
            sample_weights = calibration_sample_weights(k_struct, n_q_heads)
        fit_kw = dict(sample_weights=sample_weights, n_q_heads=n_q_heads,
                      k_struct=k_struct)
        d = k_calib.shape[-1]
        cut = int(d * self.rotary_fraction)
        k_rope, k_pass = k_calib[..., :cut], k_calib[..., cut:]
        self._cut = cut
        self._rope_vq = ProductVQKV(self.n_sub_half, self.n_codes, self.iters,
                                     normalize=True).fit(k_rope, k_rope, **fit_kw)
        self._pass_vq = ProductVQKV(self.n_sub_half, self.n_codes, self.iters,
                                    normalize=False).fit(k_pass, k_pass, **fit_kw)
        self._v_vq = ProductVQKV(2 * self.n_sub_half, self.n_codes, self.iters,
                                 normalize=True).fit(v_calib, v_calib, **fit_kw)
        return self

    def roundtrip_k(self, k):
        kr = self._rope_vq.roundtrip_k(k[..., :self._cut])
        kp = self._pass_vq.roundtrip_k(k[..., self._cut:])
        return torch.cat([kr, kp], dim=-1)

    def roundtrip_v(self, v):
        return self._v_vq.roundtrip_v(v)

    def to(self, device):
        self._rope_vq.to(device)
        self._pass_vq.to(device)
        self._v_vq.to(device)
        return self

    def bits_per_element(self, head_dim):
        # both halves use n_sub_half codes over head_dim/2 elements
        half = head_dim / 2
        sub_dim = half / self.n_sub_half
        return math.log2(self.n_codes) / sub_dim


# ----------------------------------------------------------------------------
# Data-oblivious baseline: a simplified TurboQuant-style quantizer
# ----------------------------------------------------------------------------
@dataclass
class RandomRotationScalarKV:
    """Simplified, data-OBLIVIOUS rotation-then-scalar quantizer in the spirit
    of TurboQuant / PolarQuant (Zandieh et al., ICLR 2026).

    NOT the full method: TurboQuant adds PolarQuant's normalization-free polar
    transform and a 1-bit QJL residual correction for UNBIASED inner-product
    estimation. This stand-in captures only the core data-oblivious idea --
    apply a fixed random orthogonal rotation so coordinates concentrate (a
    Beta/Gaussian-like distribution), then scalar-quantize each coordinate with
    a fixed range. It exists so the harness can run the central scientific
    comparison of this project:

        data-OBLIVIOUS rotation+scalar   vs.   data-DEPENDENT product VQ

    on real Laguna cache statistics. If you want the real thing, drop in the
    unofficial impl (github.com/0xSero/turboquant or hackimov/turboquant-kv)
    behind this same fit/roundtrip interface. The OpenReview discussion of the
    paper is contested precisely on the oblivious-vs-data-dependent claim, so a
    clean head-to-head on a NEW model is a genuine contribution either way.
    """
    nbits: int = 3
    seed: int = 0
    _R: torch.Tensor = None        # rotation
    _Rk_range: tuple = None
    _Rv_range: tuple = None

    def _rotation(self, d):
        g = torch.Generator().manual_seed(self.seed)
        a = torch.randn(d, d, generator=g)
        q, _ = torch.linalg.qr(a)   # random orthogonal matrix
        return q

    def fit(self, k_calib, v_calib):
        d = k_calib.shape[-1]
        self._R = self._rotation(d)
        # fixed (data-oblivious-ish) ranges from a robust quantile of rotated calib
        rk = k_calib @ self._R
        rv = v_calib @ self._R
        self._Rk_range = (rk.quantile(0.001), rk.quantile(0.999))
        self._Rv_range = (rv.quantile(0.001), rv.quantile(0.999))
        return self

    def _rt(self, x, rng):
        xr = x @ self._R
        lo, hi = rng
        qmax = (1 << self.nbits) - 1
        scale = (hi - lo).clamp_min(1e-8) / qmax
        q = torch.clamp(torch.round((xr - lo) / scale), 0, qmax)
        xr_hat = q * scale + lo
        return xr_hat @ self._R.T   # rotation is orthogonal, inverse == transpose

    def roundtrip_k(self, k):
        return self._rt(k, self._Rk_range)

    def roundtrip_v(self, v):
        return self._rt(v, self._Rv_range)

    def bits_per_element(self, head_dim):
        # rotation is a fixed matrix (no per-token overhead); ranges are global.
        return float(self.nbits)


# ----------------------------------------------------------------------------
# 1-bit baselines (the floor of the scalar family; rate-neighbors to VQ)
# ----------------------------------------------------------------------------
@dataclass
class SignScalarKV:
    """Symmetric 1-bit (sign) quantization with a per-group scale.

        x_hat = scale * sign(x),   scale = mean(|x|) over the group.

    This is the honest 1-bit floor of the KIVI/quanto scalar family: unlike the
    affine `ScalarKV(nbits=1)`, it stores NO zero-point (cache K/V are ~zero-mean
    after norm), so it is both cheaper and better-centered. Group axis matches
    KIVI conventions: keys per-channel (token axis), values per-token.
    `per_channel_key` toggles the key axis.
    """
    per_channel_key: bool = True
    group_dim_k: int = 0     # 0 = per-channel (token axis); -1 = per-token
    bits: float = 1.0

    def fit(self, k_calib, v_calib):
        return self

    def _sign_q(self, x, dim):
        scale = x.abs().mean(dim=dim, keepdim=True).clamp_min(1e-8)
        return torch.sign(x) * scale

    def roundtrip_k(self, k):
        dim = 0 if self.per_channel_key else -1
        return self._sign_q(k, dim)

    def roundtrip_v(self, v):
        return self._sign_q(v, dim=-1)

    def bits_per_element(self, head_dim):
        # 1 bit/element + one fp16 scale per group amortized over the group.
        # per-token value group = head_dim elements; per-channel key group is the
        # token axis (large), so its scale overhead is negligible. Report ~1 +
        # 16/head_dim as a conservative upper bound for the per-token case.
        return 1.0 + 16.0 / head_dim


@dataclass
class TernaryScalarKV:
    """1.58-bit ternary quantization {-1, 0, +1} with a per-group scale, in the
    style of BitNet b1.58 -- the "just go (almost) 1-bit" school that the
    Bonsai/PrismML line popularized. Zeros are assigned by a threshold at a
    fraction of the mean-abs, letting small-magnitude coordinates drop out.

        thr = alpha * mean(|x|);  x_hat = scale * {sign(x) if |x|>thr else 0}

    Rate ~ log2(3) ~ 1.58 bits/element. Included so the table spans the full
    aggressive regime and so a reviewer who knows BitNet sees the comparison.
    """
    alpha: float = 0.7
    per_channel_key: bool = True

    def fit(self, k_calib, v_calib):
        return self

    def _tern_q(self, x, dim):
        m = x.abs().mean(dim=dim, keepdim=True).clamp_min(1e-8)
        thr = self.alpha * m
        mask = (x.abs() > thr).to(x.dtype)
        scale = (x.abs() * mask).sum(dim=dim, keepdim=True) / \
                mask.sum(dim=dim, keepdim=True).clamp_min(1.0)
        return torch.sign(x) * mask * scale

    def roundtrip_k(self, k):
        dim = 0 if self.per_channel_key else -1
        return self._tern_q(k, dim)

    def roundtrip_v(self, v):
        return self._tern_q(v, dim=-1)

    def bits_per_element(self, head_dim):
        import math as _m
        return _m.log2(3) + 16.0 / head_dim