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 9, 01:15 UTC
[Minor] Some code simplification in `scheduler.py` (#33597)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 2, 2026 at 23:00 UTC
61397891ce00c6e28ca9918fab11be1b9e925a20
1 file changed
+19
-29
vllm/v1/core/sched/scheduler.py
+19
-29
@@ -440,17 +440,13 @@ class Scheduler(SchedulerInterface):
440
)
441
self.running.remove(preempted_req)
442
if preempted_req in scheduled_running_reqs:
443
+ preempted_req_id = preempted_req.request_id
444
scheduled_running_reqs.remove(preempted_req)
444
- token_budget += num_scheduled_tokens[
445
- preempted_req.request_id
446
- ]
447
- req_to_new_blocks.pop(preempted_req.request_id)
448
- num_scheduled_tokens.pop(preempted_req.request_id)
449
- scheduled_spec_decode_tokens.pop(
450
- preempted_req.request_id, None
451
- )
445
+ token_budget += num_scheduled_tokens.pop(preempted_req_id)
446
+ req_to_new_blocks.pop(preempted_req_id)
447
+ scheduled_spec_decode_tokens.pop(preempted_req_id, None)
448
preempted_encoder_inputs = scheduled_encoder_inputs.pop(
453
- preempted_req.request_id, None
449
+ preempted_req_id, None
450
)
451
if preempted_encoder_inputs:
452
# Restore encoder compute budget if the preempted
@@ -476,8 +472,9 @@ class Scheduler(SchedulerInterface):
472
473
# Schedule the request.
474
scheduled_running_reqs.append(request)
479
- req_to_new_blocks[request.request_id] = new_blocks
480
- num_scheduled_tokens[request.request_id] = num_new_tokens
475
+ request_id = request.request_id
476
+ req_to_new_blocks[request_id] = new_blocks
477
+ num_scheduled_tokens[request_id] = num_new_tokens
478
token_budget -= num_new_tokens
479
req_index += 1
480
@@ -492,18 +489,14 @@ class Scheduler(SchedulerInterface):
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:]
495
- scheduled_spec_decode_tokens[request.request_id] = (
496
- request.spec_token_ids
497
- )
492
+ scheduled_spec_decode_tokens[request_id] = request.spec_token_ids
493
# New spec tokens will be set in `update_draft_token_ids` before the
494
# next step when applicable.
495
request.spec_token_ids = []
496
497
# Encoder-related.
498
if encoder_inputs_to_schedule:
504
- scheduled_encoder_inputs[request.request_id] = (
505
- encoder_inputs_to_schedule
506
- )
499
+ scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule
500
# Allocate the encoder cache.
501
for i in encoder_inputs_to_schedule:
502
self.encoder_cache_manager.allocate(request, i)
@@ -535,6 +528,7 @@ class Scheduler(SchedulerInterface):
528
break
529
530
request = self.waiting.peek_request()
531
+ request_id = request.request_id
532
533
# KVTransfer: skip request if still waiting for remote kvs.
534
if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
@@ -549,7 +543,7 @@ class Scheduler(SchedulerInterface):
543
else:
544
logger.debug(
545
"%s is still in WAITING_FOR_REMOTE_KVS state.",
552
- request.request_id,
546
+ request_id,
547
)
548
self.waiting.pop_request()
549
skipped_waiting_requests.prepend_request(request)
@@ -729,7 +723,7 @@ class Scheduler(SchedulerInterface):
723
if self.connector is not None:
724
self.connector.update_state_after_alloc(
725
request,
732
- self.kv_cache_manager.get_blocks(request.request_id),
726
+ self.kv_cache_manager.get_blocks(request_id),
727
num_external_computed_tokens,
728
)
729
@@ -759,10 +753,10 @@ class Scheduler(SchedulerInterface):
753
754
if self.lora_config and request.lora_request:
755
scheduled_loras.add(request.lora_request.lora_int_id)
762
- req_to_new_blocks[request.request_id] = (
763
- self.kv_cache_manager.get_blocks(request.request_id)
756
+ req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks(
757
+ request_id
758
)
765
- num_scheduled_tokens[request.request_id] = num_new_tokens
759
+ num_scheduled_tokens[request_id] = num_new_tokens
760
token_budget -= num_new_tokens
761
request.status = RequestStatus.RUNNING
762
request.num_computed_tokens = num_computed_tokens
@@ -771,9 +765,7 @@ class Scheduler(SchedulerInterface):
765
request.num_cached_tokens = num_computed_tokens
766
# Encoder-related.
767
if encoder_inputs_to_schedule:
774
- scheduled_encoder_inputs[request.request_id] = (
775
- encoder_inputs_to_schedule
776
- )
768
+ scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule
769
# Allocate the encoder cache.
770
for i in encoder_inputs_to_schedule:
771
self.encoder_cache_manager.allocate(request, i)
@@ -806,11 +798,9 @@ class Scheduler(SchedulerInterface):
798
num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
799
with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"):
800
if self.running:
809
- any_request = self.running[0]
801
+ any_request_id = self.running[0].request_id
802
num_common_prefix_blocks = (
811
- self.kv_cache_manager.get_num_common_prefix_blocks(
812
- any_request.request_id
813
- )
803
+ self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id)
804
)
805
806
# Construct the scheduler output.