搶佔與 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