Move maybe_register_hicache_draft to mem_cache.kv_cache_builder (#25604)

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

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)