Skip to content

[Bugfix]Fix livelock under wcpb scheduling - #165

Open
LyxWxj wants to merge 1 commit into
InferMatrixLab:mainfrom
LyxWxj:fix/wcpb-prefill-livelock
Open

LyxWxj wants to merge 1 commit into
InferMatrixLab:mainfrom
LyxWxj:fix/wcpb-prefill-livelock

Conversation

@LyxWxj

@LyxWxj LyxWxj commented Nov 2, 2026

Copy link
Copy Markdown
Contributor

Purpose

Thanks @0z5a for reporting the original split-prefill livelock and providing a detailed reproduction in #137.
As the follow-up noted, the original workload drains on the newer mainline. However, a separate finite workload still reproduces a related livelock with priority preemption and WCPB: requests can be repeatedly preempted and have their unfinished prefill work rewound, even though the higher-priority request cannot obtain enough KV capacity to be admitted.
This PR addresses that preemption cycle while preserving request outputs and resource cleanup.

Bug Description

As discussion under pr #137 shows, when priority preemption is enabled and all active slots are occupied, the scheduler may suspend a lower-priority request to admit a higher-priority request. Suspending an in-flight staged-prefill request rewinds its unfinished chunk to its last fully prefilled frontier.

If the higher-priority request still cannot obtain enough KV capacity, it remains waiting. The suspended request can then resume, perform prefill work, and be preempted again before completing the chunk. Repeating this cycle can prevent a finite workload from draining even though engine.step() continues to return and execute work.

WCPB can yield between recurrence depths, so it does not guarantee that a chunk completes before the next scheduling decision. Default refill without WCPB instead runs a prefill chunk through all recurrence depths before returning to the scheduler.

The original workload reported in #137 drains on the newer mainline. This PR is motivated by a separate finite workload that reproduces a related preemption/replay cycle on the PR’s parent revision.

PR Description

  • Check KV capacity before preempting an active request for a priority slot. The estimate accounts for reclaimable allocations, shared and cached prefixes, active prompt growth, watermarks, and snapshot restoration.
  • Skip a victim when the available capacity is insufficient to admit the waiting request, and try another eligible victim when possible.
  • Retain a fallback guard for uncertain cases: if admission still fails after preemption, do not repeat the same priority-preemption attempt while the available budget remains unchanged.
  • Apply admission-bypass accounting to resumed requests as well as newly admitted requests.
  • Add CPU regression coverage for WCPB and mixed scheduling, including prefix caching, restoration, and insufficient-capacity scenarios.
    For the new ten-request regression workload:
Revision Result Preemptions / resumptions
Parent 88c8f96 3/10 finished after 500 engine steps 465 / 463
With this PR 10/10 finished in 121 engine steps 4 / 4

On the fixed revision, generated token IDs, exit depths, and finish reasons match a no-preemption reference; no KV blocks remain occupied after the workload drains.
The focused regression can be run with:

  python -m pytest -q tests/test_prefill_livelock.py -k 'seed-8 and refill'

This is bounded CPU scheduler evidence, not a general starvation-freedom guarantee or a GPU performance claim.

Required human contributor answers

The human contributor must answer all four questions before PR submission.
An explicit answer already recorded elsewhere may be referenced with a precise,
accessible source; a bare link without an identifiable answer is insufficient.
Agents must not invent or fill in the human's answers. Blank, placeholder, or
unresolved answers do not complete this gate. See the
contributor-intent gate
and contributing guide.

1. What outcome are you trying to achieve?

Prevent a finite workload from getting stuck in a preemption/replay cycle when priority scheduling and staged prefill compete for limited KV-cache capacity.

The intended outcome is that requests can make forward progress and finish, while preserving their generated outputs and releasing their resources when the workload drains.

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

The number of futile priority preemptions: cases where the scheduler rewinds a running request to make room for a higher-priority request, but the latter still cannot obtain enough KV capacity to be admitted.

Repeated futile preemptions erase useful prefill progress without admitting new work. Preventing that cycle directly addresses the observed failure.

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

Before preempting a request for a priority slot, the scheduler estimates whether the waiting request could be admitted after reclaiming eligible KV allocations. The estimate accounts for prompt growth, watermarks, shared prefixes, and resident or offloaded snapshots. If a candidate cannot make enough capacity available, the scheduler tries another eligible victim or leaves the active requests running.

Some information, such as pending prefix publication, is uncertain. If an optimistic preflight still leads to a failed admission, a fallback guard prevents the same request from repeatedly triggering priority preemption until the available budget improves.

The PR also adds regression tests that fail on the parent revision and pass with the fix, rather than relying only on a timeout or a change in scheduling statistics.

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

The users are operators and applications using priority scheduling with preemption and staged prefill under KV-cache pressure, including WCPB workloads.

They need finite requests to complete, not repeatedly replay the same prefill work; they also need correct outputs and KV/snapshot cleanup. This change is a scheduler-liveness fix, not a throughput optimization.

PR necessity and bugfix evidence

The reviewer’s original historical workload drains on the newer mainline, so that result alone is not evidence of an outstanding regression. A different finite, ten-request workload demonstrates that a related livelock remains possible on this PR’s parent.

Using CPU, FP32, the torch attention backend, priority scheduling, WCPB, 40 KV blocks of size 2, max_num_seqs=4, max_num_batched_tokens=6, and prefill_chunk_size=4:

Revision Result Preemptions / resumptions KV blocks at stop
Parent 88c8f96 3/10 requests finished after 500 engine steps 465 / 463 40
Fix a378969 10/10 finished in 121 engine steps 4 / 4 0

The fixed run’s token IDs, exit depths, and finish reasons match a no-preemption reference. The regression also covers mixedmode, prefix-cache accounting, restoration, and insufficient-capacity cases.
Run the focused regression on the PR head with:

CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 \
  python -m pytest -q tests/test_prefill_livelock.py -k 'seed-8 and refill'

For the parent comparison, copy the PR’s tests/test_prefill_livelock.py into the parent checkout first; that file does not exist on the parent. The selected test fails there with 3/10 requests finished and passes on the fix.
The evidence is bounded CPU scheduler validation. It does not establish a general starvation-freedom guarantee, GPU serving correctness, or a performance improvement.

Test Plan

**vllm-rlt Version / Commit:**88c8f96a45690c92a8efea876173cd73487dced7
Run the targeted CPU regression suite:

CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 \
  python -m pytest -q -m 'not gpu' \
  tests/test_prefill_livelock.py \
  tests/test_kv_cache.py \
  tests/test_prefix_growth.py \
  tests/test_scheduler_selection.py \
  tests/test_stage_engine.py \
  tests/test_engine_core_contract.py \
  tests/test_split_prefill.py \
  tests/test_mixed_execution.py

The tests cover:

  • The new ten-request workload and the reviewer’s workload in refill and mixed modes, with and without WCPB where applicable.
  • Output equivalence against a no-preemption reference, including token IDs, exit depths, and finish reasons.
  • Preflight rejection when KV capacity is insufficient, including shared-prefix and snapshot-restoration cases.
  • The fallback guard when an optimistic preflight is followed by failed admission.
  • KV-block and snapshot cleanup after the workload drains.

To reproduce the new workload against the parent revision, copy this PR’s tests/test_prefill_livelock.py into a parent checkout, then run:

CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 \
  python -m pytest -q tests/test_prefill_livelock.py -k 'seed-8 and refill

Test Result

Check Result
Parent 88c8f96, new workload 3/10 requests finished after 500 engine steps; 465 preemptions / 463 resumptions
Fix a378969, same workload 10/10 finished in 121 engine steps; 4 preemptions / 4 resumptions; 0 occupied KV blocks at drain
Targeted CPU suite 199 passed, 57 deselected
Broader CPU suite 866 passed, 35 skipped, 546 deselected; serving and profiling tests excluded
Ruff on vllm_rlt, tests, and benchmarks Passed

The fixed workload’s outputs match the no-preemption reference. A separate local exploratory scan of 120 bounded WCPB runs across refill/mixed and prefix-cache configurations also drained; that diagnostic script is not part of the repository.

Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant