Preemption and prefill throttle: self-preservation under KV cache pressure

源码版本v0.25.1

Responsibilities

KV cache is the scarcest resource in the v1 scheduler — the number of memory blocks on a GPU is fixed during profiling. When a request in the running queue needs to compute more tokens but KVCacheManager.allocate_slots cannot produce a new block, the scheduler must make a decision: kick out a request that is already running, return its KV blocks, and retry. This mechanism is called preemption, and the entry point is Scheduler._preempt_request (_preempt_request:1140).

A related but more subtle mechanism is prefill throttle — in DP (data parallel) scenarios, each rank's prefill step must be aligned, otherwise fast ranks end up waiting on slow ones. When the scheduler receives throttle_prefills=True, it defers in-progress prefill chunks and new prefills to the alignment step, letting decode fill the current step (defer_prefills:436-438). These two are discussed together because both are restrained behaviors that voluntarily give up pushing another prefill in a given step.

Design motivation

Why pre-empt instead of just rejecting new requests?

• Preserving progress is cheaper: a request that has already prefilled most of its prompt wastes all of its KV if rejected outright. Preemption returns its KV blocks, resets its status to PREEMPTED, and zeroes num_computed_tokens (preempt state:1152-1153), but the request itself is not discarded — it is put back at the head of the waiting queue (prepend_request:1161). The next step can re-prefill it, and any prefix-cache-hit blocks can be reused.
• Pluggable preemption policy: priority mode picks the running request with the largest (priority, arrival_time) to evict (priority preempt:547-552); FCFS mode just calls self.running.pop() (FCFS preempt:572). The semantics match the dequeue policy.
• A preempted step stops accepting new requests: the entry of the waiting loop in schedule() checks not preempted_reqs (waiting entry:637), so that a request kicked back to waiting is not picked up in the same step and ends up competing for resources with the running request that just preempted it.
• Prefill throttle is DP fairness: DP ranks must align their prefill steps. If one rank starts decoding while another is still prefilling, throughput suffers. Throttling makes the misaligned step run decode only, deferring prefill to the aligned step (defer_prefills:436-438).
• capacity-bound auto-disables: prefill_capacity_bound (prefill_capacity_bound:290) records whether the last aligned step drained the waiting queue. If it did, throttling is no longer needed — continuing to throttle would only leave the GPU idle.
• Async KV load has its own reservation: _inflight_prefill_reserved_blocks (_inflight_prefill_reserved_blocks:2400) tracks how many blocks all in-flight prefills still need. Requests loaded asynchronously are not preemptable and must have enough blocks reserved before they can enter, which avoids deadlock.

Key files

• _preempt_request:1140 — preemption entry; frees blocks, clears state, returns to waiting.
• _free_request_blocks:1149 — releases a request's KV blocks; with defer_block_free enabled they go to the deferred queue.
• PREEMPTED:1152-1153 — sets status to PREEMPTED and zeroes num_computed_tokens.
• prepend_request:1161-1162 — puts the request back at the head of waiting and adds it to reset_preempted_req_ids to notify the worker.
• preemption loop:534-575 — the while True preemption loop entered when allocate_slots fails.
• priority victim selection:547-552 — priority mode picks the running request with the largest (priority, arrival_time).
• FCFS victim selection:572 — FCFS mode pops the tail of running.
• defer_prefills:436-438 — DP prefill throttle master switch; defers prefills when throttling and not saturated.
• defer prefill in running:467-471 — in-progress prefill chunks are skipped on a throttle step.
• defer prefill in waiting:801-804 — new prefills in waiting directly break out on a throttle step.
• prefill_capacity_bound update:1019-1020 — updates this flag at the end of every non-throttle step.
• prefill_capacity_bound init:290 — defaults to False in __init__.
• _inflight_prefill_reserved_blocks:2400 — reserved block accounting for async KV loads to avoid deadlock.
• waiting entry guard:637 — once a preemption has happened this step, no more requests are taken from waiting.

Data flow

The trigger for preemption is inside the running loop, when allocate_slots returns 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:
        # The victim already got tokens in this step; return the budget.
        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

This lives in scheduler.py:545-582. There is a subtle detail in priority mode: if the chosen victim was already scheduled this round (present in scheduled_running_reqs), its token budget, new blocks, spec tokens, and encoder budget must all be returned (return budget:553-570) and req_index -= 1, because the running list shrank by one. preempted_req == request is the termination condition — once preemption reaches the request that triggered it, the scheduler gives up and breaks out of the loop.

The evicted request then goes through _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)

This is in _preempt_request:1140-1162. Note the three actions: ① _free_request_blocks returns the KV blocks to the block pool (with defer_block_free enabled they go to the deferred_frees queue waiting on a fence) (_free_request_blocks:2130-2143); ② encoder_cache_manager.free releases the encoder cache in lockstep; ③ num_computed_tokens = 0 zeroes the progress. But prefix-cache-hit blocks are preserved inside _free_request_blocks (hit blocks are shared and not exclusively owned by this request), so the next prefill does not necessarily recompute everything.

The prefill throttle path is independent. defer_prefills is computed once at the start of 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). All three conditions must hold to defer — the throttle flag comes from the DP engine core, prefill_capacity_bound is the flag set when waiting was drained last time, and the last condition ensures prefill is only deferred when there is decode work in running (decode has lower priority than in-progress prefill chunks). It then takes effect in two places: in-progress prefill chunks in running are skipped (running defer:467-471) and new prefills in waiting break out immediately (waiting defer:801-804). At the end of a non-throttle step, prefill_capacity_bound = bool(self.waiting) is updated (capacity_bound update:1019-1020); on the next throttle step, if waiting has already been drained, throttling is skipped.

Boundaries and failure modes

• Only RUNNING can be preempted: _preempt_request begins with assert request.status == RequestStatus.RUNNING (assert:1146-1148). A request in waiting cannot be "re-preempted" because it holds no KV blocks.
• assert: request already removed from running: the docstring explicitly says "popped from the running queue outside of this method" (NOTE:1143-1144). _preempt_request only handles state and returning to waiting; it does not touch self.running. The caller must pop first.
• spec_token_ids cleared: on preemption request.spec_token_ids = [] (spec clear:1154-1155), so the spec decoder does not hand out stale draft tokens on the next scheduling round.
• num_preemptions accumulates: each preemption does request.num_preemptions += 1 (num_preemptions:1156). This counter marks preempted=True in prefix cache stats (preempted marker:721) for observability.
• Deferred free prevents concurrent writes: when defer_block_free=True and last_sched_seq > processed_step_seq, _free_request_blocks does not free immediately. Instead it pushes to the deferred_frees queue waiting on a fence (_free_request_blocks:2130-2143), because under async scheduling another in-flight step may still be writing to those blocks.
• reset_preempted_req_ids notifies the worker: preempted request ids are added to reset_preempted_req_ids (reset_preempted:1162) and stuffed into SchedulerOutput.preempted_req_ids (SchedulerOutput.preempted_req_ids:1100) so the worker knows to clear the corresponding KV cache state.
• Throttle does not affect decode: defer_prefills only skips requests where is_prefill_chunk is true (running defer condition:467). Pure-decode running requests are still scheduled normally, so the GPU is not left idle on a throttle step.
• Async KV load is non-preemptable: requests going through load_kv_async enter the WAITING_FOR_REMOTE_KVS state (WAITING_FOR_REMOTE_KVS:950), not running. They rely on _inflight_prefill_reserved_blocks (_inflight_prefill_reserved_blocks:2400) to reserve blocks, avoiding a situation where they enter running and then get preempted, which would corrupt KV transfer state.

Summary

Preemption is self-preservation under KV cache pressure: pick a low-priority victim from running, release its blocks for a higher-priority request, and put the victim back at the head of waiting for the next scheduling round. Prefill throttle is a different kind of restraint in DP scenarios: non-aligned steps run decode only, deferring prefill to the aligned step. Both are cases of the scheduler voluntarily giving up "schedule one more prefill" in exchange for system stability. For where preemption sits inside the scheduling main loop, read Scheduler.schedule; for how a preempted request gets re-queued, read RequestQueue.

See official docs: vLLM docs · README