Implement Async Scheduling #19970

WoosukKwon · 2025-06-23T06:49:19Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

This PR implements async scheduler with minimal code modifications.

Test Plan

Test Result

(Optional) Documentation Update

Signed-off-by: Woosuk Kwon <[email protected]>

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.

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:

A FIXME comment in vllm/config.py could use more context.
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.

gemini-code-assist · 2025-06-23T06:51:08Z

vllm/v1/core/sched/async_scheduler.py

+            self._free_encoder_inputs(request)
+


The _free_encoder_inputs method is called here in update_before_output. However, the base Scheduler class already calls this method in update_from_output. Since AsyncScheduler inherits from Scheduler and doesn't override update_from_output, this will result in _free_encoder_inputs being called twice for each request.

While the method might be idempotent, this is inefficient and can be confusing. It seems this call should be removed from update_before_output to avoid the redundant call.

vllm/config.py

github-actions · 2025-06-23T06:54:45Z

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

🚀

lgeiger · 2025-06-23T12:14:26Z

vllm/v1/core/sched/scheduler.py

+        # OPTIMIZATION: Avoid list(set) if the set is empty.
+        if cached_encoder_input_ids:
+            for input_id in list(cached_encoder_input_ids):
+                mm_positions = request.mm_positions[input_id]
+                start_pos = mm_positions.offset
+                num_tokens = mm_positions.length
+                if start_pos + num_tokens <= request.num_computed_tokens:
+                    # The encoder output is already processed and stored
+                    # in the decoder's KV cache.
+                    self.encoder_cache_manager.free_encoder_input(
+                        request, input_id)


Not really related to this PR, but I think we can completely remove the list(set) call so which makes it a bit more readable and remove the need for the check

Suggested change

# OPTIMIZATION: Avoid list(set) if the set is empty.

if cached_encoder_input_ids:

for input_id in list(cached_encoder_input_ids):

mm_positions = request.mm_positions[input_id]

start_pos = mm_positions.offset

num_tokens = mm_positions.length

if start_pos + num_tokens <= request.num_computed_tokens:

# The encoder output is already processed and stored

# in the decoder's KV cache.

self.encoder_cache_manager.free_encoder_input(

request, input_id)

for input_id in cached_encoder_input_ids:

mm_positions = request.mm_positions[input_id]

start_pos = mm_positions.offset

num_tokens = mm_positions.length

if start_pos + num_tokens <= request.num_computed_tokens:

# The encoder output is already processed and stored

# in the decoder's KV cache.

self.encoder_cache_manager.free_encoder_input(

request, input_id)

vllm/v1/core/sched/async_scheduler.py

LucasWilkinson · 2025-06-23T20:24:09Z

vllm/v1/core/sched/scheduler.py

@@ -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,


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?

Signed-off-by: Woosuk Kwon <[email protected]>

WoosukKwon added 2 commits June 22, 2025 23:46

Implement Async Scheduler

487efdc

Signed-off-by: Woosuk Kwon <[email protected]>

Merge branch 'main' into woosuk/async-sched

a3c320d

mergify bot added the v1 label Jun 23, 2025

gemini-code-assist bot reviewed Jun 23, 2025

View reviewed changes

lgeiger reviewed Jun 23, 2025

View reviewed changes

LucasWilkinson reviewed Jun 23, 2025

View reviewed changes

vllm/v1/core/sched/async_scheduler.py Outdated Show resolved Hide resolved

LucasWilkinson reviewed Jun 23, 2025

View reviewed changes

optimization

bde64a4

Signed-off-by: Woosuk Kwon <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

Implement Async Scheduling #19970

Implement Async Scheduling #19970

WoosukKwon commented Jun 23, 2025 •

edited by github-actions bot

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jun 23, 2025

Uh oh!

Uh oh!

github-actions bot commented Jun 23, 2025

Uh oh!

lgeiger Jun 23, 2025

Uh oh!

Uh oh!

LucasWilkinson Jun 23, 2025

Uh oh!

Uh oh!

Uh oh!

Implement Async Scheduling #19970

Are you sure you want to change the base?

Implement Async Scheduling #19970

Conversation

WoosukKwon commented Jun 23, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jun 23, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Jun 23, 2025

Uh oh!

lgeiger Jun 23, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

LucasWilkinson Jun 23, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

WoosukKwon commented Jun 23, 2025 •

edited by github-actions bot

Loading