chore: import upstream snapshot with attribution

python/sglang/srt/mem_cache/registry.py

@@ -0,0 +1,235 @@
+"""Registry for pluggable RadixCache factories.
+
+If `--radix-cache-backend` is unset (by default), the built-in selection
+chain is used to pick a cache implementation.
+
+To plug in a custom backend, register it under a string name via
+`register_radix_cache_backend(name, factory)`, then select it with
+`--radix-cache-backend <name>` (the flag accepts only registered names).
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+from sglang.srt.environ import envs
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.utils.tensor_bridge import use_mlx
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TreeCacheBuildContext:
+    """Radix Cache construction arguments."""
+
+    server_args: ServerArgs
+    params: CacheInitParams
+    is_hybrid_swa: bool
+    is_hybrid_ssm: bool
+    enable_hierarchical_cache: bool
+    disable_radix_cache: bool
+    effective_chunked_prefill_size: Optional[int]
+    tp_worker: Any
+    model_config: ModelConfig
+    tp_size: int
+    tp_rank: int
+    tp_group: Any
+    full_tokens_per_layer: Optional[int] = None
+
+
+RadixCacheFactory = Callable[[TreeCacheBuildContext], BasePrefixCache]
+
+_RADIX_CACHE_REGISTRY: dict[str, RadixCacheFactory] = {}
+
+
+def register_radix_cache_backend(name: str, factory: RadixCacheFactory) -> None:
+    """Register a radix-cache factory under `name`.
+
+    Raises ValueError if `name` is empty/whitespace-only or already
+    registered.
+    """
+    if not name.strip():
+        raise ValueError(
+            f"register_radix_cache_backend: name must be non-empty, got {name!r}"
+        )
+    if name in _RADIX_CACHE_REGISTRY:
+        raise ValueError(
+            f"register_radix_cache_backend: {name!r} is already registered"
+        )
+    _RADIX_CACHE_REGISTRY[name] = factory
+
+
+def get_radix_cache_factory(name: str) -> Optional[RadixCacheFactory]:
+    return _RADIX_CACHE_REGISTRY.get(name)
+
+
+def registered_radix_cache_backends() -> list[str]:
+    return list(_RADIX_CACHE_REGISTRY.keys())
+
+
+def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
+    """Built-in Radix Cache selection chain."""
+    server_args = ctx.server_args
+    params = ctx.params
+
+    if ctx.effective_chunked_prefill_size is not None and ctx.disable_radix_cache:
+        if not ctx.is_hybrid_swa:
+            from sglang.srt.mem_cache.chunk_cache import ChunkCache
+
+            return ChunkCache(params)
+        if ctx.full_tokens_per_layer == 0:
+            from sglang.srt.mem_cache.chunk_cache import PureSWAChunkCache
+
+            return PureSWAChunkCache(params)
+        from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
+
+        return SWAChunkCache(params)
+
+    if envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get():
+        # lazy import to avoid JIT overhead
+        from sglang.srt.mem_cache.radix_cache_cpp import RadixCacheCpp
+
+        logger.info("Using experimental C++ radix tree implementation.")
+        return RadixCacheCpp(params=params, server_args=server_args)
+
+    if envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get() or use_mlx():
+        return _create_unified_radix_cache(ctx, server_args, params)
+
+    if ctx.enable_hierarchical_cache:
+        if ctx.is_hybrid_ssm or ctx.is_hybrid_swa:
+            # HybridModel launches HiCache via UnifiedRadixCache by default.
+            return _create_unified_radix_cache(ctx, server_args, params)
+        else:
+            from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
+
+            cache = HiRadixCache(params=params, server_args=server_args)
+        ctx.tp_worker.register_hicache_layer_transfer_counter(
+            cache.cache_controller.layer_done_counter
+        )
+        return cache
+
+    if ctx.is_hybrid_swa:
+        if ctx.full_tokens_per_layer == 0:
+            from sglang.srt.mem_cache.pure_swa_radix_cache import PureSWARadixCache
+
+            return PureSWARadixCache(params=params)
+        from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
+
+        return SWARadixCache(params=params)
+
+    if ctx.is_hybrid_ssm:
+        from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
+
+        return MambaRadixCache(params)
+
+    if server_args.enable_lmcache:
+        from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
+            LMCRadixCache,
+        )
+
+        return LMCRadixCache(
+            params=params,
+            model_config=ctx.model_config,
+            tp_size=ctx.tp_size,
+            rank=ctx.tp_rank,
+            tp_group=ctx.tp_group,
+        )
+
+    if server_args.enable_flexkv:
+        # Importing the package side-effect registers the explicit
+        # ``--radix-cache-backend=flexkv`` factory; we then call the
+        # factory directly so --enable-flexkv stands on its own.
+        import os
+
+        from sglang.srt.mem_cache.storage.flexkv import _flexkv_factory
+
+        # Honor a CLI --flexkv-config-file by forwarding it via the env
+        # var that FlexKV's config loader actually reads.
+        if server_args.flexkv_config_file and not os.environ.get("FLEXKV_CONFIG_PATH"):
+            os.environ["FLEXKV_CONFIG_PATH"] = server_args.flexkv_config_file
+        return _flexkv_factory(ctx)
+
+    from sglang.srt.mem_cache.radix_cache import RadixCache
+
+    return RadixCache(params)
+
+
+def _create_unified_radix_cache(
+    ctx: TreeCacheBuildContext,
+    server_args: ServerArgs,
+    params: CacheInitParams,
+) -> BasePrefixCache:
+    """Initialize a UnifiedRadixCache with proper components and optional HiCache."""
+    from sglang.srt.mem_cache.unified_cache_components import ComponentType
+    from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
+
+    tree_components = [ComponentType.FULL]
+    if ctx.is_hybrid_swa:
+        tree_components.append(ComponentType.SWA)
+    if ctx.is_hybrid_ssm:
+        tree_components.append(ComponentType.MAMBA)
+
+    params.tree_components = tuple(tree_components)
+    if use_mlx() and ctx.is_hybrid_ssm:
+        from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import (
+            MlxAuxiliaryStateComponent,
+        )
+
+        params.component_registry_override = {
+            ComponentType.MAMBA: MlxAuxiliaryStateComponent,
+        }
+    cache = UnifiedRadixCache(params)
+    if ctx.enable_hierarchical_cache:
+        cache.init_hicache(server_args, params)
+        ctx.tp_worker.register_hicache_layer_transfer_counter(
+            cache.cache_controller.layer_done_counter
+        )
+    return cache
+
+
+def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache:
+    """Route to the matching factory to construct Radix Cache."""
+    name = ctx.server_args.radix_cache_backend
+    if name:
+        factory = get_radix_cache_factory(name)
+        if factory is None:
+            raise ValueError(
+                f"--radix-cache-backend={name!r} is not registered. "
+                f"Registered backends: {registered_radix_cache_backends()}. "
+                "External backends must call register_radix_cache_backend(...) at import time."
+            )
+        cache = factory(ctx)
+        source = f"registered({name!r})"
+    else:
+        cache = default_radix_cache_factory(ctx)
+        source = "default"
+
+    streaming_wrapped = False
+    if (
+        ctx.server_args.enable_streaming_session
+        and not cache.supports_streaming_session()
+    ):
+        from sglang.srt.session.streaming_session import StreamingSession
+
+        cache = StreamingSession(cache)
+        streaming_wrapped = True
+
+    logger.info(
+        "Tree cache initialized: source=%s impl=%s hybrid_swa=%s hybrid_ssm=%s "
+        "hierarchical=%s streaming_wrapped=%s",
+        source,
+        type(cache).__name__,
+        ctx.is_hybrid_swa,
+        ctx.is_hybrid_ssm,
+        ctx.enable_hierarchical_cache,
+        streaming_wrapped,
+    )
+    return cache