[LyxWxj]

Description

What

This PR refactors the scheduler's wavefront selection and split-prefill progress paths.

The changes keep the existing request lifecycle, token-budget accounting, capacity checks,
prefix publication ordering, and asynchronous event propagation while removing duplicated
control flow between submission-time and output-time updates.

Changes

• Unify PRELUDE and RECURRENT traversal in _take_wavefront_phase.
  • Both phases use the same depth-major candidate traversal.
  • Capacity checks remain interleaved with task acceptance.
  • Already-selected requests remain protected from later preemption.
  • Token-budget truncation still happens during selection.
  • Post-selection task ordering is kept separate from queue ordering and fairness.
• Share prefill progress application between update_from_output and
  update_after_submission through _apply_prefill_phase_progress.
  • Wavefront rows are grouped and consumed once per request/phase.
  • Submission events are preserved for prefix publication and downstream consumers.
• Extract _advance_prefill_depth_frontier for the common one-depth frontier mutation.
  • _apply_split_prefill_progress retains its phase-specific validation and lifecycle rules.
  • _apply_wavefront_prefill_batch_progress retains its ordered-span and chunk-retirement rules.
  • Atomic PREFILL remains a separate path.
• Simplify _apply_prefill_progress dispatching.
• Merge contiguous PRELUDE spans using their ordered token ranges.
• Add regression coverage for multiple recurrent depths submitted in one split-prefill update,
  including frontier propagation, chunk completion, event forwarding, and prefix publication.

Scheduling semantics

The refactor does not move capacity checks or token-budget accounting out of the selection loop. In particular, PRELUDE selection remains depth-major so that a bounded budget can advance different requests at the same depth instead of allowing one request to consume all available depth work first.

Sorting of selected tasks is performed only after selection. It affects the execution-row order, not the request queue order, candidate requeue order, token-budget decisions, or fairness during selection.

Required human contributor answers

1. What outcome are you trying to achieve?

Make the scheduler's wavefront and split-prefill progress logic easier to reason about and maintain without changing the intended scheduling and lifecycle contracts.

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

The highest-leverage variable is the number of independent implementations of the same frontier and lifecycle transition. Reducing duplicated transition logic lowers the risk that submission and output paths diverge.

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

The PR centralizes the shared one-depth frontier mutation and the common prefill update routing, while keeping phase-specific validation, chunk retirement, prefix publication, and queue transitions in their respective owners. The wavefront selection refactor also makes the depth-major fairness rule explicit without moving resource checks out of the selection loop.

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

The primary users are maintainers and runtime users of split-prefill and wavefront-prefill
scheduling. They need predictable budget/fairness behavior, correct depth frontiers, correct prefix-publication timing, and equivalent behavior across synchronous submission and asynchronous output-update paths.

PR necessity

Yes. This is a core scheduler refactor involving request selection, frontier ownership,asynchronous lifecycle updates, and prefix publication. A PR review is needed to verify that the reduced duplication does not weaken scheduling fairness or request-state invariants.

Test plan

CPU Tests

• OMP_NUM_THREADS=1 python -m pytest -q tests/test_split_prefill.py tests/test_scheduler_selection.py tests/test_mixed_execution.py
  • 81 passed, 40 skipped
• OMP_NUM_THREADS=1 python -m pytest -q tests/test_engine.py tests/test_serving.py tests/test_stage_engine.py tests/test_submission_queue.py tests/test_engine_core_contract.py
  • 150 passed, 12 skipped
• git diff --check

GPU Tests

On 2 * A100

• Focused GPU scheduler tests:
  • 120 passed, 1 skipped
• Full repository Ruff check:
  • passed
• Full GPU test suite:
  • 1228 passed, 70 failed, 32 skipped

The remaining full-suite failures are environment or pre-existing test issues rather than failures
in the changed scheduler path:

• Explicit FlashAttention-4 tests are unsupported on the remote A100/SM80 environment.
• Some FlashAttention-2 tests use a paged-KV block size unsupported by the installed backend.
• The CDB cuda_graph_max_batch_size assertion reproduces on a clean 6b164d3 baseline.
• PD tests encounter remote worker-startup timeouts; a representative three-case PD test reproduced the timeout independently of the focused scheduler tests.

vLLM-rlt Version/commit 6b164d3

[LyxWxj]

cc @prettygirlisnotme @MinhaoLi0318 I would appreciate any advice you can offer to improve this implementation and adapt it for future features, such as #156. The fastloop algorithm seems to contradict the assumption that the earlier chunk will not be completed before the later chunk of the same prompt. Since I havn't perused the paper, I may have no much idea for implementing that but I'll catch up.

[MinhaoLi0318]

I reviewed this at b4a848a61 against 6b164d384, mostly from the token-sparse prefill side (#156). The active-row view there (window id, absolute position, active indices) will sit on the wavefront selection and prefill-progress code that this PR reshapes.

What I checked, CPU only (macOS, Python 3.12, torch 2.14.0, OMP_NUM_THREADS=1):

• Your two test-plan commands give the same counts here: 81 passed / 40 skipped, and 150 passed / 12 skipped. On base they give 80 / 40 and 150 / 12; the difference is the new test.
• Full pytest tests: 899 passed / 430 skipped on base, 900 passed / 430 skipped on this head.
• The new test_split_prefill_progress_advances_multiple_recurrent_depths_per_submission also passes on base. It pins the existing non-WCPB multi-depth path; it does not cover the WCPB selection change.
• I traced two 10-token prompts through WCPB (chunk size 2, budgets 3, 5 and 8) and compared every scheduled task (request, start, count, depth) on base and head. Selection, budget truncation and depth order are identical. The only difference is the row order inside PRELUDE batches (inline 1).
• Reading the two versions side by side:
  • Capacity checks still run as each task is accepted, and selected is updated before the next candidate's _ensure_execution_capacity, so already-selected requests are still protected from preemption.
  • The event passed to publish_prefix and the WCPB grouping are the same in update_from_output and update_after_submission.
  • _advance_prefill_depth_frontier reproduces both old frontier updates. The loop_count <= 0 early return keeps final_depth_reached equivalent to the old depth + loop_count == total_depths.

From the active-row side the refactor helps. After it, _advance_prefill_depth_frontier is the only per-step writer of prefill_depth_ranges; the other writers are resets at admission and in stage_engine.py. accept_task is the only place _take_wavefront_phase charges budget and KV capacity. Here is what I expect token-sparse prefill to add on top. None of it is for this PR:

1. Readiness. A sparse depth may start only on a whole window (inline 2).
2. Cost. accept_task, and take_phases L761 for non-WCPB split prefill, would charge the active row count for a sparse depth. That count is retention × window length, which the host knows, so no readback is needed.
3. Row identity. The view is keyed by (request generation, absolute window, depth), not by batch row; inline 1 shows why. Mixed batches are already decode-first, because take_phases walks (RECURRENT, PREFILLING_RECURRENT) in that order. So the sort here only orders rows inside the prompt part of the batch, which fits the decode-first partition the RFC describes.
4. Window origin. _make_wavefront_task derives chunk_start from prefill_window_origin (L621). That origin is the prefix-cache hit, len(prefix) * block_size (L390, L457), while #156 anchors windows to absolute prompt positions. The two agree once prefix hits stop at window boundaries (the prefix-cache part of #156). Until then, the window id cannot be read off chunk_start.

Inline 3 answers your question in this thread about chunk completion order.

Inline 1 is the only change I'd ask for here: one line in the description and a test for the PRELUDE row order. Inline 2 and 3 are questions about direction.

AI assistance: drafted with an AI assistant; reviewed by me before posting.

[MinhaoLi0318]

This sort now applies to PRELUDE too, which changes PRELUDE row order. Before this PR, PRELUDE tasks were appended in selection order, so rows from different requests interleaved. Now they are grouped by request. Two 10-token prompts, chunk size 2, budget 8, first batch:

• base 6b164d384: a:0, b:0, a:2, b:2
• this head: a:0, a:2, b:0, b:2

The selected spans are the same; only the order differs. I don't see anything that depends on it. build_execution_batch (worker/executor.py:26-79) and _prefill_segments (worker/model_runner.py:396-423) give each span its own segment, and prefill_submitted (pd/connector.py:186) works per request. The description says sorting affects execution-row order. Could it also say explicitly that PRELUDE rows were not sorted before? A test would pin the new order. This one fails on base and passes here:

def test_wavefront_prelude_rows_are_grouped_by_request_after_depth_major_selection():
    e = make_prefill_engine(
        scheduler_config=SchedulerConfig(
            max_num_seqs=2,
            max_num_batched_tokens=5,
            prefill_chunk_size=2,
            wavefront_chunked_prefill=True,
        )
    )
    for request_id, prompt in (("a", list(range(1, 11))), ("b", list(range(11, 21)))):
        e.add_request(request_id, prompt, SamplingParams(max_tokens=1, min_loops=1, max_loops=1))

    prelude = e.scheduler.schedule()

    assert prelude.phase is RequestPhase.PREFILLING_PRELUDE
    # Selection is still depth-major (a:0-2, b:0-2, then a:2-3 under the 5-token
    # budget); only the emitted row order is grouped by request.
    assert [
        (item.request.request_id, item.token_start, item.token_count)
        for item in prelude.requests
    ] == [("a", 0, 2), ("a", 2, 1), ("b", 0, 2)]
    e.close()

[LyxWxj]

Agreed. PRELUDE previously emitted spans in selection order; this PR changes that to request-grouped, token-position order after selection. I kept the sorting to share the emission path with RECURRENT, not to introduce a new fairness rule or claim a performance benefit.

The important distinction is that selection and budget allocation are already complete before sorting, and candidate requeueing still follows the original candidate order.

I’ll explicitly document the PRELUDE ordering change and add your proposed regression test. The budget-5 case is particularly useful because it pins both the selected spans under truncation and their emitted order. You’re also right that the existing new test covers the pre-existing multi-depth progress path rather than this ordering change.

[MinhaoLi0318]

A question for token-sparse prefill (#156); no change needed in this PR. This line publishes the next depth's readiness. Under budget truncation it can publish a point inside a chunk, so the next depth then runs part of a chunk. With the two prompts from inline 1 and budget 5, 10 recurrent tasks at depth 2 or 3 cover only half of their chunk, for example a, depth 2, tokens 2-3 of chunk 2-4. Base and head behave the same here. With retention [1, 1, 0.25, 0.10], depth 2 is the first sparse loop, and its top-k needs the scores of the whole window from depths 0 and 1.

I'd keep this helper logical and do the gating where tasks are built: a sparse depth gets a task only once right covers the whole window, since selection needs that window's complete scoring input. Whether the selected active rows may then be split across submissions by the budget is a separate choice. The places are _make_wavefront_task (L642), and for non-WCPB split prefill, _make_scheduled_request (L585-591) and _max_contiguous_prefill_loops (L834). Clamping here instead would leave the tokens between the clamp and token_end in no depth range, and Request.prefill_depth (request.py:89-95) would raise for them. Does that match how you want these helpers to evolve?

[MinhaoLi0318]

On your question in this thread, whether FlashLoop breaks the assumption about chunk completion order: as I read #156, the out-of-order check a few lines below (L1130-1132) stays valid. A sparse depth still advances the per-depth frontier over the whole logical window; only the physical rows shrink. So the final-depth frontier still passes chunk ends in prompt order, and prefill_reclaimable_tokens advances as it does today. Two things change. One is when a sparse depth may start (inline 2). The other is that the active maps and the saved old hidden rows have to stay alive until this same chunk retirement. If you had a different case in mind, could you describe it?

[LyxWxj]

Thanks, this resolves my earlier concern. I was conflating a reduction in physical model rows with a change in logical completion order. I don’t have a separate case that requires out-of-order logical chunk completion.

Provided active updates and inactive inheritance are correctly ordered, the completed logical span still advances the contiguous per-depth frontier. Fewer—or even zero—active rows in a fragment do not by themselves invalidate the existing ordered chunk-retirement check. I would therefore keep that check unchanged.

I would interpret “over the whole logical window” as accounting for all logical positions, not requiring an atomic window-sized frontier update; execution may still advance through valid fragments.

One lifetime detail for independent window/chunk boundaries: window-level maps or saved scoring inputs may be shared across several chunks. Their lifetime must cover all dependent fragments and readers, rather than being inferred solely from one chunk’s retirement. I would give that auxiliary state explicit ownership in #156, without changing this PR’s existing progress and retirement rules.

[prettygirlisnotme]

Thanks, I think we can keep the current progress and chunk retirement rules. Skipped tokens still inherit hidden/KV and count toward the logical span; only the model rows are compacted.

Selection waits for the full window's scoring inputs, then execution can split it into budget-sized fragments. Our fixed-selection CPU prototype passed unchanged on b4a848a, including zero-active fragments, mixed batches and ring reuse. I'll adapt it to Minhao's shared view once it's up.

[bjf-frz]

prelude_cursors is only used by PREFILLING_PRELUDE, but the refactor now builds it for every RECURRENT batch too. Please guard its initialization with if phase == RequestPhase.PREFILLING_PRELUDE to avoid the unnecessary allocation and iteration.

[bjf-frz]

One suggestion on the PRELUDE span validation.

[bjf-frz]

Could we keep the original frontier loop here? It checks both the initial frontier and span continuity in one pass, then updates num_preluded_tokens after validation. The replacement splits the same check across the first span, adjacent pairs, and the last span, while adding a list slice through scheduled_requests[1:]. I don’t see a clarity or correctness benefit, and there’s no measured performance benefit to justify the change.

[LyxWxj]

Aggree. Rolled back to the original implementation.

[bjf-frz]

lgtm