[Refactor][Triton] Move reject sample triton kernels into ops/triton

[whx-sjtu]

What this PR does / why we need it?

This PR moves reject sample related triton kernels into ops/triton.

Does this PR introduce any user-facing change?

No

How was this patch tested?

CI passed with existing test.

• vLLM version: release/v0.13.0
• vLLM main: vllm-project/vllm@5fbfa8d

[gemini-code-assist]

Code Review

This pull request refactors the rejection sampling Triton kernels by moving them into a new, dedicated file vllm_ascend/ops/triton/reject_sample.py. This is a good change for code organization and modularity. The calling code in vllm_ascend/sample/rejection_sampler.py has been updated accordingly.

During the review, I identified two critical bugs within the newly added (but previously existing) Triton kernels. Both rejection_greedy_sample_triton and rejection_random_sample_kernel incorrectly use Python for loops with runtime-dependent bounds. This is not supported by Triton and will cause errors. I have provided detailed comments and code suggestions to resolve these issues by using tl.range with compile-time constant bounds and appropriate masking.

[gemini-code-assist]

The loop for i in range(num_tokens1): uses a runtime variable num_tokens1 as its bound. In Triton, Python for loops are unrolled at compile time and require their bounds to be compile-time constants. Using a runtime variable here is incorrect and can lead to errors or unexpected behavior.

To fix this, you should use tl.range with a compile-time constant bound, such as max_spec_len, and mask the operations inside the loop. For max_spec_len to be a compile-time constant, you also need to remove it from the do_not_specialize list in the @triton.jit decorator for this function (line 75).

Suggested change

	for i in range(num_tokens1):
	if not rejected:
	draft_token_id = tl.load(draft_token_ids_ptr + start_idx1 + i)
	target_argmax_id = tl.load(target_argmax_ptr + start_idx1 + i)
	tl.store(
	output_token_ids_ptr + position * (max_spec_len + 1) + i,
	target_argmax_id,
	)
	if draft_token_id != target_argmax_id:
	# Reject.
	rejected = True
	for i in tl.range(0, max_spec_len):
	if i < num_tokens1 and not rejected:
	draft_token_id = tl.load(draft_token_ids_ptr + start_idx1 + i)
	target_argmax_id = tl.load(target_argmax_ptr + start_idx1 + i)
	tl.store(
	output_token_ids_ptr + position * (max_spec_len + 1) + i,
	target_argmax_id,
	)
	if draft_token_id != target_argmax_id:
	# Reject.
	rejected = True

[gemini-code-assist]

The loop for pos in range(num_draft_tokens): uses a runtime variable num_draft_tokens as its bound. In Triton, Python for loops are unrolled at compile time and require their bounds to be compile-time constants. This is a critical issue that can lead to errors.

To fix this, you should use tl.range with a compile-time constant bound, like max_spec_len, and add a mask pos < num_draft_tokens inside the loop. You will also need to make max_spec_len a compile-time constant by removing it from the do_not_specialize list in the @triton.jit decorator for this function (line 131).

    for pos in tl.range(0, max_spec_len):
        if pos < num_draft_tokens and not rejected:
            draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
            if NO_DRAFT_PROBS:
                draft_prob = 1
            else:
                draft_prob = tl.load(draft_probs_ptr +
                                     (start_idx + pos) * vocab_size +
                                     draft_token_id)
            target_prob = tl.load(target_probs_ptr +
                                  (start_idx + pos) * vocab_size +
                                  draft_token_id)
            uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
            if draft_prob > 0 and target_prob / draft_prob >= uniform_prob:
                # Accept
                token_id = draft_token_id
            else:
                # Reject. Use recovered token
                rejected = True
                token_id = tl.load(recovered_token_ids_ptr + start_idx + pos)
            tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos,
                     token_id)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[whx-sjtu]

https://github.com/vllm-project/vllm-ascend/actions/runs/20513926312/job/58938521521?pr=5324

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.