sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

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

@@ -0,0 +1,284 @@
+from __future__ import annotations
+
+import logging
+import threading
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import MatchResult
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
+
+try:
+    from lmcache.integration.sglang.sglang_adapter import (
+        LMCacheLayerwiseConnector,
+        LoadMetadata,
+        StoreMetadata,
+    )
+except ImportError as e:
+    raise RuntimeError(
+        "LMCache is not installed. Please install it by running `pip install lmcache`"
+    ) from e
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.managers.schedule_batch import Req
+
+logger = logging.getLogger(__name__)
+
+
+class LayerTransferCounter:
+    """Minimal adapter that lets the memory pool notify LMCache per-layer.
+
+    The KV pool calls `wait_until(layer_id)` after finishing a layer, which we
+    translate into a `load_kv_layerwise(layer_id)` call on the LMCache connector
+    within the provided CUDA stream.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        load_stream: torch.cuda.Stream,
+        lmc_connector: LMCacheLayerwiseConnector,
+        printable: bool = False,
+    ):
+        self.num_layers = num_layers
+        self.load_stream = load_stream
+        self.lmc_connector = lmc_connector
+
+    def wait_until(self, layer_id: int):
+        # Ensure ordering of the async loads wrt compute stream(s).
+        self.load_stream.synchronize()
+        with self.load_stream:
+            self.lmc_connector.load_kv_layerwise(layer_id)
+
+
+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
+    """
+
+    def __init__(
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
+        page_size: int,
+        disable: bool = False,
+        enable_kv_cache_events: bool = False,
+        model_config: Optional["ModelConfig"] = None,
+        tp_size: int = 1,
+        rank: int = 0,
+        tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        eviction_policy: str = "lru",
+    ):
+        super().__init__(
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+            page_size=page_size,
+            disable=disable,
+            enable_kv_cache_events=enable_kv_cache_events,
+            eviction_policy=eviction_policy,
+        )
+
+        kvcache = self.token_to_kv_pool_allocator.get_kvcache()
+        self.lmcache_connector = LMCacheLayerwiseConnector(
+            sgl_config=model_config,
+            tp_size=tp_size,
+            rank=rank,
+            # NOTE: The original implementation accessed private buffers via
+            # `_kvcache.k_buffer` / `.v_buffer`. We prefer public accessors when
+            # available; fall back to private fields if needed.
+            k_pool=getattr(
+                kvcache,
+                "k_buffer",
+                getattr(self.token_to_kv_pool_allocator._kvcache, "k_buffer"),
+            ),
+            v_pool=getattr(
+                kvcache,
+                "v_buffer",
+                getattr(self.token_to_kv_pool_allocator._kvcache, "v_buffer"),
+            ),
+            tp_group=tp_group,
+        )
+
+        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)
+
+        self._in_flight_nodes: list[TreeNode] = []
+        self._node_lock = threading.Lock()
+
+    def reset(self):  # type: ignore[override]
+        super().reset()
+        if hasattr(self, "_in_flight_nodes"):
+            with self._node_lock:
+                self._in_flight_nodes.clear()
+
+    def match_prefix(self, key: RadixKey, **kwargs) -> MatchResult:  # type: ignore[override]
+        """Match cached prefix; if there's a tail miss, prefetch from LMCache.
+
+        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.
+        """
+        if self.disable or not key:
+            return super().match_prefix(key, **kwargs)
+
+        if self.page_size != 1:
+            aligned_len = len(key) // self.page_size * self.page_size
+            key = key[:aligned_len]
+
+        base_res = super().match_prefix(key, **kwargs)
+        value: torch.Tensor = base_res.device_indices
+        last_node: TreeNode = base_res.last_device_node
+
+        if value.numel() == len(key):
+            return base_res
+
+        uncached_len = len(key) - value.numel()
+        if uncached_len == 0:
+            return base_res
+
+        chunk_size = self.lmcache_connector.chunk_size()
+        prefix_pad = value.numel() % chunk_size
+
+        if self.token_to_kv_pool_allocator.available_size() < uncached_len:
+            self.evict(uncached_len)
+
+        token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len)
+        if token_slots is None:
+            return base_res
+
+        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),
+            ]
+        )
+
+        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
+                )
+            )
+        logger.debug("num_retrieved_tokens: %s", num_retrieved)
+
+        if num_retrieved > 0:
+            self.token_to_kv_pool_allocator.free(
+                token_slots[(num_retrieved - prefix_pad) :]
+            )
+        else:
+            self.token_to_kv_pool_allocator.free(token_slots)
+
+        if num_retrieved > 0:
+            fetched = num_retrieved - prefix_pad
+            new_node = TreeNode()
+            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[self.get_child_key_fn(new_node.key)] = new_node
+            last_node = new_node
+
+            value = torch.cat([value, token_slots[:fetched]])
+            self.evictable_size_ += fetched
+
+            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,
+            )
+
+        return base_res
+
+    def cache_finished_req(self, req: "Req") -> None:  # type: ignore[override]
+        """On request completion, insert device KV into radix and store to LMCache."""
+
+        super().cache_finished_req(req)
+
+        token_ids = (req.origin_input_ids + req.output_ids)[:-1]
+        kv_indices = self.req_to_token_pool.req_to_token[
+            req.req_pool_idx, : len(token_ids)
+        ]
+
+        _, new_last_node, _, _ = self.match_prefix(RadixKey(token_ids, req.extra_key))
+        assert new_last_node is not None
+
+        self.inc_lock_ref(new_last_node)
+        store_md = StoreMetadata(
+            last_node=new_last_node,
+            token_ids=token_ids,
+            kv_indices=kv_indices,
+            offset=0,
+        )
+        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)
+
+    def evict(self, num_tokens: int) -> None:  # type: ignore[override]
+        """Before base eviction, wait for any outstanding stores and release locks."""
+        if self.disable:
+            return
+
+        self.store_stream.synchronize()
+        with self._node_lock:
+            for node in self._in_flight_nodes:
+                self.dec_lock_ref(node)
+            self._in_flight_nodes.clear()
+
+        super().evict(num_tokens)
+
+    def pretty_print(self):  # type: ignore[override]
+        super().pretty_print()
+        try:
+            logger.debug(
+                "evictable=%d protected=%d", self.evictable_size_, self.protected_size_
+            )
+        except Exception:  # pragma: no cover
+            pass
+
+
+if __name__ == "__main__":
+    cache = LMCRadixCache(
+        req_to_token_pool=None,
+        token_to_kv_pool_allocator=None,
+        page_size=1,
+        disable=False,
+        enable_kv_cache_events=False,
+        model_config=None,
+        tp_size=1,
+        rank=0,
+        tp_group=None,
+    )
+    cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 11, 12], dtype=torch.int64))
+    cache.insert(
+        RadixKey([1, 2, 3, 4]), torch.tensor([10, 11, 12, 13], dtype=torch.int64)
+    )
+    cache.pretty_print()

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

@@ -0,0 +1,121 @@
+try:
+    from lmcache.integration.sglang.sglang_adapter import (
+        LMCacheLayerwiseConnector,
+        LoadMetadata,
+        StoreMetadata,
+    )
+except ImportError:
+    raise RuntimeError(
+        "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(
+        model_path="Qwen/Qwen3-4B",
+    )
+
+    # Generate Dummy KV Cache
+    head_num = model_config.num_key_value_heads
+    head_dim = model_config.head_dim
+    layer_num = model_config.num_hidden_layers
+    buffer_size = 256
+    input_id_len = 16
+
+    k_buffer = [
+        torch.randn(buffer_size, head_num, head_dim, dtype=torch.bfloat16).cuda()
+        for _ in range(layer_num)
+    ]
+    v_buffer = [
+        torch.randn(buffer_size, head_num, head_dim, dtype=torch.bfloat16).cuda()
+        for _ in range(layer_num)
+    ]
+
+    connector = LMCacheLayerwiseConnector(model_config, 1, 0, k_buffer, v_buffer)
+
+    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,))
+    offset = 0
+
+    store_metadata = StoreMetadata(
+        last_node=None,
+        token_ids=fake_token_ids,
+        kv_indices=fake_kv_indices,
+        offset=offset,
+    )
+
+    load_metadata = LoadMetadata(
+        token_ids=fake_token_ids,
+        slot_mapping=fake_kv_indices,
+        offset=offset,
+    )
+
+    current_stream = torch.cuda.current_stream()
+
+    retrieve_token_num = connector.start_load_kv(load_metadata)
+    assert retrieve_token_num == 0
+
+    connector.store_kv(store_metadata)
+    current_stream.synchronize()
+
+    # check retrieve
+    gt_key_buffer = [
+        torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
+        for _ in range(layer_num)
+    ]
+    gt_value_buffer = [
+        torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
+        for _ in range(layer_num)
+    ]
+
+    for i in range(layer_num):
+        gt_key_buffer[i] = k_buffer[i][fake_kv_indices]
+        gt_value_buffer[i] = v_buffer[i][fake_kv_indices]
+
+    # clear the k_buffer and v_buffer
+    for _ in range(layer_num):
+        k_buffer[i].zero_()
+        v_buffer[i].zero_()
+
+    retrieve_token_num = connector.start_load_kv(load_metadata)
+    assert retrieve_token_num == input_id_len
+
+    for i in range(layer_num):
+        current_stream.synchronize()
+        connector.load_kv_layerwise(i)
+
+    current_stream.synchronize()
+    test_key_buffer = [
+        torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
+        for _ in range(layer_num)
+    ]
+    test_value_buffer = [
+        torch.zeros(input_id_len, head_num, head_dim, dtype=torch.bfloat16).cuda()
+        for _ in range(layer_num)
+    ]
+
+    for i in range(layer_num):
+        test_key_buffer[i] = k_buffer[i][fake_kv_indices]
+        test_value_buffer[i] = v_buffer[i][fake_kv_indices]
+
+    for i in range(layer_num):
+        assert torch.allclose(test_key_buffer[i], gt_key_buffer[i])
+        assert torch.allclose(test_value_buffer[i], gt_value_buffer[i])
+
+    print("================================================")
+    print("TEST_LOAD_STORE_METADATA PASSED!")
+    print("================================================")
+    connector.close()
+
+
+if __name__ == "__main__":
+    test_load_store_metadata()

sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -7,11 +7,16 @@ from typing import Any, List, Optional
 
 import torch
 
-from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
+from sglang.srt.mem_cache.hicache_storage import (
+    HiCacheStorage,
+    HiCacheStorageConfig,
+    HiCacheStorageExtraInfo,
+)
+from sglang.srt.mem_cache.memory_pool_host import HostKVCache
 
 DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
 DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024  # 16 MB
-
+DEFAULT_MOONCAKE_CONFIG_PATH_ENV = "SGLANG_HICACHE_MOONCAKE_CONFIG_PATH"
 logger = logging.getLogger(__name__)
 
 
@@ -28,13 +33,13 @@ class MooncakeStoreConfig:
     @staticmethod
     def from_file() -> "MooncakeStoreConfig":
         """Load the config from a JSON file."""
-        file_path = os.getenv("MOONCAKE_CONFIG_PATH")
-        if file_path is None:
-            raise ValueError(
-                "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
-            )
-        with open(file_path) as fin:
-            config = json.load(fin)
+        file_path = os.getenv(DEFAULT_MOONCAKE_CONFIG_PATH_ENV)
+        try:
+            with open(file_path) as fin:
+                config = json.load(fin)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load config from {file_path}: {str(e)}")
+
         return MooncakeStoreConfig(
             local_hostname=config.get("local_hostname"),
             metadata_server=config.get("metadata_server"),
@@ -72,6 +77,26 @@ class MooncakeStoreConfig:
             master_server_address=os.getenv("MOONCAKE_MASTER"),
         )
 
+    @staticmethod
+    def load_from_extra_config(extra_config: dict) -> "MooncakeStoreConfig":
+        """Load config from extra_config dictionary."""
+        if "master_server_address" not in extra_config:
+            raise ValueError("master_server_address is required in extra_config")
+
+        return MooncakeStoreConfig(
+            local_hostname=extra_config.get("local_hostname", "localhost"),
+            metadata_server=extra_config.get("metadata_server", "P2PHANDSHAKE"),
+            global_segment_size=extra_config.get(
+                "global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE
+            ),
+            local_buffer_size=extra_config.get(
+                "local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE
+            ),
+            protocol=extra_config.get("protocol", "tcp"),
+            device_name=extra_config.get("device_name", "auto"),
+            master_server_address=extra_config["master_server_address"],
+        )
+
     def __post_init__(self):
         if self.device_name == "auto":
             os.environ["MC_MS_AUTO_DISC"] = "1"
@@ -81,6 +106,7 @@ class MooncakeStoreConfig:
 
 
 class MooncakeStore(HiCacheStorage):
+
     def __init__(self, storage_config: HiCacheStorageConfig = None):
         try:
             from mooncake.store import MooncakeDistributedStore
@@ -93,14 +119,43 @@ class MooncakeStore(HiCacheStorage):
 
         try:
             self.store = MooncakeDistributedStore()
-            self.config = MooncakeStoreConfig.load_from_env()
-            logger.info("Mooncake Configuration loaded from env successfully.")
+
+            extra_config = (
+                getattr(storage_config, "extra_config", None)
+                if storage_config
+                else None
+            )
+            # Load configuration with master_server_address prioritized from extra_config if available
+            if (
+                extra_config is not None
+                and extra_config.get("master_server_address") is not None
+            ):
+                # Load from extra_config
+                self.config = MooncakeStoreConfig.load_from_extra_config(extra_config)
+                logger.info(
+                    "Mooncake Configuration loaded from extra_config successfully."
+                )
+            elif os.getenv(DEFAULT_MOONCAKE_CONFIG_PATH_ENV):
+                # Load from config file
+                self.config = MooncakeStoreConfig.from_file()
+                logger.info("Mooncake Configuration loaded from file successfully.")
+            else:
+                # Load from environment variables
+                self.config = MooncakeStoreConfig.load_from_env()
+                logger.info("Mooncake Configuration loaded from env successfully.")
+
+            tp_scale_factor = 1 if storage_config is None else storage_config.tp_size
+
+            per_tp_global_segment_size = (
+                self.config.global_segment_size // tp_scale_factor
+            )
+            per_tp_local_buffer_size = self.config.local_buffer_size // tp_scale_factor
 
             ret_code = self.store.setup(
                 self.config.local_hostname,
                 self.config.metadata_server,
-                self.config.global_segment_size,
-                self.config.local_buffer_size,
+                per_tp_global_segment_size,
+                per_tp_local_buffer_size,
                 self.config.protocol,
                 self.config.device_name,
                 self.config.master_server_address,
@@ -133,7 +188,13 @@ class MooncakeStore(HiCacheStorage):
         assert self.store.is_exist(warmup_key) == 1
         assert self.store.get(warmup_key) == warmup_value
 
-    def register_buffer(self, buffer: torch.Tensor) -> None:
+    def register_mem_pool_host(self, mem_pool_host: HostKVCache):
+        super().register_mem_pool_host(mem_pool_host)
+        assert self.mem_pool_host.layout in [
+            "page_first",
+            "page_first_direct",
+        ], "mooncake store storage backend only support page first or page first direct layout"
+        buffer = self.mem_pool_host.kv_buffer
         try:
             buffer_ptr = buffer.data_ptr()
             buffer_size = buffer.numel() * buffer.element_size()
@@ -144,6 +205,97 @@ class MooncakeStore(HiCacheStorage):
             logger.error("Failed to register buffer to Mooncake Store: %s", err)
             raise TypeError("Mooncake Store Register Buffer Error.") from err
 
+    def _get_mha_buffer_meta(self, keys, indices):
+        ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices)
+        key_list = []
+        for key_ in keys:
+            key_list.append(f"{key_}_{self.local_rank}_k")
+            key_list.append(f"{key_}_{self.local_rank}_v")
+        assert len(key_list) == len(ptr_list)
+        return key_list, ptr_list, element_size_list
+
+    def _get_mla_buffer_meta(self, keys, indices):
+        ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices)
+        key_list = []
+        for key_ in keys:
+            key_list.append(f"{key_}_k")
+        assert len(key_list) == len(ptr_list)
+        return key_list, ptr_list, element_size_list
+
+    def _batch_preprocess(self, keys, host_indices):
+        assert len(keys) > 0
+        assert len(keys) == len(host_indices) // self.mem_pool_host.page_size
+        if self.is_mla_backend:
+            return self._get_mla_buffer_meta(keys, host_indices)
+        else:
+            return self._get_mha_buffer_meta(keys, host_indices)
+
+    def _batch_postprocess(self, results: List[int], is_set_operate=False):
+        """
+        refer to https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/include/pybind_client.h
+        for batch_get_into, results is Vector of integers,
+            where each element is the number of bytes read on success, or a negative value on error
+        for batch_put_from, results is Vector of integers,
+            where each element is 0 on success, or a negative value on error
+        """
+        if self.is_mla_backend:
+            return [k_res == 0 if is_set_operate else k_res > 0 for k_res in results]
+        else:
+            kv_pairs = zip(results[::2], results[1::2])
+            return [
+                (
+                    (k_res == 0 and v_res == 0)
+                    if is_set_operate
+                    else (k_res > 0 and v_res > 0)
+                )
+                for k_res, v_res in kv_pairs
+            ]
+
+    def batch_get_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+        key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
+        get_results = self._get_batch_zero_copy_impl(
+            key_strs, buffer_ptrs, buffer_sizes
+        )
+        return self._batch_postprocess(get_results, is_set_operate=False)
+
+    def batch_set_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+        key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
+        exist_result = self._batch_exist(key_strs)
+
+        set_keys = []
+        set_buffer_ptrs = []
+        set_buffer_sizes = []
+        set_indices = []
+        set_results = [-1] * len(key_strs)
+        for i in range(len(key_strs)):
+            if exist_result[i] != 1:
+                set_keys.append(key_strs[i])
+                set_buffer_ptrs.append(buffer_ptrs[i])
+                set_buffer_sizes.append(buffer_sizes[i])
+                set_indices.append(i)
+            else:
+                set_results[i] = 0
+
+        # Only set non-existing keys to storage
+        if len(set_keys) > 0:
+            put_results = self._put_batch_zero_copy_impl(
+                set_keys, set_buffer_ptrs, set_buffer_sizes
+            )
+            for i in range(len(set_indices)):
+                set_results[set_indices[i]] = put_results[i]
+
+        return self._batch_postprocess(set_results, is_set_operate=True)
+
     def set(
         self,
         key,
@@ -264,9 +416,6 @@ class MooncakeStore(HiCacheStorage):
                 return i // key_multiplier
         return len(query_keys) // key_multiplier
 
-    def delete(self, key) -> None:
-        raise (NotImplementedError)
-
     def close(self):
         # MooncakeDistributedStore will automatically call the destructor, so
         # it is unnecessary to close it manually.