vllm-budget

vllm/vllm/v1/core/sched/scheduler.py · ef9339e · L800–871 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

800        # First, schedule the RUNNING requests.
801        req_index = 0
802        while req_index < len(self.running) and token_budget > 0:
803            request = self.running[req_index]
804            if input_budget <= draft_slots:
805                break
806
807            if (
808                request.num_output_placeholders > 0
809                # This is (num_computed_tokens + 1) - (num_output_placeholders - 1).
810                # Since output placeholders are also included in the computed tokens
811                # count, we subtract (num_output_placeholders - 1) to remove any draft
812                # tokens, so that we can be sure no further steps are needed even if
813                # they are all rejected.
814                and request.num_computed_tokens + 2 - request.num_output_placeholders
815                >= request.num_prompt_tokens + request.max_tokens
816            ):
817                # Async scheduling: Avoid scheduling an extra step when we are sure that
818                # the previous step has reached request.max_tokens. We don't schedule
819                # partial draft tokens since this prevents uniform decode optimizations.
820                req_index += 1
821                continue
822
823            if self.current_step < request.next_decode_eligible_step:
824                # V2+PP+async: enforce `pp_size` steps between same-req decodes
825                # to match worker-side sampled-tokens broadcast slot ring cadence.
826                req_index += 1
827                continue
828
829            if defer_prefills and request.is_prefill_chunk:
830                # DP prefill balancing: defer this in-progress prefill chunk to a
831                # cadence-aligned step; decodes still run to fill this step.
832                req_index += 1
833                continue
834
835            if (
836                self.ec_connector is not None
837                and request.mm_features
838                and not self.ec_connector.ensure_cache_available(
839                    request,
840                    request.num_computed_tokens - request.num_output_placeholders,
841                )
842            ):
843                req_index += 1
844                continue
845
846            num_new_tokens = (
847                request.num_tokens_with_spec
848                + request.num_output_placeholders
849                - request.num_computed_tokens
850            )
851            if 0 < long_prefill_token_threshold < num_new_tokens:
852                num_new_tokens = long_prefill_token_threshold
853            num_new_tokens = min(
854                num_new_tokens, token_budget, input_budget - draft_slots
855            )
856
857            # Make sure the input position does not exceed the max model len.
858            # This is necessary when using spec decoding.
859            num_new_tokens = min(
860                num_new_tokens,
861                self.max_model_len
862                - request.num_computed_tokens
863                - self.num_sampled_tokens_per_step,
864            )
865
866            # Apply Mamba alignment before encoder caps.
867            if self.need_mamba_block_aligned_split:
868                num_new_tokens = self._mamba_block_aligned_split(
869                    request, num_new_tokens
870                )
871