[V1][Bugfix][Spec Decode] Fix incorrect outputs in V1 speculative decoding due to batch indexing

[benchislett]

In gpu_model_runner.py, the requests in the batch are ordered according to self.input_batch.req_ids. However, the speculative decoding code is inferring the batch order from the keys of the dictionary scheduler_output.num_scheduled_tokens which is not necessarily ordered in the same way.

When multiple requests in the batch have different speculative lengths, the grouping of output probability slices into speculative decoding requests depends on this ordering. The existing end-to-end test is insufficient to capture this effect as the batch size is not large enough and the requests are too similar. I have updated the E2E test with a large batch of mixed requests, some of which will have full ngram completion and some which will not. Running this test on the main branch leads to many failures where one completion will have tokens "leaked" from other requests.

This one-line fix resolves the issue completely.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[WoosukKwon]

@benchislett Thanks for the bug fix!

[WoosukKwon]

@LiuXiaoxuanPKU Please take a look!

[LiuXiaoxuanPKU]

LGTM, thanks for the fix!