diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 0c988838c..ed4572032 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -185,6 +185,7 @@ from sglang.srt.managers.scheduler_update_weights_mixin import ( SchedulerUpdateWeightsMixin, ) 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.radix_cache import RadixCache @@ -998,37 +999,12 @@ class Scheduler( embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get() init_mm_embedding_cache(embedding_cache_size * 1024 * 1024) - @staticmethod - def get_draft_kv_pool( - *, - draft_worker: "BaseTpWorker", - spec_algorithm: SpeculativeAlgorithm, - server_args: ServerArgs, - enable_overlap: bool, - ): - """Return (draft_token_to_kv_pool, draft_model_config) for the current - draft worker, or (None, None) when no draft KV pool is available.""" - if draft_worker is None or spec_algorithm.is_ngram(): - return None, None - - if spec_algorithm.supports_spec_v2() and enable_overlap: - if server_args.enable_multi_layer_eagle: - draft_runner = draft_worker.draft_worker.draft_runner_list[0] - else: - draft_runner = draft_worker.draft_worker.draft_runner - return draft_runner.token_to_kv_pool, draft_runner.model_config - - return ( - draft_worker.model_runner.token_to_kv_pool, - draft_worker.model_config, - ) - def _maybe_register_hicache_draft(self) -> None: """Register draft KV pool with HiCacheController for piggyback L2/L3 ops.""" if not self.enable_hierarchical_cache: return - draft_kv_pool, _ = Scheduler.get_draft_kv_pool( + draft_kv_pool, _ = kv_cache_builder.get_draft_kv_pool( draft_worker=self.draft_worker, spec_algorithm=self.spec_algorithm, server_args=self.server_args, @@ -1228,7 +1204,7 @@ class Scheduler( ) # todo: should we fix this when enabling mtp or it doesn't matter since we only enable mtp in decode node thus we don't transfer draft kvs between P and D? - draft_token_to_kv_pool, model_config = Scheduler.get_draft_kv_pool( + draft_token_to_kv_pool, model_config = kv_cache_builder.get_draft_kv_pool( draft_worker=self.draft_worker, spec_algorithm=self.spec_algorithm, server_args=self.server_args, diff --git a/python/sglang/srt/managers/scheduler_components/__init__.py b/python/sglang/srt/managers/scheduler_components/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py new file mode 100644 index 000000000..8ce44f64a --- /dev/null +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + + from sglang.srt.managers.tp_worker import BaseTpWorker + from sglang.srt.server_args import ServerArgs + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + +def get_draft_kv_pool( + *, + draft_worker: "BaseTpWorker", + spec_algorithm: SpeculativeAlgorithm, + server_args: ServerArgs, + enable_overlap: bool, +): + """Return (draft_token_to_kv_pool, draft_model_config) for the current + draft worker, or (None, None) when no draft KV pool is available.""" + if draft_worker is None or spec_algorithm.is_ngram(): + return None, None + + if spec_algorithm.supports_spec_v2() and enable_overlap: + if server_args.enable_multi_layer_eagle: + draft_runner = draft_worker.draft_worker.draft_runner_list[0] + else: + draft_runner = draft_worker.draft_worker.draft_runner + return draft_runner.token_to_kv_pool, draft_runner.model_config + + return ( + draft_worker.model_runner.token_to_kv_pool, + draft_worker.model_config, + )