vqkv/compressed_cache.py

"""
vqkv.compressed_cache — inference-ready KV cache with real memory savings.

VQQuantizedCache persists uint8 codebook indices for target layers (not bf16
reconstructions). Dequantizes transiently per layer on each attention read.
Without a fused kernel this saves memory but not wall-clock; see memory_footprint().
"""

from __future__ import annotations

import math

import torch

try:
    from vqkv.quantizers import ProductVQKV, ScalarKV, KIVIScalarKV
except ModuleNotFoundError:  # flat layout (files beside this one)
    from quantizers import ProductVQKV, ScalarKV, KIVIScalarKV


# ----------------------------------------------------------------------------
# encode / decode for a *fitted* ProductVQKV  (no change to quantizers.py)
# ----------------------------------------------------------------------------
# These reach into the already-built batched stacks (cb, mu, sd) that
# ProductVQKV._ensure_stacked() prepares, so encode(x) then decode(idx) is
# bit-identical to the existing q.roundtrip_k(x). Promote them to methods on
# ProductVQKV if you prefer; kept as free functions to leave your file untouched.

def _stacked(q: ProductVQKV, which: str):
    q._ensure_stacked()
    return q._k_stacked if which == "k" else q._v_stacked


def pvq_encode(q: ProductVQKV, x: torch.Tensor, which: str = "k") -> torch.Tensor:
    """x: (N, head_dim) -> idx: (N, n_sub) uint8 (nearest codeword per sub-block)."""
    cb, mu, sd = _stacked(q, which)            # cb: (n_sub, K, sub_dim)
    n_sub, K, sub_dim = cb.shape
    if K > 256:
        raise ValueError(f"n_codes={K} > 256 needs >1 byte/index; this path is uint8. "
                         f"Use K<=256 (all headline configs do) or add bit-packing.")
    N = x.shape[0]
    c_sq = (cb * cb).sum(-1).unsqueeze(1)      # (n_sub, 1, K)
    mu_b = mu.permute(1, 0, 2) if mu is not None else None
    sd_b = sd.permute(1, 0, 2) if sd is not None else None
    chunk = max(1, (256 * 1024 * 1024) // (n_sub * K * 4))
    out = []
    for s0 in range(0, N, chunk):
        xc = x[s0:s0 + chunk]
        c = xc.shape[0]
        xb = xc.reshape(c, n_sub, sub_dim).permute(1, 0, 2).contiguous()
        if mu_b 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)
        idx = (x_sq - 2 * cross + c_sq).argmin(-1)        # (n_sub, c)
        out.append(idx.to(torch.uint8))
    return torch.cat(out, dim=1).T.contiguous()           # (N, n_sub) uint8


def pvq_decode(q: ProductVQKV, idx: torch.Tensor, which: str = "k") -> torch.Tensor:
    """idx: (N, n_sub) uint8 -> x_hat: (N, head_dim) (codebook dtype, e.g. fp32)."""
    cb, mu, sd = _stacked(q, which)
    n_sub, K, sub_dim = cb.shape
    mu_b = mu.permute(1, 0, 2) if mu is not None else None
    sd_b = sd.permute(1, 0, 2) if sd is not None else None
    idxT = idx.T.long()                                   # (n_sub, N)
    rec = torch.gather(cb, 1, idxT.unsqueeze(-1).expand(-1, -1, sub_dim))  # (n_sub,N,sub_dim)
    if mu_b is not None:
        rec = rec * sd_b + mu_b
    N = idxT.shape[1]
    return rec.permute(1, 0, 2).reshape(N, n_sub * sub_dim)


def pvq_codebook_bytes(q: ProductVQKV) -> int:
    """Fixed per-layer codebook overhead (bytes), amortized over all tokens."""
    cb_k, _, _ = _stacked(q, "k")
    cb_v, _, _ = _stacked(q, "v")
    return cb_k.numel() * cb_k.element_size() + cb_v.numel() * cb_v.element_size()


# ----------------------------------------------------------------------------
# Inference-ready cache: compressed store for target layers, native for rest
# ----------------------------------------------------------------------------
try:
    from transformers.cache_utils import DynamicCache
    _HAVE_TF = True
except Exception:  # let the file import for the standalone memory harness
    DynamicCache = object
    _HAVE_TF = False


class VQQuantizedCache(DynamicCache):
    """Drop-in cache that persists ProductVQ indices for ``target_layers`` and
    leaves all other layers in native precision (Laguna's sliding-window layers
    are bounded at 512 tokens and don't dominate, so we don't touch them).

    Memory model: ``key_cache``/``value_cache`` for target layers are never
    populated with full tensors. We keep ``k_codes[layer]`` / ``v_codes[layer]``
    as (seq, n_kv_heads, n_sub) uint8 and dequantize the whole buffer transiently
    each time the layer runs attention.

    NOTE ON transformers VERSIONS: signatures around DynamicCache shift between
    releases. This targets the modern
    ``update(key_states, value_states, layer_idx, cache_kwargs=None) -> (k, v)``
    interface and overrides ``get_seq_length``. If your pinned version calls
    additional hooks (``reorder_cache`` for beam search, ``crop`` for assisted
    decoding), forward them to the code buffers the same way ``update`` does.
    """

    def __init__(self, per_layer_quantizers: dict, target_layers, *a, **k):
        super().__init__(*a, **k)
        self.q = per_layer_quantizers
        self.target = set(int(i) for i in target_layers)
        self.k_codes: dict[int, torch.Tensor] = {}
        self.v_codes: dict[int, torch.Tensor] = {}
        self._dtype = None
        self._device = None

    # -- the hot path --------------------------------------------------------
    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
        if layer_idx not in self.target or layer_idx not in self.q:
            # native path for sliding-window / non-targeted layers
            return super().update(key_states, value_states, layer_idx, cache_kwargs)

        self._dtype = key_states.dtype
        self._device = key_states.device
        q = self.q[layer_idx]
        b, h, s, d = key_states.shape  # b == 1 in this harness

        # encode the NEW tokens (token-major, head-minor rows -> (s, h, n_sub))
        kf = key_states[0].transpose(0, 1).reshape(-1, d).float()
        vf = value_states[0].transpose(0, 1).reshape(-1, d).float()
        kc = pvq_encode(q, kf, "k").reshape(s, h, -1)
        vc = pvq_encode(q, vf, "v").reshape(s, h, -1)

        # append to the persistent compressed buffer
        if layer_idx in self.k_codes:
            self.k_codes[layer_idx] = torch.cat([self.k_codes[layer_idx], kc], dim=0)
            self.v_codes[layer_idx] = torch.cat([self.v_codes[layer_idx], vc], dim=0)
        else:
            self.k_codes[layer_idx] = kc
            self.v_codes[layer_idx] = vc

        # transiently dequantize the FULL buffer for this layer's attention
        allk, allv = self.k_codes[layer_idx], self.v_codes[layer_idx]
        S = allk.shape[0]
        kfull = pvq_decode(q, allk.reshape(-1, allk.shape[-1]), "k").reshape(S, h, d)
        vfull = pvq_decode(q, allv.reshape(-1, allv.shape[-1]), "v").reshape(S, h, d)
        kfull = kfull.permute(1, 0, 2)[None].to(self._dtype).to(self._device)
        vfull = vfull.permute(1, 0, 2)[None].to(self._dtype).to(self._device)
        return kfull, vfull

    def get_seq_length(self, layer_idx: int = 0) -> int:
        for li in sorted(self.target):
            if li in self.k_codes:
                return self.k_codes[li].shape[0]
        return super().get_seq_length(layer_idx) if _HAVE_TF else 0

    # -- live memory readout for the demo ------------------------------------
    def memory_footprint(self) -> dict:
        """Persistent bytes actually held on device, split by source."""
        code_bytes = sum(t.numel() for t in self.k_codes.values()) \
            + sum(t.numel() for t in self.v_codes.values())          # uint8 = 1 B
        cb_bytes = sum(pvq_codebook_bytes(self.q[li]) for li in self.k_codes)
        native_bytes = 0
        if _HAVE_TF:
            for kc in getattr(self, "key_cache", []):
                if isinstance(kc, torch.Tensor):
                    native_bytes += kc.numel() * kc.element_size()
            for vc in getattr(self, "value_cache", []):
                if isinstance(vc, torch.Tensor):
                    native_bytes += vc.numel() * vc.element_size()
        return {"compressed_indices_B": code_bytes,
                "codebooks_B": cb_bytes,
                "native_layers_B": native_bytes,
                "total_B": code_bytes + cb_bytes + native_bytes}


# ----------------------------------------------------------------------------
# Honest byte-accounting (model-free) -- drives the memory demo & the harness
# ----------------------------------------------------------------------------
class LagunaGeom:
    """Laguna-XS.2 cache geometry (from the config / proposal)."""
    n_layers = 40
    full_layers = 10          # full_attention layers that hold the growing cache
    sliding_layers = 30
    sliding_window = 512
    n_kv_heads = 8
    head_dim = 128


def kv_cache_bytes(context_len: int, bits_per_elt_full: float,
                   geom: LagunaGeom = LagunaGeom(),
                   bits_per_elt_sliding: float = 16.0) -> dict:
    """Total KV-cache bytes at a context length.

    Only the full-attention layers carry the growing cache; sliding layers are
    capped at ``sliding_window`` tokens. ``bits_per_elt_full`` is the rate the
    quantizer reports for the compressed layers (use 16.0 for the fp16 baseline).
    The K and V tensors are both counted.
    """
    elts_per_token = geom.n_kv_heads * geom.head_dim * 2          # K and V
    full_tokens = context_len * geom.full_layers
    slide_tokens = min(context_len, geom.sliding_window) * geom.sliding_layers
    full_B = full_tokens * elts_per_token * bits_per_elt_full / 8
    slide_B = slide_tokens * elts_per_token * bits_per_elt_sliding / 8
    return {"full_B": full_B, "sliding_B": slide_B, "total_B": full_B + slide_B}


def compression_vs_fp16(context_len: int, bits_per_elt_full: float,
                        geom: LagunaGeom = LagunaGeom()) -> float:
    base = kv_cache_bytes(context_len, 16.0, geom)["total_B"]
    comp = kv_cache_bytes(context_len, bits_per_elt_full, geom)["total_B"]
    return base / comp


# ----------------------------------------------------------------------------
# self-test: verify encode/decode == roundtrip and show the real byte win
# ----------------------------------------------------------------------------
if __name__ == "__main__":
    torch.manual_seed(0)
    N, hd = 4096, 128
    k = torch.randn(N, hd)
    v = torch.randn(N, hd)

    q = ProductVQKV(n_sub=32, n_codes=256, iters=10).fit(k, v)   # 2 bits/elt
    idx = pvq_encode(q, k, "k")
    k_hat = pvq_decode(q, idx, "k")
    rt = q.roundtrip_k(k)

    print(f"idx dtype/shape : {idx.dtype} {tuple(idx.shape)}")
    print(f"encode->decode == roundtrip : "
          f"{torch.allclose(k_hat, rt, atol=1e-4)}  "
          f"(max abs diff {(k_hat - rt).abs().max():.2e})")

    raw_B = k.numel() * 2                       # fp16
    comp_B = idx.numel() * 1                     # uint8 indices
    print(f"stored bytes  fp16={raw_B}  vq-2b={comp_B}  ratio={raw_B / comp_B:.1f}x")
    print(f"reported bits/elt: {q.bits_per_element(hd):.3f}")

    for L in (4096, 32768, 131072):
        r = compression_vs_fp16(L, q.bits_per_element(hd))
        gb = kv_cache_bytes(L, q.bits_per_element(hd))["total_B"] / 1e9
        base = kv_cache_bytes(L, 16.0)["total_B"] / 1e9
        print(f"context {L:>7}: fp16={base:6.2f} GB  vq-2b={gb:6.2f} GB  ({r:.1f}x)")