抢占与 prefill throttle:KV cache 紧张时的自保

源码版本v0.25.1

职责

KV cache 是 v1 调度器最稀缺的资源——一块 GPU 上的显存块数量在 profiling 阶段就定死了。当 running 队列里的某个 request 需要往前算更多 token、但 KVCacheManager.allocate_slots 拿不出新块时,调度器 (scheduler) 必须做一个决定:踢掉一个已经在 running 里的 request,把它的 KV 块还回去,再重试。这套机制叫抢占 (preemption),入口是 Scheduler._preempt_request(_preempt_request:1140)。

另一个相关但更隐蔽的机制是 prefill throttle——在 DP(数据并行)场景下,各 rank 的 prefill 步必须对齐,否则快的 rank 会拖着慢的 rank 等待。调度器拿到 throttle_prefills=True 时,会把 in-progress prefill chunk 和新 prefill 推迟到对齐步,让 decode 把这步填满(defer_prefills:436-438)。这两件事一起讲,因为它们都是"在某步主动放弃往前推 prefill"的克制行为。

设计动机

为什么不直接拒绝新请求、要去做抢占?

• 保留进度更划算:一个 request 已经 prefill 了一大半,如果直接拒绝,前面算的 KV 全白费;抢占 (preemption) 把它的 KV 块还回去、状态置回 PREEMPTED、num_computed_tokens 清零(preempt 状态:1152-1153),但 request 本体不丢、放回 waiting 队首(prepend_request:1161),下一步还能重新 prefill(prefix cache 命中的部分能复用)。
• 抢占策略可换:priority 模式挑 (priority, arrival_time) 最大的 running 踢(priority preempt:547-552),FCFS 模式直接 self.running.pop()(FCFS preempt:572),语义和出队一致。
• 抢占过的步不再接新请求:schedule() 的 waiting 循环入口检查 not preempted_reqs(waiting 入口:637),避免被踢回 waiting 的 request 在同一步又被取出来和正在抢占的 running 抢资源。
• prefill throttle 是 DP 公平性:DP rank 之间 prefill 步必须对齐,否则一个 rank 已经开始 decode、另一个还在 prefill,吞吐就被拖慢。throttle 让非对齐步只跑 decode,把 prefill 推到对齐步一起做(defer_prefills:436-438)。
• capacity-bound 自动覆盖:prefill_capacity_bound(prefill_capacity_bound:290)记录上一次对齐步是否把 waiting 排空,排空了就不需要再 throttle——继续 throttle 只会让 GPU 闲着。
• async KV load 单独的 reservation:_inflight_prefill_reserved_blocks(_inflight_prefill_reserved_blocks:2400)统计所有 in-flight prefill 还需要多少块,异步加载的 request 不抢占、必须预留够块才能进,避免 deadlock。

关键文件

• _preempt_request:1140 — 抢占入口,free 块、清状态、放回 waiting。
• _free_request_blocks:1149 — 释放 request 的 KV 块,defer_block_free 开启时进 deferred 队列。
• PREEMPTED:1152-1153 — 状态置 PREEMPTED,num_computed_tokens 清零。
• prepend_request:1161-1162 — 放回 waiting 队首,加入 reset_preempted_req_ids 通知 worker。
• 抢占循环:534-575 — allocate_slots 失败时的抢占 while True 循环。
• priority 选牺牲者:547-552 — priority 模式挑 (priority, arrival_time) 最大的 running。
• FCFS 选牺牲者:572 — FCFS 模式直接 pop running 末尾。
• defer_prefills:436-438 — DP prefill throttle 主开关,throttle 且非饱和时推迟 prefill。
• running 里推迟 prefill:467-471 — in-progress prefill chunk 在 throttle 步被跳过。
• waiting 里推迟 prefill:801-804 — waiting 的新 prefill 在 throttle 步直接 break。
• prefill_capacity_bound 更新:1019-1020 — 每个非 throttle 步结束更新这个 flag。
• prefill_capacity_bound 初始化:290 — __init__ 里默认 False。
• _inflight_prefill_reserved_blocks:2400 — 异步 KV load 的预留块统计,避免死锁。
• waiting 入口保护:637 — 本步发生过抢占就不再接 waiting。

数据流

抢占的触发点在 running 循环里,allocate_slots 返回 None 时:

# The request cannot be scheduled.
# Preempt the lowest-priority request.
if self.policy == SchedulingPolicy.PRIORITY:
    preempted_req = max(
        self.running,
        key=lambda r: (r.priority, r.arrival_time),
    )
    self.running.remove(preempted_req)
    if preempted_req in scheduled_running_reqs:
        # 已经在本步分到 token 的牺牲者,把预算还回去
        preempted_req_id = preempted_req.request_id
        scheduled_running_reqs.remove(preempted_req)
        token_budget += num_scheduled_tokens.pop(preempted_req_id)
        req_to_new_blocks.pop(preempted_req_id)
        ...
        req_index -= 1
else:
    preempted_req = self.running.pop()

self._preempt_request(preempted_req, scheduled_timestamp)
preempted_reqs.append(preempted_req)
if preempted_req == request:
    # No more request to preempt. Cannot schedule this request.
    break

这段在 scheduler.py:545-582。priority 模式有个细节:如果选中的牺牲者本轮已经被排过(scheduled_running_reqs 里),要把它的 token 预算、新块、spec token、encoder 预算全部还回去(还预算:553-570)并 req_index -= 1,因为 running 列表少了一个元素。preempted_req == request 是终止条件——抢占到自己头上就承认失败,break 出循环。

被踢的 request 走 _preempt_request:

def _preempt_request(self, request: Request, timestamp: float) -> None:
    """Preempt a request and put it back to the waiting queue.

    NOTE: The request should be popped from the running queue outside of this
    method.
    """
    assert request.status == RequestStatus.RUNNING, (
        "Only running requests can be preempted"
    )
    self._free_request_blocks(request)
    self.encoder_cache_manager.free(request)
    self._inflight_prefills.discard(request)
    request.status = RequestStatus.PREEMPTED
    request.num_computed_tokens = 0
    if request.spec_token_ids:
        request.spec_token_ids = []
    request.num_preemptions += 1
    ...
    # Put the request back to the waiting queue.
    self.waiting.prepend_request(request)
    self.reset_preempted_req_ids.add(request.request_id)

这段在 _preempt_request:1140-1162。注意三个动作:① _free_request_blocks 把 KV 块还给 block pool(开 defer_block_free 时进 deferred_frees 队列等 fence)(_free_request_blocks:2130-2143);② encoder_cache_manager.free 同步释放 encoder 缓存;③ num_computed_tokens = 0 把进度清零——但 prefix cache 命中的块在 _free_request_blocks 内部会保留(因为命中块是共享的、不属于这个 request 独占),所以下次重新 prefill 时不一定真重算。

prefill throttle 的路径独立。defer_prefills 在 schedule() 开头算一次(defer_prefills:436-438):throttle_prefills and not self.prefill_capacity_bound and any(not r.is_prefill_chunk for r in self.running)。三个条件全满足才推迟——throttle 标志来自 DP engine core、prefill_capacity_bound 是上次排空了 waiting 的 flag、最后一个条件保证只在 running 里有 decode 时才推迟 prefill(decode 优先级低于 in-progress prefill chunk)。然后两处生效:running 里的 in-progress prefill chunk 被跳过(running 推迟:467-471)、waiting 里的新 prefill 直接 break(waiting 推迟:801-804)。非 throttle 步结束时会更新 prefill_capacity_bound = bool(self.waiting)(capacity_bound 更新:1019-1020),下一轮 throttle 步如果发现 waiting 已经排空,就不再 throttle。

边界与失败

• 只能 preempt RUNNING:_preempt_request 开头断言 request.status == RequestStatus.RUNNING(assert:1146-1148)。waiting 里的 request 不能被"再抢占",因为它们没占 KV 块。
• assert: request 已从 running 移除:docstring 明确说"popped from the running queue outside of this method"(NOTE:1143-1144)。_preempt_request 内部只负责状态和放回 waiting,不碰 self.running 列表,调用方必须先 pop。
• spec_token_ids 清空:被抢占时 request.spec_token_ids = [](spec 清空:1154-1155),避免下一次重排时 spec decoder 给出过期 draft token。
• num_preemptions 累加:每次抢占 request.num_preemptions += 1(num_preemptions:1156),这个计数会在 prefix cache stats 里标记 preempted=True(preempted 标记:721)供可观测性使用。
• deferred free 防并发写:_free_request_blocks 在 defer_block_free=True 且 last_sched_seq > processed_step_seq 时,不立即 free,而是 push 到 deferred_frees 队列等 fence(_free_request_blocks:2130-2143),因为 async scheduling 下另一个 in-flight step 可能还在写这些块。
• reset_preempted_req_ids 通知 worker:被抢占的 request id 进入 reset_preempted_req_ids(reset_preempted:1162)并塞进 SchedulerOutput.preempted_req_ids(SchedulerOutput.preempted_req_ids:1100),让 worker 知道要清掉对应的 KV cache 状态。
• throttle 不影响 decode:defer_prefills 只跳过 is_prefill_chunk 的 request(running 推迟条件:467),纯 decode 的 running request 仍然正常排,保证 GPU 在 throttle 步不闲着。
• async KV load 不可抢占:走 load_kv_async 的 request 进 WAITING_FOR_REMOTE_KVS 状态(WAITING_FOR_REMOTE_KVS:950),不进 running,靠 _inflight_prefill_reserved_blocks(_inflight_prefill_reserved_blocks:2400)预留块,避免进 running 后又被抢占导致 KV 传输状态混乱。

小结

抢占是 KV cache 紧张时的自保:从 running 里挑一个低优先级的踢掉,释放块给高优先级 request,被踢的放回 waiting 队首等下一步重排。prefill throttle 是 DP 场景下的另一种克制:非对齐步只跑 decode,把 prefill 推到对齐步一起做。两者都是调度器主动放弃"多排一个 prefill"换取系统稳定性。要看抢占在调度主循环里的位置,读 Scheduler.schedule;要看被踢的 request 怎么重新排队,读 RequestQueue。

对照官方资料:vLLM 文档 · README。