Move build_kv_cache to mem_cache.kv_cache_builder (#25607)

python/sglang/srt/managers/scheduler.py

@@ -150,7 +150,6 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.mm_utils import (
     has_shm_features,
-    init_mm_embedding_cache,
     unwrap_shm_features,
 )
 from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
@@ -186,10 +185,7 @@ from sglang.srt.managers.scheduler_update_weights_mixin import (
 )
 from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
 from sglang.srt.mem_cache import kv_cache_builder
-from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.common import maybe_cache_unfinished_req, release_kv_cache
-from sglang.srt.mem_cache.kv_cache_builder import KVCacheBuildResult
-from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
 from sglang.srt.model_loader.utils import get_resolved_model_impl
 from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin
@@ -210,7 +206,6 @@ from sglang.srt.plugins import load_plugins
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
 from sglang.srt.session.session_controller import SessionController
-from sglang.srt.session.streaming_session import StreamingSession
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
     DynamicGradMode,
@@ -452,7 +447,7 @@ class Scheduler(
             time.sleep(t)
 
         # Init cache and memory pool
-        result = Scheduler.build_kv_cache(
+        result = kv_cache_builder.build_kv_cache(
             server_args=self.server_args,
             model_config=self.model_config,
             tp_worker=self.tp_worker,
@@ -866,199 +861,6 @@ class Scheduler(
                 startup_available_gpu_memory_gb=avail_mem,
             )
 
-    @staticmethod
-    def build_kv_cache(
-        *,
-        server_args: "ServerArgs",
-        model_config: "ModelConfig",
-        tp_worker: "BaseTpWorker",
-        page_size: int,
-        spec_algorithm: "SpeculativeAlgorithm",
-        attn_tp_cpu_group: "ProcessGroup",
-        tp_cpu_group: "ProcessGroup",
-        attn_cp_cpu_group: "ProcessGroup",
-        enable_metrics: bool,
-        enable_kv_cache_events: bool,
-        ps: "ParallelState",
-        tp_group: "GroupCoordinator",
-        enable_hierarchical_cache: bool,
-    ) -> "KVCacheBuildResult":
-        sliding_window_size: Optional[int] = None
-        full_tokens_per_layer: Optional[int] = None
-        swa_tokens_per_layer: Optional[int] = None
-        uses_transformers_backend = (
-            get_resolved_model_impl(model_config) == ModelImpl.TRANSFORMERS
-        )
-
-        # Hybrid memory pool
-        is_hybrid_swa = tp_worker.is_hybrid_swa
-        _spec = tp_worker.model_runner.linear_attn_model_spec
-        _registry_needs_mamba = (
-            _spec.uses_mamba_radix_cache if _spec is not None else False
-        )
-        is_hybrid_ssm = (
-            tp_worker.model_runner.hybrid_gdn_config is not None
-            or tp_worker.model_runner.mamba2_config is not None
-            or _registry_needs_mamba
-        )
-
-        sliding_window_size = None
-        if is_hybrid_swa:
-            sliding_window_size = tp_worker.sliding_window_size
-            full_tokens_per_layer, swa_tokens_per_layer = (
-                tp_worker.get_tokens_per_layer_info()
-            )
-
-        req_to_token_pool, token_to_kv_pool_allocator = tp_worker.get_memory_pool()
-
-        disable_radix_cache = server_args.disable_radix_cache or (
-            model_config.is_multimodal and uses_transformers_backend
-        )
-        if disable_radix_cache and not server_args.disable_radix_cache:
-            logger.warning(
-                "Radix cache is disabled for multimodal models with the "
-                "Transformers backend to avoid multimodal prefix-cache mismatches."
-            )
-
-        # Decode radix cache is unsupported with hybrid SWA/SSM models —
-        # these use specialized memory pools incompatible with the
-        # prefix-match-and-lock allocation path.
-        if (
-            server_args.disaggregation_decode_enable_radix_cache
-            and server_args.disaggregation_mode == "decode"
-        ):
-            if is_hybrid_swa:
-                raise ValueError(
-                    "--disaggregation-decode-enable-radix-cache is incompatible "
-                    "with sliding window attention (SWA) models"
-                )
-            if is_hybrid_ssm:
-                raise ValueError(
-                    "--disaggregation-decode-enable-radix-cache is incompatible "
-                    "with Mamba/SSM models"
-                )
-
-        effective_chunked_prefill_size = server_args.chunked_prefill_size
-        if model_config.is_multimodal and uses_transformers_backend:
-            effective_chunked_prefill_size = None
-
-        params = CacheInitParams(
-            disable=disable_radix_cache,
-            req_to_token_pool=req_to_token_pool,
-            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
-            page_size=page_size,
-            is_eagle=spec_algorithm.is_eagle(),
-            tp_cache_group=(
-                attn_tp_cpu_group if server_args.enable_dp_attention else tp_cpu_group
-            ),
-            attn_cp_cache_group=attn_cp_cpu_group,
-            attn_tp_cache_group=attn_tp_cpu_group,
-            eviction_policy=server_args.radix_eviction_policy,
-            enable_metrics=enable_metrics,
-            enable_kv_cache_events=enable_kv_cache_events,
-            enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(),
-            pp_rank=ps.pp_rank,
-            pp_size=ps.pp_size,
-            chunked_prefill_size=effective_chunked_prefill_size,
-            sliding_window_size=sliding_window_size,
-        )
-
-        if effective_chunked_prefill_size is not None and disable_radix_cache:
-            if not is_hybrid_swa:
-                from sglang.srt.mem_cache.chunk_cache import ChunkCache
-
-                tree_cache = ChunkCache(params)
-            else:
-                from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
-
-                tree_cache = SWAChunkCache(params)
-        else:
-            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.")
-                tree_cache = RadixCacheCpp(params=params, server_args=server_args)
-            elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get():
-                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 is_hybrid_swa or is_hybrid_ssm:
-                    tree_components.append(
-                        ComponentType.SWA if is_hybrid_swa else ComponentType.MAMBA
-                    )
-                params.tree_components = tuple(tree_components)
-                tree_cache = UnifiedRadixCache(params)
-                if enable_hierarchical_cache:
-                    tree_cache.init_hicache(server_args, params)
-                    tp_worker.register_hicache_layer_transfer_counter(
-                        tree_cache.cache_controller.layer_done_counter
-                    )
-            elif enable_hierarchical_cache:
-                if is_hybrid_ssm:
-                    from sglang.srt.mem_cache.hi_mamba_radix_cache import (
-                        HiMambaRadixCache,
-                    )
-
-                    tree_cache = HiMambaRadixCache(
-                        params=params, server_args=server_args
-                    )
-                else:
-                    from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
-
-                    tree_cache = HiRadixCache(params=params, server_args=server_args)
-                tp_worker.register_hicache_layer_transfer_counter(
-                    tree_cache.cache_controller.layer_done_counter
-                )
-            elif is_hybrid_swa:
-                from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
-
-                tree_cache = SWARadixCache(params=params)
-            elif is_hybrid_ssm:
-                from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
-
-                tree_cache = MambaRadixCache(params)
-            elif server_args.enable_lmcache:
-                from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
-                    LMCRadixCache,
-                )
-
-                tree_cache = LMCRadixCache(
-                    params=params,
-                    model_config=model_config,
-                    tp_size=ps.tp_size,
-                    rank=ps.tp_rank,
-                    tp_group=tp_group,
-                )
-            else:
-                tree_cache = RadixCache(params)
-
-        if (
-            server_args.enable_streaming_session
-            and not tree_cache.supports_streaming_session()
-        ):
-            tree_cache = StreamingSession(tree_cache)
-
-        embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
-        init_mm_embedding_cache(embedding_cache_size * 1024 * 1024)
-
-        return KVCacheBuildResult(
-            is_hybrid_swa=is_hybrid_swa,
-            is_hybrid_ssm=is_hybrid_ssm,
-            sliding_window_size=sliding_window_size,
-            full_tokens_per_layer=full_tokens_per_layer,
-            swa_tokens_per_layer=swa_tokens_per_layer,
-            req_to_token_pool=req_to_token_pool,
-            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
-            disable_radix_cache=disable_radix_cache,
-            tree_cache=tree_cache,
-        )
-
     def init_running_status(self):
         self.waiting_queue: List[Req] = []
         # The running decoding batch for continuous batching

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -23,9 +23,23 @@ class KVCacheBuildResult:
 
 from typing import TYPE_CHECKING
 
+from sglang.srt.configs.model_config import ModelImpl
+from sglang.srt.environ import envs
+from sglang.srt.managers.mm_utils import init_mm_embedding_cache
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.mem_cache.radix_cache import RadixCache
+from sglang.srt.model_loader.utils import get_resolved_model_impl
+from sglang.srt.session.streaming_session import StreamingSession
+
 if TYPE_CHECKING:
 
+    from torch.distributed import ProcessGroup
+
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+    from sglang.srt.distributed.parallel_state_wrapper import ParallelState
     from sglang.srt.managers.tp_worker import BaseTpWorker
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
     from sglang.srt.server_args import ServerArgs
     from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
@@ -113,3 +127,192 @@ def maybe_register_hicache_draft(
         return
 
     tree_cache.cache_controller.set_draft_kv_pool(pool, draft_host_pool)
+
+
+def build_kv_cache(
+    *,
+    server_args: "ServerArgs",
+    model_config: "ModelConfig",
+    tp_worker: "BaseTpWorker",
+    page_size: int,
+    spec_algorithm: "SpeculativeAlgorithm",
+    attn_tp_cpu_group: "ProcessGroup",
+    tp_cpu_group: "ProcessGroup",
+    attn_cp_cpu_group: "ProcessGroup",
+    enable_metrics: bool,
+    enable_kv_cache_events: bool,
+    ps: "ParallelState",
+    tp_group: "GroupCoordinator",
+    enable_hierarchical_cache: bool,
+) -> "KVCacheBuildResult":
+    sliding_window_size: Optional[int] = None
+    full_tokens_per_layer: Optional[int] = None
+    swa_tokens_per_layer: Optional[int] = None
+    uses_transformers_backend = (
+        get_resolved_model_impl(model_config) == ModelImpl.TRANSFORMERS
+    )
+
+    # Hybrid memory pool
+    is_hybrid_swa = tp_worker.is_hybrid_swa
+    _spec = tp_worker.model_runner.linear_attn_model_spec
+    _registry_needs_mamba = _spec.uses_mamba_radix_cache if _spec is not None else False
+    is_hybrid_ssm = (
+        tp_worker.model_runner.hybrid_gdn_config is not None
+        or tp_worker.model_runner.mamba2_config is not None
+        or _registry_needs_mamba
+    )
+
+    sliding_window_size = None
+    if is_hybrid_swa:
+        sliding_window_size = tp_worker.sliding_window_size
+        full_tokens_per_layer, swa_tokens_per_layer = (
+            tp_worker.get_tokens_per_layer_info()
+        )
+
+    req_to_token_pool, token_to_kv_pool_allocator = tp_worker.get_memory_pool()
+
+    disable_radix_cache = server_args.disable_radix_cache or (
+        model_config.is_multimodal and uses_transformers_backend
+    )
+    if disable_radix_cache and not server_args.disable_radix_cache:
+        logger.warning(
+            "Radix cache is disabled for multimodal models with the "
+            "Transformers backend to avoid multimodal prefix-cache mismatches."
+        )
+
+    # Decode radix cache is unsupported with hybrid SWA/SSM models —
+    # these use specialized memory pools incompatible with the
+    # prefix-match-and-lock allocation path.
+    if (
+        server_args.disaggregation_decode_enable_radix_cache
+        and server_args.disaggregation_mode == "decode"
+    ):
+        if is_hybrid_swa:
+            raise ValueError(
+                "--disaggregation-decode-enable-radix-cache is incompatible "
+                "with sliding window attention (SWA) models"
+            )
+        if is_hybrid_ssm:
+            raise ValueError(
+                "--disaggregation-decode-enable-radix-cache is incompatible "
+                "with Mamba/SSM models"
+            )
+
+    effective_chunked_prefill_size = server_args.chunked_prefill_size
+    if model_config.is_multimodal and uses_transformers_backend:
+        effective_chunked_prefill_size = None
+
+    params = CacheInitParams(
+        disable=disable_radix_cache,
+        req_to_token_pool=req_to_token_pool,
+        token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        page_size=page_size,
+        is_eagle=spec_algorithm.is_eagle(),
+        tp_cache_group=(
+            attn_tp_cpu_group if server_args.enable_dp_attention else tp_cpu_group
+        ),
+        attn_cp_cache_group=attn_cp_cpu_group,
+        attn_tp_cache_group=attn_tp_cpu_group,
+        eviction_policy=server_args.radix_eviction_policy,
+        enable_metrics=enable_metrics,
+        enable_kv_cache_events=enable_kv_cache_events,
+        enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(),
+        pp_rank=ps.pp_rank,
+        pp_size=ps.pp_size,
+        chunked_prefill_size=effective_chunked_prefill_size,
+        sliding_window_size=sliding_window_size,
+    )
+
+    if effective_chunked_prefill_size is not None and disable_radix_cache:
+        if not is_hybrid_swa:
+            from sglang.srt.mem_cache.chunk_cache import ChunkCache
+
+            tree_cache = ChunkCache(params)
+        else:
+            from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
+
+            tree_cache = SWAChunkCache(params)
+    else:
+        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.")
+            tree_cache = RadixCacheCpp(params=params, server_args=server_args)
+        elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get():
+            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 is_hybrid_swa or is_hybrid_ssm:
+                tree_components.append(
+                    ComponentType.SWA if is_hybrid_swa else ComponentType.MAMBA
+                )
+            params.tree_components = tuple(tree_components)
+            tree_cache = UnifiedRadixCache(params)
+            if enable_hierarchical_cache:
+                tree_cache.init_hicache(server_args, params)
+                tp_worker.register_hicache_layer_transfer_counter(
+                    tree_cache.cache_controller.layer_done_counter
+                )
+        elif enable_hierarchical_cache:
+            if is_hybrid_ssm:
+                from sglang.srt.mem_cache.hi_mamba_radix_cache import (
+                    HiMambaRadixCache,
+                )
+
+                tree_cache = HiMambaRadixCache(params=params, server_args=server_args)
+            else:
+                from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
+
+                tree_cache = HiRadixCache(params=params, server_args=server_args)
+            tp_worker.register_hicache_layer_transfer_counter(
+                tree_cache.cache_controller.layer_done_counter
+            )
+        elif is_hybrid_swa:
+            from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
+
+            tree_cache = SWARadixCache(params=params)
+        elif is_hybrid_ssm:
+            from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
+
+            tree_cache = MambaRadixCache(params)
+        elif server_args.enable_lmcache:
+            from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
+                LMCRadixCache,
+            )
+
+            tree_cache = LMCRadixCache(
+                params=params,
+                model_config=model_config,
+                tp_size=ps.tp_size,
+                rank=ps.tp_rank,
+                tp_group=tp_group,
+            )
+        else:
+            tree_cache = RadixCache(params)
+
+    if (
+        server_args.enable_streaming_session
+        and not tree_cache.supports_streaming_session()
+    ):
+        tree_cache = StreamingSession(tree_cache)
+
+    embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
+    init_mm_embedding_cache(embedding_cache_size * 1024 * 1024)
+
+    return KVCacheBuildResult(
+        is_hybrid_swa=is_hybrid_swa,
+        is_hybrid_ssm=is_hybrid_ssm,
+        sliding_window_size=sliding_window_size,
+        full_tokens_per_layer=full_tokens_per_layer,
+        swa_tokens_per_layer=swa_tokens_per_layer,
+        req_to_token_pool=req_to_token_pool,
+        token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        disable_radix_cache=disable_radix_cache,
+        tree_cache=tree_cache,
+    )