diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1cf897f2d162..7123800f5162 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -518,9 +518,30 @@ def maybe_init_draft_worker(self): self.draft_worker = None return + if self.pp_size > 1 and self.spec_algorithm.is_eagle(): + arch = self.model_config.hf_config.architectures[0] + if not arch.startswith("Deepseek"): + raise ValueError( + "PP + EAGLE speculative decoding is currently only fully supported " + f"for Deepseek architectures. Got: {arch}" + ) + + # Under PP, speculative draft worker only needs to run on the last stage. + if self.pp_size > 1 and self.pp_rank != self.pp_size - 1: + self.draft_worker = None + return + + draft_server_args = self.server_args + if self.pp_size > 1: + import copy + + draft_server_args = copy.deepcopy(self.server_args) + # Keep draft/MTP model local to one PP stage. + draft_server_args.pp_size = 1 + # Launch a draft worker for speculative decoding draft_worker_kwargs = dict( - server_args=self.server_args, + server_args=draft_server_args, gpu_id=self.gpu_id, tp_rank=self.tp_rank, moe_ep_rank=self.moe_ep_rank, @@ -529,15 +550,15 @@ def maybe_init_draft_worker(self): dp_rank=self.dp_rank, ) - if self.server_args.speculative_draft_load_format is not None: - self.server_args.load_format = ( - self.server_args.speculative_draft_load_format + if draft_server_args.speculative_draft_load_format is not None: + draft_server_args.load_format = ( + draft_server_args.speculative_draft_load_format ) logger.info( - f"Using draft model load_format: '{self.server_args.speculative_draft_load_format}'" + f"Using draft model load_format: '{draft_server_args.speculative_draft_load_format}'" ) - DraftWorkerClass = self.spec_algorithm.create_worker(self.server_args) + DraftWorkerClass = self.spec_algorithm.create_worker(draft_server_args) self.draft_worker = DraftWorkerClass(**draft_worker_kwargs) def init_model_worker(self): @@ -545,7 +566,7 @@ def init_model_worker(self): self.maybe_init_draft_worker() # Dispatch the model worker - if self.spec_algorithm.is_none(): + if self.spec_algorithm.is_none() or self.draft_worker is None: self.model_worker = self.tp_worker else: self.model_worker = self.draft_worker diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ac34bf171028..611ca718d0c5 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -511,14 +511,12 @@ def initialize(self, min_per_gpu_memory: float): if loop_num > 1: self.num_effective_layers = self.num_effective_layers * loop_num - assert ( - (not model_has_mtp_layers) - or (self.spec_algorithm.is_none()) - or ( - (not self.spec_algorithm.is_none()) - and (self.num_effective_layers == model_num_layers) + # MTP speculative decoding can run with PP on the target model as long as + # the draft side keeps full MTP layers local to the last PP stage. + if self.is_draft_worker and model_has_mtp_layers and (not self.spec_algorithm.is_none()): + assert self.num_effective_layers == model_num_layers, ( + "Draft worker must own all MTP layers when speculative decoding is enabled." ) - ), "PP is not compatible with MTP models." # Consider PP, so use start_layer and end_layer. full_attention_layer_ids = [ diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index deda4f97e889..a020b5ab5468 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -45,6 +45,8 @@ should_deepgemm_weight_requant_ue8m0, ) from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.server_args import get_global_server_args +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.models.deepseek_common.utils import ( _is_cpu, _is_cpu_amx_available, @@ -258,8 +260,18 @@ def do_load_weights( # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue - # Skip loading embed_tokens if not first rank in pipeline parallelism - if ".embed_tokens." in name and not self.pp_group.is_first_rank: + # For EAGLE + PP, keep a copy of embed tokens on last stage so + # draft worker running on last PP stage can share embed/head from target. + keep_embed_on_last_pp_stage = ( + self.pp_group.world_size > 1 + and self.pp_group.is_last_rank + and SpeculativeAlgorithm.from_string( + get_global_server_args().speculative_algorithm + ).is_eagle() + ) + if ".embed_tokens." in name and not ( + self.pp_group.is_first_rank or keep_embed_on_last_pp_stage + ): continue # Skip loading norm if not last rank in pipeline parallelism if ".norm." in name and not self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 988ed91e831d..094ff444f7fb 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -2507,7 +2507,16 @@ def __init__( else: self.cp_size = None - if self.pp_group.is_first_rank: + speculative_algo = SpeculativeAlgorithm.from_string( + get_global_server_args().speculative_algorithm + ) + keep_embed_on_last_pp_stage = ( + self.pp_group.world_size > 1 + and self.pp_group.is_last_rank + and speculative_algo.is_eagle() + ) + + if self.pp_group.is_first_rank or keep_embed_on_last_pp_stage: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size,