Reshape init_cache_with_memory_pool to match the future build_kv_cache signature (#25606)

python/sglang/srt/managers/scheduler.py

@@ -188,6 +188,7 @@ from sglang.srt.managers.utils import GenerationBatchResult, validate_input_leng
 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
@@ -451,7 +452,30 @@ class Scheduler(
             time.sleep(t)
 
         # Init cache and memory pool
-        self.init_cache_with_memory_pool()
+        result = Scheduler.build_kv_cache(
+            server_args=self.server_args,
+            model_config=self.model_config,
+            tp_worker=self.tp_worker,
+            page_size=self.page_size,
+            spec_algorithm=self.spec_algorithm,
+            attn_tp_cpu_group=self.attn_tp_cpu_group,
+            tp_cpu_group=self.tp_cpu_group,
+            attn_cp_cpu_group=self.attn_cp_cpu_group,
+            enable_metrics=self.enable_metrics,
+            enable_kv_cache_events=self.enable_kv_cache_events,
+            ps=self.ps,
+            tp_group=self.tp_group,
+            enable_hierarchical_cache=self.enable_hierarchical_cache,
+        )
+        self.is_hybrid_swa = result.is_hybrid_swa
+        self.is_hybrid_ssm = result.is_hybrid_ssm
+        self.sliding_window_size = result.sliding_window_size
+        self.full_tokens_per_layer = result.full_tokens_per_layer
+        self.swa_tokens_per_layer = result.swa_tokens_per_layer
+        self.req_to_token_pool = result.req_to_token_pool
+        self.token_to_kv_pool_allocator = result.token_to_kv_pool_allocator
+        self.disable_radix_cache = result.disable_radix_cache
+        self.tree_cache = result.tree_cache
 
         if self.enable_hisparse:
             # Coordinator was created inside ModelRunner.initialize() before CUDA graph capture
@@ -842,39 +866,55 @@ class Scheduler(
                 startup_available_gpu_memory_gb=avail_mem,
             )
 
-    def init_cache_with_memory_pool(self):
+    @staticmethod
-        server_args = self.server_args
+    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(self.model_config) == ModelImpl.TRANSFORMERS
+            get_resolved_model_impl(model_config) == ModelImpl.TRANSFORMERS
         )
 
         # Hybrid memory pool
-        self.is_hybrid_swa = self.tp_worker.is_hybrid_swa
+        is_hybrid_swa = tp_worker.is_hybrid_swa
-        _spec = self.tp_worker.model_runner.linear_attn_model_spec
+        _spec = tp_worker.model_runner.linear_attn_model_spec
         _registry_needs_mamba = (
             _spec.uses_mamba_radix_cache if _spec is not None else False
         )
-        self.is_hybrid_ssm = (
+        is_hybrid_ssm = (
-            self.tp_worker.model_runner.hybrid_gdn_config is not None
+            tp_worker.model_runner.hybrid_gdn_config is not None
-            or self.tp_worker.model_runner.mamba2_config is not None
+            or tp_worker.model_runner.mamba2_config is not None
             or _registry_needs_mamba
         )
 
-        self.sliding_window_size = None
+        sliding_window_size = None
-        if self.is_hybrid_swa:
+        if is_hybrid_swa:
-            self.sliding_window_size = self.tp_worker.sliding_window_size
+            sliding_window_size = tp_worker.sliding_window_size
-            self.full_tokens_per_layer, self.swa_tokens_per_layer = (
+            full_tokens_per_layer, swa_tokens_per_layer = (
-                self.tp_worker.get_tokens_per_layer_info()
+                tp_worker.get_tokens_per_layer_info()
             )
 
-        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
+        req_to_token_pool, token_to_kv_pool_allocator = tp_worker.get_memory_pool()
-            self.tp_worker.get_memory_pool()
-        )
 
-        self.disable_radix_cache = server_args.disable_radix_cache or (
+        disable_radix_cache = server_args.disable_radix_cache or (
-            self.model_config.is_multimodal and uses_transformers_backend
+            model_config.is_multimodal and uses_transformers_backend
         )
-        if self.disable_radix_cache and not server_args.disable_radix_cache:
+        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."
@@ -887,60 +927,58 @@ class Scheduler(
             server_args.disaggregation_decode_enable_radix_cache
             and server_args.disaggregation_mode == "decode"
         ):
-            if self.is_hybrid_swa:
+            if is_hybrid_swa:
                 raise ValueError(
                     "--disaggregation-decode-enable-radix-cache is incompatible "
                     "with sliding window attention (SWA) models"
                 )
-            if self.is_hybrid_ssm:
+            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 self.model_config.is_multimodal and uses_transformers_backend:
+        if model_config.is_multimodal and uses_transformers_backend:
             effective_chunked_prefill_size = None
 
         params = CacheInitParams(
-            disable=self.disable_radix_cache,
+            disable=disable_radix_cache,
-            req_to_token_pool=self.req_to_token_pool,
+            req_to_token_pool=req_to_token_pool,
-            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
-            page_size=self.page_size,
+            page_size=page_size,
-            is_eagle=self.spec_algorithm.is_eagle(),
+            is_eagle=spec_algorithm.is_eagle(),
             tp_cache_group=(
-                self.attn_tp_cpu_group
+                attn_tp_cpu_group if server_args.enable_dp_attention else tp_cpu_group
-                if self.server_args.enable_dp_attention
-                else self.tp_cpu_group
             ),
-            attn_cp_cache_group=self.attn_cp_cpu_group,
+            attn_cp_cache_group=attn_cp_cpu_group,
-            attn_tp_cache_group=self.attn_tp_cpu_group,
+            attn_tp_cache_group=attn_tp_cpu_group,
             eviction_policy=server_args.radix_eviction_policy,
-            enable_metrics=self.enable_metrics,
+            enable_metrics=enable_metrics,
-            enable_kv_cache_events=self.enable_kv_cache_events,
+            enable_kv_cache_events=enable_kv_cache_events,
             enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(),
-            pp_rank=self.ps.pp_rank,
+            pp_rank=ps.pp_rank,
-            pp_size=self.ps.pp_size,
+            pp_size=ps.pp_size,
             chunked_prefill_size=effective_chunked_prefill_size,
-            sliding_window_size=self.sliding_window_size,
+            sliding_window_size=sliding_window_size,
         )
 
-        if effective_chunked_prefill_size is not None and self.disable_radix_cache:
+        if effective_chunked_prefill_size is not None and disable_radix_cache:
-            if not self.is_hybrid_swa:
+            if not is_hybrid_swa:
                 from sglang.srt.mem_cache.chunk_cache import ChunkCache
 
-                self.tree_cache = ChunkCache(params)
+                tree_cache = ChunkCache(params)
             else:
                 from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
 
-                self.tree_cache = SWAChunkCache(params)
+                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.")
-                self.tree_cache = RadixCacheCpp(params=params, server_args=server_args)
+                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,
@@ -950,67 +988,77 @@ class Scheduler(
                 )
 
                 tree_components = [ComponentType.FULL]
-                if self.is_hybrid_swa or self.is_hybrid_ssm:
+                if is_hybrid_swa or is_hybrid_ssm:
                     tree_components.append(
-                        ComponentType.SWA if self.is_hybrid_swa else ComponentType.MAMBA
+                        ComponentType.SWA if is_hybrid_swa else ComponentType.MAMBA
                     )
                 params.tree_components = tuple(tree_components)
-                self.tree_cache = UnifiedRadixCache(params)
+                tree_cache = UnifiedRadixCache(params)
-                if self.enable_hierarchical_cache:
+                if enable_hierarchical_cache:
-                    self.tree_cache.init_hicache(server_args, params)
+                    tree_cache.init_hicache(server_args, params)
-                    self.tp_worker.register_hicache_layer_transfer_counter(
+                    tp_worker.register_hicache_layer_transfer_counter(
-                        self.tree_cache.cache_controller.layer_done_counter
+                        tree_cache.cache_controller.layer_done_counter
                     )
-            elif self.enable_hierarchical_cache:
+            elif enable_hierarchical_cache:
-                if self.is_hybrid_ssm:
+                if is_hybrid_ssm:
                     from sglang.srt.mem_cache.hi_mamba_radix_cache import (
                         HiMambaRadixCache,
                     )
 
-                    self.tree_cache = HiMambaRadixCache(
+                    tree_cache = HiMambaRadixCache(
                         params=params, server_args=server_args
                     )
                 else:
                     from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
 
-                    self.tree_cache = HiRadixCache(
+                    tree_cache = HiRadixCache(params=params, server_args=server_args)
-                        params=params, server_args=server_args
+                tp_worker.register_hicache_layer_transfer_counter(
-                    )
+                    tree_cache.cache_controller.layer_done_counter
-                self.tp_worker.register_hicache_layer_transfer_counter(
-                    self.tree_cache.cache_controller.layer_done_counter
                 )
-            elif self.is_hybrid_swa:
+            elif is_hybrid_swa:
                 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
 
-                self.tree_cache = SWARadixCache(params=params)
+                tree_cache = SWARadixCache(params=params)
-            elif self.is_hybrid_ssm:
+            elif is_hybrid_ssm:
                 from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 
-                self.tree_cache = MambaRadixCache(params)
+                tree_cache = MambaRadixCache(params)
             elif server_args.enable_lmcache:
                 from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
                     LMCRadixCache,
                 )
 
-                self.tree_cache = LMCRadixCache(
+                tree_cache = LMCRadixCache(
                     params=params,
-                    model_config=self.model_config,
+                    model_config=model_config,
-                    tp_size=self.ps.tp_size,
+                    tp_size=ps.tp_size,
-                    rank=self.ps.tp_rank,
+                    rank=ps.tp_rank,
-                    tp_group=self.tp_group,
+                    tp_group=tp_group,
                 )
             else:
-                self.tree_cache = RadixCache(params)
+                tree_cache = RadixCache(params)
 
         if (
             server_args.enable_streaming_session
-            and not self.tree_cache.supports_streaming_session()
+            and not tree_cache.supports_streaming_session()
         ):
-            self.tree_cache = StreamingSession(self.tree_cache)
+            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

@@ -4,6 +4,23 @@ import logging
 
 logger = logging.getLogger(__name__)
 
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class KVCacheBuildResult:
+    is_hybrid_swa: bool
+    is_hybrid_ssm: bool
+    sliding_window_size: Optional[int]
+    full_tokens_per_layer: Optional[int]
+    swa_tokens_per_layer: Optional[int]
+    req_to_token_pool: object
+    token_to_kv_pool_allocator: object
+    disable_radix_cache: bool
+    tree_cache: object
+
+
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING: