sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

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

@@ -1,20 +1,47 @@
 import json
 import logging
 import os
+import time
 import uuid
 from dataclasses import dataclass
 from typing import Any, List, Optional
 
+import requests
 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"
+SETUP_TIMEOUT = 600  # 10min
+DEFAULT_MASTER_METRICS_PORT = 9003
+DEFAULT_CHECK_SERVER = False
 
 logger = logging.getLogger(__name__)
 
 
+def _parse_global_segment_size(value) -> int:
+    if isinstance(value, int):
+        return value
+    if isinstance(value, str):
+        s = value.strip().lower()
+        if s.endswith("gb"):
+            num = s[:-2].strip()
+            if not num:
+                raise ValueError(
+                    "Invalid global_segment_size: missing number before 'gb'"
+                )
+            return int(num) * 1024 * 1024 * 1024
+        return int(s)
+    return int(value)
+
+
 @dataclass
 class MooncakeStoreConfig:
     local_hostname: str
@@ -24,28 +51,34 @@ class MooncakeStoreConfig:
     protocol: str
     device_name: str
     master_server_address: str
+    master_metrics_port: int
+    check_server: bool
 
     @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"),
-            global_segment_size=config.get(
-                "global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE
+            global_segment_size=_parse_global_segment_size(
+                config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_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"),
+            device_name=config.get("device_name", ""),
             master_server_address=config.get("master_server_address"),
+            master_metrics_port=config.get(
+                "master_metrics_port", DEFAULT_MASTER_METRICS_PORT
+            ),
+            check_server=config.get("check_server", DEFAULT_CHECK_SERVER),
         )
 
     @staticmethod
@@ -53,7 +86,7 @@ class MooncakeStoreConfig:
         """Load config from a file specified in the environment variable.
         export MOONCAKE_MASTER=10.13.3.232:50051
         export MOONCAKE_PROTOCOL="rdma"
-        export MOONCAKE_DEVICE="auto"
+        export MOONCAKE_DEVICE=""
         export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE"
         """
         # other required environment variables...
@@ -62,14 +95,18 @@ class MooncakeStoreConfig:
         return MooncakeStoreConfig(
             local_hostname=os.getenv("LOCAL_HOSTNAME", "localhost"),
             metadata_server=os.getenv("MOONCAKE_TE_META_DATA_SERVER", "P2PHANDSHAKE"),
-            global_segment_size=int(
+            global_segment_size=_parse_global_segment_size(
                 os.getenv("MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_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"),
+            device_name=os.getenv("MOONCAKE_DEVICE", ""),
             master_server_address=os.getenv("MOONCAKE_MASTER"),
+            master_metrics_port=int(
+                os.getenv("MOONCAKE_MASTER_METRICS_PORT", DEFAULT_GLOBAL_SEGMENT_SIZE)
+            ),
+            check_server=bool(os.getenv("MOONCAKE_CHECK_SERVER", DEFAULT_CHECK_SERVER)),
         )
 
     @staticmethod
@@ -81,26 +118,24 @@ class MooncakeStoreConfig:
         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
+            global_segment_size=_parse_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"),
+            device_name=extra_config.get("device_name", ""),
             master_server_address=extra_config["master_server_address"],
+            master_metrics_port=extra_config.get(
+                "master_metrics_port", DEFAULT_MASTER_METRICS_PORT
+            ),
+            check_server=extra_config.get("check_server", DEFAULT_CHECK_SERVER),
         )
 
-    def __post_init__(self):
-        if self.device_name == "auto":
-            os.environ["MC_MS_AUTO_DISC"] = "1"
-            os.environ["MC_MS_FILTERS"] = (
-                "mlx5_bond_0, mlx5_bond_1, mlx5_bond_2, mlx5_bond_3"
-            )
-
 
 class MooncakeStore(HiCacheStorage):
+
     def __init__(self, storage_config: HiCacheStorageConfig = None):
         try:
             from mooncake.store import MooncakeDistributedStore
@@ -129,6 +164,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()
@@ -141,6 +180,16 @@ class MooncakeStore(HiCacheStorage):
             )
             per_tp_local_buffer_size = self.config.local_buffer_size // tp_scale_factor
 
+            # Check if extra_backend_tag should be passed to MooncakeDistributedStore
+            self.extra_backend_tag = None
+            if extra_config and "extra_backend_tag" in extra_config:
+                self.extra_backend_tag = extra_config["extra_backend_tag"]
+                logger.info(f"Using extra_backend_tag: {self.extra_backend_tag}")
+
+            # Check server status
+            if self.config.check_server:
+                self.check_server()
+
             ret_code = self.store.setup(
                 self.config.local_hostname,
                 self.config.metadata_server,
@@ -171,6 +220,39 @@ class MooncakeStore(HiCacheStorage):
             logger.error("An error occurred while loading the configuration: %s", exc)
             raise
 
+    def check_server(self):
+        master_server_ip = self.config.master_server_address.split(":")[0]
+        segments_url = f"http://{master_server_ip}:{self.config.master_metrics_port}/get_all_segments"
+        start_time = time.perf_counter()
+
+        check_result = False
+        while time.perf_counter() - start_time < SETUP_TIMEOUT:
+            try:
+                check_segments_resp = requests.get(segments_url, timeout=3)
+            except Exception:
+                logger.info(
+                    "waiting mooncake store server started, cost_time: %.2f seconds.",
+                    time.perf_counter() - start_time,
+                )
+                time.sleep(3)
+                continue
+
+            if check_segments_resp.text == "":
+                logger.info(
+                    "waiting mooncake store server started, cost_time: %.2f seconds.",
+                    time.perf_counter() - start_time,
+                )
+                time.sleep(3)
+                continue
+
+            logger.info("Mooncake store server started successfully.")
+            check_result = True
+            break
+
+        if not check_result:
+            logger.error("Launch mooncake store server timeout")
+            raise ValueError("Launch mooncake store server timeout")
+
     def warmup(self):
         warmup_key = "sglang_mooncake_store_warmup_key" + uuid.uuid4().hex
         warmup_value = bytes(4 * 1024)  # 4 KB
@@ -178,7 +260,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 +277,107 @@ 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]:
+        # Apply extra_backend_tag prefix if available
+        if self.extra_backend_tag is not None:
+            prefix = self.extra_backend_tag
+            keys = [f"{prefix}_{key}" for key in keys]
+
+        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]:
+        # Apply extra_backend_tag prefix if available
+        if self.extra_backend_tag is not None:
+            prefix = self.extra_backend_tag
+            keys = [f"{prefix}_{key}" for key in keys]
+
+        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,
@@ -292,7 +481,9 @@ class MooncakeStore(HiCacheStorage):
         exist_result = self._batch_exist([key])
         return exist_result[0] == 1
 
-    def batch_exists(self, keys) -> int:
+    def batch_exists(
+        self, keys, extra_info: Optional[HiCacheStorageExtraInfo] = None
+    ) -> int:
         if self.is_mla_backend:
             query_keys = [f"{key}_k" for key in keys]
             key_multiplier = 1

sglang/srt/mem_cache/storage/nixl/hicache_nixl.py

@@ -1,13 +1,12 @@
-import hashlib
 import logging
 import os
 import time
 import uuid
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, List, Optional, Union
 
 import torch
 
-from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
 
 from .nixl_utils import NixlBackendSelection, NixlFileManager, NixlRegistration
 
@@ -26,7 +25,12 @@ logger = logging.getLogger(__name__)
 class HiCacheNixl(HiCacheStorage):
     """HiCacheNixl provides high-performance storage using NIXL plugins."""
 
-    def __init__(self, file_path: str = "/tmp/hicache_storage", plugin: str = "auto"):
+    def __init__(
+        self,
+        storage_config: HiCacheStorageConfig,
+        file_path: str = "/tmp/hicache_storage",
+        plugin: str = "auto",
+    ):
         """Initialize NIXL storage connector."""
         # Might be better to be unified across HiCache backends and moved to HiCacheController
         file_path = os.getenv("SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR", file_path)
@@ -36,6 +40,19 @@ class HiCacheNixl(HiCacheStorage):
             else None
         )
 
+        # Initialize suffix based on storage config
+        tp_rank, tp_size, model_name, is_mla_model = (
+            storage_config.tp_rank,
+            storage_config.tp_size,
+            storage_config.model_name,
+            storage_config.is_mla_model,
+        )
+        model_name = "-".join(model_name.split("/")) if model_name else ""
+        if is_mla_model:
+            self.config_suffix = f"_{model_name}"
+        else:
+            self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}"
+
         agent_config = nixl_agent_config(backends=[])
         self.agent_name = f"hicache_nixl_{str(uuid.uuid4())}"
         self.agent = nixl_agent(self.agent_name, agent_config)
@@ -46,6 +63,9 @@ class HiCacheNixl(HiCacheStorage):
 
         self.registration = NixlRegistration(self.agent)
 
+    def _get_suffixed_key(self, key: str) -> str:
+        return key + self.config_suffix
+
     def register_buffers(
         self, buffers: Union[torch.Tensor, List[torch.Tensor], List[tuple]]
     ) -> Optional[Any]:
@@ -194,11 +214,14 @@ class HiCacheNixl(HiCacheStorage):
         else:
             dest = target_locations
 
+        # Add suffix to keys
+        suffixed_keys = [self._get_suffixed_key(key) for key in keys]
+
         if self.backend_selector.mem_type == "FILE":
-            file_paths = [self.file_manager.get_file_path(key) for key in keys]
+            file_paths = [self.file_manager.get_file_path(key) for key in suffixed_keys]
             success = self._execute_transfer(dest, file_paths, "READ")
         else:
-            success = self._execute_transfer(dest, keys, "READ")
+            success = self._execute_transfer(dest, suffixed_keys, "READ")
         return target_locations if success and not target_sizes else [None] * len(keys)
 
     def set(
@@ -227,9 +250,12 @@ class HiCacheNixl(HiCacheStorage):
         if not values:
             values = list(zip(target_locations, target_sizes))
 
+        # Add suffix to keys
+        suffixed_keys = [self._get_suffixed_key(key) for key in keys]
+
         if self.backend_selector.mem_type == "FILE":
             file_paths = []
-            for key in keys:
+            for key in suffixed_keys:
                 file_path = self.file_manager.get_file_path(key)
                 # New file per set, to be updated when partial writes is added to HiCache
                 if not self.file_manager.create_file(file_path):
@@ -238,11 +264,14 @@ class HiCacheNixl(HiCacheStorage):
                 file_paths.append(file_path)
             return self._execute_transfer(values, file_paths, "WRITE")
         else:  # mem_type == "OBJ"
-            return self._execute_transfer(values, keys, "WRITE")
+            return self._execute_transfer(values, suffixed_keys, "WRITE")
 
     def exists(self, key: str) -> bool:
+        # Add suffix to key
+        suffixed_key = self._get_suffixed_key(key)
+
         tuples = self.registration.create_query_tuples(
-            key,
+            suffixed_key,
             self.backend_selector.mem_type,
             self.file_manager if self.backend_selector.mem_type == "FILE" else None,
         )

sglang/srt/mem_cache/storage/nixl/nixl_utils.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 

sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py

@@ -2,11 +2,12 @@
 
 import os
 import unittest
-from typing import List, Optional
+from typing import List
 from unittest.mock import MagicMock
 
 import torch
 
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
 from sglang.srt.mem_cache.storage.nixl.hicache_nixl import HiCacheNixl
 from sglang.srt.mem_cache.storage.nixl.nixl_utils import (
     NixlFileManager,
@@ -31,8 +32,22 @@ class TestNixlUnified(unittest.TestCase):
         # Create instances
         self.file_manager = NixlFileManager(self.test_dir)
         self.registration = NixlRegistration(self.mock_agent)
+
+        # Create storage config for testing
+        self.storage_config = HiCacheStorageConfig(
+            tp_rank=0,
+            tp_size=2,
+            is_mla_model=False,
+            is_page_first_layout=False,
+            model_name="test_model",
+        )
+
         try:
-            self.hicache = HiCacheNixl(file_path=self.test_dir, plugin="POSIX")
+            self.hicache = HiCacheNixl(
+                storage_config=self.storage_config,
+                file_path=self.test_dir,
+                plugin="POSIX",
+            )
         except ImportError:
             self.skipTest("NIXL not available, skipping NIXL storage tests")