[Bugfix] Resolve PD scheduler stall via lateral preemption

[Dao007forever]

Purpose

Fix a 1P2D deadlock observed under high concurrency with MultiConnector (Mooncake + NIXL) where the prefill rank's scheduler permanently stalls. Two requests, each holding ~half the rank's KV cache pool, end up at the head of self.skipped_waiting:

• Owner A is WAITING (promoted from WAITING_FOR_REMOTE_KVS after a Mooncake prefix-cache load) and fails its full-ISL admission check.
• Owner B is WAITING_FOR_REMOTE_KVS with recv complete (its ID is in finished_recving_kv_req_ids).

Today's outer break on allocation failure in Scheduler.schedule() (scheduler.py:762) exits the loop before peeking past A, so B is never promoted, never runs, and never frees its blocks. Throughput on the rank drops to 0 for 20-30 minutes until the router HTTP-times out the requests.

The bug is a control-flow head-of-line block:

• (a) a prefix-loader connector that parks requests in WAITING_FOR_REMOTE_KVS with blocks already held
• (b) two such requests that together exceed the pool, no free blocks for them to continue decoding.
• (c) running ≈ 0 so the existing running-side preemption valve has no target.

Approach

Treat a block-holder without forward-pass progress as preemptible the same way self.running is already preemptible (see _preempt_request). The admission loop, on allocation failure, searches the skipped queues for a least-progressed lateral candidate (recv-done WAITING_FOR_REMOTE_KVS or promoted-not-yet-run WAITING/PREEMPTED), preempts it, and retries. RUNNING is deliberately excluded — preempting running compute to admit a zero-compute target is priority inversion; the existing running-loop valve handles that tier unchanged.

Two clean tiers:

Loop	Victim pool	Beneficiary
Running loop (existing)	RUNNING	RUNNING needing more blocks
Admission loop (this PR)	non-RUNNING block-holders without forward-pass	WAITING admission target

Changes

• Request.has_executed (vllm/v1/request.py): set to True at the self.running.append site. Tracks whether the worker has local state for the request, used to choose the correct re-admission status.
• _preempt_blocked_waiting_request (vllm/v1/core/sched/scheduler.py): the lateral-preempt path. Frees blocks, drops stale finished_recving_kv_req_ids / failed_recving_kv_req_ids entries (so the request doesn't auto-promote off leftover state), bumps num_preemptions, and selects status by has_executed — PREEMPTED if the worker can resume from cached state, else WAITING so re-admission flows through scheduled_new_reqs.
• Admission loop: bounded retry (16 attempts per schedule() call) invoking lateral preempt on alloc failure.
• _try_promote_blocked_waiting_request: switch the PREEMPTED-vs-WAITING decision from num_preemptions to has_executed. Without this, a never-executed lateral victim that later auto-promotes from a fresh recv would be routed through scheduled_resumed_reqs and KeyError in gpu_model_runner._update_states.

Alternative considered

Reserve full ISL eagerly at async-load admission, instead of resolving wedges after the fact. Today full_sequence_must_fit checks the full sequence against block_pool.get_num_free_blocks() but only allocates the prefix — that gap is exactly what lets B park its prefix blocks while A's remaining blocks have no home. The simpler fix would track a per-request "pending full-ISL reserve" (full − currently-allocated) and subtract it from the effective free count for subsequent admission checks. No wedge can form because B is denied entry to WAITING_FOR_REMOTE_KVS while A's latent capacity is held.

Rejected because we want to keep the full-ISL check measured against currently free blocks, not against a deflated effective count. Reserving full ISL eagerly serializes async loads behind the request ahead of them, exposing recv latency that today is overlapped with other work. In the wedge-free common case (single owner per pool, or pool sized comfortably above ISL × concurrency) this costs throughput on every load; preempt only pays a cost when the wedge actually fires. Throughput is the priority — the wedge is a tail event, not the steady state.

Duplicate-work check

Searched vllm-project/vllm for open PRs and issues touching skipped_waiting, WAITING_FOR_REMOTE_KVS head-of-line, MultiConnector PD stalls, and 1P2D deadlocks — no duplicates found. PR #30794 (p2p_nccl async KV loading) is a different connector path; PR #40795 (token-level DP load balancing) is upstream of the scheduler entirely.

Test Plan

.venv/bin/python -m pytest tests/v1/core/test_scheduler.py -v

New tests in tests/v1/core/test_scheduler.py:

• test_lateral_preempt_candidate_predicate — _is_lateral_preempt_candidate classification across statuses (recv-done WAITING_FOR_REMOTE_KVS, promoted-with-nc WAITING/PREEMPTED, RUNNING and structured-output excluded).
• test_lateral_preempt_find_victim_excludes_target — the failing-to-admit request is never picked as its own victim.
• test_lateral_preempt_find_victim_picks_least_progressed — among candidates, the one with smallest num_computed_tokens wins.
• test_lateral_preempt_find_victim_searches_step_queue — both step_skipped_waiting and self.skipped_waiting are scanned.
• test_lateral_preempt_does_not_consider_running — RUNNING requests are off-limits.
• test_lateral_preempt_resolves_wedge_e2e — drives a real wedge through schedule()/update_from_output(), asserts the wedged request is admitted to self.running and the recv-done peer is preempted with num_preemptions=1, status WAITING (because has_executed=False), blocks freed, and dropped from finished_recving_kv_req_ids.
• test_lateral_preempt_holds_when_only_in_flight_recv_behind — in-flight (recv-not-yet-done) WAITING_FOR_REMOTE_KVS is not a candidate; A holds, no spurious preemption.

Test Result

Full tests/v1/core/test_scheduler.py suite:

================= 103 passed, 17 warnings in 142.55s (0:02:22) =================

pre-commit run --files on the three changed files: all hooks Passed (ruff check, ruff format, typos, mypy, SPDX, lazy-imports, forbidden imports, torch.cuda check, boolean-ops-in-with, suggestion).

Production: the corresponding lateral-preempt fix in vllm-mooncake resolved the 1P2D × 256 stall scenario that previously hung 20-30 min until router HTTP timeout, with no regressions in 2P2D runs.

---

AI assistance was used in producing this change. The submitter has reviewed each changed line and run the test suite.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request implements lateral preemption in the v1 scheduler to resolve head-of-line blocking (wedges). It introduces a mechanism to preempt requests that hold KV blocks but have not yet executed a forward pass, allowing higher-priority requests to proceed. Key changes include a retry loop in the schedule method, new helper methods for identifying and preempting lateral victims, and the addition of a has_executed flag to the Request class to ensure correct state transitions during re-admission. Extensive tests were added to validate these changes. I have no feedback to provide.

[mergify]

Hi @Dao007forever, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @Dao007forever, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Dao007forever.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @Dao007forever, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10