Update rejection_sampler.py

Author: ponix-j

vllm_ascend/sample/rejection_sampler.py

@@ -3,6 +3,7 @@
 
 import torch
 import torch.nn as nn
+import vllm.v1.sample.rejection_sampler as rs
 from vllm.logger import init_logger
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.rejection_sampler import (RejectionSampler, compute_probs,
@@ -450,3 +451,5 @@ def sample_recovered_tokens_pytorch(
 
             if IS_NGRAM:
                 target_probs[token_idx, draft_token_id] = orig_prob
+
+rs.expand_batch_to_tokens = expand_batch_to_tokens