diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 31dac386d..f4ba8281d 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -798,6 +798,18 @@ class PiecewiseCudaGraphRunner: **kwargs, ) if isinstance(output, LogitsProcessorOutput): + # Preserve mm_input_embeds when speculative decoding is + # enabled. The speculative draft's prefill path + # (eagle_worker_v2._draft_extend_for_prefill) reads + # mm_input_embeds off this LogitsProcessorOutput to reuse + # the target's encoder embeddings instead of re-embedding + # multimodal placeholder token ids. + mm_input_embeds = None + if ( + self.model_runner.spec_algorithm.is_speculative() + and output.mm_input_embeds is not None + ): + mm_input_embeds = output.mm_input_embeds[: self.raw_num_tokens] return LogitsProcessorOutput( next_token_logits=output.next_token_logits[ : self.raw_num_tokens @@ -807,6 +819,7 @@ class PiecewiseCudaGraphRunner: if output.hidden_states is not None else None ), + mm_input_embeds=mm_input_embeds, ) elif isinstance(output, EmbeddingPoolerOutput): return output