diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4f484fae5..fae8e4dbc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -481,7 +481,11 @@ class Scheduler( nccl_port=self.nccl_port, ) - def init_draft_worker(self): + def maybe_init_draft_worker(self): + if self.spec_algorithm.is_none(): + self.draft_worker = None + return + # Launch a draft worker for speculative decoding draft_worker_kwargs = dict( server_args=self.server_args, @@ -501,50 +505,12 @@ class Scheduler( f"Using draft model load_format: '{self.server_args.speculative_draft_load_format}'" ) - # FIXME: refactor the draft worker registration logic - if self.server_args.enable_multi_layer_eagle: - if self.enable_overlap: - from sglang.srt.speculative.multi_layer_eagle_worker_v2 import ( - MultiLayerEagleWorkerV2, - ) - - self.draft_worker = MultiLayerEagleWorkerV2( - gpu_id=self.gpu_id, - tp_rank=self.tp_rank, - moe_ep_rank=self.moe_ep_rank, - server_args=self.server_args, - nccl_port=self.nccl_port, - target_worker=self.tp_worker, - dp_rank=self.dp_rank, - ) - else: - from sglang.srt.speculative.multi_layer_eagle_worker import ( - MultiLayerEagleWorker, - ) - - self.draft_worker = MultiLayerEagleWorker( - gpu_id=self.gpu_id, - tp_rank=self.tp_rank, - moe_ep_rank=self.moe_ep_rank, - server_args=self.server_args, - nccl_port=self.nccl_port, - target_worker=self.tp_worker, - dp_rank=self.dp_rank, - ) - else: - WorkerClass = self.spec_algorithm.create_worker( - enable_overlap=self.enable_overlap - ) - - # FIXME: optimize the init draft worker code path - if WorkerClass is not None: - self.draft_worker = WorkerClass(**draft_worker_kwargs) - else: - self.draft_worker = None + DraftWorkerClass = self.spec_algorithm.create_worker(self.server_args) + self.draft_worker = DraftWorkerClass(**draft_worker_kwargs) def init_model_worker(self): self.init_tp_model_worker() - self.init_draft_worker() + self.maybe_init_draft_worker() # Dispatch the model worker if self.spec_algorithm.is_none(): diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 855e02d9d..a40a8aa0d 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import ModelWorkerBatch from sglang.srt.managers.tp_worker import TpModelWorker + from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.ngram_worker import NGRAMWorker @@ -49,12 +50,29 @@ class SpeculativeAlgorithm(Enum): return self.is_eagle() or self.is_standalone() def create_worker( - self, enable_overlap: bool = False + self, server_args: ServerArgs ) -> Optional[Union[Type[BaseSpecWorker], Type[TpModelWorker], Type[NGRAMWorker]]]: - if self.is_none(): - return None + assert ( + not self.is_none() + ), "Cannot create worker for NONE speculative algorithm." - if self.is_eagle(): + enable_overlap = not server_args.disable_overlap_schedule + if self.is_eagle() and server_args.enable_multi_layer_eagle: + # FIXME: migrate to EagleWorker + if enable_overlap: + from sglang.srt.speculative.multi_layer_eagle_worker_v2 import ( + MultiLayerEagleWorkerV2, + ) + + return MultiLayerEagleWorkerV2 + + from sglang.srt.speculative.multi_layer_eagle_worker import ( + MultiLayerEagleWorker, + ) + + return MultiLayerEagleWorker + + elif self.is_eagle(): if enable_overlap: from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2