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 13, 08:15 UTC
[BugFix] Priority scheduling and spec tokens preemption (#28558)
Signed-off-by: Andy Lo <andy@mistral.ai>
Andy Lo committed
Nov 12, 2025 at 20:29 UTC
58ce8d12b741b5bafe6bd0fb878727baea6171fe
2 files changed
+266
tests/v1/core/test_priority_scheduler_random.py
new
+252
@@ -0,0 +1,252 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import random
4
+import uuid
5
+
6
+import pytest
7
+
8
+from vllm.config import VllmConfig
9
+from vllm.multimodal.inputs import (
10
+ MultiModalFeatureSpec,
11
+ MultiModalKwargsItem,
12
+ PlaceholderRange,
13
+)
14
+from vllm.sampling_params import SamplingParams
15
+from vllm.utils.hashing import get_hash_fn_by_name
16
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
17
+from vllm.v1.core.sched.output import SchedulerOutput
18
+from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
19
+from vllm.v1.request import Request
20
+
21
+from .test_scheduler import create_scheduler_with_priority
22
+from .utils import EOS_TOKEN_ID
23
+
24
+pytestmark = pytest.mark.cpu_test
25
+
26
+
27
+def _create_random_request(
28
+ max_tokens_range: tuple[int, int],
29
+ num_tokens_range: tuple[int, int],
30
+ arrival_time_range: tuple[float, float],
31
+ priority_range: tuple[int, int],
32
+ num_mm_item_range: tuple[int, int],
33
+ vllm_config: VllmConfig,
34
+):
35
+ max_tokens = random.randint(*max_tokens_range)
36
+ num_tokens = random.randint(*num_tokens_range)
37
+ priority = random.randint(*priority_range)
38
+ arrival_time = random.uniform(*arrival_time_range)
39
+ num_mm_item = random.randint(*num_mm_item_range)
40
+
41
+ mm_positions: list[PlaceholderRange] = []
42
+ for mm_start in sorted(
43
+ random.sample(range(num_tokens), min(num_mm_item, num_tokens))
44
+ ):
45
+ if mm_start + 10 > num_tokens:
46
+ continue
47
+ mm_positions.append(PlaceholderRange(offset=mm_start, length=10))
48
+
49
+ request_id = uuid.uuid4().hex
50
+
51
+ sampling_params = SamplingParams(
52
+ ignore_eos=False,
53
+ max_tokens=max_tokens,
54
+ )
55
+ mm_features = []
56
+ for j, position in enumerate(mm_positions):
57
+ identifier = f"{request_id}_hash_{j}"
58
+ mm_feature = MultiModalFeatureSpec(
59
+ data=MultiModalKwargsItem.dummy("dummy_m"),
60
+ mm_position=position,
61
+ identifier=identifier,
62
+ modality="image",
63
+ )
64
+ mm_features.append(mm_feature)
65
+
66
+ prompt_token_ids = random.choices(range(100), k=num_tokens)
67
+
68
+ caching_hash_fn = get_hash_fn_by_name(
69
+ vllm_config.cache_config.prefix_caching_hash_algo
70
+ )
71
+ init_none_hash(caching_hash_fn)
72
+ block_hasher = get_request_block_hasher(
73
+ vllm_config.cache_config.block_size, caching_hash_fn
74
+ )
75
+
76
+ request = Request(
77
+ request_id=request_id,
78
+ prompt_token_ids=prompt_token_ids,
79
+ sampling_params=sampling_params,
80
+ pooling_params=None,
81
+ mm_features=mm_features if mm_features else None,
82
+ eos_token_id=EOS_TOKEN_ID,
83
+ arrival_time=arrival_time,
84
+ priority=priority,
85
+ block_hasher=block_hasher,
86
+ )
87
+ return request
88
+
89
+
90
+def _mock_execute_model(
91
+ scheduler_output: SchedulerOutput, num_output_tokens_range: tuple[int, int]
92
+) -> ModelRunnerOutput:
93
+ request_ids: list[str] = []
94
+ request_ids.extend(req.req_id for req in scheduler_output.scheduled_new_reqs)
95
+ request_ids.extend(scheduler_output.scheduled_cached_reqs.req_ids)
96
+ random.shuffle(request_ids)
97
+
98
+ num_output_tokens = [
99
+ random.randint(*num_output_tokens_range) for _ in range(len(request_ids))
100
+ ]
101
+ sampled_token_ids = [
102
+ [random.randint(0, 100) for _ in range(num_tokens)]
103
+ for num_tokens in num_output_tokens
104
+ ]
105
+
106
+ return ModelRunnerOutput(
107
+ req_ids=request_ids,
108
+ req_id_to_index={req_id: i for i, req_id in enumerate(request_ids)},
109
+ sampled_token_ids=sampled_token_ids,
110
+ logprobs=None,
111
+ prompt_logprobs_dict={},
112
+ pooler_output=[],
113
+ )
114
+
115
+
116
+def _mock_draft_token_ids(
117
+ scheduler_output: SchedulerOutput,
118
+ num_output_tokens_range: tuple[int, int],
119
+ seen_request_prompt_length: dict[str, int],
120
+) -> DraftTokenIds:
121
+ request_ids: list[str] = []
122
+ sampled_token_ids: list[list[int]] = []
123
+ for request in scheduler_output.scheduled_new_reqs:
124
+ assert request.req_id not in seen_request_prompt_length
125
+ seen_request_prompt_length[request.req_id] = len(request.prompt_token_ids or [])
126
+ if request.num_computed_tokens >= seen_request_prompt_length[request.req_id]:
127
+ num_tokens = random.randint(*num_output_tokens_range)
128
+ request_ids.append(request.req_id)
129
+ sampled_token_ids.append(
130
+ [random.randint(0, 100) for _ in range(num_tokens)]
131
+ )
132
+ for req_id, num_computed_tokens in zip(
133
+ scheduler_output.scheduled_cached_reqs.req_ids,
134
+ scheduler_output.scheduled_cached_reqs.num_computed_tokens,
135
+ ):
136
+ if num_computed_tokens >= seen_request_prompt_length[req_id]:
137
+ num_tokens = random.randint(*num_output_tokens_range)
138
+ request_ids.append(req_id)
139
+ sampled_token_ids.append(
140
+ [random.randint(0, 100) for _ in range(num_tokens)]
141
+ )
142
+ return DraftTokenIds(req_ids=request_ids, draft_token_ids=sampled_token_ids)
143
+
144
+
145
+def _chech_valid_scheduler_output(
146
+ scheduler_output: SchedulerOutput,
147
+ seen_request_ids: set[str],
148
+ seen_mm_hashes: set[str],
149
+):
150
+ for req in scheduler_output.scheduled_new_reqs:
151
+ assert req.req_id not in seen_request_ids
152
+ seen_request_ids.add(req.req_id)
153
+ for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
154
+ assert req_id in seen_request_ids
155
+
156
+ req_ids = set[str]()
157
+ req_ids.update(req.req_id for req in scheduler_output.scheduled_new_reqs)
158
+ req_ids.update(scheduler_output.scheduled_cached_reqs.req_ids)
159
+
160
+ assert set(scheduler_output.num_scheduled_tokens.keys()) == req_ids
161
+ assert (
162
+ sum(scheduler_output.num_scheduled_tokens.values())
163
+ == scheduler_output.total_num_scheduled_tokens
164
+ )
165
+
166
+ assert set(scheduler_output.scheduled_spec_decode_tokens.keys()) <= req_ids
167
+ assert set(scheduler_output.scheduled_encoder_inputs.keys()) <= req_ids
168
+
169
+ for req in scheduler_output.scheduled_new_reqs:
170
+ for mm_feature in req.mm_features:
171
+ seen_mm_hashes.add(mm_feature.identifier)
172
+ for mm_hash in scheduler_output.free_encoder_mm_hashes:
173
+ assert mm_hash in seen_mm_hashes
174
+
175
+ assert scheduler_output.finished_req_ids <= seen_request_ids
176
+
177
+
178
+@pytest.mark.parametrize("enable_prefix_caching", [True, False])
179
+@pytest.mark.parametrize("num_speculative_tokens", [None, 1, 5])
180
+@pytest.mark.parametrize(
181
+ ("max_input_tokens", "max_output_tokens", "max_num_seqs", "num_blocks"),
182
+ [
183
+ # Standard profile
184
+ (5000, 500, 256, 10000),
185
+ # Generation heavy + high max_num_seqs + low num_blocks -> Many preemptions
186
+ (500, 5000, 1024, 1000),
187
+ ],
188
+ ids=["standard", "preemption"],
189
+)
190
+def test_priority_scheduling_blast(
191
+ enable_prefix_caching: bool,
192
+ num_speculative_tokens: int | None,
193
+ max_input_tokens: int,
194
+ max_output_tokens: int,
195
+ max_num_seqs: int,
196
+ num_blocks: int,
197
+):
198
+ random.seed(42)
199
+ seen_request_prompt_length = dict[str, int]()
200
+ seen_request_ids = set[str]()
201
+ seen_mm_hashes = set[str]()
202
+
203
+ scheduler = create_scheduler_with_priority(
204
+ model="Qwen/Qwen2.5-VL-3B-Instruct",
205
+ max_num_seqs=max_num_seqs,
206
+ enable_prefix_caching=enable_prefix_caching,
207
+ num_blocks=num_blocks,
208
+ num_speculative_tokens=num_speculative_tokens,
209
+ )
210
+
211
+ num_initial_requests = 10
212
+ for _ in range(num_initial_requests):
213
+ req = _create_random_request(
214
+ max_tokens_range=(1, max_output_tokens),
215
+ num_tokens_range=(1, max_input_tokens),
216
+ arrival_time_range=(0, 1),
217
+ priority_range=(-3, 3),
218
+ num_mm_item_range=(0, 2),
219
+ vllm_config=scheduler.vllm_config,
220
+ )
221
+ scheduler.add_request(req)
222
+
223
+ for _ in range(20000):
224
+ if len(scheduler.waiting) == 0:
225
+ num_new_requests = random.randint(0, 2)
226
+ for _ in range(num_new_requests):
227
+ req = _create_random_request(
228
+ max_tokens_range=(1, max_output_tokens),
229
+ num_tokens_range=(1, max_input_tokens),
230
+ arrival_time_range=(0, 1),
231
+ priority_range=(-3, 3),
232
+ num_mm_item_range=(0, 2),
233
+ vllm_config=scheduler.vllm_config,
234
+ )
235
+ scheduler.add_request(req)
236
+ scheduler_output = scheduler.schedule()
237
+ _chech_valid_scheduler_output(
238
+ scheduler_output, seen_request_ids, seen_mm_hashes
239
+ )
240
+ model_output = _mock_execute_model(
241
+ scheduler_output,
242
+ num_output_tokens_range=(1, 1 + (num_speculative_tokens or 0)),
243
+ )
244
+ scheduler.update_from_output(scheduler_output, model_output)
245
+ if num_speculative_tokens is not None:
246
+ scheduler.update_draft_token_ids(
247
+ _mock_draft_token_ids(
248
+ scheduler_output,
249
+ (0, num_speculative_tokens),
250
+ seen_request_prompt_length,
251
+ )
252
+ )
vllm/v1/core/sched/scheduler.py
+14
@@ -300,6 +300,20 @@ class Scheduler(SchedulerInterface):
300
]
301
req_to_new_blocks.pop(preempted_req.request_id)
302
num_scheduled_tokens.pop(preempted_req.request_id)
303
+ scheduled_spec_decode_tokens.pop(
304
+ preempted_req.request_id, None
305
+ )
306
+ preempted_encoder_inputs = scheduled_encoder_inputs.pop(
307
+ preempted_req.request_id, None
308
+ )
309
+ if preempted_encoder_inputs:
310
+ # Restore encoder compute budget if the preempted
311
+ # request had encoder inputs scheduled in this step.
312
+ num_tokens_to_restore = sum(
313
+ preempted_req.get_num_encoder_tokens(i)
314
+ for i in preempted_encoder_inputs
315
+ )
316
+ encoder_compute_budget += num_tokens_to_restore
317
req_index -= 1
318
else:
319
preempted_req = self.running.pop()