diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 84e6ae1f1..a4c692b3d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -440,17 +440,13 @@ class Scheduler(SchedulerInterface): ) self.running.remove(preempted_req) if preempted_req in scheduled_running_reqs: + preempted_req_id = preempted_req.request_id scheduled_running_reqs.remove(preempted_req) - token_budget += num_scheduled_tokens[ - preempted_req.request_id - ] - req_to_new_blocks.pop(preempted_req.request_id) - num_scheduled_tokens.pop(preempted_req.request_id) - scheduled_spec_decode_tokens.pop( - preempted_req.request_id, None - ) + token_budget += num_scheduled_tokens.pop(preempted_req_id) + req_to_new_blocks.pop(preempted_req_id) + scheduled_spec_decode_tokens.pop(preempted_req_id, None) preempted_encoder_inputs = scheduled_encoder_inputs.pop( - preempted_req.request_id, None + preempted_req_id, None ) if preempted_encoder_inputs: # Restore encoder compute budget if the preempted @@ -476,8 +472,9 @@ class Scheduler(SchedulerInterface): # Schedule the request. scheduled_running_reqs.append(request) - req_to_new_blocks[request.request_id] = new_blocks - num_scheduled_tokens[request.request_id] = num_new_tokens + request_id = request.request_id + req_to_new_blocks[request_id] = new_blocks + num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens req_index += 1 @@ -492,18 +489,14 @@ class Scheduler(SchedulerInterface): if num_scheduled_spec_tokens > 0: # Trim spec_token_ids list to num_scheduled_spec_tokens. del request.spec_token_ids[num_scheduled_spec_tokens:] - scheduled_spec_decode_tokens[request.request_id] = ( - request.spec_token_ids - ) + scheduled_spec_decode_tokens[request_id] = request.spec_token_ids # New spec tokens will be set in `update_draft_token_ids` before the # next step when applicable. request.spec_token_ids = [] # Encoder-related. if encoder_inputs_to_schedule: - scheduled_encoder_inputs[request.request_id] = ( - encoder_inputs_to_schedule - ) + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule # Allocate the encoder cache. for i in encoder_inputs_to_schedule: self.encoder_cache_manager.allocate(request, i) @@ -535,6 +528,7 @@ class Scheduler(SchedulerInterface): break request = self.waiting.peek_request() + request_id = request.request_id # KVTransfer: skip request if still waiting for remote kvs. if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: @@ -549,7 +543,7 @@ class Scheduler(SchedulerInterface): else: logger.debug( "%s is still in WAITING_FOR_REMOTE_KVS state.", - request.request_id, + request_id, ) self.waiting.pop_request() skipped_waiting_requests.prepend_request(request) @@ -729,7 +723,7 @@ class Scheduler(SchedulerInterface): if self.connector is not None: self.connector.update_state_after_alloc( request, - self.kv_cache_manager.get_blocks(request.request_id), + self.kv_cache_manager.get_blocks(request_id), num_external_computed_tokens, ) @@ -759,10 +753,10 @@ class Scheduler(SchedulerInterface): if self.lora_config and request.lora_request: scheduled_loras.add(request.lora_request.lora_int_id) - req_to_new_blocks[request.request_id] = ( - self.kv_cache_manager.get_blocks(request.request_id) + req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks( + request_id ) - num_scheduled_tokens[request.request_id] = num_new_tokens + num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens request.status = RequestStatus.RUNNING request.num_computed_tokens = num_computed_tokens @@ -771,9 +765,7 @@ class Scheduler(SchedulerInterface): request.num_cached_tokens = num_computed_tokens # Encoder-related. if encoder_inputs_to_schedule: - scheduled_encoder_inputs[request.request_id] = ( - encoder_inputs_to_schedule - ) + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule # Allocate the encoder cache. for i in encoder_inputs_to_schedule: self.encoder_cache_manager.allocate(request, i) @@ -806,11 +798,9 @@ class Scheduler(SchedulerInterface): num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): if self.running: - any_request = self.running[0] + any_request_id = self.running[0].request_id num_common_prefix_blocks = ( - self.kv_cache_manager.get_num_common_prefix_blocks( - any_request.request_id - ) + self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) ) # Construct the scheduler output.