[Feat][LMCache] Support LMCache mp mode (#24089)
Signed-off-by: Shaoting-Feng <stfeng@uw.edu>
This commit is contained in:
@@ -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.
|
- **`--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
|
- **`--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 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).
|
- 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).
|
||||||
|
|||||||
@@ -1756,6 +1756,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
|||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
|
||||||
</tr>
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--lmcache-config-file`</td>
|
||||||
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Path to the LMCache YAML configuration file</td>
|
||||||
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
|
||||||
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: str</td>
|
||||||
|
</tr>
|
||||||
</tbody>
|
</tbody>
|
||||||
</table>
|
</table>
|
||||||
|
|
||||||
|
|||||||
@@ -1751,6 +1751,12 @@ click [Server Arguments](../../advanced_features/server_arguments).
|
|||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag<br/> (set to enable)</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag<br/> (set to enable)</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
||||||
</tr>
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--lmcache-config-file`</td>
|
||||||
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`None`</td>
|
||||||
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Type: str</td>
|
||||||
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
||||||
|
</tr>
|
||||||
</tbody>
|
</tbody>
|
||||||
</table>
|
</table>
|
||||||
|
|
||||||
|
|||||||
@@ -29,15 +29,43 @@ pip install -e . --no-build-isolation
|
|||||||
|
|
||||||
## Use LMCache
|
## 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
|
```bash
|
||||||
export LMCACHE_USE_EXPERIMENTAL=True
|
lmcache server \
|
||||||
export LMCACHE_CONFIG_FILE=example_config.yaml
|
--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 \
|
python -m sglang.launch_server \
|
||||||
--model-path MODEL \
|
--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
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -1,25 +1,31 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import enum
|
||||||
import logging
|
import logging
|
||||||
import threading
|
import threading
|
||||||
from typing import TYPE_CHECKING, Optional
|
from dataclasses import dataclass
|
||||||
|
from typing import TYPE_CHECKING, Optional, Tuple
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||||
EvictParams,
|
EvictParams,
|
||||||
EvictResult,
|
EvictResult,
|
||||||
|
InitLoadBackParams,
|
||||||
MatchPrefixParams,
|
MatchPrefixParams,
|
||||||
MatchResult,
|
MatchResult,
|
||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
|
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
|
||||||
|
from sglang.srt.server_args import get_global_server_args
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
from lmcache.integration.sglang.multi_process_adapter import LMCacheMPConnector
|
||||||
from lmcache.integration.sglang.sglang_adapter import (
|
from lmcache.integration.sglang.sglang_adapter import (
|
||||||
LMCacheLayerwiseConnector,
|
LMCacheLayerwiseConnector,
|
||||||
LoadMetadata,
|
LoadMetadata,
|
||||||
StoreMetadata,
|
StoreMetadata,
|
||||||
)
|
)
|
||||||
|
from lmcache.integration.sglang.utils import lmcache_get_config
|
||||||
except ImportError as e:
|
except ImportError as e:
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
"LMCache is not installed. Please install it by running `pip install lmcache`"
|
"LMCache is not installed. Please install it by running `pip install lmcache`"
|
||||||
@@ -34,6 +40,21 @@ if TYPE_CHECKING:
|
|||||||
logger = logging.getLogger(__name__)
|
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:
|
class LayerTransferCounter:
|
||||||
"""Minimal adapter that lets the memory pool notify LMCache per-layer.
|
"""Minimal adapter that lets the memory pool notify LMCache per-layer.
|
||||||
|
|
||||||
@@ -63,13 +84,17 @@ class LayerTransferCounter:
|
|||||||
class LMCRadixCache(RadixCache):
|
class LMCRadixCache(RadixCache):
|
||||||
"""RadixCache + LMCache IO.
|
"""RadixCache + LMCache IO.
|
||||||
|
|
||||||
This subclass adds:
|
IP mode keeps the existing layerwise connector and
|
||||||
- LMCache connector setup (device/host buffers, TP rank/size)
|
its per-layer transfer hook: ``match_prefix`` kicks off the load via
|
||||||
- Two CUDA streams for async load/store
|
``start_load_kv`` and SGLang's per-layer KV-pool hook drives subsequent
|
||||||
- Layer-wise transfer executor wiring to the KV cache
|
layers during forward.
|
||||||
- Overridden `match_prefix` to fetch missing prefix chunks from LMCache
|
|
||||||
- Extended cache_finalization paths to store back into LMCache
|
MP mode uses ``LMCacheMPConnector`` with a two-phase
|
||||||
- Eviction barrier that respects any in-flight host->device stores
|
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__(
|
def __init__(
|
||||||
@@ -82,8 +107,10 @@ class LMCRadixCache(RadixCache):
|
|||||||
):
|
):
|
||||||
super().__init__(params)
|
super().__init__(params)
|
||||||
|
|
||||||
|
cli_lmc_cfg = get_global_server_args().lmcache_config_file or ""
|
||||||
|
|
||||||
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
|
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
|
||||||
self.lmcache_connector = LMCacheLayerwiseConnector(
|
connector_kwargs = dict(
|
||||||
sgl_config=model_config,
|
sgl_config=model_config,
|
||||||
tp_size=tp_size,
|
tp_size=tp_size,
|
||||||
rank=rank,
|
rank=rank,
|
||||||
@@ -106,32 +133,54 @@ class LMCRadixCache(RadixCache):
|
|||||||
self.load_stream = torch.cuda.Stream()
|
self.load_stream = torch.cuda.Stream()
|
||||||
self.store_stream = torch.cuda.Stream()
|
self.store_stream = torch.cuda.Stream()
|
||||||
|
|
||||||
self.layer_done_executor = LayerTransferCounter(
|
# MP is the default. To use the in-process layerwise connector,
|
||||||
num_layers=(
|
# set ``self._mode = LMCacheMode.IP`` here.
|
||||||
model_config.num_hidden_layers if model_config is not None else 0
|
self._mode = LMCacheMode.MP
|
||||||
),
|
if self._mode is LMCacheMode.MP:
|
||||||
load_stream=self.load_stream,
|
if not cli_lmc_cfg:
|
||||||
lmc_connector=self.lmcache_connector,
|
raise ValueError(
|
||||||
)
|
"MP mode requires --lmcache-config-file (the YAML "
|
||||||
kvcache.register_layer_transfer_counter(self.layer_done_executor)
|
"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._in_flight_nodes: list[TreeNode] = []
|
||||||
self._node_lock = threading.Lock()
|
self._node_lock = threading.Lock()
|
||||||
|
self._mp_load_back_markers: dict[str, _LMCacheLoadBackMarker] = {}
|
||||||
|
|
||||||
def reset(self): # type: ignore[override]
|
def reset(self):
|
||||||
super().reset()
|
super().reset()
|
||||||
if hasattr(self, "_in_flight_nodes"):
|
if hasattr(self, "_in_flight_nodes"):
|
||||||
with self._node_lock:
|
with self._node_lock:
|
||||||
self._in_flight_nodes.clear()
|
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]
|
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
||||||
"""Match cached prefix; if there's a tail miss, prefetch from LMCache.
|
"""Dispatch to the mode-specific match_prefix.
|
||||||
|
|
||||||
Reuses the base matching logic to obtain (value, last_node). If there
|
MP mode → ``_mp_match_prefix`` (fires LOOKUP only).
|
||||||
remains a *page-aligned* uncached suffix and there is room (or after
|
IP mode → ``_ip_match_prefix`` (single-shot ``start_load_kv``
|
||||||
eviction), we allocate token slots and trigger an async LMCache load
|
plus per-layer hook).
|
||||||
into those slots, then materialize a new child node for the retrieved
|
|
||||||
chunk.
|
|
||||||
"""
|
"""
|
||||||
key = params.key
|
key = params.key
|
||||||
if self.disable or not key:
|
if self.disable or not key:
|
||||||
@@ -145,6 +194,59 @@ class LMCRadixCache(RadixCache):
|
|||||||
value: torch.Tensor = base_res.device_indices
|
value: torch.Tensor = base_res.device_indices
|
||||||
last_node: TreeNode = base_res.last_device_node
|
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):
|
if value.numel() == len(key):
|
||||||
return base_res
|
return base_res
|
||||||
|
|
||||||
@@ -152,31 +254,99 @@ class LMCRadixCache(RadixCache):
|
|||||||
if uncached_len == 0:
|
if uncached_len == 0:
|
||||||
return base_res
|
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()
|
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:
|
if self.token_to_kv_pool_allocator.available_size() < uncached_len:
|
||||||
self.evict(EvictParams(num_tokens=uncached_len))
|
self.evict(EvictParams(num_tokens=uncached_len))
|
||||||
|
|
||||||
token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len)
|
token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len)
|
||||||
if token_slots is None:
|
if token_slots is None:
|
||||||
return base_res
|
return None
|
||||||
|
|
||||||
slot_mapping = torch.cat(
|
slot_mapping = torch.empty(
|
||||||
[
|
value_numel + token_slots.numel(),
|
||||||
torch.full((value.numel(),), -1, dtype=torch.int64, device=self.device),
|
dtype=torch.int64,
|
||||||
token_slots.detach().clone().to(torch.int64).to(self.device),
|
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 = load_fn(slot_mapping, prefix_pad)
|
||||||
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
|
|
||||||
)
|
|
||||||
)
|
|
||||||
logger.debug("num_retrieved_tokens: %s", num_retrieved)
|
logger.debug("num_retrieved_tokens: %s", num_retrieved)
|
||||||
|
|
||||||
if num_retrieved > 0:
|
if num_retrieved > 0:
|
||||||
@@ -189,38 +359,80 @@ class LMCRadixCache(RadixCache):
|
|||||||
if num_retrieved > 0:
|
if num_retrieved > 0:
|
||||||
fetched = num_retrieved - prefix_pad
|
fetched = num_retrieved - prefix_pad
|
||||||
new_node = TreeNode(priority=last_node.priority)
|
new_node = TreeNode(priority=last_node.priority)
|
||||||
start = value.numel()
|
start = value_numel
|
||||||
end = start + fetched
|
end = start + fetched
|
||||||
new_node.key = key[start:end]
|
new_node.key = key[start:end]
|
||||||
new_node.value = token_slots[:fetched]
|
new_node.value = token_slots[:fetched]
|
||||||
new_node.parent = last_node
|
new_node.parent = last_node
|
||||||
last_node.children[new_node.key.child_key(self.page_size)] = new_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.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.parent)
|
||||||
self._record_store_event(new_node)
|
self._record_store_event(new_node)
|
||||||
|
|
||||||
return MatchResult(
|
return token_slots[:fetched], new_node
|
||||||
device_indices=value,
|
|
||||||
last_device_node=last_node,
|
return None
|
||||||
last_host_node=last_node,
|
|
||||||
best_match_node=last_node,
|
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:
|
||||||
|
|
||||||
def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: # type: ignore[override]
|
|
||||||
"""On request completion, insert device KV into radix and store to LMCache."""
|
"""On request completion, insert device KV into radix and store to LMCache."""
|
||||||
|
|
||||||
super().cache_finished_req(req, is_insert=is_insert)
|
super().cache_finished_req(req, is_insert=is_insert)
|
||||||
if not 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
|
return
|
||||||
|
|
||||||
from sglang.srt.server_args import get_global_server_args
|
|
||||||
|
|
||||||
global_server_args = get_global_server_args()
|
global_server_args = get_global_server_args()
|
||||||
topk = global_server_args.speculative_eagle_topk
|
topk = global_server_args.speculative_eagle_topk
|
||||||
enable_kv_committed_len = topk is None or topk == 1
|
enable_kv_committed_len = topk is None or topk == 1
|
||||||
@@ -236,7 +448,8 @@ class LMCRadixCache(RadixCache):
|
|||||||
req.req_pool_idx, :kv_committed_len
|
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))
|
MatchPrefixParams(key=RadixKey(token_ids, req.extra_key))
|
||||||
)
|
)
|
||||||
new_last_node = match_result.last_device_node
|
new_last_node = match_result.last_device_node
|
||||||
@@ -248,11 +461,19 @@ class LMCRadixCache(RadixCache):
|
|||||||
token_ids=token_ids,
|
token_ids=token_ids,
|
||||||
kv_indices=kv_indices,
|
kv_indices=kv_indices,
|
||||||
offset=0,
|
offset=0,
|
||||||
|
request_id=req.rid,
|
||||||
)
|
)
|
||||||
with torch.cuda.stream(self.store_stream):
|
with torch.cuda.stream(self.store_stream):
|
||||||
self.lmcache_connector.store_kv(store_md)
|
self.lmcache_connector.store_kv(store_md)
|
||||||
with self._node_lock:
|
if self._mode is LMCacheMode.MP:
|
||||||
self._in_flight_nodes.append(new_last_node)
|
# 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:
|
def evict(self, params: EvictParams) -> EvictResult:
|
||||||
"""Before base eviction, wait for any outstanding stores and release locks."""
|
"""Before base eviction, wait for any outstanding stores and release locks."""
|
||||||
@@ -267,7 +488,7 @@ class LMCRadixCache(RadixCache):
|
|||||||
|
|
||||||
return super().evict(params)
|
return super().evict(params)
|
||||||
|
|
||||||
def pretty_print(self): # type: ignore[override]
|
def pretty_print(self):
|
||||||
super().pretty_print()
|
super().pretty_print()
|
||||||
try:
|
try:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
|
|||||||
@@ -9,15 +9,10 @@ except ImportError:
|
|||||||
"LMCache is not installed. Please install it by running `pip install lmcache` in the root directory of LMCache"
|
"LMCache is not installed. Please install it by running `pip install lmcache` in the root directory of LMCache"
|
||||||
)
|
)
|
||||||
|
|
||||||
import os
|
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.configs.model_config import ModelConfig
|
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():
|
def test_load_store_metadata():
|
||||||
model_config = ModelConfig(
|
model_config = ModelConfig(
|
||||||
@@ -40,7 +35,9 @@ def test_load_store_metadata():
|
|||||||
for _ in range(layer_num)
|
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_token_ids = torch.randint(0, model_config.vocab_size, (input_id_len,)).tolist()
|
||||||
fake_kv_indices = torch.randint(0, buffer_size, (input_id_len,))
|
fake_kv_indices = torch.randint(0, buffer_size, (input_id_len,))
|
||||||
|
|||||||
@@ -673,6 +673,7 @@ class ServerArgs:
|
|||||||
|
|
||||||
# LMCache
|
# LMCache
|
||||||
enable_lmcache: bool = False
|
enable_lmcache: bool = False
|
||||||
|
lmcache_config_file: Optional[str] = None
|
||||||
|
|
||||||
# Ktransformers/AMX expert parallelism
|
# Ktransformers/AMX expert parallelism
|
||||||
kt_weight_path: Optional[str] = None
|
kt_weight_path: Optional[str] = None
|
||||||
@@ -6087,6 +6088,12 @@ class ServerArgs:
|
|||||||
action="store_true",
|
action="store_true",
|
||||||
help="Using LMCache as an alternative hierarchical cache solution",
|
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
|
# Ktransformer server args
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
|
|||||||
Reference in New Issue
Block a user