vllm-allocation
vllm/vllm/v1/core/sched/scheduler.py · ef9339e · L922–1028 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
922 request,
923 num_new_tokens,
924 num_lookahead_tokens=self.num_lookahead_tokens,
925 )
926
927 if new_blocks is not None:
928 # The request can be scheduled.
929 break
930
931 if (
932 self.connector is not None
933 and self.connector.has_pending_block_frees()
934 ):
935 break
936
937 # The request cannot be scheduled.
938 # Preempt the lowest-priority request.
939 if self.policy == SchedulingPolicy.PRIORITY:
940 preempted_req = max(
941 self.running,
942 key=lambda r: (r.priority, r.arrival_time),
943 )
944 else:
945 preempted_req = self.running[-1]
946
947 # A deferred free will not help with immediate allocation.
948 if not self._request_blocks_can_be_freed(preempted_req):
949 break
950
951 if self.policy == SchedulingPolicy.PRIORITY:
952 victim_index = self.running.index(preempted_req)
953 del self.running[victim_index]
954 if victim_index < req_index:
955 req_index -= 1
956
957 if preempted_req in scheduled_running_reqs:
958 preempted_req_id = preempted_req.request_id
959 scheduled_running_reqs.remove(preempted_req)
960 restored = num_scheduled_tokens.pop(preempted_req_id)
961 token_budget += restored
962 input_budget += restored + draft_slots
963 req_to_new_blocks.pop(preempted_req_id)
964 scheduled_spec_decode_tokens.pop(preempted_req_id, None)
965 preempted_encoder_inputs = scheduled_encoder_inputs.pop(
966 preempted_req_id, None
967 )
968 if preempted_encoder_inputs:
969 # Restore encoder compute budget if the preempted
970 # request had encoder inputs scheduled in this step.
971 num_embeds_to_restore = sum(
972 preempted_req.get_num_encoder_embeds(i)
973 for i in preempted_encoder_inputs
974 )
975 encoder_compute_budget += num_embeds_to_restore
976 else:
977 preempted_req = self.running.pop()
978
979 self._preempt_request(
980 preempted_req,
981 scheduled_timestamp,
982 drop_stale_output=self.requires_kv_delivery,
983 )
984 preempted_reqs.append(preempted_req)
985 if preempted_req == request:
986 # No more request to preempt. Cannot schedule this request.
987 break
988
989 if new_blocks is None:
990 if (
991 self.policy == SchedulingPolicy.PRIORITY
992 and request.status == RequestStatus.PREEMPTED
993 ):
994 # The request was the lowest-priority running request and
995 # preempted itself. The requests after it have higher
996 # priority and may still fit (it just freed its blocks);
997 # it was removed from `self.running`, so `req_index`
998 # already points at the next one.
999 continue
1000 # Cannot schedule this request.
1001 break
1002
1003 # Schedule the request.
1004 scheduled_running_reqs.append(request)
1005 prefill_scheduled |= request.is_prefill_chunk
1006 request_id = request.request_id
1007 req_to_new_blocks[request_id] = new_blocks
1008 num_scheduled_tokens[request_id] = num_new_tokens
1009 token_budget -= num_new_tokens
1010 input_budget -= num_new_tokens + draft_slots
1011 req_index += 1
1012 drafts_needed = drafts_needed or not self._drafts_unused(
1013 request, request.num_computed_tokens, num_new_tokens
1014 )
1015
1016 # Speculative decode related.
1017 if request.spec_token_ids:
1018 num_scheduled_spec_tokens = (
1019 num_new_tokens
1020 + request.num_computed_tokens
1021 - request.num_tokens
1022 - request.num_output_placeholders
1023 )
1024 if num_scheduled_spec_tokens > 0:
1025 spec_token_ids = request.spec_token_ids
1026 if len(spec_token_ids) > num_scheduled_spec_tokens:
1027 spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens]
1028 scheduled_spec_decode_tokens[request.request_id] = spec_token_ids