[Bugfix] Fix test_sample_recovered_tokens_pytorch_autoregressive meth…

[lio1226]

…od in test_rejection_sampler

What this PR does / why we need it?

This 'test_rejection_sampler' unit test is something wrong.

def test_sample_recovered_tokens_pytorch_autoregressive(self):
output_token_ids = torch.empty(2, dtype=torch.int32)
cu_num_draft_tokens = torch.tensor([1, 1])
draft_token_ids = torch.tensor([0, 1])

len(draft_token_ids ) = 2, cu_num_draft_tokens should be torch.tensor([1, 2]) or torch.tensor([2, 2])

I fix it and set cu_num_draft_tokens = torch.tensor([1, 2]). The methods before and after optimization can pass.

Does this PR introduce any user-facing change?

No

How was this patch tested?

NA

• vLLM version: v0.11.0rc3
• vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.0

[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.

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in the test_sample_recovered_tokens_pytorch_autoregressive unit test. The value of cu_num_draft_tokens was incorrect, causing the test to not properly check behavior for the second request in the batch. The provided fix corrects this value, making the test logic sound. I have one suggestion to improve the completeness of this test.

[gemini-code-assist]

This fix is correct and enables the test to properly process both requests in the batch. However, the test currently only asserts the output for the first request. To ensure the function behaves correctly for all cases covered by this test, an assertion for the second request's output should be added. Based on my analysis of the test data, output_token_ids[1] should be 0. Please consider adding assert output_token_ids[1].item() == 0 to make the test more robust.