diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 358f30a5d..8cccf094c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -850,9 +850,6 @@ class Scheduler( def init_model_worker(self): # Load model weights. self.init_tp_model_worker() - if self.spec_algorithm.is_frozen_kv_mtp(): - # Frozen-KV MTP draft construction needs the target KV pool. - self.init_target_memory_pool() self.maybe_init_draft_worker() # Allocate KV cache pools for all workers. diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 3f1b88a87..30fa31cb1 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -114,16 +114,11 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): f"{self.speculative_algorithm.name}." ) - # Draft attention uses target req_to_token + KV allocator (read-only). - self.req_to_token_pool, self.token_to_kv_pool_allocator = ( - target_worker.get_memory_pool() - ) - - target_cfg = target_worker.model_runner.memory_pool_config - self.draft_pool_config = MemoryPoolConfig( - max_total_num_tokens=64, # Dummy value - max_running_requests=target_cfg.max_running_requests, - ) + # Target pools (read-only) are bound in alloc_memory_pool(), not here, so + # the worker can be built before the target pool exists (see #29021). + self.req_to_token_pool = None + self.token_to_kv_pool_allocator = None + self.draft_pool_config: Optional[MemoryPoolConfig] = None self.hot_token_id = None @@ -144,9 +139,6 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): moe_dp_rank=moe_dp_rank, nccl_port=nccl_port, is_draft_worker=True, - req_to_token_pool=self.req_to_token_pool, - token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, - memory_pool_config=self.draft_pool_config, ) embed, head = self.target_worker.model_runner.model.get_embed_and_head() @@ -160,8 +152,6 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): ) self.kv_context: Optional[FrozenKVMTPContext] = None - if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"): - self._bind_kv_context() self.draft_tp_context = ( draft_tp_context if server_args.enable_dp_attention else empty_context @@ -180,23 +170,26 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): req_to_token_pool=None, token_to_kv_pool_allocator=None, ): + self.req_to_token_pool = req_to_token_pool + self.token_to_kv_pool_allocator = token_to_kv_pool_allocator + + self.draft_pool_config = MemoryPoolConfig( + max_total_num_tokens=64, # Dummy value + max_running_requests=memory_pool_config.max_running_requests, + ) + # NOTE: call TpModelWorker explicitly -- EagleDraftWorkerBase precedes it in # the MRO and its alloc_memory_pool is a no-op stub. TpModelWorker.alloc_memory_pool( self, memory_pool_config=self.draft_pool_config, - req_to_token_pool=( - req_to_token_pool - if req_to_token_pool is not None - else self.req_to_token_pool - ), - token_to_kv_pool_allocator=( - token_to_kv_pool_allocator - if token_to_kv_pool_allocator is not None - else self.token_to_kv_pool_allocator - ), + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=token_to_kv_pool_allocator, ) + if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"): + self._bind_kv_context() + def init_attention_backends(self): with ( self.draft_tp_context(self.draft_model_runner.tp_group), @@ -442,6 +435,9 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): self._set_positions(forward_batch) self._expand_for_topk_draft(forward_batch) + # Frozen draft never writes KV; None signals fill_from to skip the slot. + forward_batch.out_cache_loc = None + can_run_cuda_graph = ( self.cuda_graph_runner and self.cuda_graph_runner.can_run_graph(forward_batch)