# SPDX-License-Identifier: Apache-2.0 """Platform-agnostic cache-context factory. The concrete implementations live in their respective sub-packages: * :class:`~lmcache.v1.platform.cuda.cache_context.GPUCacheContext` -- CUDA-backed. * :class:`~lmcache.v1.platform.cpu.cache_context.CPUCacheContext` -- CPU-only fallback (POSIX-SHM-backed KV tensors). :func:`create_cache_context` keeps the dispatch out of the call site in :mod:`lmcache.v1.multiprocess.server`. Selection is data-driven: each backend sub-package ships its own :class:`~lmcache.v1.platform.base_cache_context.BaseCacheContext` subclass under ``platform//cache_context.py`` and declares the ``torch.device.type`` it handles via the :attr:`BaseCacheContext.device_type` ClassVar. The first :func:`create_cache_context` call discovers those subclasses with :func:`lmcache.v1.utils.subclass_discovery.discover_subclasses` and memoises the resulting ``device_type -> class`` map. Adding a new accelerator therefore requires *zero* edits to this module -- just drop a new ``platform//cache_context.py`` whose subclass sets ``device_type``. """ # Future from __future__ import annotations # Standard from collections.abc import Sequence from typing import TYPE_CHECKING # First Party from lmcache.logging import init_logger from lmcache.utils import EngineType from lmcache.v1.gpu_connector.utils import LayoutHints from lmcache.v1.multiprocess.custom_types import KVCache from lmcache.v1.platform.base_cache_context import BaseCacheContext from lmcache.v1.utils.subclass_discovery import discover_subclasses if TYPE_CHECKING: # First Party from lmcache.v1.multiprocess.group_view import EngineGroupInfo logger = init_logger(__name__) # ``device_type -> BaseCacheContext`` subclass. Populated lazily on # the first :func:`create_cache_context` call by scanning the # ``platform`` package for ``cache_context`` leaf modules at depth 2 # (i.e. ``platform//cache_context.py``). Tests substitute # entries via :func:`snapshot_backends` / :func:`restore_backends`. # # The value type is the loose ``type`` (rather than # ``type[BaseCacheContext]``) so callers can instantiate it with the # concrete subclass' positional ``__init__`` signature without mypy # resolving the abstract-base ``__init__`` instead. _BACKENDS: dict[str, type] = {} _BACKENDS_DISCOVERED: bool = False def _discover_backends_once() -> None: """Populate :data:`_BACKENDS` on first use. Walks ``lmcache.v1.platform`` two levels deep (``platform`` -> ```` -> ``cache_context``) and indexes every concrete :class:`BaseCacheContext` subclass by its ``device_type`` ClassVar. Subclasses with an empty ``device_type`` are skipped with a warning so a missing override surfaces loudly instead of silently shadowing a real backend. """ global _BACKENDS_DISCOVERED if _BACKENDS_DISCOVERED: return # First Party import lmcache.v1.platform as platform_pkg for cls in discover_subclasses( platform_pkg, BaseCacheContext, # type: ignore[type-abstract] module_filter=lambda short_name: short_name == "cache_context", levels=[2, 2], ): device_type = getattr(cls, "device_type", "") if not device_type: logger.warning( "Skipping %s: empty device_type ClassVar; concrete " "BaseCacheContext subclasses must override it.", cls.__name__, ) continue existing = _BACKENDS.get(device_type) if existing is not None and existing is not cls: logger.warning( "Multiple cache-context classes claim device_type=%r " "(%s vs %s); keeping the first.", device_type, existing.__name__, cls.__name__, ) continue _BACKENDS[device_type] = cls _BACKENDS_DISCOVERED = True def _resolve_backend(device_type: str) -> type: _discover_backends_once() cls = _BACKENDS.get(device_type) if cls is None: raise ValueError( "No cache-context class registered for device type %r. " "Make sure ``lmcache.v1.platform..cache_context`` " "ships a BaseCacheContext subclass with the matching " "``device_type`` ClassVar." % device_type ) return cls def snapshot_backends() -> dict[str, type]: """Return a shallow copy of the backend table. Pair with :func:`restore_backends` in test fixtures so installing fakes for one test does not leak into the next. """ _discover_backends_once() return dict(_BACKENDS) def restore_backends(state: dict[str, type]) -> None: """Replace the backend table with *state*. Marks the table as already-discovered so further calls do not re-trigger filesystem scanning and overwrite the test's fakes. """ global _BACKENDS_DISCOVERED _BACKENDS.clear() _BACKENDS.update(state) _BACKENDS_DISCOVERED = True def _detect_device_type(kv_caches: KVCache) -> str: """Return the ``torch.device.type`` describing *kv_caches*. All wrappers in *kv_caches* must materialize tensors on the same device type; mixed-device batches are not supported by any downstream cache-context implementation. """ device_types = {w.to_tensor().device.type for w in kv_caches} if len(device_types) != 1: raise ValueError( "create_cache_context requires all kv_caches to share one " "device type, got %r" % sorted(device_types) ) return next(iter(device_types)) def create_cache_context( kv_caches: KVCache, lmcache_tokens_per_chunk: int = 256, layout_hints: LayoutHints | None = None, engine_group_infos: "Sequence[EngineGroupInfo]" = (), engine_type: EngineType = EngineType.VLLM, separate_object_groups: bool = True, full_sw_kv: bool = False, ) -> BaseCacheContext: """Create the appropriate cache context for *kv_caches*. The signature mirrors :class:`GPUCacheContext` so callers can forward their kwargs verbatim and stay agnostic of the active backend. Selection is driven by ``tensor.device.type`` of *kv_caches*: on first use the platform package is scanned for ``BaseCacheContext`` subclasses and the one whose ``device_type`` ClassVar matches is instantiated. ``"cuda"``, ``"cpu"``, future ``"xpu"`` ... all resolve through the same code path -- no ``isinstance`` / ``if-elif`` chain. Args: kv_caches: KV cache tensor wrappers from the serving engine. Must be non-empty. lmcache_tokens_per_chunk: Number of tokens per LMCache chunk. layout_hints: Optional hints for KV format detection. Forwarded verbatim to the concrete context constructor. engine_group_infos: Engine-neutral KV cache group metadata. engine_type: Which serving engine produced the caches. separate_object_groups: When True (default), split kernel groups into one object group per sliding-window size; when False, a single full-attention object group. full_sw_kv: When True, sliding-window groups store/transfer full per-chunk KV (no sub-chunk window cutting) so chunks stay valid for reuse at any position; see :meth:`KVLayerGroupsManager.enable_full_sw_kv`. Returns: A concrete cache context instance. Raises: ValueError: If *kv_caches* is empty, mixes device types, or targets a device type with no registered backend. """ if not kv_caches: raise ValueError("create_cache_context requires a non-empty kv_caches list") device_type = _detect_device_type(kv_caches) cls = _resolve_backend(device_type) return cls( kv_caches, lmcache_tokens_per_chunk, layout_hints, engine_group_infos, engine_type, separate_object_groups, full_sw_kv, )