[Feat][LMCache] Support LMCache mp mode (#24089)

Signed-off-by: Shaoting-Feng <stfeng@uw.edu>

python/sglang/srt/mem_cache/storage/lmcache/README.md

@@ -29,15 +29,43 @@ pip install -e . --no-build-isolation
 
 ## Use LMCache
 
-Firstly, setup LMCache config. An example config is set at `example_config.yaml`. For more settings please refer to https://docs.lmcache.ai/api_reference/configurations.html.
+LMCache supports two transport modes. **MP (multi-process, default)** issues a single blocking retrieve over ZMQ to a standalone daemon that owns the KV store and survives SGLang restarts. **IP (in-process)** uses an embedded layerwise connector — the cache lives and dies with the SGLang process. Mode selection is currently a code-level setting in `LMCRadixCache.__init__` (`self._mode`); only MP is reachable by default.
 
-Secondly, setup SGLang serving engine with lmcache:
+### MP mode (default): multi-process daemon
+
+Uses `LMCacheMPConnector`. Daemon host/port come from the LMCache YAML config (`mp_host`, `mp_port`).
+
+Terminal 1 — start the LMCache daemon:
 
 ```bash
-export LMCACHE_USE_EXPERIMENTAL=True
-export LMCACHE_CONFIG_FILE=example_config.yaml
+lmcache server \
+  --host 127.0.0.1 --port 5556 \
+  --l1-size-gb 4 \
+  --eviction-policy LRU
+```
 
+Use the bundled `example_config_mp.yaml` (or any YAML setting `mp_host` / `mp_port`):
+
+Terminal 2 — start SGLang:
+
+```bash
 python -m sglang.launch_server \
   --model-path MODEL \
-  --enable-lmcache
+  --enable-lmcache \
+  --lmcache-config-file example_config_mp.yaml
+```
+
+For full LMCache config options see https://docs.lmcache.ai/api_reference/configurations.html.
+
+### IP mode: in-process
+
+Uses `LMCacheLayerwiseConnector`. KV transfer happens per layer inside the SGLang process; the cache lives and dies with the server. To enable, edit `LMCRadixCache.__init__` and set `self._mode = LMCacheMode.IP`.
+
+The LMCache config still controls chunk_size and storage; `mp_host` / `mp_port` are ignored on this path. Use the bundled `example_config_ip.yaml`:
+
+```bash
+python -m sglang.launch_server \
+  --model-path MODEL \
+  --enable-lmcache \
+  --lmcache-config-file example_config_ip.yaml
 ```

python/sglang/srt/mem_cache/storage/lmcache/example_config_mp.yaml

@@ -0,0 +1,3 @@
+# MP mode: SGLang dials the standalone `lmcache server` at this host/port.
+mp_host: 127.0.0.1
+mp_port: 5556

python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py

@@ -1,25 +1,31 @@
 from __future__ import annotations
 
+import enum
 import logging
 import threading
-from typing import TYPE_CHECKING, Optional
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
 
 from sglang.srt.mem_cache.base_prefix_cache import (
     EvictParams,
     EvictResult,
+    InitLoadBackParams,
     MatchPrefixParams,
     MatchResult,
 )
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
+from sglang.srt.server_args import get_global_server_args
 
 try:
+    from lmcache.integration.sglang.multi_process_adapter import LMCacheMPConnector
     from lmcache.integration.sglang.sglang_adapter import (
         LMCacheLayerwiseConnector,
         LoadMetadata,
         StoreMetadata,
     )
+    from lmcache.integration.sglang.utils import lmcache_get_config
 except ImportError as e:
     raise RuntimeError(
         "LMCache is not installed. Please install it by running `pip install lmcache`"
@@ -34,6 +40,21 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class _LMCacheLoadBackMarker:
+    """Carries the data ``init_load_back`` needs from the
+    ``match_prefix`` call in MP mode.
+    """
+
+    key: RadixKey  # page-aligned key the scheduler matched on
+    value_numel: int  # number of tokens already in radix at match time
+
+
+class LMCacheMode(enum.Enum):
+    MP = enum.auto()  # multi-process mode
+    IP = enum.auto()  # in-process mode
+
+
 class LayerTransferCounter:
     """Minimal adapter that lets the memory pool notify LMCache per-layer.
 
@@ -63,13 +84,17 @@ class LayerTransferCounter:
 class LMCRadixCache(RadixCache):
     """RadixCache + LMCache IO.
 
-    This subclass adds:
-      - LMCache connector setup (device/host buffers, TP rank/size)
-      - Two CUDA streams for async load/store
-      - Layer-wise transfer executor wiring to the KV cache
-      - Overridden `match_prefix` to fetch missing prefix chunks from LMCache
-      - Extended cache_finalization paths to store back into LMCache
-      - Eviction barrier that respects any in-flight host->device stores
+    IP mode keeps the existing layerwise connector and
+    its per-layer transfer hook: ``match_prefix`` kicks off the load via
+    ``start_load_kv`` and SGLang's per-layer KV-pool hook drives subsequent
+    layers during forward.
+
+    MP mode uses ``LMCacheMPConnector`` with a two-phase
+    load: ``match_prefix`` fires LOOKUP only (``connector.lookup_kv``) and
+    returns ``host_hit_length`` on the ``MatchResult``; the SGLang
+    scheduler then calls `init_load_back` at dispatch time,
+    which fires the actual RETRIEVE (``connector.retrieve_kv``) into
+    pre-allocated GPU slots.
     """
 
     def __init__(
@@ -82,8 +107,10 @@ class LMCRadixCache(RadixCache):
     ):
         super().__init__(params)
 
+        cli_lmc_cfg = get_global_server_args().lmcache_config_file or ""
+
         kvcache = self.token_to_kv_pool_allocator.get_kvcache()
-        self.lmcache_connector = LMCacheLayerwiseConnector(
+        connector_kwargs = dict(
             sgl_config=model_config,
             tp_size=tp_size,
             rank=rank,
@@ -106,32 +133,54 @@ class LMCRadixCache(RadixCache):
         self.load_stream = torch.cuda.Stream()
         self.store_stream = torch.cuda.Stream()
 
-        self.layer_done_executor = LayerTransferCounter(
-            num_layers=(
-                model_config.num_hidden_layers if model_config is not None else 0
-            ),
-            load_stream=self.load_stream,
-            lmc_connector=self.lmcache_connector,
-        )
-        kvcache.register_layer_transfer_counter(self.layer_done_executor)
+        # MP is the default. To use the in-process layerwise connector,
+        # set ``self._mode = LMCacheMode.IP`` here.
+        self._mode = LMCacheMode.MP
+        if self._mode is LMCacheMode.MP:
+            if not cli_lmc_cfg:
+                raise ValueError(
+                    "MP mode requires --lmcache-config-file (the YAML "
+                    "supplies mp_host / mp_port)."
+                )
+            lm_cfg = lmcache_get_config(cli_lmc_cfg)
+            self.lmcache_connector = LMCacheMPConnector(
+                page_size=params.page_size,
+                host=lm_cfg.mp_host,
+                port=lm_cfg.mp_port,
+                **connector_kwargs,
+            )
+        elif self._mode is LMCacheMode.IP:
+            self.lmcache_connector = LMCacheLayerwiseConnector(
+                config_file=cli_lmc_cfg, **connector_kwargs
+            )
+            # Per-layer hook
+            self.layer_done_executor = LayerTransferCounter(
+                num_layers=(
+                    model_config.num_hidden_layers if model_config is not None else 0
+                ),
+                load_stream=self.load_stream,
+                lmc_connector=self.lmcache_connector,
+            )
+            kvcache.register_layer_transfer_counter(self.layer_done_executor)
 
         self._in_flight_nodes: list[TreeNode] = []
         self._node_lock = threading.Lock()
+        self._mp_load_back_markers: dict[str, _LMCacheLoadBackMarker] = {}
 
-    def reset(self):  # type: ignore[override]
+    def reset(self):
         super().reset()
         if hasattr(self, "_in_flight_nodes"):
             with self._node_lock:
                 self._in_flight_nodes.clear()
+        if hasattr(self, "_mp_load_back_markers"):
+            self._mp_load_back_markers.clear()
 
-    def match_prefix(self, params: MatchPrefixParams) -> MatchResult:  # type: ignore[override]
-        """Match cached prefix; if there's a tail miss, prefetch from LMCache.
+    def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+        """Dispatch to the mode-specific match_prefix.
 
-        Reuses the base matching logic to obtain (value, last_node). If there
-        remains a *page-aligned* uncached suffix and there is room (or after
-        eviction), we allocate token slots and trigger an async LMCache load
-        into those slots, then materialize a new child node for the retrieved
-        chunk.
+        MP mode → ``_mp_match_prefix`` (fires LOOKUP only).
+        IP mode → ``_ip_match_prefix`` (single-shot ``start_load_kv``
+        plus per-layer hook).
         """
         key = params.key
         if self.disable or not key:
@@ -145,6 +194,59 @@ class LMCRadixCache(RadixCache):
         value: torch.Tensor = base_res.device_indices
         last_node: TreeNode = base_res.last_device_node
 
+        if self._mode is LMCacheMode.MP:
+            if params.req is None:
+                return base_res
+            return self._mp_match_prefix(key, base_res, value, last_node, params.req)
+        elif self._mode is LMCacheMode.IP:
+            return self._ip_match_prefix(key, base_res, value, last_node)
+        return base_res
+
+    def _mp_match_prefix(
+        self,
+        key: RadixKey,
+        base_res: MatchResult,
+        value: torch.Tensor,
+        last_node: TreeNode,
+        req: Req,
+    ) -> MatchResult:
+        """MP LOOKUP
+
+        Returns a ``MatchResult`` with ``host_hit_length`` set when
+        LMCache has tokens beyond radix. Otherwise releases
+        the held read locks and returns the radix-only result.
+        """
+        matched = self.lmcache_connector.lookup_kv(key.token_ids, req.rid)
+        if matched <= value.numel():
+            # Release the read locks; keep the pending session for end_session.
+            self.lmcache_connector.release_pending(req.rid)
+            return base_res
+
+        self._mp_load_back_markers[req.rid] = _LMCacheLoadBackMarker(
+            key=key,
+            value_numel=int(value.numel()),
+        )
+        return MatchResult(
+            device_indices=value,
+            last_device_node=last_node,
+            last_host_node=last_node,
+            best_match_node=last_node,
+            host_hit_length=matched - int(value.numel()),
+        )
+
+    def _ip_match_prefix(
+        self,
+        key: RadixKey,
+        base_res: MatchResult,
+        value: torch.Tensor,
+        last_node: TreeNode,
+    ) -> MatchResult:
+        """IP mode: ``start_load_kv`` + per-layer hook.
+
+        Allocates slots for the page-aligned uncached tail and kicks off
+        the layerwise load. Returns ``base_res`` if there's nothing to
+        fetch or alloc/load fails.
+        """
         if value.numel() == len(key):
             return base_res
 
@@ -152,31 +254,99 @@ class LMCRadixCache(RadixCache):
         if uncached_len == 0:
             return base_res
 
+        result = self._load_back(
+            key=key,
+            value_numel=int(value.numel()),
+            uncached_len=uncached_len,
+            last_node=last_node,
+            load_fn=lambda sm, pp: self._ip_load_back(
+                token_ids=key.token_ids,
+                value_numel=int(value.numel()),
+                slot_mapping=sm,
+                prefix_pad=pp,
+            ),
+        )
+        if result is None:
+            return base_res
+        new_slots, new_node = result
+        return MatchResult(
+            device_indices=torch.cat([value, new_slots]),
+            last_device_node=new_node,
+            last_host_node=new_node,
+            best_match_node=new_node,
+        )
+
+    def init_load_back(
+        self, params: InitLoadBackParams
+    ) -> Tuple[torch.Tensor, Optional[TreeNode]]:
+        """MP RETRIEVE.
+
+        Called by the scheduler when ``match_prefix`` returned
+        ``host_hit_length > 0``. Uses the cached LOOKUP result to
+        allocate slots and fire RETRIEVE, inserts the resulting
+        TreeNode into the radix tree, and returns
+        ``(new_indices, new_last_node)``.
+        """
+        req = params.req
+        marker = self._mp_load_back_markers.pop(req.rid)
+        last_node: TreeNode = params.best_match_node
+
+        result = self._load_back(
+            key=marker.key,
+            value_numel=marker.value_numel,
+            uncached_len=params.host_hit_length,
+            last_node=last_node,
+            load_fn=lambda sm, pp: self._mp_load_back(
+                marker=marker,
+                request_id=req.rid,
+                slot_mapping=sm,
+                prefix_pad=pp,
+            ),
+        )
+        if result is None:
+            # Either alloc failed (locks still held by lookup_kv) or
+            # retrieve returned nothing (locks already released by
+            # retrieve_kv). release_pending is idempotent on locks_held.
+            self.lmcache_connector.release_pending(req.rid)
+            return (
+                torch.empty((0,), dtype=torch.int64, device=self.device),
+                last_node,
+            )
+        return result
+
+    def _load_back(
+        self,
+        *,
+        key: RadixKey,
+        value_numel: int,
+        uncached_len: int,
+        last_node: TreeNode,
+        load_fn,  # Callable[[torch.Tensor, int], int] — (slot_mapping, prefix_pad) -> num_retrieved
+    ) -> Optional[Tuple[torch.Tensor, TreeNode]]:
+        """Alloc slots, run ``load_fn``, attach a TreeNode for what was loaded.
+
+        Returns ``(slots, new_node)`` on success, ``None`` if alloc fails
+        or the load returned zero (slots are freed in either case).
+        """
         chunk_size = self.lmcache_connector.chunk_size()
-        prefix_pad = value.numel() % chunk_size
+        prefix_pad = value_numel % chunk_size
 
         if self.token_to_kv_pool_allocator.available_size() < uncached_len:
             self.evict(EvictParams(num_tokens=uncached_len))
 
         token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len)
         if token_slots is None:
-            return base_res
+            return None
 
-        slot_mapping = torch.cat(
-            [
-                torch.full((value.numel(),), -1, dtype=torch.int64, device=self.device),
-                token_slots.detach().clone().to(torch.int64).to(self.device),
-            ]
+        slot_mapping = torch.empty(
+            value_numel + token_slots.numel(),
+            dtype=torch.int64,
+            device=self.device,
         )
+        slot_mapping[:value_numel].fill_(-1)
+        slot_mapping[value_numel:].copy_(token_slots)
 
-        with torch.cuda.stream(self.load_stream):
-            num_retrieved = self.lmcache_connector.start_load_kv(
-                LoadMetadata(
-                    token_ids=key.token_ids,  # full page-aligned key
-                    slot_mapping=slot_mapping,
-                    offset=value.numel() - prefix_pad,  # LMCache offset convention
-                )
-            )
+        num_retrieved = load_fn(slot_mapping, prefix_pad)
         logger.debug("num_retrieved_tokens: %s", num_retrieved)
 
         if num_retrieved > 0:
@@ -189,38 +359,80 @@ class LMCRadixCache(RadixCache):
         if num_retrieved > 0:
             fetched = num_retrieved - prefix_pad
             new_node = TreeNode(priority=last_node.priority)
-            start = value.numel()
+            start = value_numel
             end = start + fetched
             new_node.key = key[start:end]
             new_node.value = token_slots[:fetched]
             new_node.parent = last_node
             last_node.children[new_node.key.child_key(self.page_size)] = new_node
-            last_node = new_node
-
-            value = torch.cat([value, token_slots[:fetched]])
             self.evictable_size_ += fetched
+            self._update_leaf_status(last_node)
+            self._update_leaf_status(new_node)
 
             self._record_store_event(new_node.parent)
             self._record_store_event(new_node)
 
-            return MatchResult(
-                device_indices=value,
-                last_device_node=last_node,
-                last_host_node=last_node,
-                best_match_node=last_node,
+            return token_slots[:fetched], new_node
+
+        return None
+
+    def _mp_load_back(
+        self,
+        *,
+        marker: _LMCacheLoadBackMarker,
+        request_id: str,
+        slot_mapping: torch.Tensor,
+        prefix_pad: int,
+    ) -> int:
+        """MP non-layerwise loader: fire ``retrieve_kv`` and wait for the
+        load_stream so the compute stream observes the writes.
+        """
+        self.load_stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(self.load_stream):
+            n = self.lmcache_connector.retrieve_kv(
+                LoadMetadata(
+                    token_ids=marker.key.token_ids,
+                    slot_mapping=slot_mapping,
+                    offset=marker.value_numel - prefix_pad,
+                    prefix_pad=prefix_pad,
+                    request_id=request_id,
+                )
+            )
+        torch.cuda.current_stream().wait_stream(self.load_stream)
+        return n
+
+    def _ip_load_back(
+        self,
+        *,
+        token_ids: list[int],
+        value_numel: int,
+        slot_mapping: torch.Tensor,
+        prefix_pad: int,
+    ) -> int:
+        """IP layerwise loader: kick off ``start_load_kv`` on ``self.load_stream``.
+
+        ``start_load_kv`` enqueues the first layer's transfer; the
+        ``LayerTransferCounter`` hook drives the rest during forward.
+        """
+        with torch.cuda.stream(self.load_stream):
+            return self.lmcache_connector.start_load_kv(
+                LoadMetadata(
+                    token_ids=token_ids,
+                    slot_mapping=slot_mapping,
+                    offset=value_numel - prefix_pad,
+                )
             )
 
-        return base_res
-
-    def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:  # type: ignore[override]
+    def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
         """On request completion, insert device KV into radix and store to LMCache."""
 
         super().cache_finished_req(req, is_insert=is_insert)
         if not is_insert:
+            if self._mode is LMCacheMode.MP:
+                self._mp_load_back_markers.pop(req.rid, None)
+                self.lmcache_connector.end_session(req.rid)
             return
 
-        from sglang.srt.server_args import get_global_server_args
-
         global_server_args = get_global_server_args()
         topk = global_server_args.speculative_eagle_topk
         enable_kv_committed_len = topk is None or topk == 1
@@ -236,7 +448,8 @@ class LMCRadixCache(RadixCache):
             req.req_pool_idx, :kv_committed_len
         ]
 
-        match_result = self.match_prefix(
+        # Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix.
+        match_result = super().match_prefix(
             MatchPrefixParams(key=RadixKey(token_ids, req.extra_key))
         )
         new_last_node = match_result.last_device_node
@@ -248,11 +461,19 @@ class LMCRadixCache(RadixCache):
             token_ids=token_ids,
             kv_indices=kv_indices,
             offset=0,
+            request_id=req.rid,
         )
         with torch.cuda.stream(self.store_stream):
             self.lmcache_connector.store_kv(store_md)
-        with self._node_lock:
-            self._in_flight_nodes.append(new_last_node)
+        if self._mode is LMCacheMode.MP:
+            # MP store_kv blocks until the daemon's signal event fires, so the slots are safe to evict immediately.
+            self._mp_load_back_markers.pop(req.rid, None)
+            self.dec_lock_ref(new_last_node)
+            self.lmcache_connector.end_session(req.rid)
+        elif self._mode is LMCacheMode.IP:
+            # Layerwise store is async on store_stream; defer the unlock to evict()'s store_stream.synchronize().
+            with self._node_lock:
+                self._in_flight_nodes.append(new_last_node)
 
     def evict(self, params: EvictParams) -> EvictResult:
         """Before base eviction, wait for any outstanding stores and release locks."""
@@ -267,7 +488,7 @@ class LMCRadixCache(RadixCache):
 
         return super().evict(params)
 
-    def pretty_print(self):  # type: ignore[override]
+    def pretty_print(self):
         super().pretty_print()
         try:
             logger.debug(

python/sglang/srt/mem_cache/storage/lmcache/unit_test.py

@@ -9,15 +9,10 @@ except ImportError:
         "LMCache is not installed. Please install it by running `pip install lmcache` in the root directory of LMCache"
     )
 
-import os
-
 import torch
 
 from sglang.srt.configs.model_config import ModelConfig
 
-os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
-os.environ["LMCACHE_CONFIG_FILE"] = "example_config.yaml"
-
 
 def test_load_store_metadata():
     model_config = ModelConfig(
@@ -40,7 +35,9 @@ def test_load_store_metadata():
         for _ in range(layer_num)
     ]
 
-    connector = LMCacheLayerwiseConnector(model_config, 1, 0, k_buffer, v_buffer)
+    connector = LMCacheLayerwiseConnector(
+        model_config, 1, 0, k_buffer, v_buffer, config_file="example_config_ip.yaml"
+    )
 
     fake_token_ids = torch.randint(0, model_config.vocab_size, (input_id_len,)).tolist()
     fake_kv_indices = torch.randint(0, buffer_size, (input_id_len,))

python/sglang/srt/server_args.py

@@ -673,6 +673,7 @@ class ServerArgs:
 
     # LMCache
     enable_lmcache: bool = False
+    lmcache_config_file: Optional[str] = None
 
     # Ktransformers/AMX expert parallelism
     kt_weight_path: Optional[str] = None
@@ -6087,6 +6088,12 @@ class ServerArgs:
             action="store_true",
             help="Using LMCache as an alternative hierarchical cache solution",
         )
+        parser.add_argument(
+            "--lmcache-config-file",
+            type=str,
+            default=ServerArgs.lmcache_config_file,
+            help="Path to the LMCache YAML configuration file",
+        )
 
         # Ktransformer server args
         parser.add_argument(