From fed1197474c195890ed2c08d56a1839ac177672d Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 18 May 2026 18:28:58 +0800 Subject: [PATCH] Reshape init_cache_with_memory_pool to match the future build_kv_cache signature (#25606) --- python/sglang/srt/managers/scheduler.py | 188 +++++++++++------- .../sglang/srt/mem_cache/kv_cache_builder.py | 17 ++ 2 files changed, 135 insertions(+), 70 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 19288abe3..2524f09e0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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): - 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 diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 9c859bf5b..580ca4cce 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/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: