vLLM scheduler.py · policy.py 源码精读

class Scheduler:

    def __init__(
        self,
        scheduler_config: SchedulerConfig,
        cache_config: CacheConfig,
        lora_config: Optional[LoRAConfig],
    ) -> None:
        ...
        # Sequence groups in the WAITING state.
        # Contain new prefill or preempted requests.
        self.waiting: Deque[SequenceGroup] = deque()
        # Sequence groups in the RUNNING state.
        # Contain decode requests.
        self.running: Deque[SequenceGroup] = deque()
        # Sequence groups in the SWAPPED state.
        # Contain decode requests that are swapped out.
        self.swapped: Deque[SequenceGroup] = deque()

if self.scheduler_config.chunked_prefill_enabled:
    self.prompt_limit = self.scheduler_config.max_model_len
else:
    self.prompt_limit = min(
        self.scheduler_config.max_model_len,
        self.scheduler_config.max_num_batched_tokens)

BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class(
    version="v2" if self.scheduler_config.use_v2_block_manager else "v1")

self.block_manager = BlockSpaceManagerImpl(
    block_size=self.cache_config.block_size,
    num_gpu_blocks=self.cache_config.num_gpu_blocks,
    num_cpu_blocks=self.cache_config.num_cpu_blocks,
    sliding_window=self.cache_config.sliding_window,
    enable_caching=self.cache_config.enable_prefix_caching,
    ...)

@dataclass
class SchedulingBudget:
    """The available slots for scheduling."""
    token_budget: int        # 本批次最多处理多少 token
    max_num_seqs: int        # 本批次最多同时运行多少个序列
    _requeset_ids_num_batched_tokens: Set[str] = field(default_factory=set)
    _requeset_ids_num_curr_seqs: Set[str] = field(default_factory=set)
    _num_batched_tokens: int = 0
    _num_curr_seqs: int = 0

    def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int):
        assert num_new_tokens != 0
        assert num_new_seqs != 0
        return (self.num_batched_tokens + num_new_tokens <= self.token_budget
                and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs)

    def remaining_token_budget(self):
        return self.token_budget - self.num_batched_tokens

def add_num_batched_tokens(self, req_id: str, num_batched_tokens: int):
    if req_id in self._requeset_ids_num_batched_tokens:
        return  # 同一请求不重复计数

    self._requeset_ids_num_batched_tokens.add(req_id)
    self._num_batched_tokens += num_batched_tokens

def subtract_num_batched_tokens(self, req_id: str, num_batched_tokens: int):
    if req_id in self._requeset_ids_num_batched_tokens:
        self._requeset_ids_num_batched_tokens.remove(req_id)
        self._num_batched_tokens -= num_batched_tokens

def add_num_seqs(self, req_id: str, num_curr_seqs: int):
    if req_id in self._requeset_ids_num_curr_seqs:
        return  # 同一请求不重复计数
    ...

# 在 _schedule_default() 或 _schedule_chunked_prefill() 开始时
budget = SchedulingBudget(
    token_budget=self.scheduler_config.max_num_batched_tokens,
    max_num_seqs=self.scheduler_config.max_num_seqs,
)

class Policy:
    def get_priority(self, now: float, seq_group: SequenceGroup) -> float:
        raise NotImplementedError

    def sort_by_priority(
        self,
        now: float,
        seq_groups: Deque[SequenceGroup],
    ) -> Deque[SequenceGroup]:
        return deque(
            sorted(
                seq_groups,
                key=lambda seq_group: self.get_priority(now, seq_group),
                reverse=True,   # 优先级值越大越先调度
            ))


class FCFS(Policy):
    def get_priority(self, now: float, seq_group: SequenceGroup) -> float:
        return now - seq_group.metrics.arrival_time
        # 等待时间越长 → 优先级越高（值越大）


class PolicyFactory:
    _POLICY_REGISTRY = {'fcfs': FCFS}

    @classmethod
    def get_policy(cls, policy_name: str, **kwargs) -> Policy:
        return cls._POLICY_REGISTRY[policy_name](**kwargs)

policy = PolicyFactory.get_policy(policy_name="fcfs")
# waiting 队列按到达时间排序，优先处理等待最久的请求
waiting_queue = policy.sort_by_priority(now=time.time(),
                                         seq_groups=self.waiting)

def _schedule(self) -> SchedulerOutputs:
    """Schedule sequence groups.

    The current policy is designed to optimize the throughput. First,
    it batches as many prefill requests as possible. And it schedules
    decodes. If there's a pressure on GPU memory, preemption or
    swapping is performed.
    """
    if self.scheduler_config.chunked_prefill_enabled:
        return self._schedule_chunked_prefill()
    else:
        return self._schedule_default()

@nvtx_tag(message="schedule", domain="vllm")
def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
    # Schedule sequence groups.
    # NOTE: It is expected that the caller holds at least 1 GPU block
    # before calling this method. Else, the system will deadlock.
    scheduler_outputs = self._schedule()
    now = time.time()

    # Create input data structures.
    seq_group_metadata_list: List[SequenceGroupMetadata] = []
    for i, scheduled_seq_group in enumerate(
            scheduler_outputs.scheduled_seq_groups):
        seq_group = scheduled_seq_group.seq_group
        token_chunk_size = scheduled_seq_group.token_chunk_size
        seq_group.maybe_set_first_scheduled_time(now)

        # seq_id -> SequenceData
        seq_data: Dict[int, SequenceData] = {}
        slot_mapping: Dict[int, List[int]] = {}
        ...
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=seq_group.request_id,
                is_prompt=seq_group.is_prefill(),
                seq_data=seq_data,
                sampling_params=seq_group.sampling_params,
                block_tables=block_tables,
                token_chunk_size=token_chunk_size,
                ...
            ))
    return seq_group_metadata_list, scheduler_outputs

@dataclass
class SchedulerOutputs:
    """The scheduling decision made from a scheduler."""
    # Scheduled sequence groups.
    scheduled_seq_groups: Iterable[ScheduledSequenceGroup]
    # Number of prefill groups scheduled.
    num_prefill_groups: int
    # Total number of batched tokens.
    num_batched_tokens: int
    # Blocks to swap in. Dict of CPU -> GPU block number.
    blocks_to_swap_in: Dict[int, int]
    # Blocks to swap out. Dict of GPU -> CPU block number.
    blocks_to_swap_out: Dict[int, int]
    # Blocks to copy. Source to a list of dest blocks.
    blocks_to_copy: Dict[int, List[int]]
    # Sequence groups that are going to be ignored.
    ignored_seq_groups: List[SequenceGroup]
    # The number of slots for lookahead decoding.
    num_lookahead_slots: int
    # The number of requests in the running queue
    running_queue_size: int

@dataclass
class SchedulerRunningOutputs:
    """The requests scheduled from the running queue."""
    decode_seq_groups: List[SequenceGroup]    # 正常 decode 的请求
    prefill_seq_groups: List[SequenceGroup]   # Chunked Prefill 中分块执行的请求
    preempted: List[SequenceGroup]            # 被 Recompute 抢占的请求
    swapped_out: List[SequenceGroup]          # 被 Swap Out 的请求
    blocks_to_swap_out: Dict[int, int]
    blocks_to_copy: Dict[int, List[int]]
    num_lookahead_slots: int

@dataclass
class SchedulerSwappedInOutputs:
    """The requests scheduled from the swapped queue."""
    decode_seq_groups: List[SequenceGroup]
    prefill_seq_groups: List[SequenceGroup]
    blocks_to_swap_in: Dict[int, int]         # 需要换入的块
    blocks_to_copy: Dict[int, List[int]]
    num_lookahead_slots: int
    infeasible_seq_groups: List[SequenceGroup]  # 无法换入的请求（CPU 块已损坏等）

@dataclass
class SchedulerPrefillOutputs:
    """The requests scheduled from the waiting queue."""
    seq_groups: List[SequenceGroup]           # 被选中做 Prefill 的新请求
    ignored_seq_groups: List[SequenceGroup]   # prompt 过长被忽略的请求
    num_lookahead_slots: int

class PreemptionMode(enum.Enum):
    """Preemption modes.

    1. Swapping: Swap out the blocks of the preempted sequences to CPU memory
    and swap them back in when the sequences are resumed.
    2. Recomputation: Discard the blocks of the preempted sequences and
    recompute them when the sequences are resumed, treating the sequences as
    new prompts.
    """
    SWAP = enum.auto()
    RECOMPUTE = enum.auto()