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

sglang/srt/mem_cache/radix_cache_cpp.py

@@ -13,6 +13,7 @@ from sglang.srt.mem_cache.cpp_radix_tree.radix_tree import (
     TreeNodeCpp,
 )
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.radix_cache import RadixKey
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -93,9 +94,9 @@ class RadixCacheCpp(BasePrefixCache):
             raise NotImplementedError("Host cache is not supported yet")
         self.tree.reset()
 
-    def match_prefix(self, key: List[int], **kwargs) -> MatchResult:
+    def match_prefix(self, key: RadixKey, **kwargs) -> MatchResult:
         device_indices_vec, host_indices_length, node_gpu, node_cpu = (
-            self.tree.match_prefix(key)
+            self.tree.match_prefix(key.token_ids)
         )
         return MatchResult(
             device_indices=self._merge_tensor(device_indices_vec),
@@ -104,16 +105,16 @@ class RadixCacheCpp(BasePrefixCache):
             host_hit_length=host_indices_length,
         )
 
-    def _insert(self, key: List[int], value: torch.Tensor) -> int:
+    def _insert(self, key: RadixKey, value: torch.Tensor) -> int:
         """
         Insert a key-value pair into the radix tree.
         Args:
-            key (List[int]): The key to insert, represented as a list of integers.
+            key (RadixKey): The key to insert, represented as a RadixKey.
             value (torch.Tensor): The value to associate with the key.
         Returns:
             int: Number of device indices that were already present in the tree before the insertion.
         """
-        ongoing_write, length = self.tree.writing_through(key, value)
+        ongoing_write, length = self.tree.writing_through(key.token_ids, value)
         if self.cache_controller is None:
             assert len(ongoing_write) == 0, "Implementation error"
             return length
@@ -150,32 +151,37 @@ class RadixCacheCpp(BasePrefixCache):
     def total_size(self):
         return self.tree.total_size()
 
-    def cache_finished_req(self, req: Req):
+    def cache_finished_req(self, req: Req, is_insert: bool = True):
         """Cache request when it finishes."""
         assert req.req_pool_idx is not None
-        token_ids = (req.origin_input_ids + req.output_ids)[:-1]
+        all_token_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+        token_ids = (req.origin_input_ids + req.output_ids)[:all_token_len]
         overall_len = len(token_ids)  # prefill + decode
         kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, :overall_len]
 
         # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
         # it will automatically align them, but length of them should be equal
         old_prefix_len = len(req.prefix_indices) // self.page_size * self.page_size
-        new_prefix_len = self._insert(token_ids, kv_indices)
+        page_aligned_overall_len = overall_len // self.page_size * self.page_size
 
-        # NOTE: kv_indices[:old_prefix_len] == req.prefix_indices
-        assert old_prefix_len <= new_prefix_len, "Wrong prefix indices"
-
-        # KVCache between old & new is newly generated, but already exists in the pool
-        # we need to free this newly generated kv indices
-        if old_prefix_len < new_prefix_len:
-            self.token_to_kv_pool.free(kv_indices[old_prefix_len:new_prefix_len])
+        if is_insert:
+            new_prefix_len = self._insert(
+                RadixKey(token_ids, req.extra_key), kv_indices
+            )
+            # NOTE: kv_indices[:old_prefix_len] == req.prefix_indices
+            assert old_prefix_len <= new_prefix_len, "Wrong prefix indices"
+            # Free duplicates that were already in the pool
+            if old_prefix_len < new_prefix_len:
+                self.token_to_kv_pool.free(kv_indices[old_prefix_len:new_prefix_len])
+        else:
+            self.token_to_kv_pool.free(
+                kv_indices[old_prefix_len:page_aligned_overall_len]
+            )
 
         # need to free the unaligned part, since it cannot be inserted into the radix tree
-        if self.page_size != 1 and (  # unaligned tail only exists when page_size > 1
-            (unaligned_len := overall_len % self.page_size) > 0
-        ):
+        if page_aligned_overall_len < overall_len:
             # NOTE: sglang PagedAllocator support unaligned free (which will automatically align it)
-            self.token_to_kv_pool.free(kv_indices[overall_len - unaligned_len :])
+            self.token_to_kv_pool.free(kv_indices[page_aligned_overall_len:])
 
         # Remove req slot release the cache lock
         self.dec_lock_ref(req.last_node)
@@ -191,14 +197,16 @@ class RadixCacheCpp(BasePrefixCache):
         # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
         # it will automatically align them, but length of them should be equal
         old_prefix_len = len(req.prefix_indices) // self.page_size * self.page_size
-        new_prefix_len = self._insert(token_ids, kv_indices)
+        new_prefix_len = self._insert(RadixKey(token_ids, req.extra_key), kv_indices)
 
         # NOTE: kv_indices[:old_prefix_len] == req.prefix_indices
         assert old_prefix_len <= new_prefix_len, "Wrong prefix indices"
 
         # TODO(dark): optimize the `insert` and `match` (e.g. merge into 1 function)
         # The prefix indices need to updated to reuse the kv indices in the pool
-        new_indices_vec, _, new_last_node, _ = self.tree.match_prefix(token_ids)
+        new_indices_vec, _, new_last_node, _ = self.tree.match_prefix(
+            RadixKey(token_ids, req.extra_key).token_ids
+        )
         new_indices = self._merge_tensor(new_indices_vec)
         assert new_prefix_len <= len(new_indices)
 

sglang/srt/mem_cache/storage/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to SGLang project
+
+"""Storage backend module for SGLang HiCache."""
+
+from .backend_factory import StorageBackendFactory
+
+__all__ = [
+    "StorageBackendFactory",
+]

sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py

@@ -0,0 +1,157 @@
+import logging
+from typing import Any, List, Optional
+
+import torch
+from aibrix_kvcache import (
+    BaseKVCacheManager,
+    BlockHashes,
+    KVCacheBlockLayout,
+    KVCacheBlockSpec,
+    KVCacheConfig,
+    KVCacheTensorSpec,
+    ModelSpec,
+)
+from aibrix_kvcache.common.absl_logging import log_every_n_seconds
+
+from sglang.srt.mem_cache.hicache_storage import (
+    HiCacheStorage,
+    HiCacheStorageConfig,
+    HiCacheStorageExtraInfo,
+)
+from sglang.srt.mem_cache.memory_pool_host import HostKVCache
+
+logger = logging.getLogger(__name__)
+
+
+class AibrixKVCacheStorage(HiCacheStorage):
+    def __init__(self, storage_config: HiCacheStorageConfig, mem_pool: HostKVCache):
+        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
+        kv_cache = mem_pool.device_pool
+        self.page_size = mem_pool.page_size
+        self.kv_cache_dtype = kv_cache.dtype
+        self.layer_num = kv_cache.layer_num
+        self.kv_head_ids = [
+            self.local_rank * kv_cache.head_num + i for i in range(kv_cache.head_num)
+        ]
+        if not self.is_mla_backend:
+            self.layer_ids = range(
+                kv_cache.start_layer, kv_cache.end_layer
+            )  # for pipeline parallel
+
+            self.block_spec = KVCacheBlockSpec(
+                block_ntokens=self.page_size,
+                block_dtype=self.kv_cache_dtype,
+                block_layout=KVCacheBlockLayout(KVCacheBlockLayout.NCLD),
+                tensor_spec=KVCacheTensorSpec(
+                    heads=self.kv_head_ids,
+                    layers=self.layer_ids,
+                    head_size=kv_cache.head_dim,
+                ),
+            )
+            logger.info(self.block_spec)
+            config = KVCacheConfig(
+                block_spec=self.block_spec, model_spec=ModelSpec(102400)
+            )
+            self.kv_cache_manager = BaseKVCacheManager(config)
+        else:
+            raise NotImplementedError(
+                "MLA is not supported by AibrixKVCacheStorage yet."
+            )
+
+    def _aibrix_kvcache_metrics_report(self):
+        self.kv_cache_manager.metrics.summary()
+        self.kv_cache_manager.metrics.reset()
+
+    def batch_get(
+        self,
+        keys: List[str],
+        target_locations: List[torch.Tensor],
+        target_sizes: Optional[Any] = None,
+    ) -> List[torch.Tensor | None]:
+        block_hash = BlockHashes(keys, self.page_size)
+        status = self.kv_cache_manager.acquire(None, block_hash)
+        log_every_n_seconds(
+            logger, logging.INFO, self._aibrix_kvcache_metrics_report(), 1
+        )
+        if status.is_ok():
+            num_fetched_tokens, handle = status.value
+            kv_blocks = handle.to_tensors()
+            assert len(kv_blocks) == len(target_locations)
+            for i in range(len(kv_blocks)):
+                assert (
+                    target_locations[i].nbytes == kv_blocks[i].nbytes
+                ), f"{target_locations[i].nbytes}, {kv_blocks[i].nbytes}"
+                target_locations[i].copy_(kv_blocks[i].flatten())
+            handle.release()
+            return target_locations
+
+        return [None] * len(keys)
+
+    def get(
+        self,
+        key: str,
+        target_location: Optional[Any] = None,
+        target_size: Optional[Any] = None,
+    ) -> torch.Tensor | None:
+        return self.batch_get([key], [target_location], [target_size])[0]
+
+    def batch_set(
+        self,
+        keys: List[str],
+        values: Optional[Any] = None,
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> bool:
+        block_hash = BlockHashes(keys, self.page_size)
+        status = self.kv_cache_manager.allocate_for(None, block_hash)
+        if not status.is_ok():
+            logger.warning(
+                f"aibrix_kvcache set allocate failed, error_code {status.error_code}"
+            )
+            return False
+        handle = status.value
+        tensors = handle.to_tensors()
+        if len(tensors) != len(values):
+            logger.warning("aibrix_kvcache set allocate not enough")
+            return False
+        for i in range(len(tensors)):
+            assert (
+                tensors[i].nbytes == values[i].nbytes
+            ), f"{tensors[i].nbytes}, {values[i].nbytes}"
+            tensors[i].reshape(values[i].shape).copy_(values[i]).reshape(
+                tensors[i].shape
+            )
+        status = self.kv_cache_manager.put(None, block_hash, handle)
+        if not status.is_ok():
+            logger.info(
+                f"AIBrix KVCache Storage set failed, error_code {status.error_code}"
+            )
+            return False
+        completed = status.value
+        return completed == len(keys) * self.page_size
+
+    def set(
+        self,
+        key: str,
+        value: Optional[Any] = None,
+        target_location: Optional[Any] = None,
+        target_size: Optional[Any] = None,
+    ) -> bool:
+        return self.batch_set([key], [value], [target_location], [target_size])
+
+    def batch_exists(
+        self, keys: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None
+    ) -> int:
+        block_hash = BlockHashes(keys, self.page_size)
+        status = self.kv_cache_manager.exists(None, block_hash)
+        if status.is_ok():
+            return status.value // self.page_size
+        return 0
+
+    def exists(self, key: str) -> bool | dict:
+        return self.batch_exists([key]) > 0

sglang/srt/mem_cache/storage/aibrix_kvcache/unit_test.py

@@ -0,0 +1,97 @@
+import logging
+import os
+
+import torch
+import torch.distributed
+from aibrix_kvcache.common.absl_logging import log_every_n_seconds
+from aibrix_kvcache_storage import AibrixKVCacheStorage
+
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
+from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+
+
+def setup():
+    os.environ["RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = "63886"
+
+
+class AIBrixKVCacheStorageTest:
+    def test_with_page_size(self):
+        config = HiCacheStorageConfig(
+            tp_rank=0,
+            tp_size=1,
+            is_mla_model=False,
+            is_page_first_layout=True,
+            model_name="test",
+        )
+        for page_size in range(1, 3):
+            logger.info(f"page_size: {page_size}")
+            batch_size = 2
+            head_num = 1
+            layer_num = 64
+            head_dim = 128
+            kv_cache = MHATokenToKVPool(
+                1024,
+                page_size,
+                torch.float16,
+                head_num,
+                head_dim,
+                layer_num,
+                "cpu",
+                False,
+                0,
+                layer_num,
+            )
+            mem_pool = MHATokenToKVPoolHost(kv_cache, 2, 0, page_size, "layer_first")
+            query_length = batch_size * 2
+            partial = batch_size
+            self.aibrix_kvcache = AibrixKVCacheStorage(config, mem_pool)
+            target_shape = (2, layer_num, page_size, head_num, head_dim)
+            rand_tensor = [
+                torch.rand(target_shape, dtype=torch.float16)
+                for _ in range(query_length)
+            ]
+            keys = ["hash" + str(i) for i in range(query_length)]
+            partial_keys = keys[batch_size:query_length]
+            assert self.aibrix_kvcache.batch_exists(keys) == 0
+            assert self.aibrix_kvcache.batch_set(keys, rand_tensor)
+            get_tensor = [
+                torch.rand(target_shape, dtype=torch.float16).flatten()
+                for _ in range(query_length)
+            ]
+            self.aibrix_kvcache.batch_get(keys, get_tensor)
+            for i in range(query_length):
+                assert torch.equal(get_tensor[i], rand_tensor[i].flatten())
+            ret = self.aibrix_kvcache.batch_exists(keys)
+            assert self.aibrix_kvcache.batch_exists(keys) == query_length
+            assert self.aibrix_kvcache.batch_exists(partial_keys) == partial
+            partial_get_tensor = [
+                torch.rand(target_shape, dtype=torch.float16).flatten()
+                for _ in range(partial)
+            ]
+            self.aibrix_kvcache.batch_get(partial_keys, partial_get_tensor)
+            for i in range(partial):
+                assert torch.equal(
+                    partial_get_tensor[i], rand_tensor[i + partial].flatten()
+                )
+            log_every_n_seconds(
+                logger,
+                logging.INFO,
+                self.aibrix_kvcache.kv_cache_manager.metrics.summary(),
+                1,
+            )
+
+
+if __name__ == "__main__":
+    setup()
+    test = AIBrixKVCacheStorageTest()
+    test.test_with_page_size()

sglang/srt/mem_cache/storage/backend_factory.py

@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to SGLang project
+
+import importlib
+import logging
+from typing import TYPE_CHECKING, Any, Dict
+
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class StorageBackendFactory:
+    """Factory for creating storage backend instances with support for dynamic loading."""
+
+    _registry: Dict[str, Dict[str, Any]] = {}
+
+    @staticmethod
+    def _load_backend_class(
+        module_path: str, class_name: str, backend_name: str
+    ) -> type[HiCacheStorage]:
+        """Load and validate a backend class from module path."""
+        try:
+            module = importlib.import_module(module_path)
+            backend_class = getattr(module, class_name)
+            if not issubclass(backend_class, HiCacheStorage):
+                raise TypeError(
+                    f"Backend class {class_name} must inherit from HiCacheStorage"
+                )
+            return backend_class
+        except ImportError as e:
+            raise ImportError(
+                f"Failed to import backend '{backend_name}' from '{module_path}': {e}"
+            ) from e
+        except AttributeError as e:
+            raise AttributeError(
+                f"Class '{class_name}' not found in module '{module_path}': {e}"
+            ) from e
+
+    @classmethod
+    def register_backend(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a storage backend with lazy loading.
+
+        Args:
+            name: Backend identifier
+            module_path: Python module path containing the backend class
+            class_name: Name of the backend class
+        """
+        if name in cls._registry:
+            logger.warning(f"Backend '{name}' is already registered, overwriting")
+
+        def loader() -> type[HiCacheStorage]:
+            """Lazy loader function to import the backend class."""
+            return cls._load_backend_class(module_path, class_name, name)
+
+        cls._registry[name] = {
+            "loader": loader,
+            "module_path": module_path,
+            "class_name": class_name,
+        }
+
+    @classmethod
+    def create_backend(
+        cls,
+        backend_name: str,
+        storage_config: HiCacheStorageConfig,
+        mem_pool_host: Any,
+        **kwargs,
+    ) -> HiCacheStorage:
+        """Create a storage backend instance.
+        Args:
+            backend_name: Name of the backend to create
+            storage_config: Storage configuration
+            mem_pool_host: Memory pool host object
+            **kwargs: Additional arguments passed to external backends
+        Returns:
+            Initialized storage backend instance
+        Raises:
+            ValueError: If backend is not registered and cannot be dynamically loaded
+            ImportError: If backend module cannot be imported
+            Exception: If backend initialization fails
+        """
+        # First check if backend is already registered
+        if backend_name in cls._registry:
+            registry_entry = cls._registry[backend_name]
+            backend_class = registry_entry["loader"]()
+            logger.info(
+                f"Creating storage backend '{backend_name}' "
+                f"({registry_entry['module_path']}.{registry_entry['class_name']})"
+            )
+            return cls._create_builtin_backend(
+                backend_name, backend_class, storage_config, mem_pool_host
+            )
+
+        # Try to dynamically load backend from extra_config
+        if backend_name == "dynamic" and storage_config.extra_config is not None:
+            backend_config = storage_config.extra_config
+            return cls._create_dynamic_backend(
+                backend_config, storage_config, mem_pool_host, **kwargs
+            )
+
+        # Backend not found
+        available_backends = list(cls._registry.keys())
+
+        raise ValueError(
+            f"Unknown storage backend '{backend_name}'. "
+            f"Registered backends: {available_backends}. "
+        )
+
+    @classmethod
+    def _create_dynamic_backend(
+        cls,
+        backend_config: Dict[str, Any],
+        storage_config: HiCacheStorageConfig,
+        mem_pool_host: Any,
+        **kwargs,
+    ) -> HiCacheStorage:
+        """Create a backend dynamically from configuration."""
+        required_fields = ["backend_name", "module_path", "class_name"]
+        for field in required_fields:
+            if field not in backend_config:
+                raise ValueError(
+                    f"Missing required field '{field}' in backend config for 'dynamic' backend"
+                )
+
+        backend_name = backend_config["backend_name"]
+        module_path = backend_config["module_path"]
+        class_name = backend_config["class_name"]
+
+        try:
+            # Import the backend class
+            backend_class = cls._load_backend_class(
+                module_path, class_name, backend_name
+            )
+
+            logger.info(
+                f"Creating dynamic storage backend '{backend_name}' "
+                f"({module_path}.{class_name})"
+            )
+
+            # Create the backend instance with storage_config
+            return backend_class(storage_config, kwargs)
+        except Exception as e:
+            logger.error(
+                f"Failed to create dynamic storage backend '{backend_name}': {e}"
+            )
+            raise
+
+    @classmethod
+    def _create_builtin_backend(
+        cls,
+        backend_name: str,
+        backend_class: type[HiCacheStorage],
+        storage_config: HiCacheStorageConfig,
+        mem_pool_host: Any,
+    ) -> HiCacheStorage:
+        """Create built-in backend with original initialization logic."""
+        if backend_name == "file":
+            return backend_class(storage_config)
+        elif backend_name == "nixl":
+            return backend_class(storage_config)
+        elif backend_name == "mooncake":
+            backend = backend_class(storage_config)
+            return backend
+        elif backend_name == "aibrix":
+            backend = backend_class(storage_config, mem_pool_host)
+            return backend
+        elif backend_name == "hf3fs":
+            # Calculate bytes_per_page based on memory pool layout
+            if mem_pool_host.layout in ["page_first", "page_first_direct"]:
+                bytes_per_page = (
+                    mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size
+                )
+            elif mem_pool_host.layout == "layer_first":
+                bytes_per_page = (
+                    mem_pool_host.get_size_per_token() * mem_pool_host.page_size
+                )
+
+            dtype = mem_pool_host.dtype
+            return backend_class.from_env_config(bytes_per_page, dtype, storage_config)
+        elif backend_name == "eic":
+            return backend_class(storage_config, mem_pool_host)
+        else:
+            raise ValueError(f"Unknown built-in backend: {backend_name}")
+
+
+# Register built-in storage backends
+StorageBackendFactory.register_backend(
+    "file", "sglang.srt.mem_cache.hicache_storage", "HiCacheFile"
+)
+
+StorageBackendFactory.register_backend(
+    "nixl",
+    "sglang.srt.mem_cache.storage.nixl.hicache_nixl",
+    "HiCacheNixl",
+)
+
+StorageBackendFactory.register_backend(
+    "mooncake",
+    "sglang.srt.mem_cache.storage.mooncake_store.mooncake_store",
+    "MooncakeStore",
+)
+
+StorageBackendFactory.register_backend(
+    "hf3fs",
+    "sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs",
+    "HiCacheHF3FS",
+)
+
+StorageBackendFactory.register_backend(
+    "aibrix",
+    "sglang.srt.mem_cache.storage.aibrix_kvcache.aibrix_kvcache_storage",
+    "AibrixKVCacheStorage",
+)
+
+StorageBackendFactory.register_backend(
+    "eic",
+    "sglang.srt.mem_cache.storage.eic.eic_storage",
+    "EICStorage",
+)