diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 20c670d62..074019379 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -429,6 +429,16 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: forward_batch, ) + # Chain-style MTP: overwrite self.hidden_states with the draft model's + # output (hidden_states_before_norm) so that assign_new_state_triton + # propagates each MTP layer's own output to the next MTP layer, + # rather than always feeding the target model's hidden states. + if ( + self.eagle_worker.chain_mtp_hidden_states + and ret.hidden_states is not None + ): + self.hidden_states[:num_tokens].copy_(ret.hidden_states[:num_tokens]) + select_index = ( torch.arange(bs, device=self.model_runner.device) * (self.speculative_num_draft_tokens + self.step) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 7ef46c5df..0d5f08556 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -131,6 +131,11 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): # Alias for better readability self.draft_runner_list: List[ModelRunner] = self.draft_worker.model_runner_list + # Chain-style MTP: each step propagates its own output hidden states to the + # next step. Non-chain: each step uses the target model's hidden states. + draft_arch = self.draft_worker.model_config.hf_config.architectures[0] + self.chain_mtp_hidden_states = draft_arch in ["Step3p5MTP"] + self.init_lm_head() # Used for KV Cache reversion @@ -386,6 +391,15 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) topk_p_list.append(topk_p) topk_index_list.append(topk_index) + # Chain-style: use this step's output hidden_states as next step's input + if ( + self.chain_mtp_hidden_states + and step < self.speculative_num_steps - 1 + and output.logits_output.hidden_states is not None + ): + forward_batch.spec_info.hidden_states = ( + output.logits_output.hidden_states + ) if forward_batch.extend_seq_lens is not None: rotate_input_ids_triton( forward_batch.input_ids, @@ -605,8 +619,13 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): model_worker_batch ) - # Draft prefill - model_worker_batch.capture_hidden_mode = CaptureHiddenMode.LAST + # Chain-style MTP needs FULL to get all-token hidden states; + # non-chain only needs LAST (the target model's hidden states). + model_worker_batch.capture_hidden_mode = ( + CaptureHiddenMode.FULL + if self.draft_worker.chain_mtp_hidden_states + else CaptureHiddenMode.LAST + ) batch_output.next_draft_input = self.draft_worker._draft_extend_for_prefill( model_worker_batch, batch_output.logits_output.hidden_states,