chore: import upstream snapshot with attribution

python/tokenspeed/runtime/cache/embedding_cache.py

@@ -0,0 +1,224 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Vision-embedding cache for the EPD encode stage.
+
+The encode server runs the vision tower only; a cache lets duplicate images
+(same ``MultimodalDataItem.hash``) reuse a previously computed embedding instead
+of re-running the tower.
+
+Two interchangeable implementations share a ``get`` / ``put`` surface:
+:class:`EmbeddingCache` (single-tier, bytes-bounded LRU) and
+:class:`TieredEmbeddingCache` (L1 GPU VRAM + L2 host DRAM). Both are
+framework-agnostic and unit-testable without a GPU -- the caller supplies byte
+sizes and the device<->host copies are injectable.
+"""
+
+from __future__ import annotations
+
+import collections
+from collections.abc import Callable, Hashable
+
+
+class EmbeddingCache:
+    """Bytes-bounded LRU cache of vision-tower outputs, keyed by content hash.
+
+    Values are opaque (typically a ``torch.Tensor``); the caller supplies the
+    byte size so this stays framework-agnostic and testable without a GPU.
+
+    ``on_evict`` is an optional ``(key, value, nbytes)`` callback fired only for
+    an entry dropped by capacity overflow (not an update to an existing key,
+    which is an in-place replace, nor an explicit :meth:`pop`). It lets a second
+    tier capture LRU victims and demote them (see :class:`TieredEmbeddingCache`).
+    """
+
+    def __init__(
+        self,
+        capacity_bytes: int,
+        on_evict: Callable[[Hashable, object, int], None] | None = None,
+    ):
+        if capacity_bytes < 0:
+            raise ValueError(f"capacity_bytes must be >= 0, got {capacity_bytes}")
+        self.capacity_bytes = capacity_bytes
+        self._on_evict = on_evict
+        # key -> (value, nbytes); ordered by recency (oldest first).
+        self._store: collections.OrderedDict[Hashable, tuple[object, int]] = (
+            collections.OrderedDict()
+        )
+        self._cur_bytes = 0
+        self.hits = 0
+        self.misses = 0
+
+    def get(self, key: Hashable) -> object | None:
+        entry = self._store.get(key)
+        if entry is None:
+            self.misses += 1
+            return None
+        self._store.move_to_end(key)
+        self.hits += 1
+        return entry[0]
+
+    def put(self, key: Hashable, value: object, nbytes: int) -> None:
+        if nbytes < 0:
+            raise ValueError(f"nbytes must be >= 0, got {nbytes}")
+        # An item larger than the whole cache can never be retained; skip it
+        # rather than evicting everything for a value dropped immediately after.
+        if nbytes > self.capacity_bytes:
+            return
+        existing = self._store.pop(key, None)
+        if existing is not None:
+            self._cur_bytes -= existing[1]
+        self._store[key] = (value, nbytes)
+        self._cur_bytes += nbytes
+        self._evict()
+
+    def pop(self, key: Hashable) -> tuple[object, int] | None:
+        """Remove ``key`` and return its ``(value, nbytes)``, or ``None`` if
+        absent. A structural removal: it neither counts as a hit/miss nor fires
+        ``on_evict`` (the caller is taking ownership of the value, e.g. to
+        promote it to another tier)."""
+        existing = self._store.pop(key, None)
+        if existing is None:
+            return None
+        self._cur_bytes -= existing[1]
+        return existing
+
+    def _evict(self) -> None:
+        while self._cur_bytes > self.capacity_bytes and self._store:
+            key, (value, nbytes) = self._store.popitem(last=False)
+            self._cur_bytes -= nbytes
+            if self._on_evict is not None:
+                self._on_evict(key, value, nbytes)
+
+    def __contains__(self, key: Hashable) -> bool:
+        return key in self._store
+
+    def __len__(self) -> int:
+        return len(self._store)
+
+    @property
+    def current_bytes(self) -> int:
+        return self._cur_bytes
+
+
+def _embedding_to_host(value: object) -> object:
+    """Default L1->L2 demote: copy a vision embedding from GPU to host memory.
+
+    A cache value is a ``(main, deepstack)`` tuple (deepstack half may be
+    ``None`` for non-deepstack models) or a bare tensor.
+    """
+    if isinstance(value, tuple):
+        return tuple(None if t is None else t.cpu() for t in value)
+    return value.cpu()
+
+
+def _embedding_to_device(value: object, device) -> object:
+    """Default L2->L1 promote: copy a host-resident embedding back to ``device``
+    so the executor can stage it into the GPU send ring."""
+    if isinstance(value, tuple):
+        return tuple(None if t is None else t.to(device) for t in value)
+    return value.to(device)
+
+
+class TieredEmbeddingCache:
+    """Two-tier vision-embedding cache: L1 in GPU VRAM, L2 in host DRAM.
+
+    Exposes the same ``get`` / ``put`` surface as :class:`EmbeddingCache`, so the
+    encode worker uses either interchangeably. L2 catches L1's LRU victims in
+    cheaper host DRAM: a hit there skips the (much more expensive) ViT and only
+    pays a host->device copy.
+
+    Tiers are kept *exclusive* -- a key lives in exactly one. An L1 eviction
+    demotes the victim to L2 (device->host copy); an L2 hit promotes the entry
+    back to L1 (host->device copy) and removes it from L2; ``put`` always lands
+    in L1 and drops any stale L2 duplicate. No distributed L3: image-hash routing
+    pins each image to one worker, so a per-worker local cache captures the reuse.
+
+    The device<->host copies are injectable (``to_host`` / ``to_device``) so the
+    tiering logic is unit-testable without a GPU; the defaults copy real tensors
+    to/from ``device``.
+    """
+
+    def __init__(
+        self,
+        l1_capacity_bytes: int,
+        l2_capacity_bytes: int,
+        *,
+        device: object = None,
+        to_host: Callable[[object], object] | None = None,
+        to_device: Callable[[object], object] | None = None,
+    ):
+        # L1 demotes its evictions into L2 via the on_evict hook; L2 is the
+        # bottom tier (no hook), so its evictions are true drops -- no recursion.
+        self.l1 = EmbeddingCache(l1_capacity_bytes, on_evict=self._demote)
+        self.l2 = EmbeddingCache(l2_capacity_bytes)
+        self._device = device
+        self._to_host = to_host or _embedding_to_host
+        self._to_device = to_device or (lambda v: _embedding_to_device(v, self._device))
+        self.l1_hits = 0
+        self.l2_hits = 0
+        self.misses = 0
+        self.promotions = 0
+        self.demotions = 0
+
+    def get(self, key: Hashable) -> object | None:
+        value = self.l1.get(key)
+        if value is not None:
+            self.l1_hits += 1
+            return value
+        promoted = self.l2.pop(key)
+        if promoted is None:
+            self.misses += 1
+            return None
+        host_value, nbytes = promoted
+        device_value = self._to_device(host_value)
+        self.l2_hits += 1
+        self.promotions += 1
+        # Re-home as MRU in L1. May evict colder L1 entries, which demote to L2;
+        # the just-promoted key is MRU so it is never the victim.
+        self.l1.put(key, device_value, nbytes)
+        return device_value
+
+    def put(self, key: Hashable, value: object, nbytes: int) -> None:
+        # L1 is the write tier. Drop any stale L2 copy first so the tiers stay
+        # exclusive (e.g. a key demoted earlier and now re-encoded on a miss).
+        self.l2.pop(key)
+        self.l1.put(key, value, nbytes)
+
+    def _demote(self, key: Hashable, value: object, nbytes: int) -> None:
+        """L1 eviction hook: stash the victim in host DRAM instead of dropping
+        it. host nbytes == device nbytes (same dtype/numel). A victim larger than
+        all of L2 is silently dropped by ``L2.put`` (same as having no L2)."""
+        # Skip the wasted device->host copy for a victim L2 would just drop.
+        if nbytes > self.l2.capacity_bytes:
+            return
+        self.l2.put(key, self._to_host(value), nbytes)
+        if key in self.l2:
+            self.demotions += 1
+
+    @property
+    def hits(self) -> int:
+        return self.l1_hits + self.l2_hits
+
+    def __contains__(self, key: Hashable) -> bool:
+        return key in self.l1 or key in self.l2
+
+    def __len__(self) -> int:
+        return len(self.l1) + len(self.l2)