@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d4d309409f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 02:15 UTC

Implement Async Scheduling (#19970)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Jul 14, 2025 at 23:01 UTC d4d309409f2396e68e4b5a67ede194913502388b
11 files changed +508 -148
tests/v1/core/__init__.py
tests/v1/core/test_async_scheduler.py new
+228
@@ -0,0 +1,228 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections import deque
4 +
5 +import pytest
6 +
7 +from vllm.v1.core.sched.output import SchedulerOutput
8 +from vllm.v1.outputs import ModelRunnerOutput
9 +from vllm.v1.request import RequestStatus
10 +
11 +from .utils import create_requests, create_scheduler
12 +
13 +
14 +def _make_model_runner_output(
15 + scheduler_output: SchedulerOutput, ) -> ModelRunnerOutput:
16 + req_ids = list(scheduler_output.num_scheduled_tokens.keys())
17 + return ModelRunnerOutput(
18 + req_ids=req_ids,
19 + req_id_to_index={
20 + req_id: i
21 + for i, req_id in enumerate(req_ids)
22 + },
23 + sampled_token_ids=[[i] for i in range(len(req_ids))],
24 + spec_token_ids=None,
25 + logprobs=None,
26 + prompt_logprobs_dict={},
27 + pooler_output=[],
28 + )
29 +
30 +
31 +@pytest.mark.parametrize("max_tokens", [1, 2, 3, 5])
32 +def test_stop_by_max_tokens(max_tokens: int):
33 + scheduler = create_scheduler(async_scheduling=True)
34 + requests = create_requests(num_requests=2, max_tokens=max_tokens)
35 + req0, req1 = requests
36 +
37 + sched_outputs: deque[SchedulerOutput] = deque()
38 + scheduler.add_request(req0)
39 + sched_outputs.append(scheduler.schedule())
40 +
41 + scheduler.add_request(req1)
42 + sched_outputs.append(scheduler.schedule())
43 +
44 + while sched_outputs:
45 + sched_output = sched_outputs.popleft()
46 + model_runner_output = _make_model_runner_output(sched_output)
47 + scheduler.update_from_output(sched_output, model_runner_output)
48 +
49 + sched_output = scheduler.schedule()
50 + if sched_output.num_scheduled_tokens:
51 + sched_outputs.append(sched_output)
52 +
53 + assert scheduler.get_num_unfinished_requests() == 0
54 + assert req0.num_output_tokens == max_tokens
55 + assert req1.num_output_tokens == max_tokens
56 +
57 +
58 +def test_abort():
59 + scheduler = create_scheduler(async_scheduling=True)
60 + requests = create_requests(num_requests=10, max_tokens=20)
61 +
62 + for req in requests:
63 + scheduler.add_request(req)
64 +
65 + sched_outputs: deque[SchedulerOutput] = deque()
66 + sched_outputs.append(scheduler.schedule())
67 + sched_outputs.append(scheduler.schedule())
68 +
69 + abort_order = [0, 8, 3, 1, 6, 4, 2, 5, 7, 9]
70 + abort_order_copy = abort_order.copy()
71 +
72 + def abort_request():
73 + if not abort_order:
74 + return
75 + req = requests[abort_order.pop(0)]
76 + scheduler.finish_requests(req.request_id,
77 + RequestStatus.FINISHED_ABORTED)
78 +
79 + while sched_outputs:
80 + # Abort a scheduled request.
81 + abort_request()
82 + sched_output = sched_outputs.popleft()
83 + model_runner_output = _make_model_runner_output(sched_output)
84 + scheduler.update_from_output(sched_output, model_runner_output)
85 +
86 + sched_output = scheduler.schedule()
87 + if sched_output.num_scheduled_tokens:
88 + sched_outputs.append(sched_output)
89 +
90 + for i, req in enumerate(requests):
91 + assert req.status == RequestStatus.FINISHED_ABORTED
92 + assert req.num_output_tokens == abort_order_copy.index(i)
93 +
94 +
95 +def test_preempt():
96 + scheduler = create_scheduler(async_scheduling=True)
97 + requests = create_requests(num_requests=10, max_tokens=20)
98 +
99 + for req in requests:
100 + scheduler.add_request(req)
101 +
102 + sched_outputs: deque[SchedulerOutput] = deque()
103 + sched_outputs.append(scheduler.schedule())
104 + sched_outputs.append(scheduler.schedule())
105 +
106 + abort_order = [0, 8, 3, 1, 6, 4, 2, 5, 7, 9]
107 + abort_order_copy = abort_order.copy()
108 +
109 + def abort_request():
110 + if not abort_order:
111 + return
112 + req = requests[abort_order.pop(0)]
113 + scheduler.finish_requests(req.request_id,
114 + RequestStatus.FINISHED_ABORTED)
115 +
116 + while sched_outputs:
117 + # Abort a scheduled request.
118 + abort_request()
119 + sched_output = sched_outputs.popleft()
120 + model_runner_output = _make_model_runner_output(sched_output)
121 + scheduler.update_from_output(sched_output, model_runner_output)
122 +
123 + sched_output = scheduler.schedule()
124 + if sched_output.num_scheduled_tokens:
125 + sched_outputs.append(sched_output)
126 +
127 + for i, req in enumerate(requests):
128 + assert req.status == RequestStatus.FINISHED_ABORTED
129 + assert req.num_output_tokens == abort_order_copy.index(i)
130 +
131 +
132 +def test_prefix_caching_for_prefill_dedup():
133 + CHUNK_SIZE = 1000
134 + BLOCK_SIZE = 16
135 + num_prompt_tokens = 100
136 + scheduler = create_scheduler(async_scheduling=True,
137 + max_num_batched_tokens=CHUNK_SIZE,
138 + enable_prefix_caching=True,
139 + block_size=BLOCK_SIZE)
140 + requests = create_requests(num_requests=5,
141 + num_tokens=num_prompt_tokens,
142 + max_tokens=3,
143 + same_prompt=True)
144 + requests_copy = requests.copy()
145 +
146 + # Two requests with the same prompt.
147 + req0 = requests.pop(0)
148 + req1 = requests.pop(0)
149 + scheduler.add_request(req0)
150 + scheduler.add_request(req1)
151 +
152 + sched_outputs: deque[SchedulerOutput] = deque()
153 + sched_output = scheduler.schedule()
154 + sched_outputs.append(sched_output)
155 + # Make sure prefix caching de-duplicates the prompts in the same step,
156 + # so all the blocks except the last are shared between the two requests.
157 + assert len(sched_output.num_scheduled_tokens) == 2
158 + num_blocks = num_prompt_tokens // BLOCK_SIZE
159 + assert req0.num_cached_tokens == 0
160 + assert req1.num_cached_tokens >= num_blocks * BLOCK_SIZE
161 +
162 + sched_outputs.append(scheduler.schedule())
163 + while sched_outputs:
164 + if requests:
165 + scheduler.add_request(requests.pop(0))
166 + sched_output = sched_outputs.popleft()
167 + model_runner_output = _make_model_runner_output(sched_output)
168 + scheduler.update_from_output(sched_output, model_runner_output)
169 + sched_output = scheduler.schedule()
170 + if sched_output.num_scheduled_tokens:
171 + sched_outputs.append(sched_output)
172 +
173 + # Other requests scheduled after the two requests should also get
174 + # prefix cache hit.
175 + assert scheduler.get_num_unfinished_requests() == 0
176 + for req in requests_copy[1:]:
177 + assert req.num_cached_tokens >= num_blocks * BLOCK_SIZE
178 +
179 +
180 +def test_prefix_caching_for_multi_turn():
181 + CHUNK_SIZE = 1000
182 + BLOCK_SIZE = 16
183 + num_prompt_tokens = 100
184 + num_output_tokens = 200
185 + scheduler = create_scheduler(async_scheduling=True,
186 + max_num_batched_tokens=CHUNK_SIZE,
187 + enable_prefix_caching=True,
188 + block_size=BLOCK_SIZE)
189 + requests = create_requests(num_requests=5,
190 + num_tokens=num_prompt_tokens,
191 + max_tokens=num_output_tokens)
192 +
193 + for req in requests:
194 + scheduler.add_request(req)
195 + sched_outputs: deque[SchedulerOutput] = deque()
196 + sched_outputs.append(scheduler.schedule())
197 + sched_outputs.append(scheduler.schedule())
198 +
199 + # Process the requests.
200 + while sched_outputs:
201 + sched_output = sched_outputs.popleft()
202 + model_runner_output = _make_model_runner_output(sched_output)
203 + scheduler.update_from_output(sched_output, model_runner_output)
204 + sched_output = scheduler.schedule()
205 + if sched_output.num_scheduled_tokens:
206 + sched_outputs.append(sched_output)
207 + assert scheduler.get_num_unfinished_requests() == 0
208 +
209 + # Create next-turn requests whose prompts are the full output of the
210 + # previous turn.
211 + next_turn_requests = create_requests(
212 + num_requests=5,
213 + num_tokens=num_prompt_tokens + num_output_tokens,
214 + max_tokens=num_output_tokens,
215 + )
216 + for i, req in enumerate(next_turn_requests):
217 + req.prompt_token_ids = (requests[i].prompt_token_ids +
218 + list(requests[i].output_token_ids))
219 + # Schedule the next-turn requests.
220 + for req in next_turn_requests:
221 + scheduler.add_request(req)
222 + sched_outputs.append(scheduler.schedule())
223 +
224 + # Make sure the next-turn requests get prefix cache hit by the previous
225 + # requests.
226 + for req in next_turn_requests:
227 + assert (req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE *
228 + BLOCK_SIZE)
tests/v1/core/test_scheduler.py
+1 -127
@@ -19,133 +19,7 @@ from vllm.v1.request import Request, RequestStatus
19 from vllm.v1.structured_output import StructuredOutputManager
20 from vllm.v1.structured_output.request import StructuredOutputRequest
21
22 -EOS_TOKEN_ID = 50256
23 -
24 -
25 -def create_scheduler(
26 - model: str = "facebook/opt-125m",
27 - max_num_seqs: int = 16,
28 - max_num_batched_tokens: int = 8192,
29 - enable_prefix_caching: Optional[bool] = None,
30 - long_prefill_token_threshold: int = 0,
31 - disable_chunked_mm_input: bool = False,
32 - use_kv_connector: bool = False,
33 - num_blocks: int = 10000,
34 - block_size: int = 16,
35 - max_model_len: Optional[int] = None,
36 - num_speculative_tokens: Optional[int] = None,
37 - skip_tokenizer_init: bool = False,
38 -) -> Scheduler:
39 - '''Create scheduler under test.
40 -
41 - Args:
42 - model: model under test
43 - max_num_seqs: max sequences to schedule
44 - max_num_batch_tokens: max num tokens to batch
45 - enable_prefix_caching: optionally force APC config
46 - (True/False) or use default
47 - (None)
48 -
49 - Returns:
50 - {class}`Scheduler` instance
51 - '''
52 - if max_model_len is None:
53 - max_model_len = max_num_batched_tokens
54 - scheduler_config = SchedulerConfig(
55 - max_num_seqs=max_num_seqs,
56 - max_num_batched_tokens=max_num_batched_tokens,
57 - max_model_len=max_model_len,
58 - long_prefill_token_threshold=long_prefill_token_threshold,
59 - disable_chunked_mm_input=disable_chunked_mm_input,
60 - enable_chunked_prefill=True,
61 - )
62 - model_config = ModelConfig(
63 - model=model,
64 - task="auto",
65 - tokenizer=model,
66 - tokenizer_mode="auto",
67 - trust_remote_code=True,
68 - dtype="float16",
69 - seed=42,
70 - skip_tokenizer_init=skip_tokenizer_init,
71 - )
72 - # Cache config, optionally force APC
73 - kwargs_cache = ({} if enable_prefix_caching is None else {
74 - 'enable_prefix_caching': enable_prefix_caching
75 - })
76 - cache_config = CacheConfig(
77 - block_size=block_size,
78 - gpu_memory_utilization=0.9,
79 - swap_space=0,
80 - cache_dtype="auto",
81 - **kwargs_cache,
82 - )
83 - kv_transfer_config = KVTransferConfig(
84 - kv_connector="SharedStorageConnector",
85 - kv_role="kv_both",
86 - kv_connector_extra_config={"shared_storage_path": "local_storage"},
87 - ) if use_kv_connector else None
88 -
89 - speculative_config: Optional[SpeculativeConfig] = None
90 - if num_speculative_tokens is not None:
91 - speculative_config = SpeculativeConfig(
92 - model="ngram", num_speculative_tokens=num_speculative_tokens)
93 -
94 - vllm_config = VllmConfig(
95 - scheduler_config=scheduler_config,
96 - model_config=model_config,
97 - cache_config=cache_config,
98 - kv_transfer_config=kv_transfer_config,
99 - speculative_config=speculative_config,
100 - )
101 - kv_cache_config = KVCacheConfig(
102 - num_blocks=num_blocks, # A large number of blocks to hold all requests
103 - kv_cache_tensors=[],
104 - kv_cache_groups=[
105 - KVCacheGroupSpec(['layer'],
106 - FullAttentionSpec(block_size, 1, 1, torch.float32,
107 - False))
108 - ],
109 - )
110 - cache_config.num_gpu_blocks = num_blocks
111 - return Scheduler(
112 - vllm_config=vllm_config,
113 - kv_cache_config=kv_cache_config,
114 - log_stats=True,
115 - structured_output_manager=StructuredOutputManager(vllm_config),
116 - )
117 -
118 -
119 -def create_requests(num_requests: int,
120 - num_tokens: int = 10,
121 - mm_positions: Optional[list[PlaceholderRange]] = None,
122 - max_tokens: int = 16,
123 - stop_token_ids: Optional[list[int]] = None,
124 - prompt_logprobs: Optional[int] = None):
125 - sampling_params = SamplingParams(ignore_eos=False,
126 - max_tokens=max_tokens,
127 - stop_token_ids=stop_token_ids,
128 - prompt_logprobs=prompt_logprobs)
129 - requests = []
130 - for i in range(num_requests):
131 - if mm_positions is not None:
132 - mm_position = mm_positions[i]
133 - mm_inputs = [MultiModalKwargs({})] * len(mm_position)
134 - else:
135 - mm_position = None
136 - mm_inputs = None
137 - request = Request(
138 - request_id=f"{i}",
139 - prompt_token_ids=[i] * num_tokens,
140 - sampling_params=sampling_params,
141 - pooling_params=None,
142 - multi_modal_inputs=mm_inputs,
143 - multi_modal_placeholders=mm_position,
144 - multi_modal_hashes=None,
145 - eos_token_id=EOS_TOKEN_ID,
146 - )
147 - requests.append(request)
148 - return requests
22 +from .utils import EOS_TOKEN_ID, create_requests, create_scheduler
23
24
25 def test_add_requests():
tests/v1/core/utils.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Optional, Union
4 +
5 +import torch
6 +
7 +from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,
8 + SchedulerConfig, SpeculativeConfig, VllmConfig)
9 +from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
10 +from vllm.sampling_params import SamplingParams
11 +from vllm.v1.core.sched.async_scheduler import AsyncScheduler
12 +from vllm.v1.core.sched.scheduler import Scheduler
13 +from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
14 + KVCacheGroupSpec)
15 +from vllm.v1.request import Request
16 +from vllm.v1.structured_output import StructuredOutputManager
17 +
18 +EOS_TOKEN_ID = 50256
19 +
20 +
21 +def create_scheduler(
22 + model: str = "facebook/opt-125m",
23 + max_num_seqs: int = 16,
24 + max_num_batched_tokens: int = 8192,
25 + enable_prefix_caching: Optional[bool] = None,
26 + long_prefill_token_threshold: int = 0,
27 + disable_chunked_mm_input: bool = False,
28 + use_kv_connector: bool = False,
29 + num_blocks: int = 10000,
30 + block_size: int = 16,
31 + max_model_len: Optional[int] = None,
32 + num_speculative_tokens: Optional[int] = None,
33 + skip_tokenizer_init: bool = False,
34 + async_scheduling: bool = False,
35 +) -> Union[Scheduler, AsyncScheduler]:
36 + '''Create scheduler under test.
37 +
38 + Args:
39 + model: model under test
40 + max_num_seqs: max sequences to schedule
41 + max_num_batch_tokens: max num tokens to batch
42 + enable_prefix_caching: optionally force APC config
43 + (True/False) or use default
44 + (None)
45 +
46 + Returns:
47 + {class}`Scheduler` instance
48 + '''
49 + if max_model_len is None:
50 + max_model_len = max_num_batched_tokens
51 + scheduler_config = SchedulerConfig(
52 + max_num_seqs=max_num_seqs,
53 + max_num_batched_tokens=max_num_batched_tokens,
54 + max_model_len=max_model_len,
55 + long_prefill_token_threshold=long_prefill_token_threshold,
56 + disable_chunked_mm_input=disable_chunked_mm_input,
57 + enable_chunked_prefill=True,
58 + async_scheduling=async_scheduling,
59 + )
60 + model_config = ModelConfig(
61 + model=model,
62 + task="auto",
63 + tokenizer=model,
64 + tokenizer_mode="auto",
65 + trust_remote_code=True,
66 + dtype="float16",
67 + seed=42,
68 + skip_tokenizer_init=skip_tokenizer_init,
69 + )
70 + # Cache config, optionally force APC
71 + kwargs_cache = ({} if enable_prefix_caching is None else {
72 + 'enable_prefix_caching': enable_prefix_caching
73 + })
74 + cache_config = CacheConfig(
75 + block_size=block_size,
76 + gpu_memory_utilization=0.9,
77 + swap_space=0,
78 + cache_dtype="auto",
79 + **kwargs_cache,
80 + )
81 + kv_transfer_config = KVTransferConfig(
82 + kv_connector="SharedStorageConnector",
83 + kv_role="kv_both",
84 + kv_connector_extra_config={"shared_storage_path": "local_storage"},
85 + ) if use_kv_connector else None
86 +
87 + speculative_config: Optional[SpeculativeConfig] = None
88 + if num_speculative_tokens is not None:
89 + speculative_config = SpeculativeConfig(
90 + model="ngram", num_speculative_tokens=num_speculative_tokens)
91 +
92 + vllm_config = VllmConfig(
93 + scheduler_config=scheduler_config,
94 + model_config=model_config,
95 + cache_config=cache_config,
96 + kv_transfer_config=kv_transfer_config,
97 + speculative_config=speculative_config,
98 + )
99 + kv_cache_config = KVCacheConfig(
100 + num_blocks=num_blocks, # A large number of blocks to hold all requests
101 + kv_cache_tensors=[],
102 + kv_cache_groups=[
103 + KVCacheGroupSpec(['layer'],
104 + FullAttentionSpec(block_size, 1, 1, torch.float32,
105 + False))
106 + ],
107 + )
108 + cache_config.num_gpu_blocks = num_blocks
109 + scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
110 + return scheduler_cls(
111 + vllm_config=vllm_config,
112 + kv_cache_config=kv_cache_config,
113 + log_stats=True,
114 + structured_output_manager=StructuredOutputManager(vllm_config),
115 + )
116 +
117 +
118 +def create_requests(
119 + num_requests: int,
120 + num_tokens: int = 10,
121 + mm_positions: Optional[list[PlaceholderRange]] = None,
122 + max_tokens: int = 16,
123 + stop_token_ids: Optional[list[int]] = None,
124 + prompt_logprobs: Optional[int] = None,
125 + same_prompt: bool = False,
126 +) -> list[Request]:
127 + sampling_params = SamplingParams(ignore_eos=False,
128 + max_tokens=max_tokens,
129 + stop_token_ids=stop_token_ids,
130 + prompt_logprobs=prompt_logprobs)
131 + requests = []
132 + for i in range(num_requests):
133 + if mm_positions is not None:
134 + mm_position = mm_positions[i]
135 + mm_inputs = [MultiModalKwargs({})] * len(mm_position)
136 + else:
137 + mm_position = None
138 + mm_inputs = None
139 + prompt_token_ids = ([0] * num_tokens if same_prompt else [i] *
140 + num_tokens)
141 + request = Request(
142 + request_id=f"{i}",
143 + prompt_token_ids=prompt_token_ids,
144 + sampling_params=sampling_params,
145 + pooling_params=None,
146 + multi_modal_inputs=mm_inputs,
147 + multi_modal_placeholders=mm_position,
148 + multi_modal_hashes=None,
149 + eos_token_id=EOS_TOKEN_ID,
150 + )
151 + requests.append(request)
152 + return requests
vllm/config.py
+11
@@ -2308,6 +2308,13 @@ class SchedulerConfig:
2308 like full attention and sliding window attention.
2309 """
2310
2311 + async_scheduling: bool = False
2312 + """EXPERIMENTAL: If set to True, perform async scheduling. This may help
2313 + reduce the CPU overheads, leading to better latency and throughput. However,
2314 + async scheduling is currently not supported with some features such as
2315 + structured outputs, speculative decoding, and pipeline parallelism.
2316 + """
2317 +
2318 def compute_hash(self) -> str:
2319 """
2320 WARNING: Whenever a new field is added to this config,
@@ -2401,6 +2408,10 @@ class SchedulerConfig:
2408 if not self.cuda_graph_sizes:
2409 self.cuda_graph_sizes = [min(self.max_num_seqs * 2, 512)]
2410
2411 + if self.async_scheduling:
2412 + self.scheduler_cls = (
2413 + "vllm.v1.core.sched.async_scheduler.AsyncScheduler")
2414 +
2415 @model_validator(mode='after')
2416 def _verify_args(self) -> Self:
2417 if (self.max_num_batched_tokens < self.max_model_len
vllm/engine/arg_utils.py
+25
@@ -484,6 +484,8 @@ class EngineArgs:
484 enable_multimodal_encoder_data_parallel: bool = \
485 ParallelConfig.enable_multimodal_encoder_data_parallel
486
487 + async_scheduling: bool = SchedulerConfig.async_scheduling
488 +
489 def __post_init__(self):
490 # support `EngineArgs(compilation_config={...})`
491 # without having to manually construct a
@@ -921,6 +923,8 @@ class EngineArgs:
923 scheduler_group.add_argument(
924 "--disable-hybrid-kv-cache-manager",
925 **scheduler_kwargs["disable_hybrid_kv_cache_manager"])
926 + scheduler_group.add_argument("--async-scheduling",
927 + **scheduler_kwargs["async_scheduling"])
928
929 # vLLM arguments
930 vllm_kwargs = get_kwargs(VllmConfig)
@@ -1206,6 +1210,26 @@ class EngineArgs:
1210 self.data_parallel_rpc_port
1211 is not None) else ParallelConfig.data_parallel_rpc_port
1212
1213 + if self.async_scheduling:
1214 + # Async scheduling does not work with the uniprocess backend.
1215 + if self.distributed_executor_backend is None:
1216 + self.distributed_executor_backend = "mp"
1217 + logger.info("Using mp-based distributed executor backend "
1218 + "for async scheduling.")
1219 + if self.distributed_executor_backend == "uni":
1220 + raise ValueError("Async scheduling is not supported with "
1221 + "uni-process backend.")
1222 + if self.pipeline_parallel_size > 1:
1223 + raise ValueError("Async scheduling is not supported with "
1224 + "pipeline-parallel-size > 1.")
1225 +
1226 + # Currently, async scheduling does not support speculative decoding.
1227 + # TODO(woosuk): Support it.
1228 + if self.speculative_config is not None:
1229 + raise ValueError(
1230 + "Currently, speculative decoding is not supported with "
1231 + "async scheduling.")
1232 +
1233 parallel_config = ParallelConfig(
1234 pipeline_parallel_size=self.pipeline_parallel_size,
1235 tensor_parallel_size=self.tensor_parallel_size,
@@ -1286,6 +1310,7 @@ class EngineArgs:
1310 long_prefill_token_threshold=self.long_prefill_token_threshold,
1311 disable_hybrid_kv_cache_manager=self.
1312 disable_hybrid_kv_cache_manager,
1313 + async_scheduling=self.async_scheduling,
1314 )
1315
1316 if not model_config.is_multimodal_model and self.default_mm_loras:
vllm/v1/core/sched/async_scheduler.py new
+47
@@ -0,0 +1,47 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +from vllm.logger import init_logger
7 +from vllm.v1.core.sched.output import SchedulerOutput
8 +from vllm.v1.core.sched.scheduler import Scheduler
9 +from vllm.v1.request import Request, RequestStatus
10 +
11 +logger = init_logger(__name__)
12 +
13 +
14 +class AsyncScheduler(Scheduler):
15 +
16 + def _update_after_schedule(
17 + self,
18 + scheduler_output: SchedulerOutput,
19 + ) -> None:
20 + super()._update_after_schedule(scheduler_output)
21 + for req_id in scheduler_output.num_scheduled_tokens:
22 + request = self.requests[req_id]
23 + if (request.num_computed_tokens == request.num_tokens +
24 + request.num_output_placeholders):
25 + # The request will generate a new token in this scheduling step.
26 + # TODO(woosuk): Support speculative decoding.
27 + request.num_output_placeholders += 1
28 +
29 + def _update_request_with_output(
30 + self,
31 + request: Request,
32 + new_token_ids: list[int],
33 + ) -> tuple[list[int], bool]:
34 + status_before_update = request.status
35 + new_token_ids, stopped = super()._update_request_with_output(
36 + request, new_token_ids)
37 +
38 + # Update the number of output placeholders.
39 + request.num_output_placeholders -= len(new_token_ids)
40 + assert request.num_output_placeholders >= 0
41 +
42 + # Cache the new tokens. Preempted requests should be skipped.
43 + if status_before_update == RequestStatus.RUNNING:
44 + self.kv_cache_manager.cache_blocks(
45 + request,
46 + request.num_computed_tokens - request.num_output_placeholders)
47 + return new_token_ids, stopped
vllm/v1/core/sched/scheduler.py
+39 -21
@@ -204,7 +204,8 @@ class Scheduler(SchedulerInterface):
204 while req_index < len(self.running) and token_budget > 0:
205 request = self.running[req_index]
206
207 - num_new_tokens = (request.num_tokens_with_spec -
207 + num_new_tokens = (request.num_tokens_with_spec +
208 + request.num_output_placeholders -
209 request.num_computed_tokens)
210 if (0 < self.scheduler_config.long_prefill_token_threshold <
211 num_new_tokens):
@@ -230,9 +231,11 @@ class Scheduler(SchedulerInterface):
231 if num_new_tokens == 0:
232 # The request cannot be scheduled because one of the following
233 # reasons:
233 - # 1. No new tokens to schedule. This may happen when PP>1 and
234 - # we have already scheduled all prompt tokens but they are
235 - # not finished yet.
234 + # 1. No new tokens to schedule. This may happen when
235 + # (1) PP>1 and we have already scheduled all prompt tokens
236 + # but they are not finished yet.
237 + # (2) Async scheduling and the request has reached to either
238 + # its max_total_tokens or max_model_len.
239 # 2. The encoder budget is exhausted.
240 # 3. The encoder cache is exhausted.
241 # NOTE(woosuk): Here, by doing `continue` instead of `break`,
@@ -598,6 +601,14 @@ class Scheduler(SchedulerInterface):
601 request = self.requests[req_id]
602 request.num_computed_tokens += num_scheduled_token
603
604 + # NOTE: _free_encoder_inputs relies on num_computed_tokens, which
605 + # may be updated again in _update_from_output for speculative
606 + # decoding. However, it is safe to call the method here because
607 + # encoder inputs are always part of the prompt, not the output,
608 + # and thus are unaffected by speculative decoding.
609 + if request.has_encoder_inputs:
610 + self._free_encoder_inputs(request)
611 +
612 # Clear the finished request IDs.
613 # NOTE: We shouldn't do self.finished_req_ids.clear() here because
614 # it will also affect the scheduler output.
@@ -785,29 +796,16 @@ class Scheduler(SchedulerInterface):
796 num_draft_tokens=len(scheduled_spec_token_ids),
797 num_accepted_tokens=len(generated_token_ids) - 1)
798
788 - # NOTE(woosuk): This has to be executed after updating
789 - # `request.num_computed_tokens`.
790 - if request.has_encoder_inputs:
791 - self._free_encoder_inputs(request)
792 -
799 stopped = False
800 new_logprobs = None
801 new_token_ids = generated_token_ids
802 kv_transfer_params = None
803 status_before_stop = request.status
804
799 - # Append generated tokens and check for stop. Note that if
800 - # a request is still being prefilled, we expect the model runner
801 - # to return empty token ids for the request.
802 - for num_new, output_token_id in enumerate(new_token_ids, 1):
803 - request.append_output_token_ids(output_token_id)
804 -
805 - # Check for stop and update request state.
806 - # This must be called before we make the EngineCoreOutput.
807 - stopped = check_stop(request, self.max_model_len)
808 - if stopped:
809 - del new_token_ids[num_new:] # Trim new tokens if needed.
810 - break
805 + # Check for stop and update request status.
806 + if new_token_ids:
807 + new_token_ids, stopped = self._update_request_with_output(
808 + request, new_token_ids)
809
810 # Stop checking for pooler models.
811 pooler_output = None
@@ -915,6 +913,26 @@ class Scheduler(SchedulerInterface):
913
914 return engine_core_outputs
915
916 + def _update_request_with_output(
917 + self,
918 + request: Request,
919 + new_token_ids: list[int],
920 + ) -> tuple[list[int], bool]:
921 + # Append generated tokens and check for stop. Note that if
922 + # a request is still being prefilled, we expect the model runner
923 + # to return empty token ids for the request.
924 + stopped = False
925 + for num_new, output_token_id in enumerate(new_token_ids, 1):
926 + request.append_output_token_ids(output_token_id)
927 +
928 + # Check for stop and update request state.
929 + # This must be called before we make the EngineCoreOutput.
930 + stopped = check_stop(request, self.max_model_len)
931 + if stopped:
932 + del new_token_ids[num_new:] # Trim new tokens if needed.
933 + break
934 + return new_token_ids, stopped
935 +
936 def _free_encoder_inputs(self, request: Request) -> None:
937 cached_encoder_input_ids = (
938 self.encoder_cache_manager.get_cached_input_ids(request))
vllm/v1/executor/multiproc_executor.py
+2
@@ -367,6 +367,8 @@ class MultiprocExecutor(Executor):
367
368 @property
369 def max_concurrent_batches(self) -> int:
370 + if self.scheduler_config.async_scheduling:
371 + return 2
372 return self.parallel_config.pipeline_parallel_size
373
374 def _get_output_rank(self) -> int:
vllm/v1/executor/ray_distributed_executor.py
+2
@@ -33,6 +33,8 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
33 """Ray distributed executor supports pipeline parallelism,
34 meaning that it allows PP size batches to be executed concurrently.
35 """
36 + if self.scheduler_config.async_scheduling:
37 + return 2
38 return self.parallel_config.pipeline_parallel_size
39
40 def execute_model(
vllm/v1/request.py
+1
@@ -77,6 +77,7 @@ class Request:
77 self.num_prompt_tokens = len(self.prompt_token_ids)
78 self._output_token_ids: list[int] = []
79 self._all_token_ids: list[int] = self.prompt_token_ids.copy()
80 + self.num_output_placeholders = 0 # Used in async scheduling.
81 self.spec_token_ids: list[int] = []
82 self.num_computed_tokens = 0
83 self.cache_salt: Optional[str] = cache_salt