[Core][Speculative Decoding] Add multi-query verifier for speculative decoding without batch expansion

[sighingnow]

FIX P1 part of #4630 (link existing issues this PR will resolve)

BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE

---

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

• [Bugfix] for bug fixes.
• [CI/Build] for build or continuous integration improvements.
• [Doc] for documentation fixes and improvements.
• [Model] for adding a new model or improving an existing model. Model name should appear in the title.
• [Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
• [Kernel] for changes affecting CUDA kernels or other compute kernels.
• [Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
• [Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
• [Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

• We adhere to Google Python style guide and Google C++ style guide.
• Pass all linter checks. Please use format.sh to format your code.
• The code need to be well-documented to ensure future contributors can easily understand the code.
• Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
• Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

• After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
• After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
• After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
• Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

[mpjlu]

Have you tested this PR:#5691, which performance is better?

[sighingnow]

Have you tested this PR:#5691, which performance is better?

These two PRs shouldn't have too many differences in implementation, so the performance should be roughly the same.

One comment for #5691: it must be done carefully for prompt tokens processing in sampling_metadata.py, as the prompt tokens of sequences would affect the applying penalities logic in sampler.

[mpjlu]

@sighingnow Thanks for your kindly reply. Does this PR support cuda graph now.

[sighingnow]

@sighingnow Thanks for your kindly reply. Does this PR support cuda graph now.

Cuda graph has been supported and tested.

Hi @mpjlu I have just tried some correctness test and found that this PR cannot work as expected with cuda graph.

Apologize for the misinformation.

[cadedaniel]

I will take a pass this week. also cc @LiuXiaoxuanPKU

one major challenge is cuda graph support. it will be necessary since the pytorch dispatch and scheduling overhead is large for small batch sizes.

btw if you have performance numbers to prove ^ wrong, please share

[sighingnow]

@cadedaniel @mpjlu @LiuXiaoxuanPKU The cuda graph of both draft model and target model (with multi-query verifier) has been added.

[cadedaniel]

Hi @sighingnow , thanks for adding this. Feedback:

• The MultiQueryTop1Scorer is really good
• For the cuda graph support, can we instead integrate with @LiuXiaoxuanPKU 's work in [Kernel] Unify the kernel used in flash attention backend #6052 ? More details here. The reason is that we can use the same backend for chunked prefill / prefix caching / other MQA applications. With this PR, the speculation logic must spill into the model runner, which couples the two together.

One way forward is to merge the MultiQueryTop1Scorer without cuda graph support, then once #6052 is done we can integrate the two.

What do you think?

[sighingnow]

Hi @cadedaniel,

If the PR #6052 will be merged soon, I prefer waiting it been merged then rebasing this PR, to avoid conflicts and merging effort.

[cadedaniel]

I am not sure how soon it will be merged. cc @LiuXiaoxuanPKU

[LiuXiaoxuanPKU]

Hi @sighingnow, sorry for the very late reply.

After digging into #6052, we find there is correctness issue with cudagraph + flash_attn_varlen_func. We will look into the support. However, at the same time, we can move forward with this PR. Could you clean a bit and let us know when it's good to review. Thanks for the contribution and sorry for the delayed reply!

[cadedaniel]

cc @comaniac any thoughts on how this cuda graph approach works with model runner / prepare inputs ?

[comaniac]

Took a brief look. The approach is ok in general, but I could see that some more overheads are introduced due to more tensors/logic being processed. I could review the PR once it's rebased and ready to eliminate the overheads as possible.

[cadedaniel]

One alternative is to move this to a custom model runner, just for spec decode. Do you think that's better or worse than the current approach?

[comaniac]

One alternative is to move this to a custom model runner, just for spec decode. Do you think that's better or worse than the current approach?

In general I imagine that would have lots of code duplications, and this logic (without CUDA graph) should be very similar as prefix caching. Another way is maybe we quickly support a version without CUDA graph first and see how good/bad is?

[cadedaniel]

SGTM. @sighingnow we can merge this PR as is if we remove the cuda graph stuff, or we can add cuda graph stuff (using one of the two approaches @comaniac mentions above)

[sighingnow]

I will rebase and fixes the cudagraph issue, add testsing cases for this PR during this weekend, thanks all for your careful review.

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @sighingnow.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork