lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/v1/kv_layer_groups.py

@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Optional
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class KVLayerGroupInfo:
+    """Information about a group of layers with the same KV cache structure.
+
+    Layers within the same group have identical shape and dtype for their KV cache.
+    Different groups may have different shapes (especially head_size) and/or dtypes.
+    """
+
+    """ List of layer names belonging to this group """
+    layer_names: list[str]
+    """ List of layer indices (0-based) belonging to this group """
+    layer_indices: list[int]
+    """ Shape of the KV cache tensor for layers in this group """
+    """ For MHA: typically [2, num_blocks, block_size, num_heads, head_size] """
+    """ For MLA: typically [num_blocks, block_size, head_size] """
+    shape: torch.Size
+    """ Data type of the KV cache tensor for layers in this group """
+    dtype: torch.dtype
+
+    # Internal sets for fast membership checking
+    _layer_indices_set: set[int] = field(init=False, repr=False)
+    _layer_names_set: set[str] = field(init=False, repr=False)
+
+    def __post_init__(self):
+        """Initialize sets for fast membership checking."""
+        self._layer_indices_set = set(self.layer_indices)
+        self._layer_names_set = set(self.layer_names)
+
+    def __repr__(self) -> str:
+        if not self.layer_indices:
+            indices_repr = "[]"
+        else:
+            indices_repr = f"{self.layer_indices[0]}-{self.layer_indices[-1]}"
+        return (
+            f"KVLayerGroupInfo(layers={len(self.layer_names)}, "
+            f"indices={indices_repr}, "
+            f"shape={self.shape}, dtype={self.dtype})"
+        )
+
+    @property
+    def num_layers(self) -> int:
+        """Return the number of layers in this group."""
+        return len(self.layer_names)
+
+    @property
+    def hidden_dim_size(self) -> int:
+        """Return the size of the hidden dimension in this group."""
+        # hidden_dim_size = num_heads * head_size
+        if len(self.shape) == 5:
+            # MHA
+            return self.shape[3] * self.shape[4]
+        elif len(self.shape) == 3:
+            # MLA
+            return self.shape[2]
+        else:
+            raise ValueError(f"Invalid shape: {self.shape}")
+
+    def contains_layer(self, layer_idx: int) -> bool:
+        """Check if a layer index belongs to this group."""
+        return layer_idx in self._layer_indices_set
+
+    def contains_layer_name(self, layer_name: str) -> bool:
+        """Check if a layer name belongs to this group."""
+        return layer_name in self._layer_names_set
+
+
+@dataclass
+class KVLayerGroupsManager:
+    """Manager for KV layer groups with the same structure.
+
+    This class encapsulates the functionality for managing groups of layers
+    that have identical KV cache structure (shape and dtype).
+    """
+
+    kv_layer_groups: list[KVLayerGroupInfo] = field(default_factory=list)
+
+    @property
+    def num_groups(self) -> int:
+        """Return the number of KV layer groups."""
+        return len(self.kv_layer_groups)
+
+    def get_group_by_layer_idx(self, layer_idx: int) -> Optional[KVLayerGroupInfo]:
+        """Get the KVLayerGroupInfo for a given layer index.
+
+        Args:
+            layer_idx: The 0-based index of the layer.
+
+        Returns:
+            The KVLayerGroupInfo containing this layer, or None if not found.
+        """
+        for group in self.kv_layer_groups:
+            if group.contains_layer(layer_idx):
+                return group
+        return None
+
+    def get_group_by_layer_name(self, layer_name: str) -> Optional[KVLayerGroupInfo]:
+        """Get the KVLayerGroupInfo for a given layer name.
+
+        Args:
+            layer_name: The name of the layer.
+
+        Returns:
+            The KVLayerGroupInfo containing this layer, or None if not found.
+        """
+        for group in self.kv_layer_groups:
+            if group.contains_layer_name(layer_name):
+                return group
+        return None
+
+    def get_layer_shape(self, layer_idx: int) -> Optional[torch.Size]:
+        """Get the shape of the KV cache for a given layer index.
+
+        Args:
+            layer_idx: The 0-based index of the layer.
+
+        Returns:
+            The shape, or None if layer not found.
+        """
+        group = self.get_group_by_layer_idx(layer_idx)
+        return group.shape if group else None
+
+    def get_layer_dtype(self, layer_idx: int) -> Optional[torch.dtype]:
+        """Get the dtype of the KV cache for a given layer index.
+
+        Args:
+            layer_idx: The 0-based index of the layer.
+
+        Returns:
+            The dtype, or None if layer not found.
+        """
+        group = self.get_group_by_layer_idx(layer_idx)
+        return group.dtype if group else None
+
+    def build_kv_layer_groups(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        """Build KV layer groups structure by analyzing each layer's shape and dtype.
+
+        Layers with the same shape and dtype are grouped together. This is useful
+        because different layers may have different structures (especially the
+        last dimension head_size may differ between groups), and different groups
+        may have different dtypes.
+
+        If layer groups are already built (non-empty list), this method does nothing.
+
+        Args:
+            kv_caches: Dictionary mapping layer names to KV cache tensors.
+        """
+        # Skip if already built (non-empty list)
+        if len(self.kv_layer_groups) > 0:
+            return
+
+        if len(kv_caches) == 0:
+            logger.debug("No KV caches available, skipping KV layer groups building")
+            return
+
+        # Group layers by (shape, dtype) in a single loop
+        groups_dict: dict[tuple[torch.Size, torch.dtype], list[tuple[str, int]]] = (
+            defaultdict(list)
+        )
+
+        for idx, (layer_name, kv_cache) in enumerate(kv_caches.items()):
+            # Supports two KV cache formats:
+            # - Single-tensor format: a single tensor with shape
+            #   [2, num_blocks, block_size, num_heads, head_size].
+            # - List/tuple format (e.g., TPU/HPU): [k_tensor, v_tensor],
+            #   where each tensor has shape
+            #   [num_blocks, block_size, num_heads, head_size].
+            if isinstance(kv_cache, (tuple, list)):
+                if len(kv_cache) != 2:
+                    raise ValueError(
+                        f"Expected 2 tensors (k, v) for layer {layer_name}, "
+                        f"got {len(kv_cache)}"
+                    )
+                # Prepend the count as a leading dimension to produce the
+                # same canonical shape as the single-tensor format
+                # (e.g., [2, num_blocks, ...] for k+v), so downstream
+                # indexing (e.g., hidden_dim_size) is unaffected.
+                shape = torch.Size([len(kv_cache)] + list(kv_cache[0].shape))
+                dtype = kv_cache[0].dtype
+            else:
+                shape = kv_cache.shape
+                dtype = kv_cache.dtype
+            key = (shape, dtype)
+            groups_dict[key].append((layer_name, idx))
+
+        # Build KVLayerGroupInfo list
+        # Sort groups by the first layer index to maintain order
+        def _get_first_layer_index(shape_dtype_key):
+            """Get the index of the first layer in a layer group."""
+            layer_group = groups_dict[
+                shape_dtype_key
+            ]  # list of (layer_name, layer_index) tuples
+            first_layer_info = layer_group[0]  # first (layer_name, layer_index) tuple
+            layer_index = first_layer_info[1]  # extract the layer index
+            return layer_index
+
+        sorted_keys = sorted(groups_dict.keys(), key=_get_first_layer_index)
+
+        kv_layer_groups: list[KVLayerGroupInfo] = []
+        for shape, dtype in sorted_keys:
+            layers = groups_dict[(shape, dtype)]
+            layer_names, layer_indices = zip(*layers, strict=False)
+
+            group_info = KVLayerGroupInfo(
+                layer_names=list(layer_names),
+                layer_indices=list(layer_indices),
+                shape=shape,
+                dtype=dtype,
+            )
+            kv_layer_groups.append(group_info)
+
+        # Store the built groups
+        self.kv_layer_groups = kv_layer_groups
+
+        # Print the group structure
+        logger.info("KV layer groups: %s", kv_layer_groups)
+
+    def build_kv_layer_groups_from_list(self, kv_caches: list[torch.Tensor]) -> None:
+        """Build KV layer groups from a plain list of tensors (one per layer).
+
+        Layers with the same shape and dtype are grouped together.
+        If layer groups are already built (non-empty list), this method does nothing.
+
+        Args:
+            kv_caches: List of KV cache tensors, one per layer.
+        """
+        if len(self.kv_layer_groups) > 0:
+            return
+
+        if not kv_caches:
+            logger.debug("No KV caches available, skipping KV layer groups building")
+            return
+
+        groups_dict: dict[tuple[torch.Size, torch.dtype], list[int]] = defaultdict(list)
+        for idx, tensor in enumerate(kv_caches):
+            key = (tensor.shape, tensor.dtype)
+            groups_dict[key].append(idx)
+
+        sorted_keys = sorted(groups_dict.keys(), key=lambda k: groups_dict[k][0])
+
+        kv_layer_groups: list[KVLayerGroupInfo] = []
+        for shape, dtype in sorted_keys:
+            indices = groups_dict[(shape, dtype)]
+            group_info = KVLayerGroupInfo(
+                layer_names=[str(i) for i in indices],
+                layer_indices=indices,
+                shape=shape,
+                dtype=dtype,
+            )
+            kv_layer_groups.append(group_info)
+
+        self.kv_layer_groups = kv_layer_groups
+        logger.info("KV layer groups (from list): %s", kv_layer_groups)

lmcache/v1/lazy_memory_allocator.py

@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import List, Optional, Union
+import ctypes
+import threading
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.memory_management import (
+    AddressManager,
+    MemoryAllocatorInterface,
+    MemoryFormat,
+    MemoryObj,
+    TensorMemoryAllocator,
+)
+from lmcache.v1.system_detection import NUMAMapping
+import lmcache.c_ops as lmc_ops
+
+logger = init_logger(__name__)
+
+
+# Helper functions
+def get_numa_id(numa_mapping: NUMAMapping) -> int:
+    """
+    Get the NUMA ID for the current GPU
+
+    Args:
+        numa_mapping (NUMAMapping): The NUMA mapping object.
+
+    Returns:
+        int: The NUMA ID for the current GPU.
+
+    Raises:
+        KeyError: If GPU id is not detected in the numa mapping.
+    """
+    gpu_id = torch.cuda.current_device() if torch.cuda.is_available() else 0
+    return numa_mapping.gpu_to_numa_mapping[gpu_id]
+
+
+def align_to(size: int, align_size: int) -> int:
+    """
+    Align the given size to the nearest multiple of align_size.
+
+    Args:
+        size (int): The size to align.
+        align_size (int): The alignment size, MUST BE a power of two.
+
+    Returns:
+        int: The aligned size.
+    """
+    return (size + align_size - 1) & (~(align_size - 1))
+
+
+# Main class
+class LazyMemoryAllocator(MemoryAllocatorInterface):
+    """
+    Allocates CPU (numa) pinned memory with a initial size and expand
+    the size to the required size in the background.
+
+    Background expansion logic:
+    - After registering X GB memory, we call sbrk and updates _curr_size
+    - Once everything is registered, the background thread stops
+    """
+
+    PIN_CHUNK_SIZE = 1 << 26  # 64 MB pin chunk
+    COMMIT_SIZE = 1 << 30  # Do a commit every 1 GB
+    LOG_INTERVAL = 10 << 30  # Log expansion progress every 10 GB
+
+    def __init__(
+        self,
+        init_size: int,
+        final_size: int,
+        align_bytes: int = AddressManager.ALIGN_BYTES,
+        numa_mapping: NUMAMapping | None = None,
+    ):
+        """
+        Args:
+            init_size (int): Initial size of the memory allocation in bytes.
+            final_size (int): Final size of the memory allocation in bytes.
+            align_bytes (int, optional): Alignment in for the underlying allocations
+        """
+        # Whether using NUMA allocation
+        self._use_numa = numa_mapping is not None
+        # Currently pinned size, only accessed by the expansion thread
+        self._curr_size = align_to(init_size, self.PIN_CHUNK_SIZE)
+        # Final size of the allocation, only accessed by the expansion thread
+        self._final_size = align_to(final_size, self.PIN_CHUNK_SIZE)
+        # Underlying buffer for the memory allocation
+        self._buffer: torch.Tensor
+        # CUDA runtime API
+        self._cudart = torch.cuda.cudart()
+
+        # List of (ptr, size) for pinned memory chunks
+        self._pin_record: list[tuple[int, int]] = []
+
+        # Detect numa mapping
+        if numa_mapping is not None:
+            numa_id = get_numa_id(numa_mapping)
+            ptr = lmc_ops.alloc_numa_ptr(self._final_size, numa_id)
+            arr_type = ctypes.c_uint8 * self._final_size
+            buf = arr_type.from_address(ptr)
+            self._buffer = torch.frombuffer(buf, dtype=torch.uint8)
+        else:
+            self._buffer = torch.empty(
+                self._final_size, dtype=torch.uint8, device="cpu", pin_memory=False
+            )
+
+        # Pin the first `curr_size` bytes (aligned to the internal chunk size)
+        self._pin_memory_chunk(0, self._curr_size)
+
+        # Create the tensor memory allocator
+        self._allocator = TensorMemoryAllocator(
+            tensor=self._buffer,
+            align_bytes=align_bytes,
+            init_address_space=self._curr_size,
+        )
+
+        # Get the address manager
+        # NOTE(ApostaC): this assumes the tensor memory allocator owns the address
+        # manager, which creates extra coupling in the code.
+        # NOTE(ApostaC): this also assumes that the behavior of the allocation is
+        # completely determined by the address manager.
+        self._address_manager = self._allocator.address_manager
+
+        # Launch the background expansion thread
+        self._stop_expand = threading.Event()
+        self._expand_thread = threading.Thread(
+            target=self._expand_worker, daemon=True, name="lazy-mem-expand-thread"
+        )
+        self._expand_thread.start()
+
+    # Public methods
+    def allocate(
+        self,
+        shapes: Union[torch.Size, list[torch.Size]],
+        dtypes: Union[torch.dtype, list[torch.dtype]],
+        fmt: MemoryFormat = MemoryFormat.UNDEFINED,
+        allocator_type: Optional[str] = None,
+    ) -> Optional[MemoryObj]:
+        obj = self._allocator.allocate(shapes, dtypes, fmt, allocator_type)
+        # HACK(ApostaC): reset the parent allocator to this lazy allocator
+        # There should be a cleaner way to decouple lazy allocator and
+        # tensor memory allocator
+        if obj is not None:
+            obj.parent_allocator = self
+        return obj
+
+    def batched_allocate(
+        self,
+        shapes: Union[torch.Size, list[torch.Size]],
+        dtypes: Union[torch.dtype, list[torch.dtype]],
+        batch_size: int,
+        fmt: MemoryFormat = MemoryFormat.UNDEFINED,
+        allocator_type: Optional[str] = None,
+    ) -> Optional[List[MemoryObj]]:
+        # HACK(ApostaC): reset the parent allocator to this lazy allocator
+        # There should be a cleaner way to decouple lazy allocator and
+        # tensor memory allocator
+        ret = self._allocator.batched_allocate(
+            shapes, dtypes, batch_size, fmt, allocator_type
+        )
+
+        if ret is None:
+            return ret
+
+        for obj in ret:
+            obj.parent_allocator = self
+        return ret
+
+    def free(
+        self,
+        memory_obj: MemoryObj,
+        allocator_type: Optional[str] = None,
+    ):
+        self._allocator.free(memory_obj, allocator_type)
+
+    def batched_free(
+        self,
+        memory_objs: List[MemoryObj],
+        allocator_type: Optional[str] = None,
+        update_stats: bool = True,
+    ):
+        self._allocator.batched_free(memory_objs, allocator_type, update_stats)
+
+    def close(self):
+        # Stop the background expansion thread
+        self._stop_expand.set()
+        self._expand_thread.join()
+
+        # Unpin all pinned memory chunks
+        for ptr, size in self._pin_record:
+            self._cudart.cudaHostUnregister(ptr)
+        self._pin_record.clear()
+
+        # Free the underlying buffer if using NUMA allocation
+        if self._use_numa:
+            lmc_ops.free_numa_ptr(self._buffer.data_ptr(), self._final_size)
+
+    def memcheck(self) -> bool:
+        return self._allocator.memcheck()
+
+    def get_underlying_buffer(self) -> torch.Tensor:
+        """
+        Get the underlying buffer tensor. Will be used by RDMA registrations.
+        """
+        return self._buffer
+
+    def get_address_manager(self) -> AddressManager:
+        """
+        Get the address manager used by this allocator.
+        """
+        return self._address_manager
+
+    # Helper functions
+    def _pin_memory_chunk(self, offset: int, size: int):
+        """
+        Pin a chunk of memory.
+
+        Args:
+            offset (int): Offset in the buffer to pin.
+            size (int): Size of the memory chunk in bytes.
+        """
+        assert offset & (self.PIN_CHUNK_SIZE - 1) == 0, (
+            "Offset must be aligned to PIN_CHUNK_SIZE"
+        )
+        assert size & (self.PIN_CHUNK_SIZE - 1) == 0, (
+            "Size must be aligned to PIN_CHUNK_SIZE"
+        )
+        assert offset + size <= self._final_size, "Pinning exceeds buffer size"
+
+        ptr = self._buffer.data_ptr() + offset
+        # Use flag: cudaHostRegisterMapped (0x02)
+        self._cudart.cudaHostRegister(ptr, size, 2)
+        self._pin_record.append((ptr, size))
+
+    def _commit_expansion(self, expand_size: int):
+        """
+        Call sbrk in the address manager to commit the expansion.
+        """
+        self._address_manager.sbrk(expand_size)
+
+    def _log_expansion_progress(self, expanded_since_last_log: int):
+        """
+        Log the cumulative expansion progress since the last log.
+        """
+        percent = 100.0 * self._curr_size / self._final_size
+        logger.info(
+            "LazyMemoryAllocator: Expanded %s MB pinned memory, "
+            "now total is %s MB / %s MB (%.1f%%)",
+            expanded_since_last_log >> 20,
+            self._curr_size >> 20,
+            self._final_size >> 20,
+            percent,
+        )
+
+    def _expand_worker(self):
+        """
+        Background worker to expand the pinned memory.
+        """
+        last_commit_size = self._curr_size
+        last_log_size = self._curr_size
+        while self._curr_size < self._final_size and not self._stop_expand.is_set():
+            # Expand chunk by chunk and commit
+            for i in range(self.COMMIT_SIZE // self.PIN_CHUNK_SIZE):
+                if self._curr_size >= self._final_size:
+                    break
+                self._pin_memory_chunk(self._curr_size, self.PIN_CHUNK_SIZE)
+                self._curr_size += self.PIN_CHUNK_SIZE
+
+            expand_size = self._curr_size - last_commit_size
+            self._commit_expansion(expand_size)
+            last_commit_size = self._curr_size
+
+            # Log every LOG_INTERVAL bytes, and always on the final commit.
+            expanded_since_last_log = self._curr_size - last_log_size
+            if (
+                expanded_since_last_log >= self.LOG_INTERVAL
+                or self._curr_size >= self._final_size
+            ):
+                self._log_expansion_progress(expanded_since_last_log)
+                last_log_size = self._curr_size

lmcache/v1/lookup_client/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.v1.lookup_client.abstract_client import LookupClientInterface
+from lmcache.v1.lookup_client.chunk_statistics_lookup_client import (
+    ChunkStatisticsLookupClient,
+)
+from lmcache.v1.lookup_client.factory import LookupClientFactory
+from lmcache.v1.lookup_client.lmcache_lookup_client import (
+    LMCacheLookupClient,
+    LMCacheLookupServer,
+)
+from lmcache.v1.lookup_client.lmcache_lookup_client_bypass import (
+    LMCacheBypassLookupClient,
+)
+from lmcache.v1.lookup_client.mooncake_lookup_client import MooncakeLookupClient
+
+__all__ = [
+    "LookupClientInterface",
+    "LookupClientFactory",
+    "MooncakeLookupClient",
+    "LMCacheBypassLookupClient",
+    "LMCacheLookupClient",
+    "LMCacheLookupServer",
+    "ChunkStatisticsLookupClient",
+]

lmcache/v1/lookup_client/abstract_client.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Union
+import abc
+
+# Third Party
+import torch
+
+
+class LookupClientInterface(metaclass=abc.ABCMeta):
+    """Abstract interface for lookup clients."""
+
+    def lookup_cache(self, lookup_id: str) -> Optional[int]:
+        """
+        Lookup the cache for the given lookup ID.
+
+        Args:
+            lookup_id: The lookup ID to lookup
+
+        Returns:
+            -1 means not found;
+            None means ongoing; (this semantic is not supported in sync lookup clients)
+            int >= 0 means number of hit tokens
+        """
+        return None
+
+    @abc.abstractmethod
+    def lookup(
+        self,
+        token_ids: Union[torch.Tensor, list[int]],
+        lookup_id: str,
+        request_configs: Optional[dict] = None,
+    ) -> Optional[int]:
+        """
+        Perform lookup for the given token IDs.
+        Should be called for first lookup and pinning. Subsequent lookups for the same
+        request should call lookup_cache instead.
+
+        Caller should handle overlaps between tokens that exist in LMCache
+        and tokens that are already computed by the caller.
+
+        Args:
+            token_ids: The token IDs to lookup
+
+            lookup_id: The lookup ID to associate with the lookup
+
+            request_configs: The configs of the request,
+            includes tags and the other configs
+
+        Returns:
+            The number of tokens that exist inside LMCache.
+            None indicates the lookup/prefetch is in progress.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """Close the lookup client and clean up resources."""
+        raise NotImplementedError
+
+    def supports_producer_reuse(self) -> bool:
+        """
+        Return whether this lookup client supports producer KV cache reuse.
+
+        Returns:
+            True if producer reuse is supported, False otherwise
+        """
+        return False
+
+    def clear_lookup_status(self, lookup_id: str) -> None:
+        """
+        Clear temporary lookup status for a given lookup ID.
+
+        Args:
+            lookup_id: The lookup ID whose status needs to be cleared.
+        """
+        return

lmcache/v1/lookup_client/async_lookup_message.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Dict, Optional
+
+# Third Party
+import msgspec
+
+
+class AsyncLookupMsg(msgspec.Struct, tag=True):  # type: ignore
+    """Base class for async lookup messages"""
+
+    def describe(self) -> str:
+        return ""
+
+
+class LookupRequestMsg(AsyncLookupMsg):
+    """Async lookup request message from scheduler to worker"""
+
+    lookup_id: str
+    hashes: list[int]
+    offsets: list[int]
+    request_configs: Optional[Dict[str, str]] = None
+
+    def describe(self) -> str:
+        return (
+            f"Async lookup request for lookup_id={self.lookup_id} "
+            f"with {len(self.hashes)} hashes"
+        )
+
+
+class LookupResponseMsg(AsyncLookupMsg):
+    """Async lookup response message from worker to scheduler"""
+
+    lookup_id: str
+    num_hit_tokens: int
+
+    def describe(self) -> str:
+        return (
+            f"Async lookup response for lookup_id={self.lookup_id} "
+            f"with {self.num_hit_tokens} hit tokens"
+        )
+
+
+class LookupCleanupMsg(AsyncLookupMsg):
+    """Cleanup message from scheduler to worker to release memory objects"""
+
+    lookup_id: str
+
+    def describe(self) -> str:
+        return f"Cleanup memory for lookup_id={self.lookup_id}"