[V1][Bugfix][Spec Decode] Fix incorrect outputs in V1 speculative decoding due to batch indexing #14645
Conversation
Signed-off-by: Benjamin Chislett <[email protected]>
Signed-off-by: Benjamin Chislett <[email protected]>
benchislett
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
|
👋 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 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 🚀 |
WoosukKwon
left a comment
WoosukKwon
left a comment
There was a problem hiding this comment.
@benchislett Thanks for the bug fix!
|
@LiuXiaoxuanPKU Please take a look! |
LiuXiaoxuanPKU
left a comment
LiuXiaoxuanPKU
left a comment
There was a problem hiding this comment.
LGTM, thanks for the fix!
…oding due to batch indexing (vllm-project#14645) Signed-off-by: Benjamin Chislett <[email protected]> Signed-off-by: Richard Liu <[email protected]>
…oding due to batch indexing (vllm-project#14645) Signed-off-by: Benjamin Chislett <[email protected]> Signed-off-by: Louis Ulmer <[email protected]>
…oding due to batch indexing (vllm-project#14645) Signed-off-by: Benjamin Chislett <[email protected]>
…oding due to batch indexing (vllm-project#14645) Signed-off-by: Benjamin Chislett <[email protected]> Signed-off-by: Mu Huai <[email protected]>
3 participants
In
gpu_model_runner.py, the requests in the batch are ordered according toself.input_batch.req_ids. However, the speculative decoding code is inferring the batch order from the keys of the dictionaryscheduler_output.num_scheduled_tokenswhich 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.