[LyxWxj]

Required human contributor answers

The human contributor must answer all four questions before PR submission.

1. What outcome are you trying to achieve?

2. What is the highest-leverage variable for that outcome?

3. How will this work move that variable rather than create busywork?

4. Who is the customer or user, and what do they actually need?

PR necessity and bugfix evidence

Observed failure, reproduced locally against upstream/main a87f6c8:

tests/test_prefill_livelock.py::test_priority_staged_prefill_drains_under_cache_pressure[refill-True-seed-8]
finishes only 3/10 requests in 500 engine steps with 460 preemptions. With the
fix it finishes 10/10 in 124 steps with 11 preemptions.

Purpose

Fix a scheduler livelock in priority admission under wavefront chunked prefill
and staged prefill: restore-heavy workloads could keep bypassing a blocked
request without ever letting active requests drain.

Scheduler._admit() already bounds bypasses with admission_bypasses: each
successful fresh admission increments the counters of earlier deferred
requests, and the scan stops at max_admission_bypasses. The restoration paths
did not participate:

• ResumeResult.RESTORED returned before the common accounting, so successful
  restorations never incremented deferred bypass counters.
• ResumeResult.BLOCKED deferred the request and continued without checking
  its own counter against the limit.

As a result, a restore-heavy workload could bypass a blocked request
indefinitely and keep preempting/restoring without making progress. The fix
makes both paths honor the existing mechanism:

resume_result = self._try_resume(request)
if resume_result == ResumeResult.BLOCKED:
    deferred.append(request_id)
    if request.admission_bypasses >= self.config.max_admission_bypasses:
        break
    continue

if resume_result == ResumeResult.NOT_PREEMPTED:
    # fresh allocation path

# RESTORED falls through to the common accounting below.
active_count += 1
for blocked_id in deferred:
    self.requests[blocked_id].admission_bypasses += 1

Revision note

This revision replaces the earlier KV-capacity preflight and retry-state
approach with the smaller fairness fix requested in review. The preflight and
retry state are removed:

• Scheduler._priority_preemption_can_admit
• Scheduler._available_admission_budget
• Scheduler._record_failed_priority_preemption
• Scheduler._blocked_priority_preemptions
• Scheduler.resume_blocks_callback
• PreemptionManager.resume_required_blocks
• KVCacheManager.prefix_admission_bounds
• KVCacheManager.reclaimable_blocks

The corresponding tests and documentation were removed or reverted as well.
Those mechanisms may be a useful separate optimization, but they are not
required for this bugfix and duplicated the admission logic.

Scope and remaining limitation

The fix bounds repeated bypasses and restores forward progress. It does not
eliminate every ineffective preemption; reducing those can be a separate
optimization with its own benchmarks.

Test Plan

vllm-rlt Version / Commit: 0.2.0 / cd67fc8 on fix/wcpb-prefill-livelock,
with a net diff against upstream/main a87f6c8.

OMP_NUM_THREADS=1 python -m pytest -q tests/test_prefill_livelock.py

OMP_NUM_THREADS=1 python -m pytest -q \
  tests/test_prefill_livelock.py tests/test_cdb_runtime.py tests/test_kv_cache.py \
  tests/test_scheduler_selection.py tests/test_stage_engine.py \
  tests/test_engine_core_contract.py

OMP_NUM_THREADS=1 python -m pytest -q \
  tests/test_engine.py tests/test_split_prefill.py tests/test_prefix_growth.py \
  tests/test_mixed_execution.py tests/test_async_state.py

python -m ruff check vllm_rlt/core/scheduler.py tests/test_prefill_livelock.py
python -m ruff format --check vllm_rlt/core/scheduler.py tests/test_prefill_livelock.py

Test Result

• Refill seed-8 reproduction: 3/10 finished in 500 steps with 460 preemptions
  before the fix; 10/10 finished in 124 steps with 11 preemptions after it.
• tests/test_prefill_livelock.py: 9 passed. Covers two arrival seeds across
  refill/mixed modes with wavefront on/off, asynchronous scheduling, resumption
  bypass accounting, and the blocked-restoration limit.
• Scheduler/KV/engine CPU suites: 266 passed, 89 skipped (GPU-only cases).
  Workload outputs match the no-preemption reference; no KV allocations or
  snapshots remain after each run.
• Reviewer validation on main a87f6c8: 178 CPU regression/reproduction
  tests passed and 96/96 bounded workload cases passed for the same minimal
  patch.
• ruff check and ruff format --check pass on the changed files.

[bjf-frz]

This bugfix can be much smaller. Successful resumptions currently skip the existing admission-bypass accounting, and blocked resumptions skip its limit check. Making both paths honor the existing fairness mechanism is enough to resolve the reproduced workloads; we do not need the new KV-capacity preflight APIs and retry-state machinery for this fix.

Proposed change in Scheduler._admit() (successful resumptions fall through to the existing common accounting):

@@ -515,25 +515,25 @@ class Scheduler:
                 self._enter_prefill(request)
             else:
                 resume_result = self._try_resume(request)
-                if resume_result == ResumeResult.RESTORED:
-                    active_count += 1
-                    continue
                 if resume_result == ResumeResult.BLOCKED:
                     deferred.append(request_id)
+                    if request.admission_bypasses >= self.config.max_admission_bypasses:
+                        break
                     continue
 
-                plan = self._plan_admission(request, active_count)
-                admitted = self._commit_admission(request, plan)
-                while not admitted and self.preempt(request, waiting_only=True):
-                    # Recompute prefix claims and growth reservations after
-                    # releasing a resident allocation; the previous plan is stale.
+                if resume_result == ResumeResult.NOT_PREEMPTED:
                     plan = self._plan_admission(request, active_count)
                     admitted = self._commit_admission(request, plan)
-                if not admitted:
-                    deferred.append(request_id)
-                    if request.admission_bypasses >= self.config.max_admission_bypasses:
-                        break
-                    continue
+                    while not admitted and self.preempt(request, waiting_only=True):
+                        # Recompute prefix claims and growth reservations after
+                        # releasing a resident allocation; the previous plan is stale.
+                        plan = self._plan_admission(request, active_count)
+                        admitted = self._commit_admission(request, plan)
+                    if not admitted:
+                        deferred.append(request_id)
+                        if request.admission_bypasses >= self.config.max_admission_bypasses:
+                            break
+                        continue
 
             active_count += 1
             for blocked_id in deferred:

Validated on main a87f6c8: 178 CPU regression/reproduction tests passed. A separate bounded scan passed 96/96 workload cases across refill, mixed and asynchronous scheduling, with outputs matching the no-preemption reference and no remaining KV allocations or snapshots. This does not eliminate every ineffective preemption; reducing those can be a separate optimization.

[LyxWxj]

@bjf-frz Thanks — agreed, and I've reduced the change to the patch you proposed.

The root cause was the admission fairness accounting, not missing capacity
prediction. RESTORED returned before the common accounting, so successful
restorations never incremented admission_bypasses for earlier deferred
requests; BLOCKED deferred without checking its own counter against
max_admission_bypasses. A restore-heavy workload could therefore bypass a
blocked request indefinitely and keep preempting/restoring without draining.
The earlier capacity-preflight and "allow one failure" ideas were treating the
symptom; the existing bypass limit is the right retry bound.

I applied your _admit() change directly: RESTORED now falls through to the
common active_count / bypass accounting, BLOCKED checks the limit before
continuing, and the fresh allocation path is under NOT_PREEMPTED. I removed
the KV-capacity preflight and retry-state machinery from the previous revision
(_priority_preemption_can_admit, resume_blocks_callback,
prefix_admission_bounds / reclaimable_blocks,
_blocked_priority_preemptions, etc.) and their tests/docs.

Validation on head cd67fc8:

• refill seed-8 reproduction: 3/10 finished in 500 steps with 460 preemptions
  before the fix; 10/10 in 124 steps with 11 preemptions after it.
• 266 CPU regression tests passed (including the 9 targeted livelock tests),
  89 skipped.
• ruff check/format clean on the changed files.

I agree that eliminating the remaining ineffective preemptions is a separate
optimization; this PR now only restores forward progress. I've updated the PR
description accordingly. Thanks for the independent 178-test / 96-workload
validation on a87f6c8 — that matches what I see locally.

Could you take another look at cd67fc8 when you have a chance?