""" 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)")