186 lines
6.4 KiB
Python
186 lines
6.4 KiB
Python
"""Per-class cache pools.
|
|
|
|
No single unified block pool: cache classes differ by 150-500x in natural granularity, so
|
|
each gets its own statically budgeted pool behind one ``CacheHandle``. Four classes:
|
|
* ``FeatureCache`` — content-hash keyed, partitioned by adapter+weights (text/vision encoders)
|
|
* ``ResidualCache`` — cache-dit residuals, scoped per request AND per CFG branch
|
|
* ``SlabKVCache`` — chunk-KV slabs for causal/world models; can disable recycle
|
|
* ``PagedKVCache`` — paged text-KV stub for ar_decode (phase-2 omni); minority case, lazy
|
|
|
|
KV is the minority case — a pure bidirectional deployment (Wan/LTX T2V) allocates none of it.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections import OrderedDict
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from v2.runtime.cache.keys import CacheKey, CachePolicy
|
|
|
|
|
|
def _nbytes(value: Any) -> int:
|
|
if hasattr(value, "nbytes"):
|
|
return int(value.nbytes)
|
|
return 0
|
|
|
|
|
|
class _Pool:
|
|
|
|
def __init__(self, policy: CachePolicy):
|
|
self.policy = policy
|
|
self.used_bytes = 0
|
|
self.hits = 0
|
|
self.misses = 0
|
|
|
|
|
|
class FeatureCache(_Pool):
|
|
"""Content-hash keyed, budget-aware FIFO/LRU.
|
|
|
|
Partitioned by ``adapter_versions``/``weights_version`` through the CacheKey, so two
|
|
workflows sharing a prompt but differing in te-LoRA stack never serve stale embeddings.
|
|
"""
|
|
|
|
def __init__(self, policy: CachePolicy):
|
|
super().__init__(policy)
|
|
self._store: OrderedDict[str, tuple[Any, int, CacheKey]] = OrderedDict()
|
|
|
|
def get(self, key: CacheKey) -> Any | None:
|
|
if not self.policy.reuse_across_requests:
|
|
return None
|
|
h = key.hash
|
|
if h in self._store:
|
|
self.hits += 1
|
|
self._store.move_to_end(h) # LRU
|
|
return self._store[h][0]
|
|
self.misses += 1
|
|
return None
|
|
|
|
def put(self, key: CacheKey, value: Any) -> None:
|
|
nb = _nbytes(value)
|
|
h = key.hash
|
|
if h in self._store:
|
|
self.used_bytes -= self._store[h][1]
|
|
self._store[h] = (value, nb, key)
|
|
self._store.move_to_end(h)
|
|
self.used_bytes += nb
|
|
self._evict()
|
|
|
|
def _evict(self) -> None:
|
|
while self.used_bytes > self.policy.max_bytes and self._store:
|
|
_h, (_v, nb, _k) = self._store.popitem(last=False) # FIFO/LRU oldest
|
|
self.used_bytes -= nb
|
|
|
|
def invalidate_weights(self, version: str) -> None:
|
|
"""Drop entries from older resident-weight epochs."""
|
|
drop = [h for h, (_v, _nb, k) in self._store.items() if k.weights_version != version]
|
|
for h in drop:
|
|
_v, nb, _k = self._store.pop(h)
|
|
self.used_bytes -= nb
|
|
|
|
def invalidate_components(self, components: set[str]) -> None:
|
|
"""Drop only entries produced by the changed components (partition, not flush):
|
|
a transformer-only weight sync must NOT evict text-encoder embeddings."""
|
|
drop = [h for h, (_v, _nb, k) in self._store.items() if k.component_id in components]
|
|
for h in drop:
|
|
_v, nb, _k = self._store.pop(h)
|
|
self.used_bytes -= nb
|
|
|
|
|
|
class ResidualCache(_Pool):
|
|
"""cache-dit residual store, scoped per ``LoopState`` AND per CFG branch.
|
|
|
|
Keyed by (namespace, branch, name) where namespace is the request/loop id. This is the
|
|
structural fix for the module-global residual state that corrupts cache-dit forks under
|
|
concurrency: two interleaved requests have disjoint namespaces.
|
|
"""
|
|
|
|
def __init__(self, policy: CachePolicy):
|
|
super().__init__(policy)
|
|
self._store: dict[tuple[str, str, str], Any] = {}
|
|
|
|
def put(self, namespace: str, branch: str, name: str, value: Any) -> None:
|
|
self._store[(namespace, branch, name)] = value
|
|
|
|
def get(self, namespace: str, branch: str, name: str) -> Any | None:
|
|
v = self._store.get((namespace, branch, name))
|
|
if v is not None:
|
|
self.hits += 1
|
|
else:
|
|
self.misses += 1
|
|
return v
|
|
|
|
def clear_namespace(self, namespace: str) -> None:
|
|
for k in [k for k in self._store if k[0] == namespace]:
|
|
del self._store[k]
|
|
|
|
|
|
@dataclass
|
|
class Slab:
|
|
chunk_index: int
|
|
k: Any
|
|
v: Any
|
|
|
|
|
|
class SlabKVCache(_Pool):
|
|
"""Chunk-KV slabs for causal/world-model rollout.
|
|
|
|
``disable_recycle`` preserves the full request-local context. Otherwise the
|
|
cache behaves as a sliding window with the declared ``window`` size.
|
|
"""
|
|
|
|
def __init__(self, policy: CachePolicy):
|
|
super().__init__(policy)
|
|
self._store: dict[str, list[Slab]] = {}
|
|
self.window = max(1, policy.per_component.get("window", 1 << 30))
|
|
self.disable_recycle = policy.disable_recycle
|
|
|
|
def append(self, namespace: str, slab: Slab) -> None:
|
|
slabs = self._store.setdefault(namespace, [])
|
|
slabs.append(slab)
|
|
self.used_bytes += _nbytes(slab.k) + _nbytes(slab.v)
|
|
if not self.disable_recycle and len(slabs) > self.window:
|
|
dropped = slabs.pop(0) # sliding-window recycle (inference only)
|
|
self.used_bytes -= _nbytes(dropped.k) + _nbytes(dropped.v)
|
|
|
|
def get(self, namespace: str) -> list[Slab]:
|
|
return self._store.get(namespace, [])
|
|
|
|
def clear_namespace(self, namespace: str) -> None:
|
|
for slab in self._store.pop(namespace, []):
|
|
self.used_bytes -= _nbytes(slab.k) + _nbytes(slab.v)
|
|
|
|
|
|
class PagedKVCache(_Pool):
|
|
"""Paged text-KV stub for ar_decode (phase-2 omni). Minority case, materialized lazily."""
|
|
|
|
def __init__(self, policy: CachePolicy):
|
|
super().__init__(policy)
|
|
self.total_blocks = max(1, policy.max_bytes // max(policy.block_bytes, 1))
|
|
self.free = self.total_blocks
|
|
self._alloc: dict[str, int] = {}
|
|
|
|
def allocate(self, namespace: str, n_blocks: int) -> bool:
|
|
if n_blocks > self.free:
|
|
return False
|
|
self._alloc[namespace] = self._alloc.get(namespace, 0) + n_blocks
|
|
self.free -= n_blocks
|
|
return True
|
|
|
|
def free_namespace(self, namespace: str) -> None:
|
|
self.free += self._alloc.pop(namespace, 0)
|
|
|
|
|
|
_CLASS_REGISTRY = {
|
|
"feature": FeatureCache,
|
|
"residual": ResidualCache,
|
|
"slab_kv": SlabKVCache,
|
|
"paged_kv": PagedKVCache,
|
|
}
|
|
|
|
|
|
def make_pool(policy: CachePolicy) -> _Pool:
|
|
cls = _CLASS_REGISTRY.get(policy.class_name)
|
|
if cls is None:
|
|
raise KeyError(f"unknown cache class {policy.class_name!r} (have {list(_CLASS_REGISTRY)})")
|
|
return cls(policy)
|