[LyxWxj]

Required human contributor answers

1. What outcome are you trying to achieve?

I want to separate request lifecycle, scheduled work, and model execution. StagePhase currently carries several of these meanings, which makes scheduling and execution ownership harder to follow and extend.

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

The highest-leverage variable is how the scheduler represents and selects pending work. Explicit WorkKey values let the scheduler identify work independently from request lifecycle (RequestStatus) and model operations (Operation).

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

This change introduces lifecycle-only RequestStatus and seven valid WorkKey values while keeping Operation as the model-call boundary. Scheduler queues and immutable scheduled work use WorkKey; compatible prompt and decode work can be selected into the same single-operation batch. Admission slots separately track execution-resource ownership, including asynchronous work that has logically finished but has not retired.

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

The primary users are contributors maintaining or extending the vllm-rlt scheduler and engine. They need clear ownership of request lifecycle, pending work, and occupied execution resources, plus a reliable way to batch compatible prompt and decode work.

PR necessity and bugfix evidence

A separate PR is needed because the splite-prefill and prefill-decode mixed batching are already in main, while the Stage/WorkKey redesign is not. This PR isolates that scheduler and execution-model change for review as a follow up for #118 .

This is primarily an architecture and batching capability change, not a report of one production incident. The tests cover scheduler selection, mixed execution, preemption, asynchronous speculative work, and PD paths; no single production failure is claimed as the PR's justification.

Purpose

This PR separates request lifecycle from scheduled work and model execution:

• RequestStatus represents WAITING, RECEIVING, TRANSFERRING, ACTIVE, and FINISHED.
• WorkKey represents the seven supported pending work types: atomic prompt, split prompt prelude/recurrent, decode prelude/recurrent/coda, and speculative work.
• Operation remains the model execution boundary: PRELUDE, RECURRENT, or CODA.
• Scheduler queues are keyed by WorkKey. A request can have at most one pending work item, and scheduled batches and asynchronous tickets preserve the selected work identity.
• admission_slots tracks execution-resource ownership independently of lifecycle status. activate_request() centralizes activation and slot reservation for admission, PD activation, and preemption resume.
• Mixed scheduling can combine compatible prompt and decode prelude work, or prompt and decode recurrent work, while keeping each model batch bound to one Operation. Decode readiness and fairness decisions use exact decode work queues.
• Preemption, asynchronous speculative execution, PD, engine progress, serving tests, benchmark traces, and scheduler documentation are migrated to the new work identity.

Atomic prefill remains the default. Mixed scheduling requires split prefill and does not support speculative work. Each request still has at most one queued WorkKey; this change provides the scheduler basis for WCPB but does not add multiple concurrent prompt tasks per request.

Test Plan

vllm-rlt Version / Commit: base main 4a350344027960a48384626aa646f1765fda6d4f;

• Run the A100-supported GPU-marked suite with two GPUs reserved through canhazgpu. FlashAttention-specific tests and the listed FA4/UVA patterns are excluded from this run; FA4 is unsupported on the available A100 SM80 device.
• Run PD prefix-reuse tests with two GPUs, the split-prefill decode-graph test, and the KV memory-plan recreation test separately.
• Run Ruff and git diff --check on the changed files.
• GSM8K accuracy and performance comparisons were not run for this scheduler refactor.

Test Result

• A100-supported GPU suite: 270 passed, 1 skipped, 1 failed, 843 deselected in 241.91 seconds.
• The remaining suite failure is tests/test_cdb_runtime.py::test_cuda_memory_plan_reclaims_allocator_residue_on_recreation: within the suite, the recreated KV budget differed by 190,840,832 bytes, exceeding the test's 16 MiB tolerance. The same test passed in isolation (1 passed in 1.70 seconds); the suite-level variance remains unresolved.
• PD prefix reuse with two reserved GPUs: 6 passed in 40.26 seconds.
• tests/test_split_prefill.py::test_cuda_split_prefill_survives_decode_graph_capture_and_replay: 1 passed in 1.82 seconds.
• python -m ruff check tests/test_split_prefill.py and git diff --check passed.
• The full CPU regression suite and pre-commit run --all-files were not run on this head. The excluded FlashAttention/FA4 paths were not validated on this A100.

[LyxWxj]

cc @bjf-frz @Carlos779988

[0z5a]

Thanks — separating lifecycle (RequestStatus), pending work (WorkKey) and execution ownership (admission_slots) makes the scheduler much easier to follow. I checked the refactor against its base 4a35034 with CPU probes and on an RTX 5090.

Equivalence inside the diff looks good. For refill-atomic, no-refill, split-refill, priority + preemption (atomic and split) and synchronous speculation (atomic and split), the per-step batch composition is identical to the base, and so are the outputs. On every step of sync, CPU-async, mixed, preemption and random-abort runs, these invariants held:

• len(admission_slots) <= max_num_seqs
• every ACTIVE request holds a slot
• every slot belongs to a live ACTIVE, TRANSFERRING or FINISHED request
• at most one pending WorkKey per request
• WAITING requests are only in waiting
• no slots or blocks remain at drain.

On the RTX 5090 (Triton), the GPU-marked suite at this head gives 245 passed, 53 failed, 24 skipped. The failures are 46 FlashAttention cases (both flash_attn and flash_attn_4 report that paged FA4 is not supported on SM120), 6 PD cases that need a second GPU, and test_cuda_memory_plan_reclaims_allocator_residue_on_recreation. That is the same environmental set I see on main-based code without this PR. CPU-async in the probes does not create AsyncState and does not use UVA or graphs, so the GPU async paths are only covered by the GPU suite.

Blocker: rebase onto main (#52, #128, #141)

vllm_rlt/core/scheduler.py and vllm_rlt/engine/preemption.py conflict with #52 ("yield execution slots without offloading resident KV"). The two designs fit together: #52's "yield the execution slot, keep KV" is exactly what admission_slots + park_waiting can express. However, three assumptions in this PR contradict main and will break after a naive resolution:

1. engine/preemption.py suspend() raises when pending_work(victim) is None. On main, capacity preemption may pick a victim that is already WAITING with resident KV (a priority victim that kept its KV) and only offloads it. PreemptionManager.suspend then reuses the existing snapshot without syncing, rewinding, suspending the runner or counting a preemption. So:

   • when a snapshot exists, take work, hidden and token from it;
   • call pending_work / _rewind_split_prefill / model_runner.suspend / preemptions += 1 only on the first suspension;
   • port main's kv is None resume path.

2. _is_preemption_candidate requires ACTIVE + slot + pending work. That rejects main's WAITING-but-resident victims. Here is the failure:

   • A high-priority arrival yields a low-priority victim, which keeps its KV.
   • The arrival's allocation then fails.
   • preempt(..., waiting_only=True) finds no candidate, so the arrival is deferred.
   • The resident victim resumes into the freed slot: priority inversion, and possibly ping-pong.

   The predicate needs a second case — WAITING with an allocation, a snapshot and not priority_only — plus main's waiting_only filter.

3. preempt() always calls park_waiting. For an already-WAITING victim this moves it to the tail of waiting; main returns early to preserve its queue position.

Main also changed the following, which need porting:

• the _admit early returns, expressed with has_admission_capacity()
• the retry loops in _try_resume, in the admission commit and in _ensure_execution_capacity
• the relaxed speculative + preemption check in stage_engine.

Two places auto-merge without a conflict marker but break:

• tests/test_speculative.py (main's new test) asserts suspended.phase == RequestPhase.WAITING.
• benchmarks/preemption.py counts request.phase.value == "waiting".

Main's new speculative-preemption tests need to pass after the rebase:

• test_priority_preempts_only_between_speculative_rounds_and_resumes
• test_speculative_preemption_preserves_request_rng
• test_abort_suspended_speculative_request_allows_id_reuse

The A100 results in the description predate #52/#128/#141, so they do not cover the merged state.

Should fix

1. MixedPolicy now co-batches PRELUDE work. PRELUDE_WORKS groups DECODE_PRELUDE with PROMPT_PRELUDE, so prompt chunks join top-priority decode-prelude batches. That happens before admission and before the decode_due check, and record_batch counts these batches as prefill.
   • On CPU the outputs are correct.
   • No test asserts that such a batch forms: the mixed test only requires len(works) == 2 somewhere.
   • Under async scheduling, mixed prelude batches skip the AsyncState.prepare routing in model_runner.py, which before this PR only ran for RECURRENT.
   • Please add a test that asserts a DECODE_PRELUDE + PROMPT_PRELUDE batch, covering CUDA, async and graphs, plus a boundary-KV model if mixed mode is allowed there.
2. MixedPolicy ordering changed. The old Mixed policy took the mixed recurrent batch first whenever both recurrent queues were ready. The new one inherits refill order: DECODE_PRELUDE, then CODA, then decode_due, and only then mixed recurrent.
   • The batch composition differs from the base. In one of four seeds, mixed recurrent batches went from 39 to 35, CODA batches from 26 to 36, and steps from 119 to 131; the other seeds improved slightly.
   • tests/test_mixed_execution.py now clears the other queues to get the old selection, which suggests the test was adapted to the new order rather than kept.
   • prefer_recurrent in stage_engine.py also now fires for mixed in-flight batches.
   • Please either restore the old priority, or document it and add a throughput/GSM8K comparison for mixed mode.
3. Docs. docs/scheduler.md still says Mixed combines only DECODE_RECURRENT + PROMPT_RECURRENT and that "PRELUDE and CODA are selected separately". The MixedPolicy row still describes the old rule. The SVGs (loop-scheduler.svg, stage-components.svg and others) and docs/unified_execution.md still say "phase queues", "take_phases" or "saved phase".
4. A behavior change is folded into a test migration. tests/test_transfer_connector.py now calls _admit() before the export is suspended or acknowledged.
   • Previously, a never-admitted export cancelled into TRANSFERRING counted against max_num_seqs. Now it does not.
   • That is probably a fix, but the original "unsubmitted export" scenario is no longer tested and the description doesn't mention it.
   • Related: pd/connector.py uses not has_admission_slot for both "released after the final staging read" and "never admitted". This is not reachable today, but it is fragile.

Nits

• Quadratic per-step cost. pending_work / remove_pending_work scan all seven deques with in, and enqueue_work calls them on every enqueue. Re-enqueueing a full batch costs about 0.22 ms at 256 requests and 2.7 ms at 1024. A dict[request_id, WorkKey] index makes it O(1).
• Selection. In the selection loop, if request.request_id in selected: continue drops the queue entry without re-appending it. It is dead under the uniqueness invariant, but it would strand a request silently if the invariant ever broke. Assert or re-append instead.
• Weakened assertion. In tests/test_scheduler_selection.py, batch.requests[0].work == WorkKey.PROMPT_ATOMIC replaced the meaningful "a was not moved to WAITING" check. Assert status == ACTIVE.
• Public API. vllm_rlt.request no longer exports RequestPhase, Operation, ExecutionBatch or StateSpan, with no aliases. Also, benchmarks/cdb_runtime.py trace keys changed from stage to works. Worth a note for downstream users.
• tests/test_engine.py. drain() uses next(iter(batch.works)), which is arbitrary for mixed batches.

Pre-existing, not from this PR

Split prefill + policy="priority" + enable_preemption can livelock. I saw about 1958 consecutive PROMPT_PRELUDE batches with four requests stuck in prompt recurrent. The behaviour is identical on the base, this PR and main.

Repro: CacheConfig(40, 2, incremental_allocation=True), max_num_seqs=4, max_num_batched_tokens=6, prefill_chunk_size=4, atomic_prefill=False. I can open a separate issue if that helps.

[LyxWxj]

@0z5a
Thanks for the detailed review. I addressed the remaining findings in this update.

• Preemption now preserves and reuses an existing resident snapshot when a WAITING request is offloaded later. The candidate check accepts WAITING requests with resident KV and a snapshot for capacity reclamation, and re-preempting a WAITING request preserves its queue position. I also ported the admission and capacity retry paths and added regression coverage.
• MixedPolicy again prioritizes compatible recurrent work when both queues are ready. I added coverage for mixed PRELUDE batches, including async scheduling and CUDA Graph execution, and updated the scheduler docs. Huginn rejects staged prefill, and its existing rejection test covers that boundary.
• The unsubmitted-export case now has its own regression test. An unadmitted export does not consume a compute slot, and acknowledgement alone does not retire it.
• Pending work lookup now uses a request-to-WorkKey index. Duplicate selection raises explicitly, and the weakened ACTIVE assertion and nondeterministic mixed-batch trace helper were corrected.
• The Python API and CDB trace schema changes are documented, including the stage to works trace-key change.

Validation on the rebased revision:

• Focused CPU regression tests: 227 passed, 86 skipped; Huginn tests: 16 passed, 3 skipped.
• Focused GPU tests on an A100: 53 passed. FA4 cases were excluded because this SM80 device does not support FA4.
• Ruff check and formatting pass for the changed Python files.

The priority + split-prefill livelock remains outside this change.