delete AscendRejectionSampler

Signed-off-by: realliujiaxu <realliujiaxu@163.com>

Author: realliujiaxu

vllm_ascend/patch/worker/patch_rejection_sampler.py

@@ -0,0 +1,5 @@
+import vllm.v1.sample.rejection_sampler as rs
+from vllm_ascend.sample.rejection_sampler import expand_batch_to_tokens, rejection_sample
+
+rs.expand_batch_to_tokens = expand_batch_to_tokens
+rs.rejection_sample = rejection_sample

vllm_ascend/sample/rejection_sampler.py

@@ -18,10 +18,6 @@
 MAX_SPEC_LEN = 32
 
 
-class AscendRejectionSampler(RejectionSampler, nn.Module):
-    pass
-
-
 def rejection_sample(
     # [num_tokens]
     draft_token_ids: torch.Tensor,
@@ -695,7 +691,3 @@ def sample_recovered_tokens_kernel(
         tl.store(
             target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
             orig_prob)
-
-
-rs.expand_batch_to_tokens = expand_batch_to_tokens
-rs.rejection_sample = rejection_sample

vllm_ascend/worker/model_runner_v1.py

@@ -97,6 +97,7 @@
                              make_empty_encoder_model_runner_output)
 from vllm.v1.pool.metadata import PoolingMetadata
 from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.rejection_sampler import RejectionSampler
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
 from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
@@ -137,7 +138,6 @@
 from vllm_ascend.ops.weight_prefetch import WeightPrefetchMethod
 from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.sample.logits_processor import build_logitsprocs
-from vllm_ascend.sample.rejection_sampler import AscendRejectionSampler
 from vllm_ascend.spec_decode import get_spec_decode_method
 from vllm_ascend.spec_decode.eagle_proposer import EagleProposer
 from vllm_ascend.spec_decode.interface import SpecDcodeType
@@ -634,7 +634,7 @@ def _set_up_drafter(self):
                                              diagonal=1).to(self.device)
             if get_pp_group().is_last_rank:
                 self.drafter = self._get_drafter()
-                self.rejection_sampler = AscendRejectionSampler(self.sampler)
+                self.rejection_sampler = RejectionSampler(self.sampler)
             self.actual_seq_lengths_q = list(
                 range(self.decode_token_per_req, self.max_num_tokens + 1,
                       self.decode_token_per_req))