Update test_rejection_sampler.py

Author: ponix-j

tests/sample/test_rejection_sampler.py

@@ -8,14 +8,14 @@
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 
 from vllm_ascend.sample.rejection_sampler import (PLACEHOLDER_TOKEN_ID,
-                                                  RejectionSampler)
+                                                  AscendRejectionSampler)
 
 DEVICE = "npu"
 
 
 @pytest.fixture
 def rejection_sampler():
-    return RejectionSampler()
+    return AscendRejectionSampler()
 
 
 def create_logits_tensor(output_token_ids: list[list[int]],
@@ -423,7 +423,7 @@ def estimate_rejection_sampling_pdf(
     Returns:
         Estimated probability distribution of the output tokens.
     """
-    rejection_sampler = RejectionSampler()
+    rejection_sampler = AscendRejectionSampler()
     num_tokens = num_samples * k
     # Repeat draft probs num_samples * k times.
     draft_probs = draft_probs.reshape(1, 1,