aither-kvcache 2.1.0__tar.gz → 2.2.0__tar.gz

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aither-kvcache
-Version: 2.1.0
+Version: 2.2.0
 Summary: Near-optimal KV cache compression for LLM inference — TurboQuant (vector quantization) + TriAttention (spectral compression)
 Author: Aitherium
 License-Expression: CC-BY-4.0

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/aither_kvcache/vllm/plugin.py

@@ -1,5 +1,5 @@
 """
-vLLM plugin entrypoint.
+vLLM plugin entrypoint for the aither-kvcache wheel.
 
 Registered via pyproject.toml entry_points under vllm.general_plugins.
 vLLM loads this at startup in every process (API server + engine workers).
@@ -7,6 +7,11 @@ vLLM loads this at startup in every process (API server + engine workers).
 Registers:
   1. TurboQuant CUSTOM attention backend (--attention-backend CUSTOM)
   2. Graph-aware eviction (replaces LRU with semantic scoring)
+  3. Per-request tenant isolation (default OFF; AITHER_REQUEST_TENANT_ISOLATION=1)
+
+CANONICAL SOURCE. Synced by .github/workflows/sync-kvcache.yml into the public
+Aitherium/aitherkvcache repo as aither_kvcache/vllm/plugin.py, alongside
+tenant_isolation.py. Edit here, not in the public repo.
 """
 
 import logging
@@ -24,8 +29,8 @@ def register():
     # 1. Register TurboQuant attention backend
     try:
         from vllm.v1.attention.backends.registry import (
-            register_backend,
             AttentionBackendEnum,
+            register_backend,
         )
         register_backend(
             AttentionBackendEnum.CUSTOM,
@@ -51,3 +56,16 @@ def register():
         except Exception as e:
             print(f"[aither-kvcache] Eviction install failed: {e}",
                   file=sys.stderr, flush=True)
+
+    # 3. Per-request tenant isolation (default OFF; AITHER_REQUEST_TENANT_ISOLATION=1).
+    # Rides this plugin so it loads in every process (API server + EngineCore
+    # workers). Non-fatal: a failure leaves the process on the instance tenant.
+    if os.environ.get("AITHER_REQUEST_TENANT_ISOLATION") == "1":
+        try:
+            from .tenant_isolation import install_tenant_isolation
+            install_tenant_isolation()
+        except ImportError:
+            pass
+        except Exception as e:
+            print(f"[aither-kvcache] Tenant isolation install failed: {e}",
+                  file=sys.stderr, flush=True)

aither_kvcache-2.2.0/aither_kvcache/vllm/tenant_isolation.py

@@ -0,0 +1,364 @@
+"""
+Per-request tenant isolation for the aither-kvcache vLLM plugin.
+
+This is the SELF-CONTAINED, in-wheel port of AitherOS's per-request KV-cache
+tenant isolation. It rides the ACTIVE path on the orchestrator: the aither-kvcache
+`vllm.general_plugins` entrypoint (`plugin.register()`) loads in every vLLM
+process (API server + EngineCore workers), so a hook installed here actually runs
+— unlike the AitherOS `lib.gpu.turboquant` import hooks, which are not mounted
+into the served image.
+
+Three seams, all fail-closed and non-fatal (a wrong symbol degrades to the
+instance tenant; it never crashes vLLM startup):
+
+  1. Serving layer (API-server process): read the gateway-signed OpenAI ``user``
+     field, verify its HMAC, and publish the tenant into a ContextVar.
+  2. Engine ``generate`` (API-server process): embed the verified tenant into the
+     real vLLM ``request_id`` (positional index 2 in vLLM 0.19.x), so it crosses
+     the frontend -> EngineCore process boundary verbatim.
+  3. Block manager (EngineCore worker): on ``KVCacheManager.allocate_slots`` the
+     ``request.request_id`` carries the embedded tenant — tag the request's blocks
+     with it; untag on ``free``.
+
+Self-contained by necessity: the wheel cannot import ``lib.gpu.turboquant``. The
+crypto/transport helpers mirror ``request_tenant_registry.py`` byte-for-byte so
+tokens the AitherOS gateway signs verify here.
+
+Gate: AITHER_REQUEST_TENANT_ISOLATION=1 (default OFF). Key: AITHER_TENANT_SIGNING_KEY
+(injected from AitherSecrets into the image env; empty => fail-closed).
+"""
+
+from __future__ import annotations
+
+import contextvars
+import hashlib
+import hmac
+import logging
+import os
+import sys
+import threading
+from typing import Any
+
+logger = logging.getLogger("aither.kvcache.tenant")
+
+_ENABLED = os.environ.get("AITHER_REQUEST_TENANT_ISOLATION", "0") == "1"
+
+# Must match request_tenant_registry.py exactly (cross-repo token compatibility).
+_TOKEN_SEP = "|"
+_RID_PREFIX = "tqt1:"
+_RID_SEP = "~"
+
+_pending_tenant: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "tq_pending_tenant", default=None
+)
+
+# block_id -> tenant_slug. Populated at allocate_slots, cleared at free. Guarded.
+_block_tenant: dict[int, str] = {}
+_block_lock = threading.Lock()
+
+_isolation_status: dict[str, bool] = {
+    "enabled": _ENABLED,
+    "serving_patched": False,
+    "engine_patched": False,
+    "block_manager_patched": False,
+    "key_present": False,
+    "verified": False,
+}
+
+
+# ---------------------------------------------------------------------------
+# Token crypto + request_id transport (mirror of request_tenant_registry.py)
+# ---------------------------------------------------------------------------
+def _signing_key() -> bytes:
+    return os.environ.get("AITHER_TENANT_SIGNING_KEY", "").encode("utf-8")
+
+
+def verify_tenant_token(
+    token: str | None, *, now: float | None = None
+) -> str | None:
+    """Return the tenant slug iff ``token`` has a valid, unexpired gateway signature."""
+    import time
+    key = _signing_key()
+    if not token or not key:
+        return None
+    parts = token.split(_TOKEN_SEP)
+    if len(parts) != 3:
+        return None
+    tenant, exp_str, sig = parts
+    if not tenant or not exp_str or not sig:
+        return None
+    try:
+        exp = int(exp_str)
+    except (TypeError, ValueError):
+        return None
+    expected = hmac.new(
+        key, f"{tenant}{_TOKEN_SEP}{exp}".encode(), hashlib.sha256
+    ).hexdigest()
+    if not hmac.compare_digest(sig, expected):
+        logger.warning("Rejected tenant token with invalid signature")
+        return None
+    if exp < (now if now is not None else time.time()):
+        logger.warning("Rejected expired tenant token")
+        return None
+    return tenant
+
+
+def embed_tenant_in_request_id(request_id: str, tenant_slug: str) -> str:
+    """Prefix a verified tenant onto a request_id (idempotent, boundary-safe)."""
+    if not request_id or not tenant_slug:
+        return request_id
+    rid = str(request_id)
+    if _RID_SEP in tenant_slug or rid.startswith(_RID_PREFIX):
+        return request_id
+    return f"{_RID_PREFIX}{tenant_slug}{_RID_SEP}{rid}"
+
+
+def extract_tenant_from_request_id(
+    request_id: str | None,
+) -> tuple[str | None, str | None]:
+    """Inverse of embed_tenant_in_request_id -> (tenant_or_None, original_id)."""
+    if not request_id:
+        return None, request_id
+    rid = str(request_id)
+    if not rid.startswith(_RID_PREFIX):
+        return None, request_id
+    tenant, sep, original = rid[len(_RID_PREFIX):].partition(_RID_SEP)
+    if not sep or not tenant:
+        return None, request_id
+    return tenant, original
+
+
+# ---------------------------------------------------------------------------
+# Block -> tenant map
+# ---------------------------------------------------------------------------
+def tag_blocks(tenant_slug: str, block_ids) -> None:
+    if not tenant_slug or not block_ids:
+        return
+    with _block_lock:
+        for bid in block_ids:
+            _block_tenant[int(bid)] = tenant_slug
+
+
+def untag_blocks(block_ids) -> None:
+    if not block_ids:
+        return
+    with _block_lock:
+        for bid in block_ids:
+            _block_tenant.pop(int(bid), None)
+
+
+def block_tenant(block_id: int) -> str | None:
+    with _block_lock:
+        return _block_tenant.get(int(block_id))
+
+
+def block_tenant_count() -> int:
+    with _block_lock:
+        return len(_block_tenant)
+
+
+# ---------------------------------------------------------------------------
+# Status
+# ---------------------------------------------------------------------------
+def get_isolation_status() -> dict[str, bool]:
+    _isolation_status["enabled"] = _ENABLED
+    _isolation_status["key_present"] = bool(_signing_key())
+    _isolation_status["verified"] = (
+        _isolation_status["enabled"]
+        and _isolation_status["serving_patched"]
+        and _isolation_status["engine_patched"]
+        and _isolation_status["block_manager_patched"]
+        and _isolation_status["key_present"]
+    )
+    return dict(_isolation_status)
+
+
+# ---------------------------------------------------------------------------
+# Seam 1: serving-layer wrap (verified tenant -> ContextVar)
+# ---------------------------------------------------------------------------
+def _wrap_serving_method(cls: Any, method_name: str) -> bool:
+    orig = getattr(cls, method_name, None)
+    if orig is None:
+        return False
+    if getattr(orig, "_tq_tenant_wrapped", False):
+        return True
+
+    async def _tenant_serving_wrapper(self, request, *args, **kwargs):
+        token = None
+        try:
+            user = getattr(request, "user", None)
+            tenant = verify_tenant_token(user) if user else None
+            if tenant:
+                token = _pending_tenant.set(tenant)
+        except (AttributeError, TypeError, ValueError) as e:
+            logger.debug("tenant capture skipped: %s", e)
+            token = None
+        try:
+            return await orig(self, request, *args, **kwargs)
+        finally:
+            if token is not None:
+                _pending_tenant.reset(token)
+
+    _tenant_serving_wrapper._tq_tenant_wrapped = True
+    setattr(cls, method_name, _tenant_serving_wrapper)
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Seam 2: engine generate wrap (embed tenant into positional request_id)
+# ---------------------------------------------------------------------------
+# vLLM 0.19.x: AsyncLLM.generate(self, prompt, sampling_params, request_id, *, ...)
+# request_id is the 3rd positional arg -> index 2 in the wrapper's *args.
+_GENERATE_RID_POS = 2
+
+
+def _wrap_engine_generate(cls: Any, rid_pos: int = _GENERATE_RID_POS) -> bool:
+    orig = getattr(cls, "generate", None)
+    if orig is None:
+        return False
+    if getattr(orig, "_tq_tenant_wrapped", False):
+        return True
+
+    async def _tenant_generate_wrapper(self, *args, **kwargs):
+        tenant = _pending_tenant.get(None)
+        if tenant:
+            rid = None
+            pos = None
+            if isinstance(kwargs.get("request_id"), str) and kwargs["request_id"]:
+                rid = kwargs["request_id"]
+            elif (len(args) > rid_pos and isinstance(args[rid_pos], str)
+                  and args[rid_pos]):
+                pos = rid_pos
+                rid = args[pos]
+            if rid:
+                embedded = embed_tenant_in_request_id(rid, tenant)
+                if pos is not None:
+                    args = args[:pos] + (embedded,) + args[pos + 1:]
+                else:
+                    kwargs = dict(kwargs)
+                    kwargs["request_id"] = embedded
+        async for out in orig(self, *args, **kwargs):
+            yield out
+
+    _tenant_generate_wrapper._tq_tenant_wrapped = True
+    cls.generate = _tenant_generate_wrapper
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Seam 3: block manager wrap (tag blocks by tenant from request.request_id)
+# ---------------------------------------------------------------------------
+def _wrap_block_manager(cls: Any) -> bool:
+    alloc = getattr(cls, "allocate_slots", None)
+    free = getattr(cls, "free", None)
+    if alloc is None:
+        return False
+    if getattr(alloc, "_tq_tenant_wrapped", False):
+        return True
+
+    def _flatten_block_ids(mgr, request_id) -> list:
+        try:
+            groups = mgr.get_block_ids(request_id)
+        except Exception:
+            return []
+        ids: list = []
+        for g in groups or ():
+            try:
+                ids.extend(int(b) for b in g)
+            except TypeError:
+                ids.append(int(g))
+        return ids
+
+    def _tenant_allocate_slots(self, request, *args, **kwargs):
+        result = alloc(self, request, *args, **kwargs)
+        if result is not None:
+            try:
+                rid = getattr(request, "request_id", None)
+                tenant, _orig = extract_tenant_from_request_id(rid)
+                if tenant:
+                    tag_blocks(tenant, _flatten_block_ids(self, rid))
+            except Exception as e:  # never break allocation for tagging
+                logger.debug("tenant block-tag skipped: %s", e)
+        return result
+
+    _tenant_allocate_slots._tq_tenant_wrapped = True
+    cls.allocate_slots = _tenant_allocate_slots
+
+    if free is not None and not getattr(free, "_tq_tenant_wrapped", False):
+        def _tenant_free(self, request, *args, **kwargs):
+            try:
+                rid = getattr(request, "request_id", None)
+                if rid:
+                    untag_blocks(_flatten_block_ids(self, rid))
+            except Exception as e:
+                logger.debug("tenant block-untag skipped: %s", e)
+            return free(self, request, *args, **kwargs)
+
+        _tenant_free._tq_tenant_wrapped = True
+        cls.free = _tenant_free
+
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Install
+# ---------------------------------------------------------------------------
+def install_tenant_isolation() -> bool:
+    """Install all three seams. Called from plugin.register() when enabled.
+
+    Returns True iff verified (serving + engine + block manager + key). Every
+    failure is non-fatal and leaves the process serving with the instance tenant.
+    """
+    if not _ENABLED:
+        logger.debug("tenant isolation disabled (AITHER_REQUEST_TENANT_ISOLATION!=1)")
+        return False
+
+    serving_targets = [
+        ("vllm.entrypoints.openai.chat_completion.serving", "OpenAIServingChat",
+         "create_chat_completion"),
+        ("vllm.entrypoints.openai.serving_chat", "OpenAIServingChat",
+         "create_chat_completion"),
+        ("vllm.entrypoints.openai.completion.serving", "OpenAIServingCompletion",
+         "create_completion"),
+        ("vllm.entrypoints.openai.serving_completion", "OpenAIServingCompletion",
+         "create_completion"),
+    ]
+    seen = set()
+    for modpath, clsname, method in serving_targets:
+        try:
+            mod = __import__(modpath, fromlist=[clsname])
+            cls = getattr(mod, clsname)
+            if cls in seen:
+                continue
+            if _wrap_serving_method(cls, method):
+                seen.add(cls)
+                _isolation_status["serving_patched"] = True
+        except Exception as e:
+            logger.debug("serving wrap skipped %s.%s: %s", clsname, method, e)
+
+    for modpath, clsname in [
+        ("vllm.v1.engine.async_llm", "AsyncLLM"),
+        ("vllm.engine.async_llm_engine", "AsyncLLMEngine"),
+    ]:
+        try:
+            mod = __import__(modpath, fromlist=[clsname])
+            if _wrap_engine_generate(getattr(mod, clsname)):
+                _isolation_status["engine_patched"] = True
+        except Exception as e:
+            logger.debug("engine wrap skipped %s: %s", clsname, e)
+
+    try:
+        from vllm.v1.core.kv_cache_manager import KVCacheManager
+        if _wrap_block_manager(KVCacheManager):
+            _isolation_status["block_manager_patched"] = True
+    except Exception as e:
+        logger.debug("block-manager wrap skipped: %s", e)
+
+    st = get_isolation_status()
+    print(
+        f"[aither-kvcache] tenant isolation: serving={st['serving_patched']} "
+        f"engine={st['engine_patched']} block_mgr={st['block_manager_patched']} "
+        f"key={st['key_present']} VERIFIED={st['verified']}",
+        file=sys.stderr, flush=True,
+    )
+    return st["verified"]

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/aither_kvcache.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aither-kvcache
-Version: 2.1.0
+Version: 2.2.0
 Summary: Near-optimal KV cache compression for LLM inference — TurboQuant (vector quantization) + TriAttention (spectral compression)
 Author: Aitherium
 License-Expression: CC-BY-4.0

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/aither_kvcache.egg-info/SOURCES.txt

@@ -36,6 +36,7 @@ aither_kvcache/vllm/eviction_plugin.py
 aither_kvcache/vllm/hooks.py
 aither_kvcache/vllm/plugin.py
 aither_kvcache/vllm/sitecustomize.py
+aither_kvcache/vllm/tenant_isolation.py
 tests/test_core.py
 tests/test_fused.py
 tests/test_hybrid_hooks.py

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "aither-kvcache"
-version = "2.1.0"
+version = "2.2.0"
 description = "Near-optimal KV cache compression for LLM inference — TurboQuant (vector quantization) + TriAttention (spectral compression)"
 readme = "README.md"
 license = "CC-BY-4.0"

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/turboquant/__init__.py

@@ -22,6 +22,7 @@ Usage:
 """
 
 from .quantizer import TurboQuant, TurboQuantConfig
+from .gb10_fp8_rescue import get_gb10_fp8_status, install_gb10_fp8_rescue
 
 __all__ = [
     "TurboQuant",
@@ -34,6 +35,8 @@ __all__ = [
     "get_tier_cache_bridge",
     "get_strata_cache_shadow",
     "get_graph_eviction_advisor",
+    "install_gb10_fp8_rescue",
+    "get_gb10_fp8_status",
 ]
 __version__ = "0.9.0"
 

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/turboquant/codebook.py

@@ -82,6 +82,66 @@ def get_codebook(d: int, bits: int) -> Tuple[np.ndarray, np.ndarray, float]:
     return centroids, boundaries_inner, mse_total
 
 
+def get_correction_factors(d: int, bits: int) -> np.ndarray:
+    """
+    Compute per-centroid inner-product bias correction factors.
+
+    For Lloyd-Max quantization on N(0, 1/d), the reconstructed inner product
+    <q, x_hat> is a biased estimator of <q, x>. The correction factor for
+    centroid i is: alpha_i = E[x | x in bin_i] / c_i.
+
+    For a perfect Lloyd-Max codebook this is ~1.0, but finite-precision
+    centroids and boundary discretization introduce small biases. Computing
+    these analytically and applying them removes the systematic IP bias.
+
+    Returns:
+        correction: array of shape [2^bits], values close to 1.0.
+    """
+    if bits not in _STANDARD_CODEBOOKS:
+        raise ValueError(f"Supported bit-widths: {list(_STANDARD_CODEBOOKS.keys())}. Got {bits}")
+
+    entry = _STANDARD_CODEBOOKS[bits]
+    std_centroids = entry["centroids"]
+    num_levels = 1 << bits
+    sigma = 1.0 / math.sqrt(d)
+
+    # Standard N(0,1) boundaries
+    boundaries = np.empty(num_levels + 1)
+    boundaries[0] = -10.0  # effectively -inf
+    boundaries[-1] = 10.0  # effectively +inf
+    for i in range(num_levels - 1):
+        boundaries[i + 1] = (std_centroids[i] + std_centroids[i + 1]) / 2
+
+    # Compute E[x | x in bin_i] for N(0,1)
+    # E[x | b_lo <= x < b_hi] = (phi(b_lo) - phi(b_hi)) / (Phi(b_hi) - Phi(b_lo))
+    # where phi = pdf, Phi = cdf of N(0,1)
+    from math import exp as _exp, pi as _pi
+
+    def _phi(x):
+        return (1.0 / math.sqrt(2 * _pi)) * _exp(-0.5 * x * x)
+
+    def _Phi(x):
+        # Abramowitz & Stegun approximation (accurate to 1e-7)
+        t = 1.0 / (1.0 + 0.2316419 * abs(x))
+        poly = t * (0.319381530 + t * (-0.356563782 + t * (
+            1.781477937 + t * (-1.821255978 + t * 1.330274429))))
+        val = 1.0 - _phi(x) * poly
+        return val if x >= 0 else 1.0 - val
+
+    corrections = np.ones(num_levels, dtype=np.float64)
+    for i in range(num_levels):
+        b_lo = boundaries[i]
+        b_hi = boundaries[i + 1]
+        prob = _Phi(b_hi) - _Phi(b_lo)
+        if prob < 1e-15 or abs(std_centroids[i]) < 1e-15:
+            corrections[i] = 1.0
+            continue
+        cond_mean = (_phi(b_lo) - _phi(b_hi)) / prob
+        corrections[i] = cond_mean / std_centroids[i]
+
+    return corrections.astype(np.float32)
+
+
 def get_theory_bounds(bits: int) -> Tuple[float, float]:
     """Return (lower_bound, upper_bound) on MSE for b-bit quantization."""
     lower = 1.0 / (4 ** bits)

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/turboquant/fused_attention.py

@@ -66,9 +66,11 @@ class TQPagedAttentionRef:
         head_dim: int,
         block_size: int,
         bits: int = 4,
+        k_centroids: Optional[torch.Tensor] = None,
     ):
         self.rotation = rotation          # [head_dim, head_dim]
-        self.centroids = centroids        # [num_levels]
+        self.centroids = centroids        # [num_levels] — used for V accumulation
+        self.k_centroids = k_centroids if k_centroids is not None else centroids  # for K scoring
         self.scale = scale
         self.num_kv_heads = num_kv_heads
         self.num_query_heads = num_query_heads
@@ -81,12 +83,14 @@ class TQPagedAttentionRef:
         from .packing import unpack_4bit, unpack_3bit, unpack_2bit
         self._unpack_fns = {4: unpack_4bit, 3: unpack_3bit, 2: unpack_2bit}
 
-    def _unpack_to_full(self, packed_vec: torch.Tensor) -> torch.Tensor:
+    def _unpack_to_full(self, packed_vec: torch.Tensor,
+                        codebook: Optional[torch.Tensor] = None) -> torch.Tensor:
         """Unpack a single packed vector to full head_dim via codebook lookup."""
         # packed_vec: [packed_dim] uint8
+        cb = codebook if codebook is not None else self.centroids
         unpack_fn = self._unpack_fns.get(self.bits, self._unpack_fns[4])
         indices = unpack_fn(packed_vec.unsqueeze(0), self.head_dim).squeeze(0)
-        return self.centroids[indices.long()]
+        return cb[indices.long()]
 
     def forward(
         self,
@@ -136,9 +140,9 @@ class TQPagedAttentionRef:
                     n_valid = end - start
 
                     for pos in range(n_valid):
-                        # Unpack K → full rotated vector via codebook
+                        # Unpack K → full rotated vector via corrected codebook
                         packed_k = k_packed[phys_block, pos, kv_head]
-                        k_full = self._unpack_to_full(packed_k)
+                        k_full = self._unpack_to_full(packed_k, self.k_centroids)
                         k_even = k_full[0::2]
                         k_odd = k_full[1::2]
 
@@ -214,8 +218,9 @@ if HAS_TRITON:
         # Paged block indirection
         block_tables_ptr,     # [num_seqs, max_blocks_per_seq] int32/int64
         context_lens_ptr,     # [num_seqs] int32/int64
-        # Codebook
-        centroids_ptr,        # [16] float32
+        # Codebooks (separate for K scoring and V accumulation)
+        k_centroids_ptr,      # [16] float32 — corrected centroids for K scoring
+        v_centroids_ptr,      # [16] float32 — raw centroids for V accumulation
         # Output accumulators (rotated domain, even/odd split)
         out_even_ptr,         # [num_seqs, num_q_heads, HALF_D] float32
         out_odd_ptr,          # [num_seqs, num_q_heads, HALF_D] float32
@@ -300,8 +305,8 @@ if HAS_TRITON:
                             k_packed_ptr + kp_base + half_offs
                         ).to(tl.int32)
 
-                        k_even = tl.load(centroids_ptr + ((kp >> 4) & 0xF))
-                        k_odd = tl.load(centroids_ptr + (kp & 0xF))
+                        k_even = tl.load(k_centroids_ptr + ((kp >> 4) & 0xF))
+                        k_odd = tl.load(k_centroids_ptr + (kp & 0xF))
 
                         # ── Score ──
                         k_norm = tl.load(
@@ -332,8 +337,8 @@ if HAS_TRITON:
                             v_packed_ptr + vp_base + half_offs
                         ).to(tl.int32)
 
-                        v_even = tl.load(centroids_ptr + ((vp >> 4) & 0xF))
-                        v_odd = tl.load(centroids_ptr + (vp & 0xF))
+                        v_even = tl.load(v_centroids_ptr + ((vp >> 4) & 0xF))
+                        v_odd = tl.load(v_centroids_ptr + (vp & 0xF))
                         v_norm = tl.load(
                             v_norms_ptr
                             + phys_block * stride_kn_block
@@ -371,7 +376,7 @@ if HAS_TRITON:
         k_packed_ptr, k_norms_ptr,
         v_packed_ptr, v_norms_ptr,
         block_tables_ptr, context_lens_ptr,
-        centroids_ptr,
+        k_centroids_ptr, v_centroids_ptr,
         # Partial outputs: [num_q_heads, num_seqs, NUM_SPLITS, HALF_D]
         part_even_ptr, part_odd_ptr,
         # Partial softmax state: [num_q_heads, num_seqs, NUM_SPLITS]
@@ -431,8 +436,8 @@ if HAS_TRITON:
                             k_packed_ptr + kp_base + half_offs
                         ).to(tl.int32)
 
-                        k_even = tl.load(centroids_ptr + ((kp >> 4) & 0xF))
-                        k_odd = tl.load(centroids_ptr + (kp & 0xF))
+                        k_even = tl.load(k_centroids_ptr + ((kp >> 4) & 0xF))
+                        k_odd = tl.load(k_centroids_ptr + (kp & 0xF))
 
                         k_norm = tl.load(
                             k_norms_ptr + phys_block * stride_kn_block
@@ -457,8 +462,8 @@ if HAS_TRITON:
                             v_packed_ptr + vp_base + half_offs
                         ).to(tl.int32)
 
-                        v_even = tl.load(centroids_ptr + ((vp >> 4) & 0xF))
-                        v_odd = tl.load(centroids_ptr + (vp & 0xF))
+                        v_even = tl.load(v_centroids_ptr + ((vp >> 4) & 0xF))
+                        v_odd = tl.load(v_centroids_ptr + (vp & 0xF))
                         v_norm = tl.load(
                             v_norms_ptr + phys_block * stride_kn_block
                             + pos * stride_kn_pos + kv_head)
@@ -560,6 +565,9 @@ class TQPagedAttention:
         self.rotation = tq.rotation
         self.rotation_T = tq.rotation.T.contiguous()  # pre-transposed, contiguous
         self.centroids = tq.centroids
+        # Corrected centroids for unbiased inner product scoring (K path only).
+        # V accumulation uses raw centroids (correction only affects score bias).
+        self.k_centroids = tq.centroids * tq.correction_factors
         self.scale = 1.0 / math.sqrt(tq.head_dim)
         self.head_dim = tq.head_dim
         self.num_levels = tq.num_levels
@@ -577,6 +585,7 @@ class TQPagedAttention:
             head_dim=tq.head_dim,
             block_size=16,
             bits=tq.bits,
+            k_centroids=self.k_centroids,
         )
 
         # Block selector for sparse attention (optional)
@@ -697,7 +706,7 @@ class TQPagedAttention:
                 q_even, q_odd,
                 k_packed, k_norms, v_packed, v_norms,
                 block_tables, context_lens,
-                self.centroids,
+                self.k_centroids, self.centroids,
                 part_even, part_odd, part_m, part_l,
                 q_even.stride(0), q_even.stride(1),
                 k_packed.stride(0), k_packed.stride(1), k_packed.stride(2),
@@ -724,7 +733,7 @@ class TQPagedAttention:
                 q_even, q_odd,
                 k_packed, k_norms, v_packed, v_norms,
                 block_tables, context_lens,
-                self.centroids,
+                self.k_centroids, self.centroids,
                 out_even, out_odd,
                 q_even.stride(0), q_even.stride(1),
                 k_packed.stride(0), k_packed.stride(1), k_packed.stride(2),

{aither_kvcache-2.1.0 → aither_kvcache-2.2.0}/turboquant/quantizer.py

@@ -23,7 +23,7 @@ import torch
 from dataclasses import dataclass
 from typing import Optional, Tuple
 
-from .codebook import get_codebook, get_theory_bounds
+from .codebook import get_codebook, get_theory_bounds, get_correction_factors
 from .rotation import random_orthogonal, randomized_hadamard_matrix
 from .packing import (
     pack_4bit, unpack_4bit,
@@ -102,6 +102,12 @@ class TurboQuant:
         # Theory bounds
         self.theory_lower, self.theory_upper = get_theory_bounds(config.bits)
 
+        # Per-centroid correction factors for unbiased inner product estimation
+        correction_np = get_correction_factors(config.head_dim, config.bits)
+        self.correction_factors = torch.tensor(
+            correction_np, dtype=torch.float32, device=config.device,
+        )
+
         # Pack/unpack dispatch
         self._pack_fn = {2: pack_2bit, 3: pack_3bit, 4: pack_4bit}[config.bits]
         self._unpack_fn = {2: unpack_2bit, 3: unpack_3bit, 4: unpack_4bit}[config.bits]
@@ -211,6 +217,12 @@ class TurboQuant:
         indices = self._unpack_fn(packed, self.head_dim)
         return self.centroids[indices.long()]
 
+    def _pytorch_decode_corrected(self, packed: torch.Tensor) -> torch.Tensor:
+        """Unpack + codebook lookup with per-centroid IP bias correction."""
+        indices = self._unpack_fn(packed, self.head_dim)
+        idx_long = indices.long()
+        return self.centroids[idx_long] * self.correction_factors[idx_long]
+
     # ================================================================
     # INTERNAL: Triton path (GPU only)
     # ================================================================
@@ -295,6 +307,16 @@ class TurboQuant:
         ip_mse = (true_ip - est_ip).pow(2).mean().item()
         ip_bias = (est_ip - true_ip).mean().item()
 
+        # Corrected decode: apply per-centroid bias correction factors
+        packed_flat = packed.reshape(-1, packed.shape[-1])
+        norms_flat = norms.reshape(-1).float()
+        y_hat_corr = self._pytorch_decode_corrected(packed_flat)
+        x_hat_corr = torch.matmul(y_hat_corr, self.rotation)
+        x_hat_corr = x_hat_corr * norms_flat.unsqueeze(-1)
+        x_hat_corr = x_hat_corr.reshape(x.shape)
+        est_ip_corr = (x_hat_corr[idx_a] * x[idx_b]).sum(dim=-1)
+        ip_bias_corrected = (est_ip_corr - true_ip).mean().item()
+
         # Max absolute error
         max_err = (x - x_hat).abs().max().item()
 
@@ -308,6 +330,7 @@ class TurboQuant:
             "mse_ratio_to_lower": mse / self.theory_lower if self.theory_lower > 0 else float("inf"),
             "ip_mse": ip_mse,
             "ip_bias": ip_bias,
+            "ip_bias_corrected": ip_bias_corrected,
             "max_abs_error": max_err,
             "compression_vs_fp16": self.compression_ratio(),
             "compression_vs_fp8": self.compression_ratio_vs_fp8(),