vLLM scheduler.py 源码精读

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()

def _schedule_default(self) -> SchedulerOutputs:
    """Schedule queues to generate sequences.

    Simple, but effective policy. First, the running queue is
    scheduled (with possible preemption), then swapped queue, finally
    the waiting queue.
    """
    # Include running requests to the budget.
    budget = SchedulingBudget(
        token_budget=self.scheduler_config.max_num_batched_tokens,
        max_num_seqs=self.scheduler_config.max_num_seqs,
    )
    # Make sure we include num_seqs from scheduled sequence groups.
    for seq_group in self.running:
        budget.add_num_seqs(seq_group.request_id,
                            seq_group.get_max_num_running_seqs())

    curr_loras: Optional[Set[int]] = set(
        ..) if self.lora_enabled else None

    # 三阶段调度（顺序调用）
    running_scheduled = self._schedule_running(budget, curr_loras)
    swapped_in = self._schedule_swapped(budget, curr_loras)
    prefills = self._schedule_prefills(budget, curr_loras)

    # 断言：Swap In 和 Swap Out 不能同时发生
    assert (not running_scheduled.blocks_to_swap_out
            or not swapped_in.blocks_to_swap_in)

    # 汇总三阶段结果
    scheduled_seq_groups = (
        [ScheduledSequenceGroup(seq_group=sg, token_chunk_size=1)
         for sg in running_scheduled.decode_seq_groups]
      + [ScheduledSequenceGroup(seq_group=sg,
                                token_chunk_size=sg.get_seqs(SequenceStatus.RUNNING)[0].data.get_num_uncomputed_tokens())
         for sg in running_scheduled.prefill_seq_groups]
      + [ScheduledSequenceGroup(seq_group=sg, token_chunk_size=1)
         for sg in swapped_in.decode_seq_groups]
      + [ScheduledSequenceGroup(seq_group=sg, token_chunk_size=sg....)
         for sg in prefills.seq_groups]
    )

    blocks_to_swap_in = dict(swapped_in.blocks_to_swap_in)
    blocks_to_swap_out = dict(running_scheduled.blocks_to_swap_out)
    blocks_to_copy = merge_dicts(running_scheduled.blocks_to_copy,
                                 swapped_in.blocks_to_copy)

    return SchedulerOutputs(
        scheduled_seq_groups=scheduled_seq_groups,
        num_prefill_groups=len(prefills.seq_groups),
        num_batched_tokens=budget.num_batched_tokens,
        blocks_to_swap_in=blocks_to_swap_in,
        blocks_to_swap_out=blocks_to_swap_out,
        blocks_to_copy=blocks_to_copy,
        ignored_seq_groups=prefills.ignored_seq_groups,
        num_lookahead_slots=running_scheduled.num_lookahead_slots,
        running_queue_size=len(self.running),
    )

def _schedule_running(
    self,
    budget: SchedulingBudget,
    curr_loras: Optional[Set[int]],
    enable_chunking: bool = False,
) -> SchedulerRunningOutputs:
    """Schedule sequence groups that are running.

    Running queue should include decode and chunked prefill requests.

    Args:
        budget: The scheduling budget. The argument is in-out, so it
            will be updated based on the scheduled sequence groups.
        curr_loras: Currently batched lora request ids. The argument is
            in-out, so it will be updated based on the scheduled
            sequence groups.
        enable_chunking: If True, seq_group with prefills may be
            chunked.

    Returns:
        SchedulerRunningOutputs.
    """
    ret = SchedulerRunningOutputs.create_empty()
    # Tracks requests that failed to schedule (and should be preempted).
    victim_queue = deque()

    running_queue = self.running  # NOTE: in-place modification below

    while running_queue:
        seq_group = running_queue[0]
        # The total number of sequences in the RUNNING state.
        num_running_seqs = seq_group.get_max_num_running_seqs()
        # How many tokens need to be scheduled this step?
        # Decode = 1 token per seq; Chunked Prefill = chunk size.
        num_new_tokens = self._get_num_new_tokens(
            seq_group, SequenceStatus.RUNNING, enable_chunking, budget)

        if not enable_chunking:
            num_new_tokens = num_running_seqs

        # Try to allocate blocks for the next decode step.
        if not self._can_append_slots(seq_group, num_new_tokens):
            # 分配失败 → 需要抢占
            # 从 running 队列尾部取受害者（最近进入 running 的请求）
            budget_exhausted = False
            while running_queue:
                victim = running_queue[-1]
                if victim == seq_group:
                    # 当前请求自身无法分配且没有可驱逐的受害者
                    budget_exhausted = True
                    break
                victim_queue.appendleft(victim)
                running_queue.pop()   # 从 running 队列移除受害者

            if budget_exhausted:
                break

            # 对受害者执行抢占
            for victim in victim_queue:
                preempted_mode = self._preempt(victim, ...)
                if preempted_mode == PreemptionMode.RECOMPUTE:
                    ret.preempted.append(victim)
                else:
                    ret.swapped_out.append(victim)

            # 重新尝试分配（受害者释放了足够显存）
            if not self._can_append_slots(seq_group, num_new_tokens):
                break  # 依然不足，放弃本请求

        # 分配成功：执行 append_slots
        self._append_slots(seq_group, ret.blocks_to_copy)
        is_prefill = seq_group.is_prefill()
        if is_prefill:
            ret.prefill_seq_groups.append(seq_group)
        else:
            ret.decode_seq_groups.append(seq_group)
        budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
        # num_seqs 已在 _schedule_default 预填充，此处不再重复
        running_queue.popleft()

    return ret

def _preempt(
    self,
    seq_group: SequenceGroup,
    blocks_to_swap_out: Dict[int, int],
    preemption_mode: Optional[PreemptionMode] = None,
) -> PreemptionMode:
    # Determine the preemption mode.
    if preemption_mode is None:
        if seq_group.get_max_num_running_seqs() == 1:
            preemption_mode = PreemptionMode.RECOMPUTE
        else:
            preemption_mode = PreemptionMode.SWAP

    if preemption_mode == PreemptionMode.RECOMPUTE:
        self._preempt_by_recompute(seq_group)
    elif preemption_mode == PreemptionMode.SWAP:
        self._preempt_by_swap(seq_group, blocks_to_swap_out)
    else:
        raise AssertionError("Unknown preemption mode.")

    return preemption_mode

def _schedule_swapped(
    self,
    budget: SchedulingBudget,
    curr_loras: Optional[Set[int]],
    enable_chunking: bool = False,
) -> SchedulerSwappedInOutputs:
    """Schedule sequence groups that are swapped out.

    It schedules swapped requests as long as it fits `budget`.
    """
    # Swapped-in requests are ordered by their arrival time
    # (same as waiting queue ordering).
    swapped_queue = self.policy.sort_by_priority(
        time.time(), self.swapped)

    ret = SchedulerSwappedInOutputs.create_empty()
    leftover_swapped: Deque[SequenceGroup] = deque()

    while swapped_queue:
        seq_group = swapped_queue[0]
        # How many tokens to schedule?
        num_new_tokens = self._get_num_new_tokens(
            seq_group, SequenceStatus.SWAPPED, enable_chunking, budget)

        # Budget checks (both token and seq dimension)
        if not budget.can_schedule(
            num_new_tokens=num_new_tokens,
            num_new_seqs=seq_group.get_max_num_running_seqs()
        ):
            break

        # LoRA constraint check
        if self.lora_enabled and not self._can_lora_fetch(seq_group, curr_loras):
            leftover_swapped.append(seq_group)
            swapped_queue.popleft()
            continue

        # Check if we can swap in (enough GPU blocks for the KV cache)
        if not self.block_manager.can_swap_in(seq_group,
                                               self._get_num_lookahead_slots()):
            # Infeasible: not enough GPU free blocks
            ret.infeasible_seq_groups.append(seq_group)
            swapped_queue.popleft()
            continue

        # Execute swap-in
        alloc_status = self.block_manager.swap_in(
            seq_group, ret.blocks_to_swap_in)
        assert alloc_status == AllocStatus.OK

        self._append_slots(seq_group, ret.blocks_to_copy)
        is_prefill = seq_group.is_prefill()
        if is_prefill:
            ret.prefill_seq_groups.append(seq_group)
        else:
            ret.decode_seq_groups.append(seq_group)

        budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
        budget.add_num_seqs(seq_group.request_id,
                            seq_group.get_max_num_running_seqs())
        swapped_queue.popleft()

    # Put back leftover (LoRA conflict) requests
    self.swapped = leftover_swapped + swapped_queue

    return ret

def _schedule_prefills(
    self,
    budget: SchedulingBudget,
    curr_loras: Optional[Set[int]],
    enable_chunking: bool = False,
) -> SchedulerPrefillOutputs:
    """Schedule sequence groups that are in prefill stage.

    Note that the current scheduler treats PREEMPTED_FOR_RECOMPUTE
    as a new seq_group in the waiting queue (restart from scratch).
    """
    ret = SchedulerPrefillOutputs.create_empty()
    waiting_queue = self.waiting

    # Important guard: in Default path, do NOT schedule prefills
    # if there are swapped-out requests waiting to be restored.
    # Mixing prefill with swap-in creates memory pressure oscillation.
    leftover_waiting_seqs: Deque[SequenceGroup] = deque()

    # Sort waiting queue by FCFS priority
    waiting_queue = self.policy.sort_by_priority(
        time.time(), waiting_queue)

    while waiting_queue:
        seq_group = waiting_queue[0]

        # Skip if there are still swapped requests
        # (guard against Prefill+Swap coexistence in Default path)
        if len(self.swapped) + len(ret.seq_groups) > 0:
            # NOTE: in Default scheduling, if there are swapped sequences,
            # we skip all prefills to prevent memory oscillation.
            break

        waiting_seqs = seq_group.get_seqs(SequenceStatus.WAITING)
        assert len(waiting_seqs) == 1, \
            "Waiting sequence group should have only one prompt sequence."

        num_new_tokens = self._get_num_new_tokens(
            seq_group, SequenceStatus.WAITING, enable_chunking, budget)

        # Guard 1: prompt exceeds model limit
        if num_new_tokens > self.prompt_limit:
            logger.warning(
                "Input prompt (%d tokens) is too long and exceeds "
                "limit of %d", num_new_tokens, self.prompt_limit)
            for seq in waiting_seqs:
                seq.status = SequenceStatus.FINISHED_IGNORED
            ret.ignored_seq_groups.append(seq_group)
            waiting_queue.popleft()
            continue  # 跳过，不阻塞后续请求

        # Guard 2: can we allocate GPU blocks for this request?
        can_allocate = self.block_manager.can_allocate(seq_group)
        if can_allocate == AllocStatus.LATER:
            break  # 当前无可用块，但以后可能有，停止本轮
        elif can_allocate == AllocStatus.NEVER:
            # 即使独占全部 GPU 也无法满足，直接放弃
            logger.warning(
                "Input prompt (%d tokens) is too long and exceeds "
                "the capacity of block_manager", num_new_tokens)
            for seq in waiting_seqs:
                seq.status = SequenceStatus.FINISHED_IGNORED
            ret.ignored_seq_groups.append(seq_group)
            waiting_queue.popleft()
            continue

        # Guard 3: budget check
        num_new_seqs = seq_group.get_max_num_running_seqs()
        if not budget.can_schedule(num_new_tokens=num_new_tokens,
                                   num_new_seqs=num_new_seqs):
            break

        # All guards passed: allocate blocks and schedule
        seq_group.maybe_set_first_scheduled_time(time.time())
        self._allocate_and_set_running(seq_group)

        budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
        budget.add_num_seqs(seq_group.request_id, num_new_seqs)
        ret.seq_groups.append(seq_group)
        waiting_queue.popleft()

    # Re-queue leftover (LoRA conflict) requests
    self.waiting = leftover_waiting_seqs + waiting_queue
    return ret

def _allocate_and_set_running(self, seq_group: SequenceGroup) -> None:
    self.block_manager.allocate(seq_group)
    for seq in seq_group.get_seqs(SequenceStatus.WAITING):
        seq.status = SequenceStatus.RUNNING

# 初始预算
budget = SchedulingBudget(token_budget=T, max_num_seqs=N)

# Phase 1：running 请求预填充 num_seqs（token 小，seq 槽位已预占）
for seq_group in self.running:
    budget.add_num_seqs(...)  # 预占 seq 槽位，几乎不消耗 token_budget

# _schedule_running：追加实际消耗（每请求约 1 token × num_seqs）
#   token_budget 消耗 ≈ len(running) * 1（极小）
#   max_num_seqs 消耗 ≈ len(running)（已预占，无额外消耗）

# _schedule_swapped：消耗剩余 token_budget 的一小部分
#   每个 swapped 请求 ≈ 1 token × num_seqs
#   max_num_seqs 消耗 = swapped_in 的序列数

# _schedule_prefills：使用剩余预算
#   每个 prefill 请求消耗 prompt_len 个 token（可能很大）
#   token_budget 通常在第一个 prefill 请求后接近耗尽

@dataclass
class SchedulerRunningOutputs:
    """The requests scheduled from a running queue.

    Could contain prefill (prefill that needs to be chunked or
    continued) or decode requests. Note that decode requests are the
    majority, and prefill requests are only added when chunked prefill
    is enabled.

    Attributes:
        decode_seq_groups: Sequence groups that are going to be decoded.
        prefill_seq_groups: Sequence groups that are going to be (continued)
            prefilled.
        preempted: Sequence groups that are preempted.
        swapped_out: Sequence groups that are swapped out.
        blocks_to_swap_out: GPU -> CPU block number to swap out.
        blocks_to_copy: GPU -> GPU block numbers to copy.
        num_lookahead_slots: a number of speculative tokens that will be
            decoded, currently used by draft models in speculative decoding.
    """
    decode_seq_groups: List[SequenceGroup]
    prefill_seq_groups: List[SequenceGroup]
    preempted: List[SequenceGroup]
    swapped_out: List[SequenceGroup]
    blocks_to_swap_out: Dict[int, int]        # GPU 块 → CPU 块
    blocks_to_copy: Dict[int, List[int]]      # Copy-on-Write 指令
    num_lookahead_slots: int

    @classmethod
    def create_empty(cls) -> "SchedulerRunningOutputs":
        return cls(
            decode_seq_groups=[],
            prefill_seq_groups=[],
            preempted=[],
            swapped_out=[],
            blocks_to_swap_out={},
            blocks_to_copy={},
            num_lookahead_slots=0,
        )

@dataclass
class SchedulerSwappedInOutputs:
    """The requests scheduled from a swapped queue.

    Could contain prefill (when enable_chunking=True) or decode requests.

    Attributes:
        decode_seq_groups: Sequence groups that are going to be decoded.
        prefill_seq_groups: Sequence groups that are going to be prefilled.
        blocks_to_swap_in: CPU -> GPU block number to swap in.
        blocks_to_copy: GPU -> GPU block numbers to copy.
        num_lookahead_slots: A number of speculative tokens that will be
            decoded.
        infeasible_seq_groups: Sequence groups that are infeasible due to
            resource constraints.
    """
    decode_seq_groups: List[SequenceGroup]
    prefill_seq_groups: List[SequenceGroup]
    blocks_to_swap_in: Dict[int, int]         # CPU 块 → GPU 块
    blocks_to_copy: Dict[int, List[int]]      # Copy-on-Write 指令
    num_lookahead_slots: int
    infeasible_seq_groups: List[SequenceGroup]  # 无法换入的请求

    @classmethod
    def create_empty(cls) -> "SchedulerSwappedInOutputs":
        return cls(
            decode_seq_groups=[],
            prefill_seq_groups=[],
            blocks_to_swap_in={},
            blocks_to_copy={},
            num_lookahead_slots=0,
            infeasible_seq_groups=[],
        )

@dataclass
class SchedulerPrefillOutputs:
    """The requests scheduled from a waiting queue.

    Attributes:
        seq_groups: Sequence groups that are going to be prefilled.
        ignored_seq_groups: Sequence groups that are going to be ignored.
        num_lookahead_slots: A number of speculative tokens that will be
            decoded.
    """
    seq_groups: List[SequenceGroup]
    ignored_seq_groups: List[SequenceGroup]
    num_lookahead_slots: int

    @classmethod
    def create_empty(cls) -> "SchedulerPrefillOutputs":
        return cls(
            seq_groups=[],
            ignored_seq_groups=[],
            num_lookahead_slots=0,
        )

# 更新队列状态
self.running.extendleft(reversed(running_scheduled.decode_seq_groups))
self.running.extendleft(reversed(swapped_in.decode_seq_groups))
for seq_group in running_scheduled.preempted:
    self.waiting.appendleft(seq_group)
for seq_group in running_scheduled.swapped_out:
    self.swapped.appendleft(seq_group)
for seq_group in prefills.seq_groups:
    self.running.appendleft(seq_group)

# 构建最终输出
return SchedulerOutputs(
    scheduled_seq_groups=(
        [ScheduledSequenceGroup(sg, token_chunk_size=1)
         for sg in running_scheduled.decode_seq_groups]
      + [ScheduledSequenceGroup(sg, token_chunk_size=1)
         for sg in swapped_in.decode_seq_groups]
      + [ScheduledSequenceGroup(sg, token_chunk_size=sg.get_seqs(...)[0].data.get_num_uncomputed_tokens())
         for sg in prefills.seq_groups]
    ),
    num_prefill_groups=len(prefills.seq_groups),
    num_batched_tokens=budget.num_batched_tokens,
    blocks_to_swap_in=dict(swapped_in.blocks_to_swap_in),
    blocks_to_swap_out=dict(running_scheduled.blocks_to_swap_out),
    blocks_to_copy=merge_dicts(running_scheduled.blocks_to_copy,
                               swapped_in.blocks_to_copy),
    ignored_seq_groups=prefills.ignored_seq_groups,
    num_lookahead_slots=running_scheduled.num_lookahead_slots,
    running_queue_size=len(self.running),
)