sglang 0.5.2rc2__py3-none-any.whl → 0.5.3.post1__py3-none-any.whl

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

@@ -0,0 +1,151 @@
+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
+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]) -> 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,109 @@
+import logging
+import os
+
+import torch
+import torch.distributed
+from aibrix_kvcache import (
+    BaseKVCacheManager,
+    GroupAwareKVCacheManager,
+    KVCacheBlockLayout,
+    KVCacheBlockSpec,
+    KVCacheConfig,
+    KVCacheMetrics,
+    KVCacheTensorSpec,
+    ModelSpec,
+    TokenListView,
+)
+from aibrix_kvcache.common.absl_logging import getLogger, log_every_n_seconds, log_if
+from aibrix_kvcache_storage import AibrixKVCacheStorage
+from torch.distributed import Backend, ProcessGroup
+
+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()
+        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 == "page_first":
+                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",
+)