Implement Async Scheduling #19970
Implement Async Scheduling #19970
Conversation
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
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 new asynchronous scheduling mechanism to the system. It involves creating a dedicated AsyncScheduler class, making it configurable via new engine arguments and CLI options, and refactoring existing scheduler and model runner components to support its unique operational flow. The changes aim to enable more flexible and potentially optimized request processing by altering how tokens are managed and KV cache blocks are allocated in an asynchronous context.
Highlights
- New Async Scheduler Implementation: Introduced a new
AsyncSchedulerclass, inheriting from the baseScheduler, specifically designed to handle asynchronous scheduling logic. This includes distinct behaviors for token management and KV cache block allocation. - Configurability and CLI Integration: Added an
async_schedulingboolean option toSchedulerConfigandEngineArgs, which is also exposed as a command-line argument (--async-scheduling), allowing users to easily enable or disable the new asynchronous scheduling mode. - Refactored Scheduler Logic: The base
Schedulerclass has been refactored to extract common logic into new, overridable methods such as_free_encoder_inputs,update_request, andupdate_before_output. This modularization enables theAsyncSchedulerto implement its specific token and KV cache handling. - Modified Token and KV Cache Management: Adjusted how token IDs are appended and KV cache blocks are managed within both the
AsyncSchedulerand thegpu_model_runner. In async mode, token updates and block caching are handled differently, often in a more deferred or single-token-at-a-time manner. - Batching Strategy Adaptation: The
multiproc_executornow setsmax_concurrent_batchesto2when asynchronous scheduling is enabled, indicating a tailored batching strategy for this new mode. - Temporary Distributed Executor Backend Fix: A temporary
# FIXMEchange was added tovllm/config.pyto default thedistributed_executor_backendtomp(multiprocessing) whenworld_sizeis1, overriding the previousuni(unified) setting.
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
This pull request introduces asynchronous scheduling, which is a significant feature. The changes are well-structured, with a new AsyncScheduler and modifications to support it throughout the codebase. The use of a feature flag async_scheduling is a good approach for introducing this functionality.
I've identified a couple of points for improvement:
- A
FIXMEcomment invllm/config.pycould use more context. - There appears to be a redundant method call in
vllm/v1/core/sched/async_scheduler.pythat could lead to inefficiency.
Overall, this is a solid implementation. Addressing these points will improve the code's clarity and correctness.
|
👋 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 🚀 |
vllm/v1/core/sched/scheduler.py
Outdated
| @@ -253,7 +261,9 @@ def schedule(self) -> SchedulerOutput: | |||
| request, | |||
| num_new_tokens, | |||
| num_draft_tokens=num_draft_tokens, | |||
| num_lookahead_tokens=self.num_lookahead_tokens) | |||
| num_lookahead_tokens=self.num_lookahead_tokens, | |||
| delay_cache_blocks=self.is_async, | |||
There was a problem hiding this comment.
do we want to preserve this behavior:
vllm/docs/design/v1/prefix_caching.md
Line 143 in e6327c9
There was a problem hiding this comment.
Good point. In the current version of this PR I make sure that the deduplication in between the same batch happens. Please check out the test_prefix_caching_for_prefill_dedup test.
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
…vllm into woosuk/async-sched
| if self.speculative_config is not None: | ||
| raise ValueError( | ||
| "Currently, speculative decoding is not supported with " | ||
| "async scheduling.") |
There was a problem hiding this comment.
nit: should we be checking for structured output backends here too? I assume those are also incompatible with async scheduling
There was a problem hiding this comment.
IIUC, structured outputs are runtime parameter per request, so it's tricky to detect it when initializing the server. 😅 What about printing a warning msg for now?
There was a problem hiding this comment.
sorry should have been more clear!; I guess what I meant was: could we check if the server is started with --guided-decoding-backend= set here? (my bad should have looked up what that config name was before commenting, I think this sets StructuredOutputManager.backend; very out of my depth here haha)
There was a problem hiding this comment.
Will we reject requests that use structured output?
There was a problem hiding this comment.
I want to look into this issue and try to make asynchronous scheduling work together with StructuredOutput.
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: avigny <47987522+avigny@users.noreply.github.com>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: x22x22 <wadeking@qq.com>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Diego-Castan <diego.castan@ibm.com>
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This PR implements an initial version of the asynchronous scheduler, similar to the approach described in the NanoFlow paper. The primary goal is to minimize scheduler overhead by overlapping scheduling with model execution. This is achieved by making the scheduler operate one step ahead of the current execution.
Feature Compatibility
Async scheduling is introduced as an experimental feature and can be enabled using the
--async-schedulingflag. It has currently only been validated with basic setups and has not yet been tested in the following scenarios:At this stage, async scheduling does not support:
Most of these limitations are potentially addressable in future iterations. For example, I believe we can make most speculative decoding methods (such as EAGLE) compatible with async scheduling.
Implementation
To reuse code and minimize disruption, the async scheduler is implemented as a subclass of the existing scheduler. The key addition is an "output placeholder," representing tokens scheduled but not yet generated.
One limitation of the current design is that the scheduler and GPU worker must run in separate processes for parallel execution. This setup causes extra serialization overhead for input/output data (which is unnecessary in single-gpu cases). This may impact performance, especially for multimodal models with large inputs.
Performance
Async scheduling increases throughput by 3-15%. The speedup is more pronounced with smaller models and larger batch sizes. No significant latency improvement is observed for batch size 1 scenarios, as scheduler overhead in those cases is already minimal.
In terms of serving latency, async scheduling generally reduces TPOT but slightly increases TTFT. The TTFT increases because incoming requests must wait for an additional scheduling step before it is actually processed on GPUs. Despite this, the overall end-to-end latency typically gets faster.