triton code modify rejection_sampler.py

Signed-off-by: yuxingcyx <[email protected]>

Author: yuxingcyx

vllm_ascend/sample/rejection_sampler.py

@@ -616,7 +616,7 @@ def rejection_greedy_sample_spec_len_1_triton(
     target_argmax_id = tl.load(target_argmax_ptr + offset, mask)
     tl.store(output_token_ids_ptr + offset * 2, target_argmax_id, mask)
 
-    for pos in tl.arange(0, BLOCK_SIZE):
+    for pos in tl.range(0, BLOCK_SIZE):
         draft_token_id1 = tl.get_element(draft_token_id, (pos, ))
         target_argmax1 = tl.get_element(target_argmax_id, (pos, ))
         position = block_idx * BLOCK_SIZE + pos