Preemption und Prefill-Throttle: Selbstschutz bei knapper KV-Cache

源码版本v0.25.1

Verantwortung

KV-Cache ist die knappste Ressource des v1-Schedulers — die Anzahl der Speicherblöcke auf einer GPU ist in der Profiling-Phase festgelegt. Wenn ein Request in der running-Warteschlange mehr Token vorwärts rechnen will, aber KVCacheManager.allocate_slots keine neuen Blöcke liefern kann, muss der Scheduler (scheduler) eine Entscheidung treffen: einen Request aus running herauswerfen, dessen KV-Blöcke zurückgeben und es erneut versuchen. Dieser Mechanismus heißt Preemption (preemption); der Eingang ist Scheduler._preempt_request (_preempt_request:1140).

Ein weiterer, subtilerer Mechanismus ist Prefill-Throttle (prefill throttle) — in DP-Szenarien (Datenparallel, data parallel) müssen die Prefill-Schritte der einzelnen Ranks ausgerichtet sein, sonst wartet der schnelle Rank auf den langsamen. Wenn der Scheduler throttle_prefills=True erhält, werden in-progress prefill chunks und neue Prefills auf den Alignment-Schritt verschoben, sodass Decode diesen Schritt füllt (defer_prefills:436-438). Beide Themen werden gemeinsam behandelt, weil sie jeweils ein „in einem Schritt aktiv darauf verzichten, den Prefill voranzutreiben" sind.

Entwurfsmotivation

Warum nicht einfach neue Requests ablehnen, statt Preemption durchzuführen?

Fortschritt zu bewahren lohnt sich: Ein Request, der bereits mehrheitlich im Prefill ist, würde beim direkten Ablehnen den bereits berechneten KV verlieren; Preemption gibt die KV-Blöcke zurück, setzt den Status auf PREEMPTED und num_computed_tokens auf 0 (preempt Status:1152-1153), verwirft aber den Request selbst nicht — er wird an die Spitze von waiting gelegt (prepend_request:1161) und im nächsten Schritt neu prefilled (die per Prefix-Cache getroffenen Teile können wiederverwendet werden).
Preemptions-Strategie austauschbar: Im Priority-Modus wird der running-Request mit dem größten (priority, arrival_time) gekickt (priority preempt:547-552); im FCFS-Modus wird direkt self.running.pop() ausgeführt (FCFS preempt:572) — semantisch konsistent mit der Ausleitung.
Schritt mit Preemption nimmt keine neuen Requests an: Der Eingang der waiting-Schleife von schedule() prüft not preempted_reqs (waiting Eingang:637), damit ein zurück nach waiting gekickter Request nicht im selben Schritt erneut herausgeholt wird und mit den laufenden preemptenden running-Requests um Ressourcen streitet.
Prefill-Throttle ist DP-Fairness: Zwischen DP-Ranks müssen die Prefill-Schritte ausgerichtet sein, sonst hat ein Rank bereits mit Decode begonnen, während ein anderer noch im Prefill ist, was den Durchsatz bremst. Throttle lässt im nicht ausgerichteten Schritt nur Decode laufen und verschiebt den Prefill auf den Alignment-Schritt (defer_prefills:436-438).
Capacity-bound schaltet automatisch ab: prefill_capacity_bound (prefill_capacity_bound:290) notiert, ob der letzte Alignment-Schritt waiting geleert hat; ist das der Fall, muss nicht weiter gedrosselt werden — weiteres Throttle würde die GPU nur leerlaufen lassen.
Async KV-Load mit eigener Reservierung: _inflight_prefill_reserved_blocks (_inflight_prefill_reserved_blocks:2400) zählt, wie viele Blöcke noch für alle in-flight Prefills benötigt werden; asynchron geladene Requests werden nicht preempted und müssen ausreichend Blöcke reserviert bekommen, bevor sie eintreten dürfen, um Deadlock zu vermeiden.

Schlüsseldateien

_preempt_request:1140 — Preemption-Eingang; Blöcke freigeben, Status zurücksetzen, zurück nach waiting legen.
_free_request_blocks:1149 — Gibt die KV-Blöcke des Requests frei; bei aktiviertem defer_block_free werden sie in die deferred-Warteschlange verschoben.
PREEMPTED:1152-1153 — Status auf PREEMPTED setzen, num_computed_tokens auf 0 zurücksetzen.
prepend_request:1161-1162 — An die Spitze von waiting legen und in reset_preempted_req_ids aufnehmen, um den Worker zu benachrichtigen.
Preemption-Schleife:534-575 — while True-Schleife der Preemption, wenn allocate_slots fehlschlägt.
priority wählt Opfer:547-552 — Im Priority-Modus wird der running-Request mit dem größten (priority, arrival_time) gewählt.
FCFS wählt Opfer:572 — Im FCFS-Modus wird das Ende von running gepoppt.
defer_prefills:436-438 — Hauptschalter für DP-Prefill-Throttle; bei Throttle und nicht gesättigt wird Prefill verschoben.
running verschiebt Prefill:467-471 — in-progress prefill chunks werden im Throttle-Schritt übersprungen.
waiting verschiebt Prefill:801-804 — neue Prefills in waiting brechen im Throttle-Schritt direkt ab.
prefill_capacity_bound Update:1019-1020 — Aktualisiert dieses Flag am Ende jedes Nicht-Throttle-Schritts.
prefill_capacity_bound Initialisierung:290 — In __init__ standardmäßig False.
_inflight_prefill_reserved_blocks:2400 — Statistik der reservierten Blöcke für async KV-Load, um Deadlock zu vermeiden.
waiting Eingangsschutz:637 — Tritt in diesem Schritt eine Preemption auf, werden keine neuen Requests aus waiting angenommen.

Datenfluss

Der Trigger der Preemption liegt in der running-Schleife, wenn allocate_slots None zurückgibt:

python

# 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

Dieser Abschnitt steht in scheduler.py:545-582. Im Priority-Modus gibt es ein Detail: Wenn das gewählte Opfer in dieser Runde bereits eingeplant war (in scheduled_running_reqs), müssen sein Token-Budget, seine neuen Blöcke, Spec-Token und Encoder-Budget vollständig zurückgegeben werden (Budget zurückgeben:553-570) und req_index -= 1, da die running-Liste um ein Element kürzer wurde. preempted_req == request ist die Abbruchbedingung — wenn die Preemption bei sich selbst ankommt, wird das Scheitern anerkannt und die Schleife per break verlassen.

Der gekickte Request läuft durch _preempt_request:

python

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)

Dieser Abschnitt steht in _preempt_request:1140-1162. Beachte drei Aktionen: ① _free_request_blocks gibt die KV-Blöcke an die block pool zurück (bei aktivem defer_block_free werden sie in die deferred_frees-Warteschlange geschickt, die auf ein fence wartet) (_free_request_blocks:2130-2143); ② encoder_cache_manager.free gibt den Encoder-Cache synchron frei; ③ num_computed_tokens = 0 setzt den Fortschritt auf 0 — aber die per Prefix-Cache getroffenen Blöcke werden innerhalb von _free_request_blocks beibehalten (da getroffene Blöcke geteilt sind und nicht exklusiv diesem Request gehören), sodass beim erneuten Prefill nicht zwingend wirklich neu gerechnet wird.

Der Pfad des Prefill-Throttle ist separat. defer_prefills wird zu Beginn von schedule() einmal berechnet (defer_prefills:436-438): throttle_prefills and not self.prefill_capacity_bound and any(not r.is_prefill_chunk for r in self.running). Nur wenn alle drei Bedingungen erfüllt sind, wird verschoben — das Throttle-Flag kommt vom DP-Engine-Core, prefill_capacity_bound ist das Flag, dass waiting im vorherigen Schritt geleert wurde, und die letzte Bedingung stellt sicher, dass Prefill nur verschoben wird, wenn sich in running auch Decode befindet (Decode hat niedrigere Priorität als in-progress prefill chunks). Danach wirkt es an zwei Stellen: in-progress prefill chunks in running werden übersprungen (running verschieben:467-471) und neue Prefills in waiting brechen direkt ab (waiting verschieben:801-804). Am Ende eines Nicht-Throttle-Schritts wird prefill_capacity_bound = bool(self.waiting) (capacity_bound Update:1019-1020) aktualisiert; wenn im nächsten Throttle-Schritt waiting bereits leer ist, wird nicht weiter gedrosselt.

Grenzen und Fehler

Nur RUNNING kann preempted werden: Am Anfang von _preempt_request steht die Assertion request.status == RequestStatus.RUNNING (assert:1146-1148). Requests in waiting können nicht „nochmal preempted" werden, da sie keine KV-Blöcke halten.
assert: Request wurde bereits aus running entfernt: Der docstring stellt klar „popped from the running queue outside of this method" (NOTE:1143-1144). _preempt_request ist nur für Status und das Zurücklegen in waiting zuständig und berührt die Liste self.running nicht; der Aufrufer muss vorher poppen.
spec_token_ids wird geleert: Bei Preemption wird request.spec_token_ids = [] (spec leeren:1154-1155) gesetzt, damit der Spec-Decoder bei der nächsten Neuplanung keine veralteten Draft-Token liefert.
num_preemptions wird hochgezählt: Bei jeder Preemption wird request.num_preemptions += 1 (num_preemptions:1156) ausgeführt; dieser Zähler markiert in den Prefix-Cache-Stats preempted=True (preempted Markierung:721) für die Observability.
Deferred Free gegen nebenläufige Schreibzugriffe: _free_request_blocks gibt bei defer_block_free=True und last_sched_seq > processed_step_seq nicht sofort frei, sondern pusht in die deferred_frees-Warteschlange und wartet auf ein fence (_free_request_blocks:2130-2143), da bei async Scheduling ein anderer in-flight Schritt diese Blöcke gerade noch schreiben könnte.
reset_preempted_req_ids benachrichtigt den Worker: Die preemptete Request-ID landet in reset_preempted_req_ids (reset_preempted:1162) und wird in SchedulerOutput.preempted_req_ids (SchedulerOutput.preempted_req_ids:1100) gepackt, damit der Worker den zugehörigen KV-Cache-Zustand löschen kann.
Throttle betrifft Decode nicht: defer_prefills überspringt nur Requests mit is_prefill_chunk (running Bedingung für Verschiebung:467); reine Decode-running-Requests werden weiterhin normal eingeplant, damit die GPU im Throttle-Schritt nicht leerläuft.
Async KV-Load nicht preemptbar: Ein Request, der über load_kv_async läuft, kommt in den Status WAITING_FOR_REMOTE_KVS (WAITING_FOR_REMOTE_KVS:950) und nicht nach running; über _inflight_prefill_reserved_blocks (_inflight_prefill_reserved_blocks:2400) werden Blöcke reserviert, damit er nicht nach running gelangt und anschließend preempted wird, was den KV-Transferzustand durcheinanderbringen würde.

Zusammenfassung

Preemption ist der Selbstschutz bei knapper KV-Cache: Ein niedrig priorisierter Request wird aus running gekickt, Blöcke für hoch priorisierte Requests freigegeben und der Gekickte an die Spitze von waiting gelegt, damit er im nächsten Schritt neu eingeplant wird. Prefill-Throttle ist eine andere Form der Zurückhaltung in DP-Szenarien: im nicht ausgerichteten Schritt läuft nur Decode, Prefills werden auf den Alignment-Schritt verschoben. Beides ist der Scheduler, der aktiv auf „noch einen Prefill einplanen" verzichtet, um die Systemstabilität zu wahren. Wo die Preemption in der Hauptschleife des Schedulers liegt, siehe Scheduler.schedule; wie der gekickte Request sich neu einreiht, siehe RequestQueue.

Siehe offizielle Dokumentation: vLLM-Dokumentation · README.

Preemption und Prefill-Throttle: Selbstschutz bei knapper KV-Cache ​

Verantwortung ​

Entwurfsmotivation ​

Schlüsseldateien ​

Datenfluss ​

Grenzen und Fehler ​

Zusammenfassung ​