diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 81490315a..35376ae6c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -454,7 +454,7 @@ class Scheduler( self.init_cache_with_memory_pool() # Register draft KV pool (when spec + HiCache co-enabled). - Scheduler.maybe_register_hicache_draft( + kv_cache_builder.maybe_register_hicache_draft( tree_cache=self.tree_cache, draft_worker=self.draft_worker, spec_algorithm=self.spec_algorithm, @@ -1007,66 +1007,6 @@ class Scheduler( embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get() init_mm_embedding_cache(embedding_cache_size * 1024 * 1024) - @staticmethod - def maybe_register_hicache_draft( - *, - tree_cache: "BasePrefixCache", - draft_worker: "BaseTpWorker", - spec_algorithm: SpeculativeAlgorithm, - server_args: ServerArgs, - enable_hierarchical_cache: bool, - enable_overlap: bool, - page_size: int, - ) -> None: - """Register draft KV pool with HiCacheController for piggyback L2/L3 ops.""" - if not enable_hierarchical_cache: - return - - draft_kv_pool, _ = kv_cache_builder.get_draft_kv_pool( - draft_worker=draft_worker, - spec_algorithm=spec_algorithm, - server_args=server_args, - enable_overlap=enable_overlap, - ) - if draft_kv_pool is None: - return - - from sglang.srt.mem_cache.memory_pool import ( - HybridLinearKVPool, - MHATokenToKVPool, - MLATokenToKVPool, - ) - from sglang.srt.mem_cache.memory_pool_host import ( - MHATokenToKVPoolHost, - MLATokenToKVPoolHost, - ) - - pool = draft_kv_pool - if isinstance(pool, HybridLinearKVPool): - pool = pool.full_kv_pool - - # Create host pool for draft with the same slot count as the target host pool, - # so that host indices stay 1-to-1 between target and draft KV caches. - primary = tree_cache.cache_controller.mem_pool_host - kw = dict( - host_to_device_ratio=primary.size / pool.size, - host_size=0, - page_size=page_size, - layout=server_args.hicache_mem_layout, - ) - if isinstance(pool, MHATokenToKVPool): - draft_host_pool = MHATokenToKVPoolHost(pool, **kw) - elif isinstance(pool, MLATokenToKVPool): - draft_host_pool = MLATokenToKVPoolHost(pool, **kw) - else: - logger.warning( - "Draft pool type %s not supported for HiCache, skipping.", - type(pool).__name__, - ) - return - - tree_cache.cache_controller.set_draft_kv_pool(pool, draft_host_pool) - 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 8ce44f64a..9c859bf5b 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -1,5 +1,9 @@ from __future__ import annotations +import logging + +logger = logging.getLogger(__name__) + from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -32,3 +36,63 @@ def get_draft_kv_pool( draft_worker.model_runner.token_to_kv_pool, draft_worker.model_config, ) + + +def maybe_register_hicache_draft( + *, + tree_cache: "BasePrefixCache", + draft_worker: "BaseTpWorker", + spec_algorithm: SpeculativeAlgorithm, + server_args: ServerArgs, + enable_hierarchical_cache: bool, + enable_overlap: bool, + page_size: int, +) -> None: + """Register draft KV pool with HiCacheController for piggyback L2/L3 ops.""" + if not enable_hierarchical_cache: + return + + draft_kv_pool, _ = get_draft_kv_pool( + draft_worker=draft_worker, + spec_algorithm=spec_algorithm, + server_args=server_args, + enable_overlap=enable_overlap, + ) + if draft_kv_pool is None: + return + + from sglang.srt.mem_cache.memory_pool import ( + HybridLinearKVPool, + MHATokenToKVPool, + MLATokenToKVPool, + ) + from sglang.srt.mem_cache.memory_pool_host import ( + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + ) + + pool = draft_kv_pool + if isinstance(pool, HybridLinearKVPool): + pool = pool.full_kv_pool + + # Create host pool for draft with the same slot count as the target host pool, + # so that host indices stay 1-to-1 between target and draft KV caches. + primary = tree_cache.cache_controller.mem_pool_host + kw = dict( + host_to_device_ratio=primary.size / pool.size, + host_size=0, + page_size=page_size, + layout=server_args.hicache_mem_layout, + ) + if isinstance(pool, MHATokenToKVPool): + draft_host_pool = MHATokenToKVPoolHost(pool, **kw) + elif isinstance(pool, MLATokenToKVPool): + draft_host_pool = MLATokenToKVPoolHost(pool, **kw) + else: + logger.warning( + "Draft pool type %s not supported for HiCache, skipping.", + type(pool).__name__, + ) + return + + tree_cache.cache_controller.set_draft_kv_pool(pool, draft_host_pool)