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

2026-05-18 18:28:58 +08:00
parent 60337abe24
commit fed1197474
2 changed files with 135 additions and 70 deletions
@@ -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):
-        server_args = self.server_args
+    @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(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
-        _spec = self.tp_worker.model_runner.linear_attn_model_spec
+        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
        )
-        self.is_hybrid_ssm = (
-            self.tp_worker.model_runner.hybrid_gdn_config is not None
-            or self.tp_worker.model_runner.mamba2_config is not None
+        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
        )

-        self.sliding_window_size = None
-        if self.is_hybrid_swa:
-            self.sliding_window_size = self.tp_worker.sliding_window_size
-            self.full_tokens_per_layer, self.swa_tokens_per_layer = (
-                self.tp_worker.get_tokens_per_layer_info()
+        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()
            )

-        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
-            self.tp_worker.get_memory_pool()
-        )
+        req_to_token_pool, token_to_kv_pool_allocator = tp_worker.get_memory_pool()

-        self.disable_radix_cache = server_args.disable_radix_cache or (
-            self.model_config.is_multimodal and uses_transformers_backend
+        disable_radix_cache = server_args.disable_radix_cache or (
+            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,
-            req_to_token_pool=self.req_to_token_pool,
-            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
-            page_size=self.page_size,
-            is_eagle=self.spec_algorithm.is_eagle(),
+            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=(
-                self.attn_tp_cpu_group
-                if self.server_args.enable_dp_attention
-                else self.tp_cpu_group
+                attn_tp_cpu_group if server_args.enable_dp_attention else tp_cpu_group
            ),
-            attn_cp_cache_group=self.attn_cp_cpu_group,
-            attn_tp_cache_group=self.attn_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=self.enable_metrics,
-            enable_kv_cache_events=self.enable_kv_cache_events,
+            enable_metrics=enable_metrics,
+            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_size=self.ps.pp_size,
+            pp_rank=ps.pp_rank,
+            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 not self.is_hybrid_swa:
+        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

-                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)
-                if self.enable_hierarchical_cache:
-                    self.tree_cache.init_hicache(server_args, params)
-                    self.tp_worker.register_hicache_layer_transfer_counter(
-                        self.tree_cache.cache_controller.layer_done_counter
+                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 self.enable_hierarchical_cache:
-                if self.is_hybrid_ssm:
+            elif enable_hierarchical_cache:
+                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(
-                        params=params, server_args=server_args
-                    )
-                self.tp_worker.register_hicache_layer_transfer_counter(
-                    self.tree_cache.cache_controller.layer_done_counter
+                    tree_cache = HiRadixCache(params=params, server_args=server_args)
+                tp_worker.register_hicache_layer_transfer_counter(
+                    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)
-            elif self.is_hybrid_ssm:
+                tree_cache = SWARadixCache(params=params)
+            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,
-                    tp_size=self.ps.tp_size,
-                    rank=self.ps.tp_rank,
-                    tp_group=self.tp_group,
+                    model_config=model_config,
+                    tp_size=ps.tp_size,
+                    rank=ps.tp_rank,
+                    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
@@ -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: