[Spec Decode] (1/2) Remove batch expansion (#8839)

vllm/engine/output_processor/interfaces.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Callable, List
+from typing import Callable, List, Optional
 
 from vllm.config import SchedulerConfig
 from vllm.core.scheduler import Scheduler
@@ -58,10 +58,14 @@ class SequenceGroupOutputProcessor(ABC):
     @abstractmethod
     def process_outputs(self, sequence_group: SequenceGroup,
                         outputs: List[SequenceGroupOutput],
-                        is_async: bool) -> None:
+                        is_async: bool) -> Optional[int]:
         """Process new token ids for the sequence group. Handles logic such as
         detokenization, stop checking, and freeing/forking sequences in the
         scheduler.
+        
+        Return the number of new tokens generated in the sequence group.
+        The returned value is optional because it is only used for 
+        speculative decoding mqa scorer.
         """
         pass
 

vllm/engine/output_processor/multi_step.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Callable, List
+from typing import Callable, List, Optional
 
 from vllm.core.scheduler import Scheduler
 from vllm.engine.output_processor.interfaces import (
@@ -69,7 +69,7 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
     def process_outputs(self,
                         sequence_group: SequenceGroup,
                         outputs: List[SequenceGroupOutput],
-                        is_async: bool = False) -> None:
+                        is_async: bool = False) -> Optional[int]:
         """Append new tokens in the outputs to sequences in the sequence group.
 
         This only supports sequence groups of size 1. It supports greater than
@@ -84,6 +84,10 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
             tokens from the previous step. If this is true, then
             no tokens need to be appended since it is already done
             externally (before the next schedule() call)
+            
+        Returns:
+            The number of tokens appended to the sequence. This is optional
+            because only speculative decode uses this return value.
         """
         # Sequences can be in RUNNING or FINISHED_ABORTED state
         # once scheduled, as a sequence is moved to FINSIHED_ABORTED
@@ -106,6 +110,7 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
             # was already appended, so we only need to do the rest of the
             # postprocessor: Detokenization + stopping logic
             self._process_decode_and_stop(seq, sequence_group.sampling_params)
+            return None
         else:
             # Standard multi-step case
 
@@ -121,8 +126,8 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
             ]
             assert valid_samples
 
-            self._process_seq_outputs(seq, valid_samples,
-                                      sequence_group.sampling_params)
+            return self._process_seq_outputs(seq, valid_samples,
+                                             sequence_group.sampling_params)
 
     def _process_decode_and_stop(self, seq: Sequence,
                                  sampling_params: SamplingParams) -> None:
@@ -140,7 +145,7 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
 
     def _process_seq_outputs(self, seq: Sequence,
                              valid_samples: List[SequenceOutput],
-                             sampling_params: SamplingParams) -> None:
+                             sampling_params: SamplingParams) -> int:
         output_token_ids = [sample.output_token for sample in valid_samples]
         output_logprobs = [sample.logprobs for sample in valid_samples]
 
@@ -148,7 +153,6 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
         remaining_tokens = sampling_params.max_tokens - (seq.get_output_len() +
                                                          len(output_token_ids))
         if remaining_tokens < 0:
-            valid_samples = valid_samples[:remaining_tokens]
             output_token_ids = output_token_ids[:remaining_tokens]
 
         # Truncate any tokens after EOS. This is required as spec decode
@@ -162,7 +166,6 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
             for i in range(len(output_token_ids)):
                 if output_token_ids[i] == eos_token_id:
                     output_token_ids = output_token_ids[:i + 1]
-                    valid_samples = valid_samples[:i + 1]
                     break
 
         # Incrementally append tokens to the sequence, as if we had only one new
@@ -173,9 +176,9 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
                 token_id=output_token_id,
                 logprobs=output_logprob,
             )
-            seq.data.update_num_computed_tokens(1)
 
             self._process_decode_and_stop(seq, sampling_params)
 
             if seq.is_finished():
                 break
+        return len(output_token_ids)