sglang 0.4.9.post6__py3-none-any.whl → 0.4.10.post1__py3-none-any.whl

sglang/srt/mem_cache/mooncake_store/mooncake_store.py

@@ -0,0 +1,264 @@
+import hashlib
+import json
+import logging
+import os
+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
+
+DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
+DEFAULT_LOCAL_BUFFER_SIZE = 128 * 1024 * 1024  # 128 MB
+
+logger = logging.getLogger(__name__)
+
+
+def get_hash_str_mooncake(current_page_ids: List, prefix_block_key: str):
+    local_rank = get_tensor_model_parallel_rank()
+    prefix_str = ""
+    if prefix_block_key:
+        if len(prefix_block_key):
+            prefix_str = hashlib.sha256(prefix_block_key.encode()).hexdigest()
+    current_token_ids_bytes = np.array(current_page_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)}_{local_rank}"
+
+
+@dataclass
+class MooncakeStoreConfig:
+    local_hostname: str
+    metadata_server: str
+    global_segment_size: int
+    local_buffer_size: int
+    protocol: str
+    device_name: str
+    master_server_address: str
+
+    @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)
+        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
+            ),
+            local_buffer_size=config.get(
+                "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"),
+        )
+
+    @staticmethod
+    def load_from_env() -> "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_TE_META_DATA_SERVER="P2PHANDSHAKE"
+        """
+        # other required environment variables...
+        if not os.getenv("MOONCAKE_MASTER"):
+            raise ValueError("The environment variable 'MOONCAKE_MASTER' is not set.")
+        return MooncakeStoreConfig(
+            local_hostname=os.getenv("LOCAL_HOSTNAME", "localhost"),
+            metadata_server=os.getenv("MOONCAKE_TE_META_DATA_SERVER", "P2PHANDSHAKE"),
+            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)
+            ),
+            protocol=os.getenv("MOONCAKE_PROTOCOL", "tcp"),
+            device_name=os.getenv("MOONCAKE_DEVICE", "auto"),
+            master_server_address=os.getenv("MOONCAKE_MASTER"),
+        )
+
+    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):
+        try:
+            from mooncake.store import MooncakeDistributedStore
+        except ImportError as e:
+            raise ImportError(
+                "Please install mooncake by following the instructions at "
+                "https://kvcache-ai.github.io/Mooncake/getting_started/build.html"
+                "to run SGLang with MooncakeConnector."
+            ) from e
+
+        try:
+            self.store = MooncakeDistributedStore()
+            self.config = MooncakeStoreConfig.load_from_env()
+            logger.info("Mooncake Configuration loaded from env successfully.")
+
+            ret_code = self.store.setup(
+                self.config.local_hostname,
+                self.config.metadata_server,
+                self.config.global_segment_size,
+                self.config.local_buffer_size,
+                self.config.protocol,
+                self.config.device_name,
+                self.config.master_server_address,
+            )
+            if ret_code:
+                logger.error(f"failed to setup mooncake store, error code: {ret_code}")
+
+            logger.info("Connect to Mooncake store successfully.")
+            self.warmup()
+            logger.info("Mooncake store warmup successfully.")
+
+        except ValueError as e:
+            logger.error("Configuration loading failed: %s", e)
+            raise
+        except Exception as exc:
+            logger.error("An error occurred while loading the configuration: %s", exc)
+            raise
+
+    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)
+        assert self.store.is_exist(warmup_key) == 1
+        self.store.get(warmup_key)
+        self.store.remove(warmup_key)
+
+    def register_buffer(self, buffer: torch.Tensor) -> None:
+        try:
+            buffer_ptr = buffer.data_ptr()
+            buffer_size = buffer.numel() * buffer.element_size()
+            ret_code = self.store.register_buffer(buffer_ptr, buffer_size)
+            if ret_code:
+                logger.error(f"failed to register buffer, error code: {ret_code}")
+        except TypeError as err:
+            logger.error("Failed to register buffer to Mooncake Store: %s", err)
+            raise TypeError("Mooncake Store Register Buffer Error.") from err
+
+    def set(
+        self,
+        key,
+        value: Optional[Any] = None,
+        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)
+
+    def batch_set(
+        self,
+        keys: List[str],
+        value: Optional[Any] = None,
+        target_location: Optional[List[int]] = None,
+        target_sizes: Optional[List[int]] = None,
+    ) -> bool:
+        assert len(keys) == len(target_location) == len(target_sizes)
+        if len(keys) == 0:
+            return
+
+        for i in range(len(keys)):
+            if keys[i] is None or target_location[i] is None or target_sizes[i] is None:
+                return
+
+        self._put_batch_zero_copy_impl(keys, target_location, target_sizes)
+
+    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)
+
+    def batch_get(
+        self,
+        keys: List[str],
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> torch.Tensor | None:
+        assert len(keys) == len(target_location) == len(target_sizes)
+        if len(keys) == 0:
+            return
+
+        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 = torch.cuda.current_device()
+        for key in keys:
+            if key is None:
+                return None
+            # Since mooncake store is stored in layer by layer,
+            # only the first layer is checked here.
+            _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)
+
+    def close(self):
+        # MooncakeDistributedStore will automatically call the destructor, so
+        # it is unnecessary to close it manually.
+        pass
+
+    def clear(self) -> None:
+        raise (NotImplementedError)
+
+    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
+
+    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

sglang/srt/mem_cache/mooncake_store/unit_test.py

@@ -0,0 +1,40 @@
+import torch
+from mooncake_store import MooncakeStore
+
+
+def test_init_and_warmup():
+    store = MooncakeStore()
+    assert store.store is not None
+
+
+def test_register_buffer():
+    store = MooncakeStore()
+    tensor = torch.zeros(1024, dtype=torch.float32)
+    store.register_buffer(tensor)
+
+
+def test_set_and_get():
+    store = MooncakeStore()
+
+    key = ["test_key_" + str(i) for i in range(2)]
+    tensor = torch.arange(256, dtype=torch.float32).cuda()
+    ptrs = [tensor.data_ptr(), tensor.data_ptr()]
+    sizes = [tensor.numel() * tensor.element_size()] * 2
+
+    store.set(key, target_location=ptrs, target_sizes=sizes)
+    store.get(key, target_location=ptrs, target_sizes=sizes)
+
+
+def test_exists():
+    store = MooncakeStore()
+    keys = ["test_key_0", "non_existent_key"]
+    result = store.exists(keys)
+    assert isinstance(result, dict)
+    assert "test_key_0" in result
+
+
+if __name__ == "__main__":
+    test_init_and_warmup()
+    test_register_buffer()
+    test_set_and_get()
+    test_exists()

sglang/srt/mem_cache/nixl/hicache_nixl.py

@@ -0,0 +1,163 @@
+import hashlib
+import logging
+import os
+import time
+import uuid
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
+
+from .nixl_utils import NixlBackendSelection, NixlFileManager, NixlRegistration
+
+try:
+    from nixl._api import nixl_agent, nixl_agent_config
+except ImportError as e:
+    raise ImportError(
+        "Please install NIXL by following the instructions at "
+        "https://github.com/ai-dynamo/nixl/blob/main/README.md "
+        "to use HiCacheNixl storage backend."
+    ) from e
+
+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"):
+        """Initialize NIXL storage connector."""
+        self.file_manager = (
+            NixlFileManager(file_path)
+            if plugin not in NixlBackendSelection.OBJ_PLUGINS
+            else None
+        )
+
+        agent_config = nixl_agent_config(backends=[])
+        self.agent_name = f"hicache_nixl_{str(uuid.uuid4())}"
+        self.agent = nixl_agent(self.agent_name, agent_config)
+
+        self.backend_selector = NixlBackendSelection(plugin)
+        if not self.backend_selector.create_backend(self.agent):
+            raise RuntimeError("Failed to create NIXL backend")
+
+        self.registration = NixlRegistration(self.agent)
+
+    def _execute_transfer(
+        self, tensors: List[torch.Tensor], keys: List[str], direction: str
+    ) -> bool:
+        if len(tensors) != len(keys):
+            logger.error("Mismatch between number of tensors and files/objects")
+            return False
+
+        if not self.registration.register_buffers(tensors):
+            logger.error("Failed to register tensors")
+            return False
+
+        # Get transfer tuples based on backend type
+        tensor_sizes = [tensor.element_size() * tensor.numel() for tensor in tensors]
+        if self.backend_selector.mem_type == "FILE":
+            file_tuples = self.file_manager.files_to_nixl_tuples(keys)
+            if not file_tuples or not self.registration.register_files(file_tuples):
+                logger.error("Failed to prepare files for transfer")
+                return False
+            transfer_tuples = [
+                (x[0], s, x[2]) for x, s in zip(file_tuples, tensor_sizes)
+            ]
+        else:
+            if not self.registration.register_objects(keys, tensors):
+                logger.error("Failed to register objects")
+                return False
+            transfer_tuples = [(0, s, key) for s, key in zip(tensor_sizes, keys)]
+
+        try:
+            # Get transfer descriptors
+            if (tensor_descs := self.agent.get_xfer_descs(tensors)) is None or (
+                file_descs := self.agent.get_xfer_descs(
+                    transfer_tuples, self.backend_selector.mem_type
+                )
+            ) is None:
+                logger.error("Failed to get transfer descriptors")
+                return False
+
+            # Initialize and execute transfer
+            if (
+                xfer_req := self.agent.initialize_xfer(
+                    direction, tensor_descs, file_descs, self.agent_name
+                )
+            ) is None:
+                logger.error("Failed to create transfer request")
+                return False
+
+            state = self.agent.transfer(xfer_req)
+            while state != "DONE":
+                state = self.agent.check_xfer_state(xfer_req)
+                if state == "ERR":
+                    logger.error("Transfer failed")
+                    return False
+            time.sleep(0.0001)  # Can be changed to os.sched_yield() or parametrized
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to execute transfer: {e}")
+            import traceback
+
+            logger.error(f"Traceback: {traceback.format_exc()}")
+            return False
+
+    def batch_set(self, keys: List[str], values: List[torch.Tensor]) -> bool:
+        if not keys:
+            return True
+
+        if self.backend_selector.mem_type == "FILE":
+            file_paths = []
+            for key in keys:
+                tensor_path = self.file_manager.get_file_path(key)
+                if not self.file_manager.create_file(tensor_path):
+                    logger.error(f"Failed to create file {tensor_path}")
+                    return False
+                file_paths.append(tensor_path)
+            return self._execute_transfer(values, file_paths, "WRITE")
+        else:
+            return self._execute_transfer(values, keys, "WRITE")
+
+    def set(self, key: str, value: torch.Tensor) -> bool:
+        return self.batch_set([key], [value])
+
+    def get(
+        self, key: str, dst_tensor: Optional[torch.Tensor] = None
+    ) -> torch.Tensor | None:
+        if dst_tensor is None:  # To be removed, being compatible with the current API
+            return None
+        result = self.batch_get([key], [dst_tensor])
+        return result[0] if result else None
+
+    def batch_get(
+        self, keys: List[str], dst_tensors: List[torch.Tensor]
+    ) -> List[Optional[torch.Tensor]]:
+        if not keys:
+            return []
+
+        if self.backend_selector.mem_type == "FILE":
+            file_paths = [self.file_manager.get_file_path(key) for key in keys]
+            success = self._execute_transfer(dst_tensors, file_paths, "READ")
+        else:
+            success = self._execute_transfer(dst_tensors, keys, "READ")
+        return dst_tensors if success else [None] * len(keys)
+
+    def exists(self, key: str) -> bool:
+        tuples = self.registration.create_query_tuples(
+            key,
+            self.backend_selector.mem_type,
+            self.file_manager if self.backend_selector.mem_type == "FILE" else None,
+        )
+        if not tuples:
+            return False
+
+        query_res = self.agent.query_memory(
+            tuples,
+            self.backend_selector.backend_name,
+            mem_type=self.backend_selector.mem_type,
+        )
+        return query_res[0] is not None  # can be expanded to multiple keys