[RFC]: Restructure the core loop to allow more asynchrony

[WoosukKwon]

Motivation.

Today, the core loop looks like:

while True:
    scheduler_output = self.scheduler.schedule()
    model_runner_output = self.model_executor.execute_model(scheduler_output)
    engine_core_outputs = self.scheduler.update_from_output(scheduler_output, model_runner_output)
    yield engine_core_outputs

While simple, this structure doesn't allow us to utilize CPU cycles while the GPU is running the model in execute_model.

Proposed Change.

The proposal is to carve out the sampling stage from the execute_model method. This way, execute_model will be a non-blocking call that returns without no GPU->CPU synchronization at the end. That is:

while True:
    scheduler_output = self.scheduler.schedule()

    # Prepare inputs and execute the model until the last hidden states. This is non-blocking.
    self.model_executor.execute_model(scheduler_output)

    # If structured outputs is used, produce the bitmask here. Otherwise, bitmask is None.
    bitmask = self.scheduler.get_grammar_bitmask(scheduler_output)

    # Prepare sampling metadata and sample the next token ids. This is blocking.
    model_runner_output = self.model_executor.sample(bitmask)

    engine_core_outputs = self.scheduler.update_from_output(scheduler_output, model_runner_output)
    yield engine_core_outputs

This gives two performance benefits:

1. (On the worker side) We can overlap the construction of sampling metadata (and logits processor) with execute_model.
2. (On the scheduler side) We can overlap the construction of bitmask with execute_model.

For async scheduling, the core loop will look like

# Initial step
scheduler_output = self.scheduler.schedule()
self.model_executor.prepare_inputs(scheduler_output)
self.model_executor.execute_model()
bitmask = self.scheduler.get_grammar_bitmask(scheduler_output)
prev_scheduler_output = scheduler_output

while True:
    scheduler_output = self.scheduler.schedule()
    self.model_executor.prepare_inputs(scheduler_output)
    model_output = self.model_executor.sample(bitmask)
    self.model_executor.execute_model()
    bitmask = self.scheduler.get_grammar_bitmask(scheduler_output)
    engine_core_outputs = self.scheduler.update_from_output(prev_scheduler_output, model_output)
    prev_scheduler_output = scheduler_output

Feedback Period.

No response

CC List.

No response

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[cadedaniel]

Two comments:

How to handle bitmask creation failure

If the grammar library fails during bitmask creation, vLLM needs to make a decision on whether the request continues to be scheduled in an undefined state (or aborted), or if vLLM should crash.

In the past, I've seen this in large scale structured outputs where complex schemas and/or output generations cause xgrammar to recurse deeply and exceed a production stack size limit. I have a PR I need to upstream (building upon #18780) that aborts requests in this case.

This is fixed in newer xgrammars but it's feasible that xgrammar could fail to step the state machine in any particular step, and in this case vLLM should abort the request instead of crashing and impacting all users.

How does this RFC interact with the stable "Scheduler" pluggable interface?

Can we more precisely define the Scheduler interface, so that it remains pluggable towards the goals of version 1.0.0? Relevant issue: #22932

[simon-mo]

Dumb question, how does this work with async scheduling?

[robertgshaw2-redhat]

How does this RFC interact with the stable "Scheduler" pluggable interface?

Can we more precisely define the Scheduler interface, so that it remains pluggable towards the goals of version 1.0.0? Relevant issue: #22932

My opinion is that we should not cut a 1.0 until we are ready to standardize the "Scheduler Interface" and are happy with keeping it in place until 2.0

[vadiklyutiy]

@WoosukKwon
This is a alternative to #23224, lets me comment here.
I see your point. But with my motivation examples - Qwen[2,2.5,3] small model with high throughput - seems there is not a lot of free CPU cycles. With one GPU scheduler and worker share the same CPU.
Below is nsys profile with my additional annotation and comment

The blue handwritten interval is free CPU cycles and seems they are not always enough.

[njhill]

Dumb question, how does this work with async scheduling?

I asked the same question :-)

1. (On the worker side) We can overlap the construction of sampling metadata (and logits processor) with execute_model.

I think this is actually independent of the proposal in this RFC, but it is a great idea regardless! I have opened a draft PR with a simple change for this: #23383.

[WoosukKwon]

If the grammar library fails during bitmask creation, vLLM needs to make a decision on whether the request continues to be scheduled in an undefined state (or aborted), or if vLLM should crash.

@cadedaniel I think the new structure can handle this failure? We can just abort the request when the scheduler finds the grammar invalid. We will waste a small amount of compute because we schedule the next step before we abort the request, but I think it's fine.

[WoosukKwon]

Dumb question, how does this work with async scheduling?

@simon-mo Added a section about it. #23391 includes async scheduling.

[WoosukKwon]

With one GPU scheduler and worker share the same CPU.

@vadiklyutiy I see. FYI, this is not the case for async scheduling though. They use different processes.