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 8, 22:45 UTC
[Perf] Optimize spec decoding + async scheduling, 1.5% Throughput improvement (#33612)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Wentao Ye committed
Feb 4, 2026 at 09:34 UTC
711edaf0d089a15df5fa2b99248c516e53929bd2
2 files changed
+14
-6
vllm/v1/core/sched/async_scheduler.py
+7
-2
@@ -10,6 +10,11 @@ logger = init_logger(__name__)
10
11
12
class AsyncScheduler(Scheduler):
13
+ def __init__(self, *args, **kwargs) -> None:
14
+ super().__init__(*args, **kwargs)
15
+ # reusable read-only placeholder list for speculative decoding.
16
+ self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens
17
+
18
def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
19
super()._update_after_schedule(scheduler_output)
20
has_structured_output_requests = False
@@ -31,9 +36,9 @@ class AsyncScheduler(Scheduler):
36
# The request will generate a new token plus num_spec_tokens
37
# in this scheduling step.
38
request.num_output_placeholders += 1 + cur_num_spec_tokens
34
- # Add placeholders for the new tokens in spec_token_ids.
39
+ # Add placeholders for the new draft/spec tokens.
40
# We will update the actual spec token ids in the worker process.
36
- request.spec_token_ids = [-1] * self.num_spec_tokens
41
+ request.spec_token_ids = self._spec_token_placeholders
42
43
scheduler_output.has_structured_output_requests = has_structured_output_requests
44
scheduler_output.pending_structured_output_tokens = (
vllm/v1/core/sched/scheduler.py
+7
-4
@@ -487,9 +487,11 @@ class Scheduler(SchedulerInterface):
487
- request.num_output_placeholders
488
)
489
if num_scheduled_spec_tokens > 0:
490
- # Trim spec_token_ids list to num_scheduled_spec_tokens.
491
- del request.spec_token_ids[num_scheduled_spec_tokens:]
492
- scheduled_spec_decode_tokens[request_id] = request.spec_token_ids
490
+ spec_token_ids = request.spec_token_ids
491
+ if len(spec_token_ids) > num_scheduled_spec_tokens:
492
+ spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens]
493
+ scheduled_spec_decode_tokens[request.request_id] = spec_token_ids
494
+
495
# New spec tokens will be set in `update_draft_token_ids` before the
496
# next step when applicable.
497
request.spec_token_ids = []
@@ -887,7 +889,8 @@ class Scheduler(SchedulerInterface):
889
self.encoder_cache_manager.free(request)
890
request.status = RequestStatus.PREEMPTED
891
request.num_computed_tokens = 0
890
- request.spec_token_ids.clear()
892
+ if request.spec_token_ids:
893
+ request.spec_token_ids = []
894
request.num_preemptions += 1
895
if self.log_stats:
896
request.record_event(EngineCoreEventType.PREEMPTED, timestamp)