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 12, 06:00 UTC
[BugFix] Preserve spec decoding uniform decode when scheduling (#29759)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Dec 1, 2025 at 17:15 UTC
44822d7ff22cb62856b4a107f97f00fc8e1199a0
3 files changed
+25
-17
tests/v1/e2e/test_spec_decode.py
+2
-2
@@ -191,8 +191,8 @@ def test_suffix_decoding_acceptance(
191
# Expect the acceptance rate to improve.
192
assert first_accept_rate < last_accept_rate
193
194
- # Heuristic: expect at least 85% acceptance rate at the end.
195
- assert last_accept_rate > 0.85
194
+ # Heuristic: expect at least 82.5% acceptance rate at the end.
195
+ assert last_accept_rate > 0.825
196
197
del spec_llm
198
torch.cuda.empty_cache()
vllm/v1/core/sched/async_scheduler.py
+1
-1
@@ -33,7 +33,7 @@ class AsyncScheduler(Scheduler):
33
# in this scheduling step.
34
request.num_output_placeholders += 1 + cur_num_spec_tokens
35
# Add placeholders for the new tokens in spec_token_ids.
36
- # Wwe will update the actual spec token ids in the worker process.
36
+ # We will update the actual spec token ids in the worker process.
37
request.spec_token_ids = [-1] * self.num_spec_tokens
38
39
scheduler_output.pending_structured_output_tokens = (
vllm/v1/core/sched/scheduler.py
+22
-14
@@ -236,6 +236,22 @@ class Scheduler(SchedulerInterface):
236
while req_index < len(self.running) and token_budget > 0:
237
request = self.running[req_index]
238
239
+ if (
240
+ request.num_output_placeholders > 0
241
+ # This is (num_computed_tokens + 1) - (num_output_placeholders - 1).
242
+ # Since output placeholders are also included in the computed tokens
243
+ # count, we subtract (num_output_placeholders - 1) to remove any draft
244
+ # tokens, so that we can be sure no further steps are needed even if
245
+ # they are all rejected.
246
+ and request.num_computed_tokens + 2 - request.num_output_placeholders
247
+ >= request.num_prompt_tokens + request.max_tokens
248
+ ):
249
+ # Async scheduling: Avoid scheduling an extra step when we are sure that
250
+ # the previous step has reached request.max_tokens. We don't schedule
251
+ # partial draft tokens since this prevents uniform decode optimizations.
252
+ req_index += 1
253
+ continue
254
+
255
num_new_tokens = (
256
request.num_tokens_with_spec
257
+ request.num_output_placeholders
@@ -245,18 +261,10 @@ class Scheduler(SchedulerInterface):
261
num_new_tokens = self.scheduler_config.long_prefill_token_threshold
262
num_new_tokens = min(num_new_tokens, token_budget)
263
248
- num_spec_placeholders = max(0, request.num_output_placeholders - 1)
249
- max_total_tokens = min(
250
- # Avoid scheduling tokens that we're sure won't will be needed based on
251
- # request.max_tokens. For this calculation we assume placeholder
252
- # speculated output tokens are rejected.
253
- request.num_prompt_tokens + request.max_tokens + num_spec_placeholders,
254
- # Make sure the input position does not exceed the max model len.
255
- # This is necessary when using spec decoding.
256
- self.max_model_len,
257
- )
264
+ # Make sure the input position does not exceed the max model len.
265
+ # This is necessary when using spec decoding.
266
num_new_tokens = min(
259
- num_new_tokens, max_total_tokens - 1 - request.num_computed_tokens
267
+ num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens
268
)
269
270
# Schedule encoder inputs.
@@ -799,15 +807,15 @@ class Scheduler(SchedulerInterface):
807
for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)):
808
req_id = req.request_id
809
req_ids.append(req_id)
802
- num_tokens = num_scheduled_tokens[req_id] - len(
803
- spec_decode_tokens.get(req_id, ())
804
- )
810
if self.use_pp:
811
# When using PP, the scheduler sends the sampled tokens back,
812
# because there's no direct communication between the first-
813
# stage worker and the last-stage worker. Otherwise, we don't
814
# need to send the sampled tokens back because the model runner
815
# will cache them.
816
+ num_tokens = num_scheduled_tokens[req_id] - len(
817
+ spec_decode_tokens.get(req_id, ())
818
+ )
819
token_ids = req.all_token_ids[
820
req.num_computed_tokens : req.num_computed_tokens + num_tokens
821
]