sglang 0.5.3rc0__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/mem_cache/storage/eic/test_unit.py

@@ -0,0 +1,115 @@
+import argparse
+import os
+
+import eic
+import torch
+import yaml
+
+
+def pase_args():
+    parser = argparse.ArgumentParser(description="EIC Storage Unit Test")
+    parser.add_argument(
+        "--config",
+        "-c",
+        type=str,
+        default="/sgl-workspace/config/remote-eic.yaml",
+        help="EIC yaml config",
+    )
+    args, _ = parser.parse_known_args()
+    return args
+
+
+def init_eic_client():
+    args = pase_args()
+    config_path = os.path.abspath(args.config)
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+    with open(config_path, "r") as fin:
+        config = yaml.safe_load(fin)
+
+    remote_url = config.get("remote_url", None)
+    if remote_url is None:
+        AssertionError("remote_url is None")
+    endpoint = remote_url[len("eic://") :]
+    eic_instance_id = config.get("eic_instance_id", None)
+    eic_log_dir = config.get("eic_log_dir", None)
+    eic_log_level = config.get("eic_log_level", 2)
+    eic_trans_type = config.get("eic_trans_type", 3)
+    eic_flag_file = config.get("eic_flag_file", None)
+
+    if not os.path.exists(eic_log_dir):
+        os.makedirs(eic_log_dir, exist_ok=True)
+    eic_client = eic.Client()
+    init_option = eic.InitOption()
+    init_option.log_dir = eic_log_dir
+    init_option.log_level = eic.LogLevel(eic_log_level)
+    init_option.transport_type = eic.TransportType(eic_trans_type)
+    init_option.flag_file = eic_flag_file
+    ret = eic_client.init(eic_instance_id, endpoint, init_option)
+    if ret != 0:
+        raise RuntimeError(f"EIC Client init failed with error code: {ret}")
+    return eic_client
+
+
+def test_set(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    tensors = [
+        torch.ones([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
+        for _ in range(16)
+    ]
+    data_keys = eic.StringVector()
+    data_vals = eic.IOBuffers()
+    for i in range(16):
+        data_keys.append(test_key[i])
+        data_vals.append(
+            tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
+        )
+    set_opt = eic.SetOption()
+    set_opt.ttl_second = 3
+    status_code, set_outcome = eic_client.mset(data_keys, data_vals, set_opt)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Set failed with status code: {status_code}"
+
+
+def test_get(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    tensors = [
+        torch.zeros([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
+        for _ in range(16)
+    ]
+    data_keys = eic.StringVector()
+    data_vals = eic.IOBuffers()
+    for i in range(16):
+        data_keys.append(test_key[i])
+        data_vals.append(
+            tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
+        )
+    get_opt = eic.GetOption()
+    status_code, data_vals, get_outcome = eic_client.mget(data_keys, get_opt, data_vals)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Get failed with status code: {status_code}"
+
+
+def test_exists(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    data_keys = eic.StringVector()
+    for key in test_key:
+        data_keys.append(key)
+    exists_opt = eic.ExistOption()
+    status_code, exists_outcome = eic_client.mexist(data_keys, exists_opt)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Exists failed with status code: {status_code}"
+
+
+def main():
+    eic_client = init_eic_client()
+    test_set(eic_client)
+    test_exists(eic_client)
+    test_get(eic_client)
+
+
+if __name__ == "__main__":
+    main()

sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py

@@ -12,7 +12,12 @@ from typing import Any, List, Optional, Tuple
 
 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
 from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsClient
 from sglang.srt.metrics.collector import StorageMetrics
 
@@ -178,11 +183,14 @@ class HiCacheHF3FS(HiCacheStorage):
             self.skip_backup = True
             self.rank = 0
 
+        self.is_zero_copy = False
+
         logger.info(
             f"[Rank {self.rank}] HiCacheHF3FS Client Initializing: "
             f"file_path={self.file_path}, "
             f"file_size={self.file_size / (2 ** 30):.2f} GB, "
-            f"num_pages={self.num_pages}"
+            f"num_pages={self.num_pages}, "
+            f"is_mla_model={self.is_mla_model}"
         )
 
         self.ac = AtomicCounter(self.numjobs)
@@ -323,25 +331,12 @@ class HiCacheHF3FS(HiCacheStorage):
             use_mock_client=use_mock_client,
         )
 
-    def get(
-        self,
-        key: str,
-        target_location: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> torch.Tensor | None:
-        return self.batch_get(
-            [key],
-            [target_location] if target_location is not None else None,
-            [target_sizes] if target_sizes is not None else None,
-        )[0]
-
     @synchronized()
-    def batch_get(
+    def _batch_get(
         self,
         keys: List[str],
-        target_locations: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> List[torch.Tensor | None]:
+        values: List[torch.Tensor],
+    ) -> List[bool]:
         page_indices = self.metadata_client.get_page_indices(self.rank, keys)
 
         batch_indices, file_offsets = [], []
@@ -350,15 +345,9 @@ class HiCacheHF3FS(HiCacheStorage):
                 batch_indices.append(i)
                 file_offsets.append(page_index * self.bytes_per_page)
 
-        if target_locations is not None:
-            for target_location in target_locations:
-                assert target_location.is_contiguous()
-            file_results = target_locations
-        else:
-            file_results = [
-                torch.empty(self.numel, dtype=self.dtype)
-                for _ in range(len(batch_indices))
-            ]
+        for target_location in values:
+            assert target_location.is_contiguous()
+        file_results = values
 
         start_time = time.perf_counter()
 
@@ -379,12 +368,10 @@ class HiCacheHF3FS(HiCacheStorage):
             ionum / (end_time - start_time) * self.gb_per_page
         )
 
-        results = [None] * len(keys)
-        for batch_index, file_result, read_result in zip(
-            batch_indices, file_results, read_results
-        ):
+        results = [False] * len(keys)
+        for batch_index, read_result in zip(batch_indices, read_results):
             if read_result == self.bytes_per_page:
-                results[batch_index] = file_result
+                results[batch_index] = True
             else:
                 logger.error(
                     f"[Rank {self.rank}] HiCacheHF3FS get {keys[batch_index]} failed"
@@ -392,28 +379,12 @@ class HiCacheHF3FS(HiCacheStorage):
 
         return results
 
-    def set(
-        self,
-        key: str,
-        value: Optional[Any] = None,
-        target_location: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> bool:
-        return self.batch_set(
-            [key],
-            [value] if value is not None else None,
-            [target_location] if target_location is not None else None,
-            [target_sizes] if target_sizes is not None else None,
-        )
-
     @synchronized()
-    def batch_set(
+    def _batch_set(
         self,
         keys: List[str],
         values: Optional[Any] = None,
-        target_locations: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> bool:
+    ) -> List[bool]:
         # In MLA backend, only one rank needs to backup the KV cache
         if self.skip_backup:
             return True
@@ -474,7 +445,7 @@ class HiCacheHF3FS(HiCacheStorage):
                 self.rank, written_keys_to_confirm, pages_to_release
             )
 
-        return all(results)
+        return results
 
     def delete(self, key: str) -> None:
         self.metadata_client.delete_keys(self.rank, [key])
@@ -484,21 +455,25 @@ class HiCacheHF3FS(HiCacheStorage):
         return result[0] if result else False
 
     def batch_exists(self, keys: List[str]) -> int:
+        factor = 1
+        if self.is_zero_copy and not self.is_mla_model:
+            keys = self._get_mha_zero_copy_keys(keys)
+            factor = 2
+
         results = self.metadata_client.exists(self.rank, keys)
-        for i in range(len(keys)):
-            if not results[i]:
-                return i
 
-        return len(keys)
+        i = 0
+        while i < len(keys) and results[i]:
+            i += 1
 
-    def clear(self) -> bool:
+        return i // factor
+
+    def clear(self) -> None:
         try:
             self.metadata_client.clear(self.rank)
             logger.info(f"Cleared HiCacheHF3FS for rank {self.rank}")
-            return True
         except Exception as e:
             logger.error(f"Failed to clear HiCacheHF3FS: {e}")
-            return False
 
     def close(self) -> None:
         try:
@@ -521,3 +496,143 @@ class HiCacheHF3FS(HiCacheStorage):
         self.prefetch_bandwidth.clear()
         self.backup_bandwidth.clear()
         return storage_metrics
+
+    def register_mem_pool_host(self, mem_pool_host: HostKVCache):
+        super().register_mem_pool_host(mem_pool_host)
+        self.is_zero_copy = self.mem_pool_host.layout == "page_first"
+        logger.info(f"{self.is_zero_copy=}")
+
+    def _get_mha_zero_copy_keys(self, keys: List[str]) -> List[str]:
+        _keys = []
+        for k in keys:
+            _keys.append(f"{k}-k")
+            _keys.append(f"{k}-v")
+        return _keys
+
+    def _get_mha_zero_copy_values(
+        self, values: List[torch.Tensor]
+    ) -> List[torch.Tensor]:
+        _values = []
+        for value in values:
+            _values.append(value[0])
+            _values.append(value[1])
+        return _values
+
+    def _batch_get_preprocess(self, keys, host_indices):
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+        # host_indices to kv_buffer
+        flat = not self.is_zero_copy
+        values = (
+            [
+                self.mem_pool_host.get_data_page(
+                    host_indices[i * self.mem_pool_host.page_size], flat=flat
+                )
+                for i in range(page_num)
+            ]
+            if self.is_zero_copy
+            else [
+                self.mem_pool_host.get_dummy_flat_data_page() for _ in range(page_num)
+            ]
+        )
+
+        if self.is_zero_copy and not self.is_mla_model:
+            keys = self._get_mha_zero_copy_keys(keys)
+            values = self._get_mha_zero_copy_values(values)
+
+        return keys, values
+
+    def _batch_get_postprocess(self, host_indices, values, results):
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+
+        if self.is_zero_copy:
+            if not self.is_mla_model:
+                results = [
+                    (results[2 * i] and results[2 * i + 1]) for i in range(page_num)
+                ]
+                results = results[:page_num]
+            return results
+
+        for i in range(page_num):
+            if not results[i]:
+                break
+            self.mem_pool_host.set_from_flat_data_page(
+                host_indices[i * self.mem_pool_host.page_size], values[i]
+            )
+
+        return results
+
+    def batch_get_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+        keys, values = self._batch_get_preprocess(keys, host_indices)
+        results = self._batch_get(keys, values)
+        return self._batch_get_postprocess(host_indices, values, results)
+
+    def _batch_set_preprocess(self, keys, host_indices):
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+        # host_indices to kv_buffer
+        flat = not self.is_zero_copy
+        values = [
+            self.mem_pool_host.get_data_page(
+                host_indices[i * self.mem_pool_host.page_size], flat=flat
+            )
+            for i in range(page_num)
+        ]
+
+        if self.is_zero_copy and not self.is_mla_model:
+            keys = self._get_mha_zero_copy_keys(keys)
+            values = self._get_mha_zero_copy_values(values)
+
+        return keys, values
+
+    def batch_set_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+        len_keys = len(keys)
+        keys, values = self._batch_set_preprocess(keys, host_indices)
+        results = self._batch_set(keys, values)
+        return results
+
+    # Deprecated
+    def get(
+        self,
+        key: str,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> torch.Tensor | None:
+        pass
+
+    # Deprecated
+    def batch_get(
+        self,
+        keys: List[str],
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> List[torch.Tensor | None] | int:
+        pass
+
+    # Deprecated
+    def set(
+        self,
+        key: str,
+        value: Optional[Any] = None,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> bool:
+        pass
+
+    # Deprecated
+    def batch_set(
+        self,
+        keys: List[str],
+        values: Optional[Any] = None,
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> bool:
+        pass

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

@@ -9,7 +9,7 @@ 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, TreeNode
+from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
 
 try:
     from lmcache.integration.sglang.sglang_adapter import (
@@ -78,6 +78,7 @@ class LMCRadixCache(RadixCache):
         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,
@@ -85,6 +86,7 @@ class LMCRadixCache(RadixCache):
             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()
@@ -129,7 +131,7 @@ class LMCRadixCache(RadixCache):
             with self._node_lock:
                 self._in_flight_nodes.clear()
 
-    def match_prefix(self, key: List[int], **kwargs) -> MatchResult:  # type: ignore[override]
+    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
@@ -176,7 +178,7 @@ class LMCRadixCache(RadixCache):
         with torch.cuda.stream(self.load_stream):
             num_retrieved = self.lmcache_connector.start_load_kv(
                 LoadMetadata(
-                    token_ids=key,  # full page-aligned key
+                    token_ids=key.token_ids,  # full page-aligned key
                     slot_mapping=slot_mapping,
                     offset=value.numel() - prefix_pad,  # LMCache offset convention
                 )
@@ -225,7 +227,7 @@ class LMCRadixCache(RadixCache):
             req.req_pool_idx, : len(token_ids)
         ]
 
-        _, new_last_node, _, _ = self.match_prefix(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)
@@ -275,6 +277,8 @@ if __name__ == "__main__":
         rank=0,
         tp_group=None,
     )
-    cache.insert([1, 2, 3], torch.tensor([10, 11, 12], dtype=torch.int64))
-    cache.insert([1, 2, 3, 4], torch.tensor([10, 11, 12, 13], dtype=torch.int64))
+    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/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"),
@@ -101,6 +106,7 @@ class MooncakeStoreConfig:
 
 
 class MooncakeStore(HiCacheStorage):
+
     def __init__(self, storage_config: HiCacheStorageConfig = None):
         try:
             from mooncake.store import MooncakeDistributedStore
@@ -129,6 +135,10 @@ class MooncakeStore(HiCacheStorage):
                 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()
@@ -178,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()
@@ -189,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,