# SPDX-License-Identifier: Apache-2.0 """CPU-only KV-cache IPC wrapper backed by POSIX shared memory. Mirrors the GPU-mode CUDA-IPC zero-copy semantics for hosts without an accelerator: client and LMCache mp server map the **same** physical pages so transfers are pointer-shuffles rather than memcpys. Self-registers a ``"cpu"`` factory with :mod:`lmcache.v1.platform._registry` at import time, so the multiprocess adapter can dispatch by ``tensor.device.type`` without any if/elif chain. """ # Future from __future__ import annotations # Standard from typing import ClassVar import ctypes import itertools import os import threading import weakref # Third Party import torch # First Party from lmcache.logging import init_logger from lmcache.v1.multiprocess.posix_shm import ( shm_create_readwrite, shm_map_readwrite, shm_munmap, shm_unlink, ) from lmcache.v1.platform.base_ipc_wrapper import DeviceIPCWrapper logger = init_logger(__name__) # Re-export POSIX-SHM primitives so existing callers keep working. # The canonical home is :mod:`lmcache.v1.multiprocess.posix_shm`; new # code (e.g. the MP non-GPU SHM transport) should import from there. __all__ = [ "CpuShmTensorWrapper", "inject_stale_cache_entry_for_test", "migrate_to_shm_and_wrap", "shm_create_readwrite", "shm_map_readwrite", "shm_munmap", "shm_unlink", ] # --------------------------------------------------------------------------- # Wrapper class # # --------------------------------------------------------------------------- class CpuShmTensorWrapper(DeviceIPCWrapper): """IPC wrapper for CPU tensors backed by POSIX shared memory. Used by the ``lmcache bench kvcache --mode cpu`` path and the vLLM CPU integration so that the client and the LMCache mp server map the **same** physical pages for the KV cache, mirroring the GPU-mode CUDA-IPC zero-copy semantics. Subclassing :class:`DeviceIPCWrapper` is load-bearing for the same reason :class:`RawCudaIPCWrapper` does it: msgspec does not support unions of custom ext-encoded types, so all wire-level KV-cache wrappers must share the single ext code (1) registered for ``DeviceIPCWrapper``. Pickle preserves the subclass identity so ``to_tensor`` dispatches correctly on both sides. """ #: ``torch.device.type`` this wrapper handles (used by auto-discovery #: in :func:`~lmcache.v1.platform._registry._discover_wrappers_once`). device_type: ClassVar[str] = "cpu" #: Marked ``True`` so auto-discovery picks this as the default #: factory for ``"cpu"``. _is_default_wrapper: ClassVar[bool] = True # POSIX shared-memory name (``/lmcache_...``) -- leading ``/`` is # required by ``shm_open(3)`` on both Linux and macOS. SHM_NAME_PREFIX = "/lmcache_kv_" @classmethod def wrap(cls, tensor: torch.Tensor) -> "CpuShmTensorWrapper": """Factory used by :func:`~lmcache.v1.platform._registry._discover_wrappers_once`. Delegates to :func:`migrate_to_shm_and_wrap`, which migrates the tensor's storage to a POSIX SHM segment so the LMCache mp server can map the same physical pages. Args: tensor: A contiguous CPU tensor to migrate and wrap. Returns: A new :class:`CpuShmTensorWrapper` referencing the SHM segment that now backs ``tensor``. """ return migrate_to_shm_and_wrap(tensor) def __init__(self, tensor: torch.Tensor, shm_name: str) -> None: if tensor.device.type != "cpu": raise ValueError( "CpuShmTensorWrapper requires a CPU tensor, got %s" % tensor.device ) if not tensor.is_contiguous(): raise ValueError("CpuShmTensorWrapper requires a contiguous tensor") self.shm_name = shm_name # ``numel * element_size`` is the correct logical byte size; the # underlying storage may be larger when the tensor is a view. self.nbytes = tensor.numel() * tensor.element_size() # DeviceIPCWrapper interface fields. ``handle`` / ``device_uuid`` # are unused on the CPU path but kept to satisfy the base # contract used by equality checks. self.handle = None self.dtype = tensor.dtype self.shape = tuple(tensor.shape) self.stride = tuple(tensor.stride()) self.storage_offset = int(tensor.storage_offset()) self.device_uuid = "cpu" def to_tensor(self) -> torch.Tensor: """Reconstruct the tensor by mapping the same SHM segment. The returned tensor owns the mmap: a ``weakref.finalize`` hook runs ``munmap`` once the tensor (and any views derived from it) is garbage-collected, so the per-process virtual address space does not leak across repeated ``to_tensor`` calls. We rebuild the view through ``as_strided`` so the original memory layout (stride / storage_offset / memory_format) is replayed faithfully on the receiving side; reshape would silently re-coalesce strides and lose, e.g., channels_last. """ # Empty tensors carry no SHM segment (mmap with length 0 is # undefined / EINVAL on POSIX); rebuild the empty view in-process. if self.nbytes == 0: return torch.empty(self.shape, dtype=self.dtype) addr = shm_map_readwrite(self.shm_name, self.nbytes) # ``torch.frombuffer`` requires a writable buffer; build one # via ctypes so the resulting torch tensor shares storage # with the SHM mapping (zero copy across processes). buf_type = ctypes.c_uint8 * self.nbytes buf = buf_type.from_address(addr) flat = torch.frombuffer(buf, dtype=torch.uint8) typed = flat.view(self.dtype) out = torch.as_strided(typed, self.shape, self.stride, self.storage_offset) # Pin the mmap to the *storage*, not the outer tensor: views # (reshape / slicing) create new tensor objects that share the # storage but do not inherit Python attributes, so a finalizer # attached to ``out`` would munmap as soon as ``out`` is GC'd # even when a view is still reading the SHM segment. # ``UntypedStorage`` is shared across views, so finalizing on it # only fires once every view is also dropped. storage = out.untyped_storage() _CPU_SHM_KEEP_ALIVE[id(storage)] = flat weakref.finalize(storage, _release_shm_segment, id(storage), addr, self.nbytes) return out # --------------------------------------------------------------------------- # Migrate-and-wrap factory (used by the multiprocess adapter) # # --------------------------------------------------------------------------- # Per-process registry of SHM segments we have created, so the same # tensor object is only migrated to SHM once even if the factory is # called multiple times. # # Keyed by ``id(tensor)`` for cheap O(1) lookup, but each entry also # holds a ``weakref.ref`` to the original tensor and we *verify the # referent is still that exact object* before reusing the cached SHM # name. CPython recycles object IDs, so a fresh tensor allocated at # the same address as a previously migrated (now garbage-collected) # one would otherwise inherit a stale name -- and because # :func:`shm_create_readwrite` uses ``O_EXCL``, the next migration # would crash with ``EEXIST`` ("File exists"). The weakref-validated # lookup below makes that race impossible: a stale entry can only # point at a dead referent, which we treat as a miss. _CPU_SHM_NAMES: dict[int, tuple["weakref.ReferenceType[torch.Tensor]", str]] = {} _CPU_SHM_LOCK = threading.Lock() _CPU_SHM_COUNTER = itertools.count() # Process-level registry that pins the base ``flat`` buffer of every live # ``to_tensor()`` mmap until its storage is finalized. Keyed by ``id(storage)``, # which is stable across views because PyTorch caches the storage Python # wrapper (so reshape / slicing returns the same ``UntypedStorage`` object). _CPU_SHM_KEEP_ALIVE: dict[int, torch.Tensor] = {} def _release_shm_segment(storage_id: int, addr: int, nbytes: int) -> None: """Drop the pinned base buffer and ``munmap`` the mapping. Invoked by ``weakref.finalize`` on the tensor's ``UntypedStorage`` once every view of the mapping is gone, so views (e.g. ``reshape`` returning a new tensor without ``_lmcache_shm_buf``) cannot trigger a premature unmap that would turn into a use-after-free in the next read. """ _CPU_SHM_KEEP_ALIVE.pop(storage_id, None) shm_munmap(addr, nbytes) def _cleanup_shm_segment(tid: int, shm_name: str, addr: int, nbytes: int) -> None: """Release the mmap, unlink, and forget the cached SHM name.""" with _CPU_SHM_LOCK: # Only drop the entry if it still points at *this* segment; # a future tensor reusing ``tid`` may already have replaced it. cached = _CPU_SHM_NAMES.get(tid) if cached is not None and cached[1] == shm_name: _CPU_SHM_NAMES.pop(tid, None) shm_munmap(addr, nbytes) shm_unlink(shm_name) def migrate_to_shm_and_wrap(tensor: torch.Tensor) -> CpuShmTensorWrapper: """Re-point ``tensor``'s storage at a POSIX SHM segment, then wrap. Used as the registered ``"cpu"`` KV-wrapper factory: the LMCache mp server can mmap the same physical pages on the receiving side. Idempotent per tensor identity (validated via a stored weakref so Python's id-recycling cannot produce a stale-name hit). The SHM segment is released (``munmap`` + ``shm_unlink``) automatically when the migrated tensor is garbage-collected. """ # First Party from lmcache.v1.gpu_connector.kv_format.contiguity import ( attempt_permute_to_contiguous_view, ) # Validate and normalise the tensor *before* touching the registry # or mutating storage, so a bad input never leaves things half-done. tensor = attempt_permute_to_contiguous_view(tensor) if tensor.device.type != "cpu": raise ValueError( "migrate_to_shm_and_wrap requires a CPU tensor, got %s" % tensor.device ) if not tensor.is_contiguous(): raise ValueError("migrate_to_shm_and_wrap requires a contiguous tensor") tid = id(tensor) # Fast path: check the registry under the lock, return early if the # tensor has already been migrated. with _CPU_SHM_LOCK: cached = _CPU_SHM_NAMES.get(tid) if cached is not None: ref, cached_name = cached if ref() is tensor: return CpuShmTensorWrapper(tensor, cached_name) # Stale entry from a GC'd tensor whose id has been # reused; drop it and fall through to allocate fresh. _CPU_SHM_NAMES.pop(tid, None) nbytes = tensor.numel() * tensor.element_size() assert tensor.storage_offset() == 0, ( "migrate_to_shm_and_wrap: SHM segment is sized to " "numel*elem_size; a nonzero storage_offset would cause " "OOB access. Got offset=%d" % tensor.storage_offset() ) if nbytes == 0: # No SHM segment for empty tensors: ``mmap`` with length 0 # is undefined / EINVAL on POSIX. ``to_tensor`` rebuilds an # empty view directly when ``shm_name`` is empty. return CpuShmTensorWrapper(tensor, "") shm_name = "%s%d_%d" % ( CpuShmTensorWrapper.SHM_NAME_PREFIX, os.getpid(), next(_CPU_SHM_COUNTER), ) # Perform the heavy work (syscall + tensor mutation) outside the lock # to keep the critical section small. addr = shm_create_readwrite(shm_name, nbytes) try: buf_type = ctypes.c_uint8 * nbytes buf = buf_type.from_address(addr) shm_storage = torch.frombuffer(buf, dtype=torch.uint8).untyped_storage() tensor.set_( shm_storage, tensor.storage_offset(), tensor.shape, tensor.stride(), ) except Exception: # Make sure the SHM resources don't leak if migration fails # part-way (e.g. ``set_`` rejects an unusual stride). shm_munmap(addr, nbytes) shm_unlink(shm_name) raise with _CPU_SHM_LOCK: _CPU_SHM_NAMES[tid] = (weakref.ref(tensor), shm_name) weakref.finalize(tensor, _cleanup_shm_segment, tid, shm_name, addr, nbytes) logger.info( "Migrated CPU KV cache tensor (nbytes=%d) to SHM %s", nbytes, shm_name, ) return CpuShmTensorWrapper(tensor, shm_name) def inject_stale_cache_entry_for_test( tensor: torch.Tensor, dead_ref: "weakref.ReferenceType[torch.Tensor]", stale_shm_name: str, ) -> None: """Test-only hook: pre-seed the registry with a stale entry. Lets unit tests reproduce the CPython id-reuse race -- where a fresh tensor lands on the same id as a previously migrated and garbage-collected one -- without the per-test global-state surgery that would otherwise have to reach into the module's private dict / lock. """ with _CPU_SHM_LOCK: _CPU_SHM_NAMES[id(tensor)] = (dead_ref, stale_shm_name)