[fix]解决宽松mtp1报错

vllm/v1/spec_decode/eagle.py

@@ -239,12 +239,17 @@ class EagleProposer:
 
         if envs.VLLM_REJECT_SAMPLE_OPT:
             draft_prob = logits.softmax(dim=-1, dtype=torch.float32)
-            draft_probs_list = [draft_prob]
 
         # Early exit if there is only one draft token to be generated.
         if self.num_speculative_tokens == 1:
             # [batch_size, 1]
+            if envs.VLLM_REJECT_SAMPLE_OPT:
+                return draft_token_ids.view(-1, 1), draft_prob.view(-1, 1, logits.shape[-1])
+
             return draft_token_ids.view(-1, 1)
+        
+        if envs.VLLM_REJECT_SAMPLE_OPT:
+            draft_probs_list = [draft_prob]
 
         # TODO: Currently, MTP module released by deepseek only has
         # one layer. Adapt this code to support multiple layers once

vllm/zero_overhead/v1/eagle.py

@@ -166,12 +166,17 @@ class V1ZeroEagleProposer(EagleProposer):
 
         if envs.VLLM_REJECT_SAMPLE_OPT:
             draft_prob = logits.softmax(dim=-1, dtype=torch.float32)
-            draft_probs_list = [draft_prob]
 
         # Early exit if there is only one draft token to be generated.
         if self.num_speculative_tokens == 1:
             # [batch_size, 1]
+            if envs.VLLM_REJECT_SAMPLE_OPT:
+                return draft_token_ids.view(-1, 1), draft_prob.view(-1, 1, logits.shape[-1])
+
             return draft_token_ids.view(-1, 1)
+        
+        if envs.VLLM_REJECT_SAMPLE_OPT:
+            draft_probs_list = [draft_prob]
 
         # TODO: Currently, MTP module released by deepseek only has
         # one layer. Adapt this code to support multiple layers once