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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aither-kvcache
-Version: 2.0.1
+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.0.1 → aither_kvcache-2.2.0}/aither_kvcache/vllm/backend.py

@@ -306,7 +306,7 @@ def _make_backend_class():
 
         @staticmethod
         def get_supported_head_sizes():
-            return [64, 96, 128, 256]
+            return [64, 96, 128, 256, 512]
 
     return TurboQuantBackend
 

{aither_kvcache-2.0.1 → aither_kvcache-2.2.0}/aither_kvcache/vllm/engine.py

@@ -133,6 +133,27 @@ def _tq_dim_for_head(head_size: int) -> int:
         return packed_size(head_size, _TQ_BITS) + 4
 
 
+def _align_page_size(raw_bytes: int) -> int:
+    """Round page size up to the next power of 2 for cross-layer divisibility.
+
+    vLLM's unify_kv_cache_spec_page_size() requires all layer page sizes to be
+    divisible by the max page size. With heterogeneous head dims (e.g. Gemma 4:
+    256 + 512), the smaller layers produce page sizes that don't divide the
+    larger ones. Padding to next power-of-2 ensures any layer's padded page
+    size divides any other's (both are powers of 2 -> larger divides by smaller).
+
+    For uniform models, this is a no-op if page size is already a power of 2,
+    and at worst wastes a few bytes per block (acceptable for correctness).
+    """
+    if raw_bytes <= 0:
+        return raw_bytes
+    # Next power of 2 >= raw_bytes
+    p = 1
+    while p < raw_bytes:
+        p <<= 1
+    return p
+
+
 def _patch_page_size() -> bool:
     """
     Patch real_page_size_bytes on both AttentionSpec and FullAttentionSpec.
@@ -148,8 +169,11 @@ def _patch_page_size() -> bool:
 
     @property
     def _tq_real_page_size_bytes(self):
-        """Return TQ-compressed page size (handles both uniform and hybrid)."""
-        return _tq_page_size_bytes(self.block_size, self.num_kv_heads, self.head_size)
+        """Return TQ-compressed page size (handles both uniform and hybrid).
+        Pads to ensure cross-layer divisibility for heterogeneous head dims
+        (e.g. Gemma 4: head_dim=256 local + head_dim=512 global)."""
+        raw = _tq_page_size_bytes(self.block_size, self.num_kv_heads, self.head_size)
+        return _align_page_size(raw)
 
     # Patch BOTH classes -- FullAttentionSpec shadows parent's property
     AttentionSpec.real_page_size_bytes = _tq_real_page_size_bytes

{aither_kvcache-2.0.1 → aither_kvcache-2.2.0}/aither_kvcache/vllm/hooks.py

@@ -29,7 +29,7 @@ Set AITHER_TQ_DEBUG_STEPS=N to log the first N forward calls per layer-0
 import logging
 import os
 import sys
-from typing import Optional
+from typing import Dict, Optional, Tuple
 
 import torch
 
@@ -83,8 +83,9 @@ _hooks_applied = False
 _original_forward = None  # Saved reference to TritonAttentionImpl.forward
 
 # TQ state shared across all layers (class-level in the patched impl)
-_tq_quantizer: Optional[object] = None
-_tq_packed_dim: int = 0     # packed bytes per head (derived from quantizer)
+# Per-head-dim quantizers for heterogeneous models (e.g. Gemma 4: 256 + 512)
+_tq_quantizers: Dict[int, object] = {}    # head_dim -> TurboQuant instance
+_tq_packed_dims: Dict[int, int] = {}      # head_dim -> packed bytes per head
 _tq_fused_attn: dict = {}   # layer_idx -> TQPagedAttention instance
 _primary_k_norms: Optional[torch.Tensor] = None  # None for hybrid (norms embedded)
 _primary_v_norms: Optional[torch.Tensor] = None
@@ -96,6 +97,18 @@ _prefill_k_buf: Optional[torch.Tensor] = None
 _prefill_v_buf: Optional[torch.Tensor] = None
 _PREFILL_BUF_BLOCKS = 512
 
+def _get_tq(head_size: int) -> Tuple[object, int]:
+    """Get quantizer and packed_dim for a given head_size."""
+    return _tq_quantizers[head_size], _tq_packed_dims[head_size]
+
+
+def _any_tq() -> Optional[object]:
+    """Get any initialized quantizer (for operations that don't care about head_dim)."""
+    if _tq_quantizers:
+        return next(iter(_tq_quantizers.values()))
+    return None
+
+
 # Custom op flag -- set by _register_custom_ops()
 _USE_CUSTOM_OP = False
 
@@ -135,10 +148,9 @@ def _register_custom_ops():
             num_actual: int,
         ) -> torch.Tensor:
             """Encode new token + fused TQ attention. No graph break.
-            Accesses module globals (_tq_quantizer, _tq_fused_attn) in body --
+            Accesses module globals (_tq_quantizers, _tq_fused_attn) in body --
             these are initialized during vLLM warmup before CUDA graph capture."""
-            tq = _tq_quantizer
-            packed_dim = head_size // 2
+            tq, packed_dim = _get_tq(head_size)
             key_cache = kv_cache[:, 0]
             value_cache = kv_cache[:, 1]
             block_size = key_cache.shape[1]
@@ -260,8 +272,7 @@ def _register_hybrid_custom_ops():
             """
             global _hybrid_dk_buf, _hybrid_dv_buf
 
-            tq = _tq_quantizer
-            packed_dim = _tq_packed_dim
+            tq, packed_dim = _get_tq(head_size)
             key_cache = kv_cache[:, 0]
             value_cache = kv_cache[:, 1]
             block_size = key_cache.shape[1]
@@ -369,32 +380,38 @@ def _tq_init_layer(self_impl):
 
 @torch.compiler.disable
 def _ensure_quantizer(device, head_size, num_heads, num_kv_heads):
-    """Lazy-init TQ quantizer and fused attention on first real forward."""
-    global _tq_quantizer, _tq_packed_dim, _num_layers
-    if _tq_quantizer is not None:
-        return
+    """Lazy-init TQ quantizer for a given head_size. Creates/caches per unique head_dim.
+    Supports heterogeneous head dimensions (e.g. Gemma 4: 256 local + 512 global)."""
+    global _num_layers
+    if head_size in _tq_quantizers:
+        return  # Already initialized for this head_dim
 
-    _num_layers = _layer_counter
+    if not _tq_quantizers:  # First quantizer -- snapshot layer count
+        _num_layers = _layer_counter
 
     if _TQ_IS_HYBRID:
         from ..hybrid_quantizer import HybridTurboQuant
-        _tq_quantizer = HybridTurboQuant(
+        tq = HybridTurboQuant(
             head_dim=head_size, mode=_TQ_MODE, device=str(device),
             num_kv_heads=num_kv_heads)
-        _tq_quantizer.calibrate_uniform()
-        _tq_packed_dim = _tq_quantizer.layout.packed_dim
+        tq.calibrate_uniform()
+        pd = tq.layout.packed_dim
+        _tq_quantizers[head_size] = tq
+        _tq_packed_dims[head_size] = pd
         logger.info("TQ hybrid quantizer initialized: head_dim=%d, mode=%s, "
-                    "packed_dim=%d, %d layers",
-                    head_size, _TQ_MODE, _tq_packed_dim, _num_layers)
+                    "packed_dim=%d, %d layers (total quantizers: %d)",
+                    head_size, _TQ_MODE, pd, _num_layers, len(_tq_quantizers))
     else:
         from ..quantizer import TurboQuant
         from ..packing import packed_size
-        _tq_quantizer = TurboQuant(
+        tq = TurboQuant(
             head_dim=head_size, bits=_TQ_BITS, device=str(device))
-        _tq_packed_dim = packed_size(head_size, _TQ_BITS)
+        pd = packed_size(head_size, _TQ_BITS)
+        _tq_quantizers[head_size] = tq
+        _tq_packed_dims[head_size] = pd
         logger.info("TQ quantizer initialized: head_dim=%d, bits=%d, "
-                    "packed_dim=%d, %d layers",
-                    head_size, _TQ_BITS, _tq_packed_dim, _num_layers)
+                    "packed_dim=%d, %d layers (total quantizers: %d)",
+                    head_size, _TQ_BITS, pd, _num_layers, len(_tq_quantizers))
 
 
 @torch.compiler.disable
@@ -420,12 +437,11 @@ def _ensure_norms(num_layers, num_blocks, block_size, num_kv_heads, device):
 def _tq_encode_phase(layer_idx, key, value, kv_cache, slot_mapping):
     """Encode new K,V tokens into TQ primary cache. Runs eagerly.
     Supports both uniform (separate norms) and hybrid (norms embedded) modes."""
-    global _tq_quantizer, _primary_k_norms, _primary_v_norms
-    tq = _tq_quantizer
-    if tq is None:
+    global _primary_k_norms, _primary_v_norms
+    head_size = key.shape[-1]
+    if head_size not in _tq_quantizers:
         return
-
-    packed_dim = _tq_packed_dim
+    tq, packed_dim = _get_tq(head_size)
 
     key_cache = kv_cache[:, 0]    # [blocks, block_size, kv_heads, tq_dim]
     value_cache = kv_cache[:, 1]
@@ -508,18 +524,17 @@ def _tq_decompress_active(layer_idx, kv_cache, block_table, num_kv_heads,
     """Decompress active TQ blocks into a temporary fp8/bf16 buffer.
     Returns a standard-shaped kv_cache tensor for TritonAttention.
     Runs eagerly."""
-    global _tq_quantizer, _primary_k_norms, _primary_v_norms
+    global _primary_k_norms, _primary_v_norms
     global _prefill_k_buf, _prefill_v_buf
 
-    tq = _tq_quantizer
-    if tq is None:
+    if head_size not in _tq_quantizers:
         return kv_cache  # Fallback: return original (will fail but shouldn't happen)
+    tq, packed_dim = _get_tq(head_size)
 
     key_cache = kv_cache[:, 0]
     value_cache = kv_cache[:, 1]
     num_blocks = key_cache.shape[0]
     block_size = key_cache.shape[1]
-    packed_dim = head_size // 2
 
     # Allocate or reuse decompress buffers
     need_realloc = (
@@ -576,15 +591,16 @@ def _tq_fused_decode(layer_idx, query, kv_cache, attn_metadata, output,
                      num_heads, num_kv_heads, head_size, scale):
     """Fused TQ decode using TQPagedAttention kernel. No decompress needed.
     Runs eagerly (graph break)."""
-    global _tq_fused_attn, _tq_quantizer, _primary_k_norms, _primary_v_norms
+    global _tq_fused_attn, _primary_k_norms, _primary_v_norms
+
+    tq, packed_dim = _get_tq(head_size)
 
     # Lazy-init fused attention for this layer
     if layer_idx not in _tq_fused_attn:
         from ..fused_attention import TQPagedAttention
-        _tq_fused_attn[layer_idx] = TQPagedAttention(_tq_quantizer, num_heads)
+        _tq_fused_attn[layer_idx] = TQPagedAttention(tq, num_heads)
 
     fused = _tq_fused_attn[layer_idx]
-    packed_dim = head_size // 2
 
     key_cache = kv_cache[:, 0]
     value_cache = kv_cache[:, 1]
@@ -609,8 +625,7 @@ def _encode_inline(layer_idx, key, value, kv_cache, slot_mapping):
     """Encode K/V into TQ cache. No decorator -- called from _tq_decode_step.
     Branchless for decode: no .any() CPU-GPU syncs. Works for any N tokens.
     Supports both uniform and hybrid modes."""
-    tq = _tq_quantizer
-    packed_dim = _tq_packed_dim
+    tq, packed_dim = _get_tq(key.shape[-1])
     key_cache = kv_cache[:, 0]
     value_cache = kv_cache[:, 1]
     block_size = key_cache.shape[1]
@@ -658,14 +673,14 @@ def _tq_decode_step(layer_idx, query, key, value, kv_cache, attn_metadata,
         _encode_inline(layer_idx, key, value, kv_cache,
                        attn_metadata.slot_mapping)
 
+    tq, packed_dim = _get_tq(head_size)
+
     # Fused attention -- dict lookup for cached TQPagedAttention
     fused = _tq_fused_attn.get(layer_idx)
     if fused is None:
         from ..fused_attention import TQPagedAttention
-        fused = TQPagedAttention(_tq_quantizer, num_heads)
+        fused = TQPagedAttention(tq, num_heads)
         _tq_fused_attn[layer_idx] = fused
-
-    packed_dim = _tq_packed_dim
     key_cache = kv_cache[:, 0]
     value_cache = kv_cache[:, 1]
 
@@ -715,7 +730,12 @@ def _tq_prefill_sdpa(layer_idx, query, key, value, kv_cache, attn_metadata,
     Manual GQA expansion. is_causal=True for prefill, is_causal=False for decode
     (single query attends to all prior positions -- no mask needed)."""
     import torch.nn.functional as F
-    global _tq_quantizer, _primary_k_norms, _primary_v_norms
+    global _primary_k_norms, _primary_v_norms
+
+    if head_size not in _tq_quantizers:
+        output.fill_(0)
+        return output
+    tq_inst, packed_dim = _get_tq(head_size)
 
     num_actual = attn_metadata.num_actual_tokens
     cu_seqlens = attn_metadata.query_start_loc
@@ -729,7 +749,6 @@ def _tq_prefill_sdpa(layer_idx, query, key, value, kv_cache, attn_metadata,
     key_cache = kv_cache[:, 0]
     value_cache = kv_cache[:, 1]
     block_size = key_cache.shape[1]
-    packed_dim = _tq_packed_dim
     block_table = attn_metadata.block_table
 
     for i in range(num_seqs):
@@ -746,7 +765,7 @@ def _tq_prefill_sdpa(layer_idx, query, key, value, kv_cache, attn_metadata,
             # First prefill: all context is in current K/V
             ki = k_raw[q_start:q_end]
             vi = v_raw[q_start:q_end]
-        elif k_raw is not None and _tq_quantizer is not None:
+        elif k_raw is not None and tq_inst is not None:
             # Continuation: decompress prior blocks + append current
             prior_len = ctx_len - q_len
             prior_blocks = (prior_len + block_size - 1) // block_size
@@ -754,7 +773,7 @@ def _tq_prefill_sdpa(layer_idx, query, key, value, kv_cache, attn_metadata,
             bt = bt[bt >= 0]
 
             if bt.numel() > 0:
-                tq = _tq_quantizer
+                tq = tq_inst
                 H = num_kv_heads
 
                 kp = key_cache[bt, :, :, :packed_dim]
@@ -843,7 +862,7 @@ def _make_tq_forward(original_fwd):
         # Runs during BOTH normal execution AND CUDA graph capture.
         # Falls through to slow path for: init, prefill, missing key/value.
         if (hasattr(self, '_tq_layer_idx')
-                and _tq_quantizer is not None
+                and self.head_size in _tq_quantizers
                 and attn_metadata is not None
                 and output is not None
                 and key is not None and value is not None
@@ -892,7 +911,7 @@ def _make_tq_forward(original_fwd):
                 output_block_scale=output_block_scale)
 
         # --- Init quantizer + norms (guarded: no graph break after first call) ---
-        if _tq_quantizer is None:
+        if self.head_size not in _tq_quantizers:
             _ensure_quantizer(query.device, self.head_size,
                              self.num_heads, self.num_kv_heads)
         if not _TQ_IS_HYBRID and _primary_k_norms is None:
@@ -946,7 +965,7 @@ def _make_tq_forward(original_fwd):
         # --- Phase 2a: Decode (fused TQ attention for uniform, custom op for hybrid) ---
         # Hybrid mode uses tq::hybrid_decode_step custom op (zero graph breaks),
         # or falls through to encode+SDPA path if custom op not available.
-        if is_decode and _tq_quantizer is not None and _TQ_IS_HYBRID:
+        if is_decode and self.head_size in _tq_quantizers and _TQ_IS_HYBRID:
             if _USE_HYBRID_CUSTOM_OP and key is not None and value is not None:
                 num_actual = attn_metadata.num_actual_tokens
                 try:
@@ -965,7 +984,7 @@ def _make_tq_forward(original_fwd):
                                  self._tq_layer_idx, e)
                     # Fall through to graph-break SDPA
 
-        if is_decode and _tq_quantizer is not None and not _TQ_IS_HYBRID:
+        if is_decode and self.head_size in _tq_quantizers and not _TQ_IS_HYBRID:
             num_actual = attn_metadata.num_actual_tokens
             try:
                 if _TQ_DEBUG and self._tq_layer_idx == 0:
@@ -1119,8 +1138,7 @@ def _debug_prefill_sdpa(layer_idx, query, key, value, kv_cache, attn_metadata,
                           file=sys.stderr, flush=True)
 
                 # Decode first few vectors and check
-                tq = _tq_quantizer
-                packed_dim = _tq_packed_dim
+                tq, packed_dim = _get_tq(head_size)
                 kc = kv_cache[:, 0]
                 kp0 = kc[block_ids[0], 0, 0, :packed_dim]  # first slot, first head
                 kn0 = _primary_k_norms[0, block_ids[0], 0, 0]  # first slot, first head

{aither_kvcache-2.0.1 → 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.0.1 → aither_kvcache-2.2.0}/aither_kvcache.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aither-kvcache
-Version: 2.0.1
+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.0.1 → 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.0.1 → aither_kvcache-2.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "aither-kvcache"
-version = "2.0.1"
+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.0.1 → 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.0.1 → 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.0.1 → 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.0.1 → 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(),