[Feature] implement the standard multi-layer MTP for step3p5 (#18564)

Co-authored-by: mei ran <meiran0528@gmail.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>

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)

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,