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

sglang/srt/mem_cache/hiradix_cache.py

@@ -79,7 +79,9 @@ class HiRadixCache(RadixCache):
         self.write_through_threshold = (
             1 if hicache_write_policy == "write_through" else 3
         )
-        self.write_through_threshold_storage = 3
+        self.write_through_threshold_storage = (
+            1 if hicache_write_policy == "write_through" else 3
+        )
         self.load_back_threshold = 10
         super().__init__(
             req_to_token_pool, token_to_kv_pool_allocator, page_size, disable=False
@@ -111,6 +113,7 @@ class HiRadixCache(RadixCache):
             )
         if host_indices is not None:
             node.host_value = host_indices
+            assert len(node.host_value) > 0
             self.ongoing_write_through[node.id] = node
             if not write_back:
                 # no need to lock nodes if write back
@@ -388,10 +391,14 @@ class HiRadixCache(RadixCache):
                 self.cache_controller.ack_backup_queue.get()
             )
             host_node = self.ongoing_backup[ack_id]
-            if completed_tokens < len(host_node.key):
+            if completed_tokens == 0:
+                host_node.hash_value = None
+            elif completed_tokens < len(host_node.key):
                 # backup is only partially successful, split the node
                 new_node = self._split_node(host_node.key, host_node, completed_tokens)
                 new_node.hash_value = hash_value
+            else:
+                host_node.hash_value = hash_value
             host_node.release_host()
             del self.ongoing_backup[ack_id]
 
@@ -431,6 +438,8 @@ class HiRadixCache(RadixCache):
             written_indices,
             hash_value[:min_completed_tokens],
         )
+        if len(written_indices):
+            self.cache_controller.mem_pool_host.update_prefetch(written_indices)
 
         self.cache_controller.mem_pool_host.free(host_indices[:matched_length])
         self.cache_controller.mem_pool_host.free(
@@ -551,13 +560,11 @@ class HiRadixCache(RadixCache):
             prefix_len = self.key_match_fn(child.key, key)
             if prefix_len < len(child.key):
                 new_node = self._split_node(child.key, child, prefix_len)
-                self.inc_hit_count(new_node)
                 if not new_node.evicted:
                     value.append(new_node.value)
                 node = new_node
                 break
             else:
-                self.inc_hit_count(child)
                 if not child.evicted:
                     value.append(child.value)
                 node = child
@@ -587,6 +594,10 @@ class HiRadixCache(RadixCache):
         if child.backuped:
             new_node.host_value = child.host_value[:split_len]
             child.host_value = child.host_value[split_len:]
+
+        if child.hash_value:
+            new_node.hash_value = child.hash_value[: split_len // self.page_size]
+            child.hash_value = child.hash_value[split_len // self.page_size :]
         child.parent = new_node
         child.key = child.key[split_len:]
         new_node.parent.children[self.get_child_key_fn(key)] = new_node

sglang/srt/mem_cache/memory_pool_host.py

@@ -25,7 +25,6 @@ def synchronized(debug_only=False):
         @wraps(func)
         def wrapper(self, *args, **kwargs):
             if (not debug_only) or self.debug:
-                return func(self, *args, **kwargs)
                 with self.lock:
                     return func(self, *args, **kwargs)
             else:
@@ -181,6 +180,15 @@ class HostKVCache(abc.ABC):
             )
         self.mem_state[indices] = MemoryStateInt.BACKUP
 
+    @synchronized(debug_only=True)
+    def update_prefetch(self, indices: torch.Tensor):
+        if not self.is_reserved(indices):
+            raise ValueError(
+                f"The host memory slots should be in RESERVED state before turning into BACKUP. "
+                f"Current state: {self.get_state(indices)}"
+            )
+        self.mem_state[indices] = MemoryStateInt.BACKUP
+
     @synchronized(debug_only=True)
     def update_synced(self, indices: torch.Tensor):
         self.mem_state[indices] = MemoryStateInt.SYNCED
@@ -257,6 +265,43 @@ class MHATokenToKVPoolHost(HostKVCache):
             self.head_dim,
         )
 
+    def get_buffer_meta(self, keys, indices):
+        ptr_list = []
+        key_list = []
+        kv_buffer_data_ptr = self.kv_buffer.data_ptr()
+        v_offset = (
+            self.layer_num
+            * self.size
+            * self.head_num
+            * self.head_dim
+            * self.dtype.itemsize
+        )
+        for index in range(0, len(indices), self.page_size):
+            for layer_id in range(self.layer_num):
+                k_ptr = (
+                    kv_buffer_data_ptr
+                    + indices[index]
+                    * self.head_num
+                    * self.head_dim
+                    * self.dtype.itemsize
+                    + layer_id
+                    * self.size
+                    * self.head_num
+                    * self.head_dim
+                    * self.dtype.itemsize
+                )
+                v_ptr = k_ptr + v_offset
+                ptr_list.append(k_ptr)
+                ptr_list.append(v_ptr)
+                key_ = keys[index // self.page_size]
+                key_list.append(f"{key_}_{layer_id}_k")
+                key_list.append(f"{key_}_{layer_id}_v")
+        element_size = (
+            self.dtype.itemsize * self.page_size * self.head_num * self.head_dim
+        )
+        element_size_list = [element_size] * len(key_list)
+        return key_list, ptr_list, element_size_list
+
     @property
     def k_buffer(self):
         return self.kv_buffer[0]
@@ -317,3 +362,30 @@ class MLATokenToKVPoolHost(HostKVCache):
             1,
             self.kv_lora_rank + self.qk_rope_head_dim,
         )
+
+    def get_buffer_meta(self, keys, indices):
+        ptr_list = []
+        key_list = []
+        kv_buffer_data_ptr = self.kv_buffer.data_ptr()
+        for index in range(0, len(indices), self.page_size):
+            for layer_id in range(self.layer_num):
+                k_ptr = (
+                    kv_buffer_data_ptr
+                    + indices[index]
+                    * (self.kv_lora_rank + self.qk_rope_head_dim)
+                    * self.dtype.itemsize
+                    + layer_id
+                    * self.size
+                    * (self.kv_lora_rank + self.qk_rope_head_dim)
+                    * self.dtype.itemsize
+                )
+                ptr_list.append(k_ptr)
+                key_ = keys[index // self.page_size]
+                key_list.append(f"{key_}_{layer_id}_k")
+        element_size = (
+            self.dtype.itemsize
+            * self.page_size
+            * (self.kv_lora_rank + self.qk_rope_head_dim)
+        )
+        element_size_list = [element_size] * len(key_list)
+        return key_list, ptr_list, element_size_list

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/storage/hf3fs/client_hf3fs.py

@@ -0,0 +1,177 @@
+import logging
+import multiprocessing
+import os
+import threading
+from functools import wraps
+from pathlib import Path
+from typing import List
+
+import torch
+from torch.utils.cpp_extension import load
+
+root = Path(__file__).parent.resolve()
+hf3fs_utils = load(name="hf3fs_utils", sources=[f"{root}/hf3fs_utils.cpp"])
+
+logger = logging.getLogger(__name__)
+
+try:
+    from hf3fs_fuse.io import (
+        deregister_fd,
+        extract_mount_point,
+        make_ioring,
+        make_iovec,
+        register_fd,
+    )
+except ImportError as e:
+    logger.warning(f"hf3fs_fuse.io is not available: {e}")
+
+
+def rsynchronized():
+    def _decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            with self.rlock:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return _decorator
+
+
+def wsynchronized():
+    def _decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            with self.wlock:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return _decorator
+
+
+class Hf3fsClient:
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        self.path = path
+        self.size = size
+        self.bytes_per_page = bytes_per_page
+        self.entries = entries
+
+        self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
+        os.ftruncate(self.file, size)
+        register_fd(self.file)
+
+        self.hf3fs_mount_point = extract_mount_point(path)
+        self.bs = self.bytes_per_page
+        self.shm_r = multiprocessing.shared_memory.SharedMemory(
+            size=self.bs * self.entries, create=True
+        )
+        self.shm_w = multiprocessing.shared_memory.SharedMemory(
+            size=self.bs * self.entries, create=True
+        )
+
+        self.shm_r_tensor = torch.frombuffer(self.shm_r.buf, dtype=torch.uint8)
+        self.shm_w_tensor = torch.frombuffer(self.shm_w.buf, dtype=torch.uint8)
+
+        self.numa = -1
+        self.ior_r = make_ioring(
+            self.hf3fs_mount_point,
+            self.entries,
+            for_read=True,
+            timeout=1,
+            numa=self.numa,
+        )
+        self.ior_w = make_ioring(
+            self.hf3fs_mount_point,
+            self.entries,
+            for_read=False,
+            timeout=1,
+            numa=self.numa,
+        )
+        self.iov_r = make_iovec(self.shm_r, self.hf3fs_mount_point)
+        self.iov_w = make_iovec(self.shm_w, self.hf3fs_mount_point)
+
+        self.rlock = threading.RLock()
+        self.wlock = threading.RLock()
+
+    @rsynchronized()
+    def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        self.check(offsets, tensors)
+
+        # prepare
+        current = 0
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+            self.ior_r.prepare(
+                self.iov_r[current : current + size], True, self.file, offset
+            )
+            current += size
+
+        # submit
+        ionum = len(offsets)
+        resv = self.ior_r.submit().wait(min_results=ionum)
+
+        # results
+        hf3fs_utils.read_shm(self.shm_r_tensor, tensors)
+        results = [res.result for res in resv]
+
+        return results
+
+    @wsynchronized()
+    def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        self.check(offsets, tensors)
+
+        # prepare
+        hf3fs_utils.write_shm(tensors, self.shm_w_tensor)
+        current = 0
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+            self.ior_w.prepare(
+                self.iov_w[current : current + size], False, self.file, offset
+            )
+            current += size
+
+        # submit
+        ionum = len(offsets)
+        resv = self.ior_w.submit().wait(min_results=ionum)
+
+        # results
+        results = [res.result for res in resv]
+
+        return results
+
+    def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
+        sizes = [t.numel() * t.itemsize for t in tensors]
+        if any(
+            [
+                len(offsets) > self.entries,
+                len(offsets) != len(sizes),
+                all(
+                    [
+                        offset < 0 or offset + size > self.size
+                        for offset, size in zip(offsets, sizes)
+                    ]
+                ),
+                all([size > self.bytes_per_page for size in sizes]),
+            ]
+        ):
+            self.close()
+            raise ValueError(f"Hf3fsClient.check: {offsets=}, {sizes=}")
+
+    def get_size(self) -> int:
+        return self.size
+
+    def close(self) -> None:
+        deregister_fd(self.file)
+        os.close(self.file)
+        del self.ior_r
+        del self.ior_w
+        del self.iov_r
+        del self.iov_w
+        self.shm_r.close()
+        self.shm_w.close()
+        self.shm_r.unlink()
+        self.shm_w.unlink()
+
+    def flush(self) -> None:
+        os.fsync(self.file)