[Spec Decode] Make speculative decoding compatible with pipeline parallelism

Signed-off-by: Xin Yang <xyangx@amazon.com>

Author: xyang16

vllm/attention/layer.py

@@ -153,7 +153,7 @@ def __init__(
         # this variable will not be accessed if use_direct_call is True
         self.kv_cache = [
             torch.tensor([]) for _ in range(
-                get_current_vllm_config().parallel_config.virtual_engine_size)
+                get_current_vllm_config().parallel_config.num_virtual_engine)
         ]
 
         self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)

vllm/config.py

@@ -1417,7 +1417,7 @@ class ParallelConfig:
     data_parallel_master_ip: str = "127.0.0.1"
     data_parallel_master_port: int = 29500  # Port of the data parallel master.
     enable_expert_parallel: bool = False  # Use EP instead of TP for MoE layers.
-    virtual_engine_size: int = 1  # Number of virtual engine.
+    num_virtual_engine: int = 1  # Number of virtual engines.
 
     # Maximum number of multiple batches
     # when load model sequentially. To avoid RAM OOM when using tensor
@@ -1927,6 +1927,9 @@ class SpeculativeConfig:
         - draft_tensor_parallel_size (Optional[int]): The degree of the tensor
             parallelism for the draft model. Can only be 1 or the same as the
             target model's tensor parallel size.
+        - draft_pipeline_parallel_size (Optional[int]): The degree of the
+            pipeline parallelism for the draft model. Can only be 1 or the
+            same as the target model's pipeline parallel size.
         - disable_logprobs (bool): If set to True, token log probabilities are
             not returned during speculative decoding. If set to False, token
             log probabilities are returned according to the log probability
@@ -2321,7 +2324,7 @@ def create_draft_parallel_config(
         draft_parallel_config = ParallelConfig(
             pipeline_parallel_size=speculative_draft_pipeline_parallel_size,
             tensor_parallel_size=speculative_draft_tensor_parallel_size,
-            virtual_engine_size=target_parallel_config.virtual_engine_size,
+            num_virtual_engine=target_parallel_config.num_virtual_engine,
             distributed_executor_backend=target_parallel_config.
             distributed_executor_backend,
             max_parallel_loading_workers=target_parallel_config.

vllm/engine/arg_utils.py

@@ -181,7 +181,7 @@ class EngineArgs:
     guided_decoding_backend: str = 'xgrammar'
     logits_processor_pattern: Optional[str] = None
 
-    speculative_config: Optional[Union[str, Dict[str, Any]]] = None
+    speculative_config: Optional[Dict[str, Any]] = None
 
     qlora_adapter_name_or_path: Optional[str] = None
     show_hidden_metrics_for_version: Optional[str] = None
@@ -1189,7 +1189,7 @@ def create_engine_config(
         parallel_config = ParallelConfig(
             pipeline_parallel_size=self.pipeline_parallel_size,
             tensor_parallel_size=self.tensor_parallel_size,
-            virtual_engine_size=self.pipeline_parallel_size,
+            num_virtual_engine=self.pipeline_parallel_size,
             data_parallel_size=self.data_parallel_size,
             enable_expert_parallel=self.enable_expert_parallel,
             max_parallel_loading_workers=self.max_parallel_loading_workers,

vllm/engine/async_llm_engine.py

@@ -812,9 +812,9 @@ async def run_engine_loop(engine_ref: ReferenceType):
         if not engine:
             return
 
-        virtual_engine_size = \
-                engine.engine.parallel_config.virtual_engine_size
-        has_requests_in_progress = [False] * virtual_engine_size
+        num_virtual_engine = \
+                engine.engine.parallel_config.num_virtual_engine
+        has_requests_in_progress = [False] * num_virtual_engine
         while True:
             if not any(has_requests_in_progress):
                 logger.debug("Waiting for new requests...")
@@ -839,9 +839,9 @@ async def run_engine_loop(engine_ref: ReferenceType):
                 logger.debug("Got new requests!")
                 requests_in_progress = [
                     asyncio.create_task(engine.engine_step(ve))
-                    for ve in range(virtual_engine_size)
+                    for ve in range(num_virtual_engine)
                 ]
-                has_requests_in_progress = [True] * virtual_engine_size
+                has_requests_in_progress = [True] * num_virtual_engine
 
             # Abort if iteration takes too long due to unrecoverable errors
             # (eg. NCCL timeouts).
@@ -850,7 +850,7 @@ async def run_engine_loop(engine_ref: ReferenceType):
                     done, _ = await asyncio.wait(
                         requests_in_progress,
                         return_when=asyncio.FIRST_COMPLETED)
-                    for _ in range(virtual_engine_size):
+                    for _ in range(num_virtual_engine):
                         await asyncio.sleep(0)
                 for task in done:
                     result = task.result()

vllm/engine/metrics.py

@@ -266,7 +266,7 @@ def __init__(self, labelnames: List[str], vllm_config: VllmConfig):
         # Speculative decoding stats
         self.gauge_spec_decode_draft_acceptance_rate = self._gauge_cls(
             name="vllm:spec_decode_draft_acceptance_rate",
-            documentation="Speulative token acceptance rate.",
+            documentation="Speculative token acceptance rate.",
             labelnames=labelnames,
             multiprocess_mode="sum")
         self.gauge_spec_decode_efficiency = self._gauge_cls(

vllm/sequence.py

@@ -1246,6 +1246,8 @@ def update(self,
         decode steps"""
         assert len(seq_group_metadata_list) == len(hidden_states)
         self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list))
+        if self.seq_group_metadata_list is not None:
+            self.seq_group_metadata_list.extend(seq_group_metadata_list)
         self.hidden_states = torch.cat([self.hidden_states, hidden_states])
 
         if self.second_last_token_hidden_states is not None:
@@ -1270,6 +1272,10 @@ def prune(self,
             # Batch contents changed - prune removed sequences.
             index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
             self.hidden_states = self.hidden_states[index]
+            if self.seq_group_metadata_list is not None:
+                self.seq_group_metadata_list = [
+                    self.seq_group_metadata_list[i] for i in index
+                ]
             if self.second_last_token_hidden_states is not None:
                 self.second_last_token_hidden_states = self\
                     .second_last_token_hidden_states[index]
@@ -1284,12 +1290,23 @@ def expand_with_bonus_tokens(
             return
 
         index = []
-        for seq_id in self._seq_ids:
-            i = self._seq_ids.index(seq_id)
+        expanded_seq_ids = []
+        expanded_seq_group_metadata_list = []
+        for i, seq_id in enumerate(self._seq_ids):
             if seq_id in seq_with_bonus_token_in_last_step:
                 index.append(i + len(self._seq_ids))
+                expanded_seq_ids.append(seq_id)
+                if self.seq_group_metadata_list is not None:
+                    expanded_seq_group_metadata_list.append(
+                        self.seq_group_metadata_list[i])
             index.append(i)
+            expanded_seq_ids.append(seq_id)
+            if self.seq_group_metadata_list is not None:
+                expanded_seq_group_metadata_list.append(
+                    self.seq_group_metadata_list[i])
 
+        self._seq_ids = expanded_seq_ids
+        self.seq_group_metadata_list = expanded_seq_group_metadata_list
         self.hidden_states = torch.cat(
             [self.hidden_states, self.second_last_token_hidden_states])[index]
 
@@ -1370,7 +1387,7 @@ def clone(
             virtual_engine=self.virtual_engine,
             num_lookahead_slots=self.num_lookahead_slots,
             running_queue_size=self.running_queue_size,
-            previous_hidden_states=self.previous_hidden_states,
+            previous_hidden_states=copy.copy(self.previous_hidden_states),
             num_steps=self.num_steps,
             finished_requests_ids=self.finished_requests_ids,
             last_sampled_token_ids=self.last_sampled_token_ids.clone()

vllm/spec_decode/batch_expansion.py

@@ -43,7 +43,7 @@ def score_proposals(
         self,
         execute_model_req: ExecuteModelRequest,
         proposals: SpeculativeProposals,
-    ) -> Optional[SpeculativeScores]:
+    ) -> SpeculativeScores:
         """Score the proposed tokens via the scorer model.
 
         This converts each input sequence to a set of k+1 target sequences. The

vllm/spec_decode/interfaces.py

@@ -94,5 +94,5 @@ def score_proposals(
         self,
         execute_model_req: ExecuteModelRequest,
         proposals: SpeculativeProposals,
-    ) -> Optional[SpeculativeScores]:
+    ) -> SpeculativeScores:
         raise NotImplementedError

vllm/spec_decode/mqa_scorer.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional
-
+from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.sequence import (ExecuteModelRequest, SequenceData,
                            SequenceGroupMetadata, get_all_seq_ids)
 from vllm.spec_decode.interfaces import (SpeculativeProposals,
@@ -17,7 +17,7 @@ def score_proposals(
         self,
         execute_model_req: ExecuteModelRequest,
         proposals: SpeculativeProposals,
-    ) -> Optional[SpeculativeScores]:
+    ) -> SpeculativeScores:
         target_seq_group_metadata_list = []
         target_seq_id_start = max(
             get_all_seq_ids(execute_model_req.seq_group_metadata_list)) + 1
@@ -70,8 +70,40 @@ def score_proposals(
                 seq_group_metadata_list=target_seq_group_metadata_list))
 
         target_sampler_output = target_sampler_output[0]
-        if target_sampler_output is None:
-            return None
+        if get_pp_group().is_last_rank:
+            assert len(
+                target_sampler_output) == 1, "expected single-step output"
+            target_sampler_output = target_sampler_output[0]
+            # Store hidden states from target model execution, BxD.
+            sampled_token_probs = target_sampler_output.sampled_token_probs
+            logprobs = target_sampler_output.logprobs
+            sampled_token_ids = target_sampler_output.sampled_token_ids
+            hidden_states = target_sampler_output.hidden_states
+            prefill_hidden_states = target_sampler_output.prefill_hidden_states
+            tensors = {
+                "sampled_token_probs": sampled_token_probs,
+                "logprobs": logprobs,
+                "sampled_token_ids": sampled_token_ids,
+                "hidden_states": hidden_states,
+                "prefill_hidden_states": prefill_hidden_states
+            }
+            get_pp_group().broadcast_tensor_dict(
+                tensors, src=get_pp_group().world_size - 1)
+        else:
+            tensors = get_pp_group().broadcast_tensor_dict(
+                src=get_pp_group().world_size - 1)
+            sampled_token_probs = tensors["sampled_token_probs"]
+            logprobs = tensors["logprobs"]
+            sampled_token_ids = tensors["sampled_token_ids"]
+            hidden_states = tensors["hidden_states"]
+            prefill_hidden_states = tensors["prefill_hidden_states"]
+            target_sampler_output = SamplerOutput(
+                outputs=None,
+                sampled_token_probs=sampled_token_probs,
+                logprobs=logprobs,
+                sampled_token_ids=sampled_token_ids,
+                hidden_states=hidden_states,
+                prefill_hidden_states=prefill_hidden_states)
 
         k = execute_model_req.num_lookahead_slots
         bs = len(execute_model_req.seq_group_metadata_list)

vllm/spec_decode/spec_decode_worker.py

@@ -337,7 +337,7 @@ def __init__(
 
         # Hidden states from target model to pass to proposer
         # in the subsequent step.
-        self.previous_hidden_states: Optional[HiddenStates] = None
+        self.previous_hidden_states: Dict[int, Optional[HiddenStates]] = {}
         self._disable_logprobs = disable_logprobs
         self._disable_log_stats = disable_log_stats
         self._num_spec_prefill_steps = num_spec_prefill_steps
@@ -374,11 +374,13 @@ def init_device(self) -> None:
             self.proposer_worker.maybe_load_lm_head_weight(
                 target_lm_head_weight)
 
-        self._metrics.init_tensors(self.rank, device_type=self.device)
         if model_parallel_is_initialized():
+            self._metrics.init_tensors(get_tp_group().rank_in_group,
+                                       device_type=self.device)
             self.spec_decode_sampler.init_tensors(get_tp_group().local_rank,
                                                   device_type=self.device)
         else:
+            self._metrics.init_tensors(self.rank, device_type=self.device)
             self.spec_decode_sampler.init_tensors(self.rank,
                                                   device_type=self.device)
 
@@ -467,7 +469,9 @@ def execute_model(
     ) -> List[SamplerOutput]:
         """Perform speculative decoding on the input batch.
         """
-        if self.rank % self.tensor_parallel_size != self._driver_rank:
+        rank = get_tp_group().rank_in_group if model_parallel_is_initialized(
+        ) else self.rank
+        if rank != self._driver_rank:
             self._run_non_driver_rank()
             return []
 
@@ -721,14 +725,19 @@ def _run_no_spec(self, execute_model_req: ExecuteModelRequest,
                 hidden_states = hidden_states[
                     torch.where(sampler_output.sampled_token_ids -
                                 VLLM_INVALID_TOKEN_ID)[0]]
-            if self.previous_hidden_states is None and len(
-                    seq_group_meta_with_hidden):
-                self.previous_hidden_states = HiddenStates(
-                    hidden_states, seq_group_meta_with_hidden)
-            elif self.previous_hidden_states and len(
-                    seq_group_meta_with_hidden):
-                self.previous_hidden_states.update(hidden_states,
-                                                   seq_group_meta_with_hidden)
+            if execute_model_req.virtual_engine not in \
+                    self.previous_hidden_states and \
+                    len(seq_group_meta_with_hidden):
+                self.previous_hidden_states[
+                    execute_model_req.virtual_engine] = HiddenStates(
+                        hidden_states, seq_group_meta_with_hidden)
+            elif execute_model_req.virtual_engine in \
+                    self.previous_hidden_states and \
+                    len(seq_group_meta_with_hidden):
+                previous_hidden_states: HiddenStates = \
+                    self.previous_hidden_states[execute_model_req.virtual_engine]
+                previous_hidden_states.update(hidden_states,
+                                              seq_group_meta_with_hidden)
 
         if not skip_proposer:
             # We prepare the prefill hidden states here so that there no
@@ -804,17 +813,15 @@ def _run_speculative_decoding_step(
         Returns a list of SamplerOutput, each containing a single token per
         sequence.
         """
-        if self.previous_hidden_states is not None:
-            self.previous_hidden_states.seq_group_metadata_list = execute_model_req.seq_group_metadata_list
         if get_pp_group().is_first_rank:
             # With prefill chunking, expect requests to have prompts first
             # so that backend gets prefill|decode.
             assert num_lookahead_slots == execute_model_req.num_lookahead_slots
 
             # Pass last hidden states from target model to proposer
             execute_model_req.previous_hidden_states = \
-                self.previous_hidden_states
-            self.previous_hidden_states = None
+                self.previous_hidden_states[execute_model_req.virtual_engine]
+            self.previous_hidden_states.pop(execute_model_req.virtual_engine)
 
             with Timer() as proposal_timer:
                 # Generate proposals using draft worker.
@@ -883,8 +890,8 @@ def _run_speculative_decoding_step(
 
         with Timer() as verification_timer:
             accepted_token_ids, target_logprobs = self._verify_tokens(
-                execute_model_req.seq_group_metadata_list, proposal_scores,
-                proposals, execute_model_req.num_lookahead_slots)
+                execute_model_req, proposal_scores, proposals,
+                execute_model_req.num_lookahead_slots)
 
         stage_times = (proposal_execute_time, scoring_timer.elapsed_time_ms,
                        verification_timer.elapsed_time_ms)
@@ -901,7 +908,7 @@ def _run_speculative_decoding_step(
     @nvtx_range("spec_decode_worker._verify_tokens")
     def _verify_tokens(
         self,
-        seq_group_metadata_list: List[SequenceGroupMetadata],
+        execute_model_req: ExecuteModelRequest,
         proposal_scores: SpeculativeScores,
         proposals: SpeculativeProposals,
         max_proposal_len: int,
@@ -912,6 +919,7 @@ def _verify_tokens(
         Returns a tuple of Tensors, one for the accepted token ids and one for
         the logprobs according to the scoring model.
         """
+        seq_group_metadata_list = execute_model_req.seq_group_metadata_list
         proposal_lens_list = proposals.proposal_lens.tolist()
 
         # vLLM currently only supports proposal lens equal to zero or the batch
@@ -991,9 +999,10 @@ def _verify_tokens(
             second_last_token_hidden_states = hidden_states[:, -2]  # b x d
             hidden_states = hidden_states.gather(1, index).squeeze(1)  # b x d
             # Store hidden states from target model for subsequent decode step
-            self.previous_hidden_states = HiddenStates(
-                hidden_states, terminal_metadata,
-                second_last_token_hidden_states)
+            self.previous_hidden_states[
+                execute_model_req.virtual_engine] = HiddenStates(
+                    hidden_states, terminal_metadata,
+                    second_last_token_hidden_states)
         return accepted_token_ids, logprobs
 
     def _create_output_sampler_list(