lmcache-cli 0.5.6rc4.dev27__py3-none-any.whl → 0.5.7.dev0__py3-none-any.whl

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.6rc4.dev27'
-__version_tuple__ = version_tuple = (0, 5, 6, 'rc4', 'dev27')
+__version__ = version = '0.5.7.dev0'
+__version_tuple__ = version_tuple = (0, 5, 7, 'dev0')
 
-__commit_id__ = commit_id = 'g7d7ca4704'
+__commit_id__ = commit_id = 'g3e11e1081'

lmcache/cli/commands/bench/l2_adapter_bench/data.py

@@ -140,7 +140,7 @@ def make_memory_objects(
         metadata = MemoryObjMetadata(
             shape=torch.Size([data_size]),
             dtype=torch.uint8,
-            address=raw_tensor.data_ptr(),
+            address=start,
             phy_size=data_size * raw_tensor.element_size(),
             fmt=MemoryFormat.KV_2LTD,
             ref_count=1,

lmcache/cli/commands/query/_request.py

@@ -130,8 +130,23 @@ def _stream(
     chat: bool,
     max_tokens: int,
 ) -> dict[str, Any]:
-    """POST with ``stream: true``; parse SSE; return the completion text plus
-    TTFT/TPOT and token metrics."""
+    """POST with ``stream: true`` and collect text and token/latency metrics.
+
+    Args:
+        url: Endpoint URL for the streaming request.
+        body: OpenAI-compatible request payload.
+        timeout: HTTP request timeout in seconds.
+        chat: Whether the endpoint uses chat completion chunks.
+        max_tokens: Output-token estimate when the server omits its count.
+
+    Returns:
+        Completion text, server token counts, and measured latency metrics.
+        A reported completion count of zero is preserved.
+
+    Raises:
+        RuntimeError: If the HTTP request fails, the server reports an error,
+            or the stream contains neither text nor usage.
+    """
     payload = {
         **body,
         "stream": True,
@@ -199,9 +214,8 @@ def _stream(
 
     u = usage or {}
     prompt_tokens = int(u.get("prompt_tokens") or 0)
-    num_completion = int(u.get("completion_tokens") or 0)
-    # Match V2RequestSender: server count if present, else max_tokens cap.
-    num_generated = num_completion if num_completion > 0 else max_tokens
+    num_completion = u.get("completion_tokens")
+    num_generated = int(num_completion) if num_completion is not None else max_tokens
     if first_token_t is None:
         # Use total round-trip as a conservative TTFT approximation.
         ttft_s = t1 - t0
@@ -294,7 +308,9 @@ class Request:
         Returns:
             A ``(answer, metrics)`` tuple where ``answer`` is the model's
             completion text and ``metrics`` is a :data:`MetricMap` of token
-            and latency stats keyed by metric id.
+            and latency stats keyed by metric id. Output tokens use the
+            server's completion count, including zero, or ``max_tokens``
+            when the count is absent.
         """
         request_data = self.build_request(prompt)
         result = self._query_with_fallback(request_data)

lmcache/integration/sglang/multi_process_adapter.py

@@ -101,7 +101,7 @@ def _wrap_sglang_kv_caches(
     device = _validate_sglang_kv_pools(k_pool, v_pool)
     tensors = [*k_pool, *v_pool]
     device_spec = get_device_spec(device.type)
-    if device_spec is None or not device_spec.is_handle_transfer_available():
+    if device_spec is None or not device_spec.is_lmcache_driven_available():
         raise ValueError(
             "SGLang MP handle transfer is unavailable for device type "
             f"{device.type!r}: required memory IPC, event IPC, cache context, "

lmcache/integration/sglang/unified_lmcache_mp_connector.py

@@ -4,7 +4,7 @@
 This module deliberately talks to LMCache's engine-neutral MP protocol.  It
 does not use LMCache's legacy SGLang integration and it never constructs an
 in-process LMCache engine.  The registered SGLang GPU KV tensors remain owned
-by SGLang; LMCache accesses them through device-memory and event IPC handles.
+by SGLang; the selected transfer context moves their KV data.
 """
 
 # Future
@@ -62,7 +62,7 @@ class _ImmediateFuture:
 
 
 class UnifiedLMCacheMPConnector:
-    """Asynchronous, CUDA-IPC connector to a standalone LMCache server."""
+    """Asynchronous device-IPC connector to a standalone LMCache server."""
 
     def __init__(
         self,
@@ -113,8 +113,6 @@ class UnifiedLMCacheMPConnector:
         if not kv_groups or any(not group.kv_tensors for group in kv_groups):
             raise ValueError("LMCache KV group registration cannot be empty")
         kv_tensors = [tensor for group in kv_groups for tensor in group.kv_tensors]
-        if any(t.device.type != "cuda" for t in kv_tensors):
-            raise NotImplementedError("LMCache MP currently requires CUDA KV tensors")
         if any(t.device != kv_tensors[0].device for t in kv_tensors):
             raise ValueError("All LMCache-registered KV tensors must share one device")
         if any(t.dim() < 2 for t in kv_tensors):
@@ -240,6 +238,9 @@ class UnifiedLMCacheMPConnector:
                 _DEFAULT_HEARTBEAT_INTERVAL_SECONDS,
             )
         )
+        self._mp_transfer_mode = config.get_extra_config_value(  # type: ignore[attr-defined]
+            "lmcache.mp.mp_transfer_mode", None
+        )
         if any(
             group.recurrent_state or group.sliding_window_size >= 0
             for group in wire_groups
@@ -302,7 +303,6 @@ class UnifiedLMCacheMPConnector:
         ) = self._build_engine_group_info_specs()
         self._req_client = RequestClientFactory.create(self.server_url)
         self._transfer_ctx: Optional["TransferContext"] = None
-        self._event_backend: Any = None
         self._registered = False
         self._closed = False
         self._lookups: dict[str, LMCacheLookupOperation] = {}
@@ -321,9 +321,10 @@ class UnifiedLMCacheMPConnector:
                 f"LMCache chunk size {self.chunk_size} must be a positive "
                 "multiple of every SGLang group tokens_per_block"
             )
-        # LMCache-driven MP ignores this compatibility argument; per-group
-        # block counts are derived from EngineGroupInfo.tokens_per_block.
-        self.blocks_in_chunk = self.chunk_size
+        # Engine-driven currently accepts one transfer group and needs its
+        # physical block count. LMCache-driven derives per-group counts from
+        # EngineGroupInfo and ignores this compatibility argument.
+        self.blocks_in_chunk = self.chunk_size // self._kv_groups[0].tokens_per_block
         self.register_kv_cache()
         logger.info("UnifiedLMCacheMPConnector initialized succeed.")
 
@@ -525,22 +526,19 @@ class UnifiedLMCacheMPConnector:
             dist.all_reduce(tensor, op=op, group=self.pp_group)
 
     def register_kv_cache(self) -> None:
-        """Export the SGLang GPU tensors to the LMCache MP server once."""
+        """Register SGLang KV tensors with the selected transfer context."""
         # First Party
         from lmcache.utils import EngineType
         from lmcache.v1.multiprocess.group_view import EngineGroupInfo
         from lmcache.v1.multiprocess.transfer_context import create_transfer_context
-        from lmcache.v1.platform.base.event_ipc import get_event_ipc_backend
 
         if self._registered:
             raise RuntimeError("LMCache KV tensors are already registered")
-        self._event_backend = get_event_ipc_backend(self.device)
-        self._event_backend.check_event_support(self.device)
         transfer_ctx = create_transfer_context(
             self._kv_caches,
             instance_id=self.instance_id,
             req_client=self._req_client,
-            mode="lmcache_driven",
+            mode=self._mp_transfer_mode,
         )
         self._transfer_ctx = transfer_ctx
         engine_group_infos = [
@@ -599,13 +597,6 @@ class UnifiedLMCacheMPConnector:
             self._heartbeat_thread = heartbeat_thread
             logger.info("LMCache MP heartbeat thread started")
 
-    def _new_event(self, stream: Any = None) -> Any:
-        event = self._event_backend.create_event(self.device)
-        if stream is None:
-            stream = torch.get_device_module(self.device).current_stream()
-        self._event_backend.record_event(event, stream)
-        return event
-
     def _create_key(
         self,
         operation: LMCacheLookupOperation,
@@ -886,11 +877,7 @@ class UnifiedLMCacheMPConnector:
             transfer_ctx = self._transfer_ctx
             if transfer_ctx is None:
                 raise RuntimeError("LMCache KV tensors are not registered")
-            event = (
-                self._new_event()
-                if producer_stream is None
-                else self._new_event(producer_stream)
-            )
+            event = self._create_recorded_event(transfer_ctx, producer_stream)
             future = transfer_ctx.submit_retrieve(
                 operation.request_id,
                 key,
@@ -941,11 +928,11 @@ class UnifiedLMCacheMPConnector:
             otherwise ``False``.
 
         Notes:
-            Waiting for the raw MQ response imports the LMCache server's
-            completion event. Device completion is then ordered with a stream
-            wait, so the CPU does not wait for H2D. On a cross-rank failure the
-            successful ranks synchronize their local work before the caller
-            releases destination slots.
+            LMCache-driven retrieval imports the server's completion event and
+            orders it on ``stream`` without a CPU device synchronization.
+            Engine-driven retrieval completes its worker-side scatter before
+            returning its plain future. On a cross-rank failure, successful
+            ranks synchronize local work before destination slots are released.
         """
         local_success = False
         try:
@@ -1121,7 +1108,7 @@ class UnifiedLMCacheMPConnector:
                 key = self._create_key(
                     lookup, start=start, end=aligned_end, worker_id=self.kv_worker_id
                 )
-                event = self._new_event()
+                event = transfer_ctx.create_recorded_event()
                 future = transfer_ctx.submit_store(
                     request_id,
                     key,
@@ -1260,3 +1247,13 @@ class UnifiedLMCacheMPConnector:
             transfer_ctx.close()
             self._transfer_ctx = None
         self._req_client.close()
+
+    def _create_recorded_event(
+        self, transfer_ctx: "TransferContext", stream: Any = None
+    ) -> Any:
+        """Create a context-owned event on an optional SGLang producer stream."""
+        if stream is None:
+            return transfer_ctx.create_recorded_event()
+        device_module = torch.get_device_module(self.device)
+        with device_module.stream(stream):
+            return transfer_ctx.create_recorded_event()

lmcache/integration/vllm/lmcache_mp_connector.py

@@ -10,6 +10,7 @@ import time
 # Third Party
 from vllm.config import VllmConfig
 from vllm.distributed.kv_events import (
+    BlockRemoved,
     BlockStored,
     KVCacheEvent,
     KVConnectorKVEvents,
@@ -81,6 +82,7 @@ from lmcache.integration.vllm.utils import (
     mla_only,
     vllm_layout_hints,
 )
+from lmcache.utils import CacheRemoveEvent
 from lmcache.utils import init_logger as lmcache_init_logger
 
 try:
@@ -146,16 +148,23 @@ class LMCacheMPKVEvents(KVConnectorKVEvents):
 
     def __init__(self, num_workers: int) -> None:
         self._aggregator = KVEventAggregator(num_workers)
+        self._batches: list[list[KVCacheEvent]] = []
 
     def add_events(self, events: list[KVCacheEvent]) -> None:
         """Add events from one worker."""
         self._aggregator.add_events(events)
+        self._batches.append(list(events))
 
     def aggregate(self) -> "LMCacheMPKVEvents":
-        """Retain only events seen from every contributing worker."""
-        common_events = self._aggregator.get_common_events()
+        """Keep all pollers' events, preserving transitions within each batch."""
+        remaining = set(self._aggregator.get_all_events())
+        merged: list[KVCacheEvent] = []
+        for events in self._batches:
+            merged.extend(event for event in events if event in remaining)
+            remaining.difference_update(events)
+        self._batches = [merged] if merged else []
         self._aggregator.clear_events()
-        self._aggregator.add_events(common_events)
+        self._aggregator.add_events(merged)
         self._aggregator.reset_workers()
         return self
 
@@ -164,8 +173,8 @@ class LMCacheMPKVEvents(KVConnectorKVEvents):
         self._aggregator.increment_workers(count)
 
     def get_all_events(self) -> list[KVCacheEvent]:
-        """Return all buffered events."""
-        return self._aggregator.get_all_events()
+        """Return all buffered events in worker order."""
+        return [event for batch in self._batches for event in batch]
 
     def get_number_of_workers(self) -> int:
         """Return the number of contributing workers."""
@@ -175,6 +184,7 @@ class LMCacheMPKVEvents(KVConnectorKVEvents):
         """Clear buffered events and reset the worker count."""
         self._aggregator.clear_events()
         self._aggregator.reset_workers()
+        self._batches.clear()
 
 
 # Helper functions
@@ -727,7 +737,8 @@ class LMCacheMPConnector(KVConnectorBase_V1, SupportsHMA):
                 extra_config=vllm_config.kv_transfer_config.kv_connector_extra_config,
             )
             self.request_trackers: dict[str, LMCacheMPRequestTracker] = {}
-            self._kv_cache_events: LMCacheMPKVEvents | None = None
+            # Worker events aggregated per step, awaiting take_events().
+            self._kv_cache_events: list[KVCacheEvent] = []
 
             # Align-mode Mamba keeps one speculative block per draft token at
             # the tail of a request's block list.
@@ -1124,8 +1135,10 @@ class LMCacheMPConnector(KVConnectorBase_V1, SupportsHMA):
         """Return worker-side KV cache events collected since the last step.
 
         Returns:
-            A vLLM KV event container with completed LMCache store events, or
-            None when disabled or when no store completed.
+            A vLLM KV event container with ``BlockStored`` events for LMCache
+            stores (own completed stores and stores learned from the MP
+            server) and ``BlockRemoved`` events for CPU evictions, or None
+            when disabled or when nothing happened.
         """
         if not self._enable_kv_events:
             return None
@@ -1133,21 +1146,27 @@ class LMCacheMPConnector(KVConnectorBase_V1, SupportsHMA):
         if not events:
             return None
 
-        blocks = [
-            BlockStored(
-                block_hashes=[
-                    _convert_kv_event_hash(block_hash)
-                    for block_hash in event.block_hashes
-                ],
-                parent_block_hash=_convert_kv_event_hash(event.parent_block_hash),
-                token_ids=event.token_ids,
-                lora_id=event.lora_id,
-                block_size=event.block_size,
-                medium=event.medium,
-                lora_name=event.lora_name,
+        blocks: list[KVCacheEvent] = []
+        for event in events:
+            block_hashes = [
+                _convert_kv_event_hash(block_hash) for block_hash in event.block_hashes
+            ]
+            if isinstance(event, CacheRemoveEvent):
+                blocks.append(
+                    BlockRemoved(block_hashes=block_hashes, medium=event.medium)
+                )
+                continue
+            blocks.append(
+                BlockStored(
+                    block_hashes=block_hashes,
+                    parent_block_hash=_convert_kv_event_hash(event.parent_block_hash),
+                    token_ids=event.token_ids,
+                    lora_id=event.lora_id,
+                    block_size=event.block_size,
+                    medium=event.medium,
+                    lora_name=event.lora_name,
+                )
             )
-            for event in events
-        ]
         kv_events = LMCacheMPKVEvents(num_workers=1)
         kv_events.add_events(blocks)
         return kv_events
@@ -1487,13 +1506,8 @@ class LMCacheMPConnector(KVConnectorBase_V1, SupportsHMA):
         """
         kv_cache_events = connector_output.kv_cache_events
         if kv_cache_events and isinstance(kv_cache_events, LMCacheMPKVEvents):
-            if self._kv_cache_events is None:
-                self._kv_cache_events = kv_cache_events
-            else:
-                self._kv_cache_events.add_events(kv_cache_events.get_all_events())
-                self._kv_cache_events.increment_workers(
-                    kv_cache_events.get_number_of_workers()
-                )
+            # One aggregated batch per step keeps the workers' event order.
+            self._kv_cache_events.extend(kv_cache_events.aggregate().get_all_events())
 
         if not self.lazy_offload:
             return
@@ -1577,14 +1591,13 @@ class LMCacheMPConnector(KVConnectorBase_V1, SupportsHMA):
         """
         Take the KV cache events from the connector.
 
-        Yields:
-            New KV cache events since the last call.
+        Returns:
+            The KV cache events collected since the last call, in the order
+            the workers reported them. The buffer is cleared.
         """
-        if self._kv_cache_events is not None:
-            self._kv_cache_events.aggregate()
-            yield from self._kv_cache_events.get_all_events()
-            self._kv_cache_events.clear_events()
-            self._kv_cache_events = None
+        events = self._kv_cache_events
+        self._kv_cache_events = []
+        return events
 
     def has_pending_push_work(self) -> bool:
         """Return whether vLLM should keep stepping for pending push work.