[Optimization] Cache sampled token ids in model runner #20291
Conversation
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
…vllm into woosuk/async-sched
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
WoosukKwon
requested review from
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
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 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @WoosukKwon, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces a significant optimization to the VLLM system by changing how sampled token IDs are handled between the scheduler and the GPU model runner. Previously, the scheduler always sent these IDs, even when unnecessary. Now, the model runner directly caches sampled token IDs, eliminating the need for the scheduler to transmit them, except in scenarios involving Pipeline Parallelism (PP). This change reduces communication overhead and serialization, leading to general performance improvements and laying groundwork for future features.
Highlights
- Optimized Token ID Handling: The system now caches sampled token IDs directly within the GPU model runner, reducing the need for the scheduler to transmit them in most cases, thereby minimizing communication overhead.
- Conditional Communication for Pipeline Parallelism (PP): The scheduler will only send sampled token IDs to the model runner when Pipeline Parallelism (PP) is active, maintaining necessary communication for distributed setups while optimizing for non-PP configurations.
- Reduced Serialization Overhead: By minimizing the data transferred between the scheduler and model runner, this change significantly reduces serialization overhead, contributing to overall performance gains.
- Foundation for Future Features: This optimization is a prerequisite step for implementing future features, specifically referencing #19970.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
The code changes introduce the ability to cache sampled token IDs within the model runner, which reduces communication overhead with the scheduler. I've identified a critical issue that could cause an UnboundLocalError and some redundant logic that could be simplified.
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
vllm/v1/worker/gpu_model_runner.py
Outdated
|
|
||
| start_idx = self.input_batch.num_tokens_no_spec[req_idx] | ||
| end_idx = start_idx + len(sampled_ids) | ||
| if end_idx >= self.max_model_len: |
There was a problem hiding this comment.
nit: should this be an assert? shouldnt the scheduler never schedule anything greater than max_model_len? if it can do we need to handle multiple sampled tokens here (i.e. spec decode), like do end_idx = min(end_idx, self.max_model_len ).
Also since we use this as the non-inclusive end to a slice wouldnt it be ok if end_idx == self.max_model_len
There was a problem hiding this comment.
Nice catch! Added assert end_idx <= self.max_model_len and fixed a small bug in the scheduler.
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
…20291) Signed-off-by: Woosuk Kwon <[email protected]>
…20291) Signed-off-by: Woosuk Kwon <[email protected]>
…20291) Signed-off-by: Woosuk Kwon <[email protected]>
# Description This branch has a fix for: - Caching the token_ids (now the new tokens are cached in `execute_model` instead of `update_states`. This is because of vllm-project/vllm#20291. ) - Changes from the `CachedRequestData` (#273) ## Related Issues Fix for #271 --------- Signed-off-by: Prashant Gupta <[email protected]> Signed-off-by: Max de Bayser <[email protected]> Co-authored-by: Max de Bayser <[email protected]>
Signed-off-by: Chendi Xue <[email protected]>
…20291) Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: avigny <[email protected]>
Currently, the GPU model runner always fetches new token IDs from the scheduler without caching them after they are sampled. This behavior exists because, in pipeline parallelism (PP), there is no direct communication channel between the first-stage and last-stage workers, making it necessary to route tokens through the scheduler.
However, this approach is suboptimal due to the communication overhead it introduces (although they are small). This PR addresses the issue by caching the sampled token IDs within the model runner. As a result, the scheduler no longer needs to send new token IDs unless PP is being used.
While this change is a required step for implementing #19970, it also provides a general performance improvement by slightly reducing serialization overhead.