diff --git a/docs_new/docs/advanced_features/hicache_design.mdx b/docs_new/docs/advanced_features/hicache_design.mdx index 24ca4dac1..0265ea029 100644 --- a/docs_new/docs/advanced_features/hicache_design.mdx +++ b/docs_new/docs/advanced_features/hicache_design.mdx @@ -163,6 +163,8 @@ Specifically, **LMCache**, an efficient KV cache layer for enterprise-scale LLM - **`--enable-lmcache`**: Using LMCache as an alternative hierarchical cache solution. +- **`--lmcache-config-file`**: Path to the LMCache YAML configuration file. + - **`--hicache-storage-backend-extra-config HICACHE_STORAGE_BACKEND_EXTRA_CONFIG`**: the extra config can be either - a JSON string containing extra configuration for the storage backend, e.g., `--hicache-storage-backend-extra-config '{"prefetch_threshold":512, "prefetch_timeout_base": 0.5, "prefetch_timeout_per_ki_token": 0.25}' `, or - a TOML or JSON or YAML file specifying the extra configuration for the storage backend (to differentiate from the JSON string input, prepend a `@` in front of the file name), e.g., `--hicache-storage-backend-extra-config "@config.toml"` where `config.toml` is the config file containing the complex configurations. This can be useful when the configuration consists of many or complex key-value pairs (for instance, it is preferred to use a config file for NIXL backend as its configurations can be complex). diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index 68f19455f..c2a9f282e 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -1756,6 +1756,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s `False` bool flag (set to enable) + + `--lmcache-config-file` + Path to the LMCache YAML configuration file + `None` + Type: str + diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx index ded6c90db..166572491 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx @@ -1751,6 +1751,12 @@ click [Server Arguments](../../advanced_features/server_arguments). bool flag
(set to enable) Special for GPU + + `--lmcache-config-file` + `None` + Type: str + Special for GPU + diff --git a/python/sglang/srt/mem_cache/storage/lmcache/README.md b/python/sglang/srt/mem_cache/storage/lmcache/README.md index 7177e21e5..f23954d97 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/README.md +++ b/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 ``` diff --git a/python/sglang/srt/mem_cache/storage/lmcache/example_config.yaml b/python/sglang/srt/mem_cache/storage/lmcache/example_config_ip.yaml similarity index 100% rename from python/sglang/srt/mem_cache/storage/lmcache/example_config.yaml rename to python/sglang/srt/mem_cache/storage/lmcache/example_config_ip.yaml diff --git a/python/sglang/srt/mem_cache/storage/lmcache/example_config_mp.yaml b/python/sglang/srt/mem_cache/storage/lmcache/example_config_mp.yaml new file mode 100644 index 000000000..c398addc3 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 1df5ec639..88a640041 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/storage/lmcache/unit_test.py b/python/sglang/srt/mem_cache/storage/lmcache/unit_test.py index 68dfe939d..b6dae15d0 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/unit_test.py +++ b/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,)) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 35596790a..ec851bb03 100644 --- a/python/sglang/srt/server_args.py +++ b/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(