[Model] Integrate PARD into vLLM pre-commit

Signed-off-by: root <anzihao_hh@126.com>

Signed-off-by:  <anzihao_hh@126.com>

Author: zihaoanllm

vllm/spec_decode/batch_expansion.py

@@ -73,14 +73,17 @@ def score_proposals(
             if VLLM_INVALID_TOKEN_ID not in proposals
         ]
 
-        (spec_indices, non_spec_indices, target_seq_group_metadata_list, num_scoring_tokens) = self._expand_batch(
+        (spec_indices, non_spec_indices, target_seq_group_metadata_list,
+         num_scoring_tokens) = self._expand_batch(
              seq_group_metadata_list=execute_model_req.seq_group_metadata_list,
              proposal_token_ids_list=proposal_token_ids_list_without_skips,
              proposal_lens_list=proposal_lens_list,
          )
 
         if keep_index is not None:
-            target_seq_group_metadata_list = [target_seq_group_metadata_list[i] for i in keep_index]
+            target_seq_group_metadata_list = [
+                target_seq_group_metadata_list[i] for i in keep_index
+            ]
         target_sampler_output = self._scorer_worker.execute_model(
             execute_model_req=execute_model_req.clone(
                 seq_group_metadata_list=target_seq_group_metadata_list))

vllm/spec_decode/multi_step_worker.py

@@ -13,13 +13,11 @@
 from vllm.sequence import (ExecuteModelRequest, HiddenStates, SequenceData,
                            SequenceGroupMetadata)
 from vllm.spec_decode.interfaces import SpeculativeProposals
-from vllm.spec_decode.batch_expansion import BatchExpansionTop1Scorer
 
 if current_platform.is_cuda_alike():
     from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
 
-from vllm.spec_decode.interfaces import (SpeculativeProposals,
-                                         SpeculativeProposer)
+from vllm.spec_decode.interfaces import SpeculativeProposer
 from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
 from vllm.spec_decode.top1_proposer import Top1Proposer
 from vllm.worker.worker_base import DelegateWorkerBase
@@ -89,7 +87,8 @@ def sampler_output(
         model_outputs: List[SamplerOutput] = []
         if current_platform.is_cuda_alike() and isinstance(
                 self.model_runner, TP1DraftModelRunner
-        ) and self.model_runner.supports_gpu_multi_step(expanded_request) and not self.pard:
+        ) and self.model_runner.supports_gpu_multi_step(
+                expanded_request) and not self.pard:
             # Here we run the draft_model_runner with multi-step prepare
             # on the GPU directly
             expanded_request.num_steps = sample_len
@@ -107,7 +106,8 @@ def sampler_output(
                 self.worker.model_runner.return_hidden_states = True
 
             if hasattr(self, "pard") and self.pard is True:
-                filtered_model_outputs = self.pard_infer(expanded_request, sample_len)
+                filtered_model_outputs = self.pard_infer(
+                    expanded_request, sample_len)
                 return filtered_model_outputs, True
 
             for _ in range(sample_len):
@@ -124,7 +124,6 @@ def sampler_output(
                     indices_of_seq_with_bonus_tokens)
                 model_outputs.append(model_output)
 
-
         # move indices to device to avoid stream sync
         indices_of_seq_with_bonus_tokens = torch.tensor(
             indices_of_seq_with_bonus_tokens, device=self.device)
@@ -133,7 +132,7 @@ def sampler_output(
         return filtered_model_outputs, True
 
     def pard_infer(self, expanded_request: ExecuteModelRequest,
-            sample_len: int) -> List[SamplerOutput]:
+                   sample_len: int) -> List[SamplerOutput]:
         # prepare recompute kv token
         # update seq_group_metadata_list
         mask_token_id = self.pard_token
@@ -147,69 +146,92 @@ def pard_infer(self, expanded_request: ExecuteModelRequest,
         for name, tmp_request in request_by_id.items():
             seq_num_base = len(tmp_request)
             group_key = list(tmp_request[-1].seq_data.keys())[0]
-            output_token_ids = tmp_request[-1].seq_data[group_key].output_token_ids
-            rm_num = min(sample_len - 1 + seq_num_base - 1, len(output_token_ids) - 1)
-            rm_token_ids = list(output_token_ids[len(output_token_ids) - rm_num:])
+            output_token_ids = tmp_request[-1].seq_data[
+                group_key].output_token_ids
+            rm_num = min(sample_len - 1 + seq_num_base - 1,
+                         len(output_token_ids) - 1)
+            rm_token_ids = list(output_token_ids[len(output_token_ids) -
+                                                 rm_num:])
             all_rm_token_ids.append(rm_token_ids)
             tmp_new_requests = tmp_request[-1]
-            tmp_new_requests.seq_data[group_key].output_token_ids = output_token_ids[:len(output_token_ids) - rm_num]
+            tmp_new_requests.seq_data[
+                group_key].output_token_ids = output_token_ids[:len(
+                    output_token_ids) - rm_num]
             tmp_new_requests.seq_data[group_key]._num_computed_tokens -= rm_num
             new_request_list.append(tmp_new_requests)
         expanded_request.seq_group_metadata_list = new_request_list
         max_rm_num = max([len(i) for i in all_rm_token_ids])
-        min_rm_num = min([len(i) for i in all_rm_token_ids])
 
         # get proposal
         proposal = SpeculativeProposals(
-                proposal_token_ids = torch.tensor([
-                    rm_token_ids + [mask_token_id for i in range(sample_len -1 + max_rm_num - len(rm_token_ids))] 
-                    for rm_token_ids in all_rm_token_ids], device=self.device),
-                proposal_probs=torch.tensor([sample_len - 1 + max_rm_num] * len(new_request_list), device=self.device), #fake
-                proposal_lens=torch.tensor([sample_len - 1 + max_rm_num] * len(new_request_list), device=self.device)
-                )
+            proposal_token_ids=torch.tensor([
+                rm_token_ids + [
+                    mask_token_id for i in range(sample_len - 1 + max_rm_num -
+                                                 len(rm_token_ids))
+                ] for rm_token_ids in all_rm_token_ids
+            ],
+                                            device=self.device),
+            proposal_probs=torch.tensor([sample_len - 1 + max_rm_num] *
+                                        len(new_request_list),
+                                        device=self.device),  #fake
+            proposal_lens=torch.tensor([sample_len - 1 + max_rm_num] *
+                                       len(new_request_list),
+                                       device=self.device))
 
         # pard forward
         keep_index = []
         rm_token_num = []
         rm_token_num_sum = []
         for i, rm_token_ids in enumerate(all_rm_token_ids):
-            keep_index.extend([i * (sample_len + max_rm_num) + j for j in range(sample_len + len(rm_token_ids))])
+            keep_index.extend([
+                i * (sample_len + max_rm_num) + j
+                for j in range(sample_len + len(rm_token_ids))
+            ])
             rm_token_num.append(len(rm_token_ids))
             rm_token_num_sum.append(sum(rm_token_num))
 
-        pard_draft_out = self.pard_scorer.score_proposals(expanded_request, proposal, return_output=True, keep_index=keep_index)
+        pard_draft_out = self.pard_scorer.score_proposals(
+            expanded_request,
+            proposal,
+            return_output=True,
+            keep_index=keep_index)
 
         # align probs shape of target and draft model
         target_dim = self.pard_scorer._vocab_size
         if pard_draft_out.sampled_token_probs.shape[1] > target_dim:
-            pard_draft_out.sampled_token_probs = pard_draft_out.sampled_token_probs[:, :target_dim]
+            tmp_draft_probs = pard_draft_out.sampled_token_probs[:, :
+                                                                 target_dim]
+            pard_draft_out.sampled_token_probs = tmp_draft_probs
         elif pard_draft_out.sampled_token_probs.shape[1] < target_dim:
             pard_draft_out.sampled_token_probs = torch.nn.functional.pad(
-                    pard_draft_out.sampled_token_probs, (0, target_dim - pard_draft_out.sampled_token_probs.shape[1]), value=0)
+                pard_draft_out.sampled_token_probs,
+                (0, target_dim - pard_draft_out.sampled_token_probs.shape[1]),
+                value=0)
 
         # get output
-        output_indices = torch.tensor([[i + tmp_rm + j * sample_len for j, tmp_rm in enumerate(rm_token_num_sum)] 
-            for i in range(sample_len)], device=self.device)
+        output_indices = torch.tensor([[
+            i + tmp_rm + j * sample_len
+            for j, tmp_rm in enumerate(rm_token_num_sum)
+        ] for i in range(sample_len)],
+                                      device=self.device)
         filtered_model_outputs = [
             SamplerOutput(
                 outputs=[
                     pard_draft_out.outputs[i] for i in output_indices_to_retain
                 ] if len(pard_draft_out.outputs) > 0 else [],
                 sampled_token_probs=(
-                    pard_draft_out.sampled_token_probs[output_indices_to_retain]
-                    if pard_draft_out.sampled_token_probs is not None
-                    else None),
-                logprobs=(
-                     pard_draft_out.logprobs[output_indices_to_retain]
-                    if  pard_draft_out.logprobs is not None else None),
-                sampled_token_ids=(pard_draft_out.
-                                   sampled_token_ids[output_indices_to_retain]
-                                   if pard_draft_out.sampled_token_ids
-                                   is not None else None))
+                    pard_draft_out.
+                    sampled_token_probs[output_indices_to_retain] if
+                    pard_draft_out.sampled_token_probs is not None else None),
+                logprobs=(pard_draft_out.logprobs[output_indices_to_retain]
+                          if pard_draft_out.logprobs is not None else None),
+                sampled_token_ids=(
+                    pard_draft_out.sampled_token_ids[output_indices_to_retain]
+                    if pard_draft_out.sampled_token_ids is not None else None))
             for output_indices_to_retain in output_indices
-            ]
+        ]
         return filtered_model_outputs
- 
+
     @staticmethod
     def _maybe_update_previous_hidden_states(
             model_output: SamplerOutput,

vllm/spec_decode/spec_decode_worker.py

@@ -101,10 +101,8 @@ def create_spec_worker(*args, **kwargs) -> "SpecDecodeWorker":
         scorer_worker=target_worker,
         draft_worker_kwargs=draft_worker_kwargs,
         disable_mqa_scorer=speculative_config.disable_mqa_scorer,
-        disable_by_batch_size=speculative_config.
-        disable_by_batch_size,
-        draft_token_acceptance_method=speculative_config.
-        acceptance_method,
+        disable_by_batch_size=speculative_config.disable_by_batch_size,
+        draft_token_acceptance_method=speculative_config.acceptance_method,
         typical_acceptance_sampler_posterior_threshold=speculative_config.
         posterior_threshold,
         typical_acceptance_sampler_posterior_alpha=speculative_config.
@@ -201,7 +199,6 @@ def create_worker(
                 if draft_model_config.hf_config.model_type == "eagle":
                     enable_lm_head_weight_load = True
 
-
                 proposer_worker = MultiStepWorker(**draft_worker_kwargs)
 
                 if draft_model_config.hf_config.model_type == "deepseek_mtp":
@@ -210,10 +207,13 @@ def create_worker(
 
             proposer_worker = SmallerTpProposerWorker.maybe_wrap_worker(
                 proposer_worker, draft_tp, target_tp)
-            pard = draft_model_config.hf_config.__dict__.get('spd_type', None) == 'pard'
+            pard = draft_model_config.hf_config.__dict__.get('spd_type',
+                                                             None) == 'pard'
             proposer_worker.pard = pard
             if pard:
-                proposer_worker.pard_token = draft_model_config.hf_config.__dict__['pard_token']
+                pard_token = draft_model_config.hf_config.__dict__[
+                    'pard_token']
+                proposer_worker.pard_token = pard_token
 
         logger.info("Configuring SpecDecodeWorker with proposer=%s",
                     type(proposer_worker))
@@ -350,7 +350,6 @@ def __init__(
         self._disable_log_stats = disable_log_stats
         self._num_spec_prefill_steps = num_spec_prefill_steps
 
-
     def init_device(self) -> None:
         """Initialize both scorer and proposer models.
         """
@@ -396,11 +395,11 @@ def init_device(self) -> None:
                                  device=self.device,
                                  vocab_size=self._vocab_size)
 
-
         if self.proposer_worker.pard:
-            self.proposer_worker.pard_scorer = scorer_cls(scorer_worker=self.proposer_worker,
-                                 device=self.device,
-                                 vocab_size=self._vocab_size)
+            self.proposer_worker.pard_scorer = scorer_cls(
+                scorer_worker=self.proposer_worker,
+                device=self.device,
+                vocab_size=self._vocab_size)
 
         self._configure_model_sampler_for_spec_decode()
 
@@ -796,7 +795,7 @@ def _run_speculative_decoding_step(
         # Pass last hidden states from target model to proposer
         execute_model_req.previous_hidden_states = self.previous_hidden_states
         self.previous_hidden_states = None
-        
+
         with Timer() as proposal_timer:
             # Generate proposals using draft worker.
             proposals = self.proposer_worker.get_spec_proposals(
@@ -1275,7 +1274,8 @@ def _vocab_size(self) -> int:
             for worker in [self.proposer_worker, self.scorer_worker]
         ]
         if not self.proposer_worker.pard:
-            assert all(vocab_sizes[0] == vocab_size for vocab_size in vocab_sizes)
+            assert all(vocab_sizes[0] == vocab_size
+                       for vocab_size in vocab_sizes)
         return vocab_sizes[0]
 
     @property