sglang 0.5.1.post2__py3-none-any.whl → 0.5.2__py3-none-any.whl

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

@@ -1,4 +1,3 @@
-import hashlib
 import json
 import logging
 import os
@@ -6,28 +5,16 @@ import uuid
 from dataclasses import dataclass
 from typing import Any, List, Optional
 
-import numpy as np
 import torch
 
-from sglang.srt.distributed import get_tensor_model_parallel_rank
-from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
 
 DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
-DEFAULT_LOCAL_BUFFER_SIZE = 128 * 1024 * 1024  # 128 MB
+DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024  # 16 MB
 
 logger = logging.getLogger(__name__)
 
 
-def get_hash_str_mooncake(token_ids: List[int], prior_hash: str = None):
-    prefix_str = ""
-    if prior_hash:
-        prefix_str = hashlib.sha256(prior_hash.encode()).hexdigest()
-    current_token_ids_bytes = np.array(token_ids).tobytes()
-    current_hash_object = hashlib.sha256(current_token_ids_bytes)
-    current_hash_hex = current_hash_object.hexdigest()
-    return f"{prefix_str}_{int(current_hash_hex[:16], 16)}"
-
-
 @dataclass
 class MooncakeStoreConfig:
     local_hostname: str
@@ -54,9 +41,8 @@ class MooncakeStoreConfig:
             global_segment_size=config.get(
                 "global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE
             ),
-            local_buffer_size=config.get(
-                "local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE
-            ),
+            # Zero copy interface does not need local buffer
+            local_buffer_size=DEFAULT_LOCAL_BUFFER_SIZE,
             protocol=config.get("protocol", "tcp"),
             device_name=config.get("device_name", "auto"),
             master_server_address=config.get("master_server_address"),
@@ -79,14 +65,33 @@ class MooncakeStoreConfig:
             global_segment_size=int(
                 os.getenv("MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_SIZE)
             ),
-            local_buffer_size=int(
-                os.getenv("MOONCAKE_LOCAL_BUFFER_SIZE", DEFAULT_LOCAL_BUFFER_SIZE)
-            ),
+            # Zero copy interface does not need local buffer
+            local_buffer_size=DEFAULT_LOCAL_BUFFER_SIZE,
             protocol=os.getenv("MOONCAKE_PROTOCOL", "tcp"),
             device_name=os.getenv("MOONCAKE_DEVICE", "auto"),
             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"
@@ -96,7 +101,7 @@ class MooncakeStoreConfig:
 
 
 class MooncakeStore(HiCacheStorage):
-    def __init__(self, is_mla: bool = False):
+    def __init__(self, storage_config: HiCacheStorageConfig = None):
         try:
             from mooncake.store import MooncakeDistributedStore
         except ImportError as e:
@@ -108,8 +113,26 @@ 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."
+                )
+            else:
+                # Load from environment variables
+                self.config = MooncakeStoreConfig.load_from_env()
+                logger.info("Mooncake Configuration loaded from env successfully.")
 
             ret_code = self.store.setup(
                 self.config.local_hostname,
@@ -126,7 +149,13 @@ class MooncakeStore(HiCacheStorage):
             logger.info("Connect to Mooncake store successfully.")
             self.warmup()
             logger.info("Mooncake store warmup successfully.")
-            self.is_mla = is_mla
+
+            if storage_config is not None:
+                self.is_mla_backend = storage_config.is_mla_model
+                self.local_rank = storage_config.tp_rank
+            else:
+                self.is_mla_backend = False
+                self.local_rank = 0
 
         except ValueError as e:
             logger.error("Configuration loading failed: %s", e)
@@ -137,12 +166,10 @@ class MooncakeStore(HiCacheStorage):
 
     def warmup(self):
         warmup_key = "sglang_mooncake_store_warmup_key" + uuid.uuid4().hex
-        # 10 MB
-        warmup_value = bytes(10 * 1024 * 1024)
-        self.store.put(warmup_key, warmup_value)
+        warmup_value = bytes(4 * 1024)  # 4 KB
+        assert self.store.put(warmup_key, warmup_value) == 0
         assert self.store.is_exist(warmup_key) == 1
-        self.store.get(warmup_key)
-        self.store.remove(warmup_key)
+        assert self.store.get(warmup_key) == warmup_value
 
     def register_buffer(self, buffer: torch.Tensor) -> None:
         try:
@@ -162,81 +189,118 @@ class MooncakeStore(HiCacheStorage):
         target_location: Optional[List[int]] = None,
         target_sizes: Optional[List[int]] = None,
     ) -> bool:
-        assert len(key) == len(target_location) == len(target_sizes)
-        if len(key) == 0:
-            return
-
-        for i in range(len(key)):
-            if key[i] is None or target_location[i] is None or target_sizes[i] is None:
-                return
-
-        self._put_batch_zero_copy_impl(key, target_location, target_sizes)
+        # Only support zero copy set for now
+        assert target_location is not None and target_sizes is not None
+        exist_result = self._batch_exist([key])
+        if exist_result[0] == 1:
+            return True
+        put_result = self._put_batch_zero_copy_impl(
+            [key], [target_location], [target_sizes]
+        )
+        return put_result[0] == 0
 
     def batch_set(
         self,
         keys: List[str],
-        value: Optional[Any] = None,
-        target_location: Optional[List[int]] = None,
+        values: Optional[List[torch.Tensor]] = None,
+        target_locations: Optional[List[int]] = None,
         target_sizes: Optional[List[int]] = None,
     ) -> bool:
-        assert len(keys) == len(target_location) == len(target_sizes)
+        # Only support zero copy set for now
+        assert target_locations is not None and target_sizes is not None
+        assert len(keys) == len(target_locations) == len(target_sizes)
+
         if len(keys) == 0:
-            return
+            return False
 
         for i in range(len(keys)):
-            if keys[i] is None or target_location[i] is None or target_sizes[i] is None:
-                return
+            if (
+                keys[i] is None
+                or target_locations[i] is None
+                or target_sizes[i] is None
+            ):
+                return False
+
+        exist_result = self._batch_exist(keys)
+        set_keys = []
+        set_target_locations = []
+        set_target_sizes = []
+        set_indices = []
+        for i in range(len(keys)):
+            if exist_result[i] != 1:
+                set_keys.append(keys[i])
+                set_target_locations.append(target_locations[i])
+                set_target_sizes.append(target_sizes[i])
+                set_indices.append(i)
+        # Only set non-existing keys to storage
+        put_result = self._put_batch_zero_copy_impl(
+            set_keys, set_target_locations, set_target_sizes
+        )
+        for i in range(len(set_indices)):
+            if put_result[i] == 0:
+                exist_result[set_indices[i]] = 1
 
-        self._put_batch_zero_copy_impl(keys, target_location, target_sizes)
+        success_count = 0
+        for i in range(len(keys)):
+            if exist_result[i] == 0:
+                break
+            success_count += 1
+        # TODO: return the number of consecutive successful operations from the start.
+        return success_count == len(keys)
 
     def get(
         self,
         key,
         target_location: Optional[Any] = None,
         target_sizes: Optional[Any] = None,
-    ) -> torch.Tensor | None:
-        assert len(key) == len(target_location) == len(target_sizes)
-        if len(key) == 0:
-            return
-
-        for i in range(len(key)):
-            if key[i] is None or target_location[i] is None or target_sizes[i] is None:
-                return
-
-        return self._get_batch_zero_copy_impl(key, target_location, target_sizes)
+    ) -> bool:
+        assert target_location is not None and target_sizes is not None
+        get_result = self._get_batch_zero_copy_impl(
+            [key], [target_location], [target_sizes]
+        )
+        return get_result[0] >= 0
 
     def batch_get(
         self,
         keys: List[str],
-        target_location: Optional[Any] = None,
+        target_locations: Optional[Any] = None,
         target_sizes: Optional[Any] = None,
-    ) -> torch.Tensor | None:
-        assert len(keys) == len(target_location) == len(target_sizes)
+    ) -> int:
+        assert len(keys) == len(target_locations) == len(target_sizes)
         if len(keys) == 0:
-            return
-
+            return 0
+        get_result = self._get_batch_zero_copy_impl(
+            keys, target_locations, target_sizes
+        )
+        if self.is_mla_backend:
+            key_multiplier = 1
+        else:
+            key_multiplier = 2
         for i in range(len(keys)):
-            if keys[i] is None or target_location[i] is None or target_sizes[i] is None:
-                return
-
-        return self._get_batch_zero_copy_impl(keys, target_location, target_sizes)
-
-    def exists(self, keys) -> bool | dict:
-        _keys = []
-        local_rank = get_tensor_model_parallel_rank()
-        for key in keys:
-            if key is None:
-                return None
-
-            if self.is_mla:
-                _keys.append(f"{key}_k")
-            else:
-                _keys.append(f"{key}_{local_rank}_k")
-        result = {k: v for k, v in zip(keys, self.store.batch_is_exist(_keys))}
-        return result
-
-    def delete(self, key) -> None:
-        raise (NotImplementedError)
+            if get_result[i] < 0:
+                return i // key_multiplier
+        return len(keys) // key_multiplier
+
+    def exists(self, key) -> bool:
+        exist_result = self._batch_exist([key])
+        return exist_result[0] == 1
+
+    def batch_exists(self, keys) -> int:
+        if self.is_mla_backend:
+            query_keys = [f"{key}_k" for key in keys]
+            key_multiplier = 1
+        else:
+            query_keys = []
+            for key in keys:
+                query_keys.append(f"{key}_{self.local_rank}_k")
+                query_keys.append(f"{key}_{self.local_rank}_v")
+            key_multiplier = 2
+
+        exist_result = self._batch_exist(query_keys)
+        for i in range(len(query_keys)):
+            if exist_result[i] != 1:
+                return i // key_multiplier
+        return len(query_keys) // key_multiplier
 
     def close(self):
         # MooncakeDistributedStore will automatically call the destructor, so
@@ -244,22 +308,17 @@ class MooncakeStore(HiCacheStorage):
         pass
 
     def clear(self) -> None:
-        raise (NotImplementedError)
+        self.store.remove_all()
 
     def _put_batch_zero_copy_impl(
         self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int]
-    ) -> None:
-        try:
-            self.store.batch_put_from(key_strs, buffer_ptrs, buffer_sizes)
-        except TypeError as err:
-            logger.error("Failed to put value to Mooncake Store: %s", err)
-            raise TypeError("Mooncake Store Put Type Error.") from err
+    ) -> List[int]:
+        return self.store.batch_put_from(key_strs, buffer_ptrs, buffer_sizes)
 
     def _get_batch_zero_copy_impl(
         self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int]
-    ) -> None:
-        try:
-            self.store.batch_get_into(key_strs, buffer_ptrs, buffer_sizes)
-        except TypeError as err:
-            logger.error("Failed to get value from Mooncake Store: %s", err)
-            raise TypeError("Mooncake Store Get Type Error.") from err
+    ) -> List[int]:
+        return self.store.batch_get_into(key_strs, buffer_ptrs, buffer_sizes)
+
+    def _batch_exist(self, key_strs: List[str]) -> List[int]:
+        return self.store.batch_is_exist(key_strs)

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

@@ -0,0 +1,161 @@
+import logging
+import uuid
+
+import torch
+from mooncake_store import MooncakeStore
+
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+def generate_batch_query_keys(kv_num: int, config: HiCacheStorageConfig):
+    keys = []
+    for _ in range(kv_num):
+        key = "test_" + str(uuid.uuid4())
+        keys.append(key)
+    set_keys = []
+    for key in keys:
+        if config.is_mla_model:
+            set_keys.append(key + "_k")
+        else:
+            set_keys.append(key + f"_{config.tp_rank}_k")
+            set_keys.append(key + f"_{config.tp_rank}_v")
+    get_keys = set_keys
+    exist_keys = keys
+    return set_keys, get_keys, exist_keys
+
+
+def test_single_operation():
+    """Test the set API with a single key-value pair."""
+    print("=" * 100)
+    print("Testing single operation")
+
+    buffer_size = 1024 * 1024 * 16  # 16MB
+    value_elements = 1024
+    store = MooncakeStore()
+    buffer = torch.randn(buffer_size, dtype=torch.float32)
+    store.register_buffer(buffer)
+    value_size = value_elements * buffer.element_size()
+
+    key = str(uuid.uuid4())
+    set_slice = buffer[:value_elements]
+    get_slice = buffer[value_elements : 2 * value_elements]
+    set_location = set_slice.data_ptr()
+    get_location = get_slice.data_ptr()
+
+    # Test set operation
+    result = store.set(key, target_location=set_location, target_sizes=value_size)
+    assert result is True, f"❌set operation failed for key: {key}"
+
+    # Test exists operation
+    assert store.exists(key), f"❌key {key} should exist after set operation"
+
+    # Test get operation
+    result = store.get(key, target_location=get_location, target_sizes=value_size)
+    assert result is True, f"❌get operation failed for key: {key}"
+
+    # Compare the data using proper tensor indices
+    assert torch.allclose(
+        set_slice, get_slice, atol=1e-6
+    ), f"❌get operation failed for key: {key}"
+
+    logger.info(f"✅ Single operation passed")
+
+
+def test_batch_operation(config: HiCacheStorageConfig):
+    """Test the batch set/get APIs with multiple key-value pairs."""
+    print("=" * 100)
+    print(f"Testing batch operation with config: {config}")
+
+    buffer_size = 1024 * 1024 * 16  # 16MB
+    value_elements = 256
+    kv_num = 13
+    store = MooncakeStore(config)
+    buffer = torch.randn(buffer_size, dtype=torch.float32)
+    store.register_buffer(buffer)
+    value_size = value_elements * buffer.element_size()
+
+    set_keys, get_keys, exist_keys = generate_batch_query_keys(kv_num, config)
+    set_slices = [
+        buffer[i * value_elements : (i + 1) * value_elements]
+        for i in range(len(set_keys))
+    ]
+    set_locations = [set_slice.data_ptr() for set_slice in set_slices]
+    target_sizes = [value_size for _ in range(len(set_keys))]
+
+    # Test batch set operation
+    result = store.batch_set(
+        set_keys, target_locations=set_locations, target_sizes=target_sizes
+    )
+    assert result is True, f"❌batch set operation failed"
+
+    # Test batch exists operation
+    assert store.batch_exists(
+        exist_keys
+    ), f"❌keys should exist after batch set operation"
+
+    # Test batch get operation
+    get_slices = [
+        buffer[
+            (len(set_keys) + i)
+            * value_elements : (len(set_keys) + i + 1)
+            * value_elements
+        ]
+        for i in range(len(get_keys))
+    ]
+    get_locations = [get_slice.data_ptr() for get_slice in get_slices]
+    result = store.batch_get(
+        get_keys, target_locations=get_locations, target_sizes=target_sizes
+    )
+    assert result == kv_num, f"❌batch get operation failed"
+    for i in range(len(get_keys)):
+        assert torch.allclose(
+            set_slices[i], get_slices[i], atol=1e-6
+        ), f"❌batch get operation failed for key: {get_keys[i]}"
+
+    logger.info(f"✅ Batch operation passed")
+
+
+if __name__ == "__main__":
+    test_single_operation()
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=False,
+            tp_rank=0,
+            tp_size=1,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=True,
+            tp_rank=0,
+            tp_size=1,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=False,
+            tp_rank=1,
+            tp_size=4,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=True,
+            tp_rank=3,
+            tp_size=8,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    logger.info(f"✅ All tests passed")

sglang/srt/mem_cache/swa_radix_cache.py

@@ -60,8 +60,6 @@ class TreeNode:
         self.last_access_time = time.monotonic()
 
         self.hit_count = 0
-        # indicating the node is loading KV cache from host
-        self.loading = False
         # store the host indices of KV cache
         self.host_value = None
 
@@ -464,7 +462,7 @@ class SWARadixCache(BasePrefixCache):
         self.req_to_token_pool.free(req.req_pool_idx)
         self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock)
 
-    def cache_unfinished_req(self, req: Req) -> None:
+    def cache_unfinished_req(self, req: Req, chunked=False) -> None:
         """Cache request when it is unfinished."""
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[