vLLM 抢占机制源码精读

抢占的触发条件

抢占发生在 _schedule_running 方法中，当调度器发现无法为当前 decode 请求分配新的 KV Cache slot 时触发：

scheduler.py — 抢占触发点 L443-478

running_queue.popleft()
while not self._can_append_slots(seq_group):
    budget.subtract_num_batched_tokens(seq_group.request_id,
                                       num_running_tokens)
    num_running_seqs = seq_group.get_max_num_running_seqs()
    budget.subtract_num_seqs(seq_group.request_id,
                             num_running_seqs)

    if running_queue:
        # Preempt the lowest-priority sequence groups.
        victim_seq_group = running_queue.pop()
        preempted_mode = self._preempt(victim_seq_group,
                                       blocks_to_swap_out)
        if preempted_mode == PreemptionMode.RECOMPUTE:
            preempted.append(victim_seq_group)
        else:
            swapped_out.append(victim_seq_group)
    else:
        # No other sequence groups can be preempted.
        # Preempt the current sequence group.
        preempted_mode = self._preempt(seq_group,
                                       blocks_to_swap_out)
        if preempted_mode == PreemptionMode.RECOMPUTE:
            preempted.append(seq_group)
        else:
            swapped_out.append(seq_group)
        break

_can_append_slots（L1001-1018）委托给 Block Manager 检查是否有空闲 block：

scheduler.py — _can_append_slots L1001-1018

def _can_append_slots(self, seq_group: SequenceGroup) -> bool:
    # It is True only for testing case to trigger artificial preemption.
    if (self.enable_artificial_preemption
            and random.uniform(0, 1) < ARTIFICIAL_PREEMPTION_PROB
            and self.artificial_preempt_cnt > 0):
        self.artificial_preempt_cnt -= 1
        return False

    is_prefill = False
    return self.block_manager.can_append_slots(
        seq_group=seq_group,
        num_lookahead_slots=self._get_num_lookahead_slots(is_prefill),
    )

人工抢占测试

通过环境变量 VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT 可启用人工抢占，以 50% 概率强制返回 False，用于测试抢占路径。最多触发 500 次（ARTIFICIAL_PREEMPTION_MAX_CNT）。

PreemptionMode 枚举

vLLM 定义了两种抢占模式：

scheduler.py — PreemptionMode L30-41

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

模式	操作	优点	缺点
RECOMPUTE	释放 GPU block，序列回到 WAITING 队列，重新计算	实现简单，不需要 CPU 内存	需要重新做 prefill，浪费计算
SWAP	KV Cache 从 GPU 搬到 CPU，序列进入 SWAPPED 队列	恢复时不需要重新计算	需要 CPU 内存，有 PCIe 传输开销

_preempt 决策入口

_preempt 方法（L1155-1188）是抢占的主入口，负责选择抢占模式并执行：

scheduler.py — _preempt L1155-1188

def _preempt(
    self,
    seq_group: SequenceGroup,
    blocks_to_swap_out: Dict[int, int],
    preemption_mode: Optional[PreemptionMode] = None,
) -> PreemptionMode:
    # We use recomputation by default since it incurs lower overhead than
    # swapping. However, when the sequence group has multiple sequences
    # (e.g., beam search), recomputation is not currently supported.
    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("Invalid preemption mode.")
    return preemption_mode

graph TD A["_preempt(seq_group)"] --> B{"指定了 preemption_mode?"} B -->|"否"| C{"seq_group 有多个序列?
(beam search)"} C -->|"单序列"| D["RECOMPUTE"] C -->|"多序列"| E["SWAP"] B -->|"是"| F["使用指定模式"] D --> G["_preempt_by_recompute"] E --> H["_preempt_by_swap"] F --> I{"RECOMPUTE?"} I -->|"是"| G I -->|"否"| H

为什么单序列默认 Recompute

源码注释明确指出：Recompute 的开销通常低于 Swap（避免了 PCIe 数据传输）。但 Recompute 目前不支持多序列（beam search），因为需要更复杂的 CUDA kernel 来处理多序列的重新计算。这也导致了一个"奇怪"的副作用：使用 beam search 的请求会被隐式地优先于单序列请求（因为 SWAP 请求进入 swapped 队列，而 swapped 队列的调度优先级高于 waiting 队列）。

Recompute 模式

_preempt_by_recompute（L1190-1202）是最简单的抢占实现：

scheduler.py — _preempt_by_recompute L1190-1202

def _preempt_by_recompute(
    self,
    seq_group: SequenceGroup,
) -> None:
    seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)
    assert len(seqs) == 1
    for seq in seqs:
        seq.status = SequenceStatus.WAITING
        self.free_seq(seq)
        seq.reset_state_for_recompute()

三步操作：

状态变更：RUNNING → WAITING
释放 block：self.free_seq(seq) → block_manager.free(seq)，释放所有 GPU block
重置计算状态：seq.reset_state_for_recompute()，将 num_computed_tokens 重置为 0，但保留已生成的 output token

reset_state_for_recompute 的巧妙设计

被 recompute 抢占的序列重新进入 waiting 队列后，其 prompt 长度 = 原始 prompt + 已生成 output token。这意味着重新 prefill 时会把之前生成的 token 也重新计算 KV Cache。这保证了 KV Cache 的正确性——相当于从头开始，但"prompt"变长了。

Swap 模式

Swap 分为两个操作：swap out（GPU → CPU）和 swap in（CPU → GPU）。

Swap Out

scheduler.py — _swap_out L1229-1247

def _swap_out(
    self,
    seq_group: SequenceGroup,
    blocks_to_swap_out: Dict[int, int],
) -> None:
    if not self.block_manager.can_swap_out(seq_group):
        raise RuntimeError(
            "Aborted due to the lack of CPU swap space. Please increase "
            "the swap space to avoid this error.")
    mapping = self.block_manager.swap_out(seq_group)
    blocks_to_swap_out.update(mapping)
    for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
        seq.status = SequenceStatus.SWAPPED

执行流程：

检查 CPU 空间：can_swap_out 确认有足够的 CPU block
获取映射：block_manager.swap_out() 返回 GPU block → CPU block 的映射关系
记录映射：累积到 blocks_to_swap_out，供后续 Worker 执行实际的数据传输
状态变更：RUNNING → SWAPPED

CPU swap 空间不足时的处理

如果 CPU 没有足够的 swap block，当前实现会直接 raise RuntimeError 终止整个引擎。源码中标注了 FIXME，建议未来改为只 abort 当前请求而不是整个引擎。这是一个需要注意的边界情况。

Swap In

scheduler.py — _swap_in L1215-1227

def _swap_in(
    self,
    seq_group: SequenceGroup,
    blocks_to_swap_in: Dict[int, int],
) -> None:
    mapping = self.block_manager.swap_in(seq_group)
    blocks_to_swap_in.update(mapping)
    for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
        seq.status = SequenceStatus.RUNNING

Swap in 是 swap out 的逆操作：将 KV Cache 从 CPU 搬回 GPU，状态从 SWAPPED → RUNNING。

_schedule_running 中的抢占循环

抢占的触发逻辑在 _schedule_running（L382-519）的内层 while 循环中。这个循环的设计非常精妙：

graph TD A["从 running_queue 取出高优先级请求 seq_group"] --> B{"_can_append_slots?"} B -->|"是"| C["_append_slots, 加入 batch"] B -->|"否"| D{"running_queue 还有其他请求?"} D -->|"是"| E["从队尾弹出最低优先级 victim"] E --> F["_preempt(victim)"] F --> G["victim 释放 block"] G --> B D -->|"否"| H["抢占自己 _preempt(seq_group)"] H --> I["break 退出循环"]

关键设计点：

优先级排序：running_queue 按 FCFS 排序，队头是最高优先级（最早到达的请求）
牺牲者选择：从队尾（最低优先级）弹出作为被抢占者
循环抢占：如果一个 victim 释放的 block 不够，继续抢占下一个最低优先级的请求
自我抢占：如果队列中只剩自己，说明即使释放所有其他请求也不够，只能牺牲自己

每次抢占后会从 budget 中扣除对应的 token 数和 seq 数，确保预算与实际调度一致。

Swap-in 恢复流程

被 swap out 的请求会进入 self.swapped 队列，在后续调度轮次中通过 _schedule_swapped（L521-635）恢复：

scheduler.py — _schedule_swapped 恢复条件 L566-607

while swapped_queue:
    seq_group = swapped_queue[0]

    # If the sequence group cannot be swapped in, stop.
    alloc_status = self.block_manager.can_swap_in(seq_group)
    if alloc_status == AllocStatus.LATER:
        break
    elif alloc_status == AllocStatus.NEVER:
        logger.warning(
            "Failing the request %s because there's not enough kv "
            "cache blocks to run the entire sequence.",
            seq_group.request_id)
        for seq in seq_group.get_seqs():
            seq.status = SequenceStatus.FINISHED_IGNORED
        infeasible_seq_groups.append(seq_group)
        swapped_queue.popleft()
        continue

    num_new_seqs = seq_group.get_max_num_running_seqs()
    num_new_tokens = self._get_num_new_tokens(seq_group,
                                              SequenceStatus.SWAPPED,
                                              enable_chunking, budget)
    if (num_new_tokens == 0
            or not budget.can_schedule(num_new_tokens=num_new_tokens,
                                       num_new_seqs=num_new_seqs)):
        break

    swapped_queue.popleft()
    self._swap_in(seq_group, blocks_to_swap_in)
    self._append_slots(seq_group, blocks_to_copy)

恢复的三个前置条件：

can_swap_in 返回 AllocStatus.OK（GPU 有足够的空闲 block）
Budget 有足够的 token 和 seq 配额
没有发生抢占（在 Default 和 Chunked Prefill 策略中都有此检查）

如果 can_swap_in 返回 AllocStatus.NEVER，说明即使把所有 GPU block 都给这个请求也放不下，直接标记为 FINISHED_IGNORED 丢弃。

与 Block Manager 的交互

抢占机制与 Block Manager 的交互接口总结：

Scheduler 调用	Block Manager 方法	作用
`_can_append_slots`	`can_append_slots(seq_group)`	检查是否有空闲 block 供 decode 生成新 token
`free_seq`	`free(seq)`	释放序列占用的所有 GPU block（Recompute 模式）
`_swap_out`	`can_swap_out(seq_group)` `swap_out(seq_group)`	检查 CPU 空间并执行 GPU→CPU 块映射
`_swap_in`	`can_swap_in(seq_group)` `swap_in(seq_group)`	检查 GPU 空间并执行 CPU→GPU 块映射

注意：swap_out 和 swap_in 返回的是块映射关系（Dict[int, int]），实际的数据传输由 Worker 的 CacheEngine 在模型执行前异步完成。Scheduler 只负责"决策"，Worker 负责"执行"。

性能影响与优化

抢占对系统性能的影响：

Recompute 的开销

需要重新做完整的 prefill（prompt + 已生成 output）
如果序列已经生成了很多 token，重新计算的代价可能很高
适合已生成 token 较少的场景

Swap 的开销

PCIe 带宽有限（通常 32-64 GB/s），大量 KV Cache 交换会成为瓶颈
需要预留 CPU 内存作为 swap 空间
swap-in/swap-out 可以与计算部分重叠（异步执行），降低实际延迟

delay_factor 调优

scheduler.py — _passed_delay L1249-1263

def _passed_delay(self, now: float) -> bool:
    if self.prev_prompt:
        self.last_prompt_latency = now - self.prev_time
    self.prev_time, self.prev_prompt = now, False
    # Delay scheduling prompts to let waiting queue fill up
    if self.scheduler_config.delay_factor > 0 and self.waiting:
        earliest_arrival_time = min(
            [e.metrics.arrival_time for e in self.waiting])
        passed_delay = (
            (now - earliest_arrival_time) >
            (self.scheduler_config.delay_factor * self.last_prompt_latency)
            or not self.running)
    else:
        passed_delay = True
    return passed_delay

delay_factor 通过延迟调度新 prefill 来减少抢占发生的频率。当设置 delay_factor > 0 时，只有等待时间超过 delay_factor × 上次 prompt 延迟 才会调度新 prefill。这给了 running 请求更多时间完成并释放 block。

减少抢占的最佳实践

增大 gpu_memory_utilization（默认 0.9）以获得更多 KV Cache block
降低 max_num_seqs 以限制并发请求数
设置合理的 max_num_batched_tokens 避免一次性调度太多 token
使用 Chunked Prefill 策略减少 prefill 对 decode 的影响
增加 CPU swap 空间（swap_space 参数，单位 GiB）以支持 Swap 模式

vLLM 抢占机制 源码精读

抢占的触发条件

PreemptionMode 枚举

_preempt 决策入口

Recompute 模式

Swap 模式

Swap Out

Swap In

_schedule_running 中的抢占循环

Swap-in 恢复流程

与 Block Manager 的交互

性能影响与优化

Recompute 的开销

Swap 的开销

delay_factor 调优

vLLM 抢占机制源码精读