Implement Async Scheduling

[WoosukKwon]

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-scheduling flag. It has currently only been validated with basic setups and has not yet been tested in the following scenarios:

• KV connector
• Data parallel

At this stage, async scheduling does not support:

• Speculative decoding
• Structured outputs
• Pipeline parallelism

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.

[gemini-code-assist]

Summary of Changes

Hello @WoosukKwon, I'm Gemini Code Assist¹! 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 AsyncScheduler class, inheriting from the base Scheduler, 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_scheduling boolean option to SchedulerConfig and EngineArgs, 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 Scheduler class has been refactored to extract common logic into new, overridable methods such as _free_encoder_inputs, update_request, and update_before_output. This modularization enables the AsyncScheduler to 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 AsyncScheduler and the gpu_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_executor now sets max_concurrent_batches to 2 when asynchronous scheduling is enabled, indicating a tailored batching strategy for this new mode.
• Temporary Distributed Executor Backend Fix: A temporary # FIXME change was added to vllm/config.py to default the distributed_executor_backend to mp (multiprocessing) when world_size is 1, overriding the previous uni (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

1. 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. ↩

[gemini-code-assist]

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:

1. A FIXME comment in vllm/config.py could use more context.
2. There appears to be a redundant method call in vllm/v1/core/sched/async_scheduler.py that could lead to inefficiency.

Overall, this is a solid implementation. Addressing these points will improve the code's clarity and correctness.

[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.

🚀

[LucasWilkinson]

do we want to preserve this behavior:

vllm/docs/design/v1/prefix_caching.md

Line 143 in e6327c9

4. If an allocated block is already full of tokens, we immediately add it to the Cache Block, so that the block can be reused by other requests in the same batch.

in the asynchronous case?

[WoosukKwon]

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.

[mergify]

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

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

[LucasWilkinson]

Clean! Looks goods! Left on nit otherwise LGTM

[LucasWilkinson]

nit: should we be checking for structured output backends here too? I assume those are also incompatible with async scheduling

[WoosukKwon]

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?

[LucasWilkinson]

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)

[mgoin]

Will we reject requests that use structured output?

[chaunceyjiang]

I want to look into this issue and try to make asynchronous scheduling work together with StructuredOutput.