update stale files

README.md

@@ -4,7 +4,7 @@ license: apache-2.0
 
 # AttnVQ — Attention-Aware KV Cache Quantization
 
-Training-free **product vector quantization** of the KV cache for long-context LLMs. AttnVQ fits small per-subspace codebooks with LBG, but scores distortion by **attention-output error** (and key cosine / inner-product bias), not cache MSE. Calibration is light: **10–15 agent traces, ~15 s on GPU** — enough to capture the **model's** K/V geometry (data-aware, not corpus-dependent).
+Training-free **product vector quantization** of the KV cache for long-context LLMs. AttnVQ fits small per-subspace codebooks with **attention-weighted batched LBG** (centroids weighted by key attention mass from GQA causal attention), but scores distortion by **attention-output error** (and key cosine / inner-product bias), not cache MSE. Calibration is light: **10–15 agent traces, ~15 s on GPU** — enough to capture the **model's** K/V geometry (data-aware, not corpus-dependent).
 
 Primary target: **Laguna-XS.2** (model-agnostic). Only the **10 full-attention layers** are compressed; 30 sliding-window layers stay fp16.
 
@@ -67,7 +67,7 @@ CODEBOOKS_PATH = "artifacts/codebooks.pt"
 codebooks = torch.load(CODEBOOKS_PATH, map_location="cuda", weights_only=False)
 
 # build cache
-quantizers, layers = codebooks["fitted"]["productvq-32x256-2b"], blob["meta"]["full_layers"]
+quantizers, layers = codebooks["fitted"]["productvq-32x256-2b"], codebooks["meta"]["full_layers"]
 cache = VQQuantizedCache(quantizers, layers)  # persists uint8 codebook indices
 
 # generate

benchmark.py

@@ -27,7 +27,8 @@ import torch
 from vqkv.quantizers import (ScalarKV, KIVIScalarKV, ProductVQKV, RoPESplitVQKV,
                             SignScalarKV, TernaryScalarKV)
 from vqkv.metrics import (key_cosine, cache_mse, inner_product_distortion,
-                          attention_output, attn_output_cosine, attn_output_error)
+                          attention_output, attn_output_cosine, attn_output_error,
+                          calibration_sample_weights)
 
 MODEL_ID = os.environ.get("LAGUNA_ID", "poolside/Laguna-XS.2")
 ARTIFACT_DIR = os.environ.get("VQKV_ARTIFACTS", "./artifacts")
@@ -296,17 +297,39 @@ def stage_fit(only: list[str] | None = None):
         print(f"[fit] GPU available ({torch.cuda.get_device_name()}) -- fitting LBG on CUDA")
     n_workers = 1 if fit_device == "cuda" else min(len(layer_ids), os.cpu_count() or 1)
 
+    attn_weighted = any(
+        isinstance(f(), (ProductVQKV, RoPESplitVQKV)) for _, f in configs
+    )
+    if attn_weighted:
+        print("[fit] ProductVQ / RoPESplit: attention-weighted LBG "
+              "(centroids weighted by key attention mass)")
+
+    n_q = meta.get("n_q_heads", 48)
     for name, factory in configs:
         t0 = time.time()
 
         def _fit_layer(i, _factory=factory, _device=fit_device):
             if calib is not None and i in calib:
-                kf = calib[i]["k"].reshape(-1, hd)[:200_000].to(_device)
-                vf = calib[i]["v"].reshape(-1, hd)[:200_000].to(_device)
+                k_struct = calib[i]["k"]
+                v_struct = calib[i]["v"]
+                if k_struct.shape[0] > 512:
+                    k_struct = k_struct[-512:]
+                    v_struct = v_struct[-512:]
+                w = calibration_sample_weights(k_struct, n_q)
+                kf = k_struct.reshape(-1, hd)[:200_000].to(_device)
+                vf = v_struct.reshape(-1, hd)[:200_000].to(_device)
+                if w is not None:
+                    w = w[:kf.shape[0]].to(_device)
             else:
                 kf = torch.zeros(1, hd, device=_device)
                 vf = torch.zeros(1, hd, device=_device)
-            q = _factory().fit(kf, vf)
+                w = k_struct = None
+            q = _factory()
+            if isinstance(q, (ProductVQKV, RoPESplitVQKV)):
+                q.fit(kf, vf, sample_weights=w, n_q_heads=n_q,
+                      k_struct=k_struct.to(_device) if k_struct is not None else None)
+            else:
+                q.fit(kf, vf)
             # Codebooks are saved to disk as CPU tensors; move back before returning.
             if _device != "cpu" and hasattr(q, "to"):
                 q.to("cpu")

generate.py

@@ -13,7 +13,7 @@ CODEBOOKS_PATH = "artifacts/codebooks.pt"
 codebooks = torch.load(CODEBOOKS_PATH, map_location="cuda", weights_only=False)
 
 # build cache
-quantizers, layers = codebooks["fitted"]["productvq-32x256-2b"], blob["meta"]["full_layers"]
+quantizers, layers = codebooks["fitted"]["productvq-32x256-2b"], codebooks["meta"]["full_layers"]
 cache = VQQuantizedCache(quantizers, layers)  # persists uint8 codebook indices
 
 # generate

vqkv/metrics.py

@@ -114,6 +114,53 @@ def key_cosine(k_ref, k_hat):
     return torch.nn.functional.cosine_similarity(a, b, dim=-1, eps=1e-8).mean().item()
 
 
+def key_attention_mass(k, n_q_heads, q=None, seed=0):
+    """Per-(token, kv-head) importance for attention-weighted codebook fitting.
+
+    k: (T, n_kv_heads, head_dim). Uses synthetic unit-norm queries (fixed seed)
+    when ``q`` is None — same idea as the cheap-metrics window.
+
+    Returns (T, n_kv_heads) non-negative masses; recent keys get more mass under
+    causal attention because more queries can attend to them.
+    """
+    T, n_kv, head_dim = k.shape
+    groups = n_q_heads // n_kv
+    if q is None:
+        g = torch.Generator(device=k.device).manual_seed(seed)
+        q = torch.randn(T, n_q_heads, head_dim, generator=g, device=k.device, dtype=k.dtype)
+        q = q / q.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+
+    k_exp = k.repeat_interleave(groups, dim=1)
+    qh = q.transpose(0, 1)
+    kh = k_exp.transpose(0, 1)
+    scores = torch.matmul(qh, kh.transpose(-1, -2)) * (head_dim ** -0.5)
+    causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=k.device))
+    scores = scores.masked_fill(~causal, float("-inf"))
+    attn = F.softmax(scores, dim=-1)  # (n_q_heads, T, T)
+
+    mass = torch.zeros(T, n_kv, device=k.device, dtype=attn.dtype)
+    for h in range(n_kv):
+        sl = slice(h * groups, (h + 1) * groups)
+        mass[:, h] = attn[sl, :, :].sum(dim=(0, 1))
+    return mass
+
+
+def calibration_sample_weights(k, n_q_heads, seed=0, max_tokens=512):
+    """Flattened row weights for ``k.reshape(-1, head_dim)`` LBG fitting.
+
+    k: (T, n_kv_heads, head_dim). Optionally truncates to the last ``max_tokens``
+    rows (matches the cheap-metrics attention window). Weights are normalized
+    to mean 1 over kept rows.
+    """
+    if k.dim() == 2:
+        return None
+    if k.shape[0] > max_tokens:
+        k = k[-max_tokens:]
+    mass = key_attention_mass(k, n_q_heads, seed=seed)
+    w = mass.reshape(-1)
+    return w / w.mean().clamp_min(1e-8)
+
+
 def inner_product_distortion(q, k_ref, k_hat, n_pairs=4096):
     """Relative error in the q.k inner products that drive attention logits.
 
@@ -135,4 +182,4 @@ def inner_product_distortion(q, k_ref, k_hat, n_pairs=4096):
     ip_hat = (qd[qi] * kh[ki]).sum(-1)
     rel = ((ip_hat - ip_ref).abs() / ip_ref.abs().clamp_min(1e-6)).mean().item()
     bias = (ip_hat - ip_ref).mean().item()
-    return {"ip_rel_err": round(rel, 5), "ip_bias": round(bias, 6)}
+    return {"ip_rel_err": round(rel, 5), "ip_bias": round(bias, 6)}

vqkv/quantizers.py

@@ -1,10 +1,11 @@
 """
 vqkv.quantizers — KV-cache quantizers for AttnVQ.
 
-ScalarKV, KIVIScalarKV, ProductVQKV (LBG product VQ), RoPESplitVQKV,
+ScalarKV, KIVIScalarKV, ProductVQKV (attention-weighted batched LBG), RoPESplitVQKV,
 SignScalarKV, TernaryScalarKV. Each exposes fit() + roundtrip_k/v() and
-bits_per_element(). Distortion is evaluated via vqkv.metrics (attention-output
-error, not cache MSE alone).
+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
@@ -88,15 +89,18 @@ class KIVIScalarKV:
 # 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) -> torch.Tensor:
-    """Linde-Buzo-Gray vector quantizer design (== Lloyd / k-means for MSE).
+                 seed: int = 0, sample_weights: torch.Tensor | None = None) -> torch.Tensor:
+    """Linde-Buzo-Gray / Lloyd design on sub-vectors.
 
-    data:    (N, d) training vectors
-    returns: (n_codes, d) codebook
+    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)
@@ -109,8 +113,9 @@ def lbg_codebook(data: torch.Tensor, n_codes: int, iters: int = 25,
         # Was: 256 Python iterations; now: 3 torch ops.
         new_cb = torch.zeros_like(cb)
         counts = torch.zeros(n_codes, dtype=data.dtype, device=dev)
-        new_cb.scatter_add_(0, assign.unsqueeze(1).expand(-1, d), data)
-        counts.scatter_add_(0, assign, ones)
+        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)
@@ -127,19 +132,24 @@ def lbg_codebook(data: torch.Tensor, n_codes: int, iters: int = 25,
 
 
 def lbg_codebook_batched(xb: torch.Tensor, n_codes: int, iters: int = 25,
-                          seed: int = 0) -> torch.Tensor:
-    """Fit n_sub independent LBG codebooks in one batched pass.
+                          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]
@@ -155,12 +165,14 @@ def lbg_codebook_batched(xb: torch.Tensor, n_codes: int, iters: int = 25,
         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
+        # 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)
-        new_cb.scatter_add_(1, assign_exp, xb)
-        counts.scatter_add_(1, assign, ones.unsqueeze(0).expand(n_sub, -1))
+        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)
@@ -236,7 +248,7 @@ class ProductVQKV:
         # x: (N, head_dim) -> list of (N, sub_dim)
         return list(torch.chunk(x, self.n_sub, dim=-1))
 
-    def _fit_one(self, x):
+    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
@@ -251,13 +263,18 @@ class ProductVQKV:
         else:
             stats = [None] * self.n_sub
 
-        cb_batched = lbg_codebook_batched(xb, self.n_codes, self.iters)
+        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):
-        self.k_codebooks, self._k_stats = self._fit_one(k_calib)
-        self.v_codebooks, self._v_stats = self._fit_one(v_calib)
+    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):
@@ -439,17 +456,23 @@ class RoPESplitVQKV:
     _pass_vq: ProductVQKV = None
     _v_vq: ProductVQKV = None
 
-    def fit(self, k_calib, v_calib):
+    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)
+                                     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)
+                                    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)
+                                 normalize=True).fit(v_calib, v_calib, **fit_kw)
         return self
 
     def roundtrip_k(self, k):