[LyxWxj]

[FEAT] Implement wavefront chunked prefill batching

Required human contributor answers

1. What outcome are you trying to achieve?

Higher throughput by improving recurrent-core batch occupancy at small prefill chunk sizes, while retaining the scheduler boundaries those small chunks provide and preserving LAST_EXITED KV semantics.

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

The number of useful rows in each recurrent-core invocation at a chosen chunk size. Ordinary small-chunk prefill lowers this row count and can leave the GPU underfilled; WCPB aims to increase it without making chunks larger.

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

Batch independent, ready (chunk, depth) tasks into the same recurrent invocation when their hidden-state and KV dependencies are satisfied. This directly increases rows per batch while keeping small chunks as scheduling units. In the cited chunk-16 A100 runs, rows per batch and throughput increased for the four-request and mixed workloads; those results are directional because each variant had only one measured run. The static tests did not include concurrent decode, so they do not establish decode-latency improvements.

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

Users especially when prompt work is sparse or uneven. They need better prompt throughput from small chunks without giving up opportunities to schedule decode or admit new prompts, with TTFT and decode latency measured as part of the trade-off.

PR necessity and bugfix evidence

A PR is needed to deliver the requested WCPB feature. This is a feature implementation, not a claimed bugfix.

Purpose

• Add opt-in WCPB through wavefront_chunked_prefill and --wavefront-chunked-prefill.
• Batch ready prefill chunks across requests and recurrence depths under the existing token and sequence budgets. Mixed scheduling can return to decode work between prefill stages.
• Allocate a bounded per-request hidden-state ring sized for the configured chunk size and recurrence depth. Advance its reusable boundary only after final-depth completion.
• Derive execution boundaries from ExecutionBatch.loop_count, keeping batches atomic when no intermediate scheduling decision is needed.
• Integrate staged prefill with prefix reuse, preemption and recovery, PD chunk transfers, speculative execution, and CUDA graph paths. Reject unsupported shared-KV wavefront configurations.

Test Plan

• CPU targeted tests:
  • tests/test_split_prefill.py tests/test_stage_engine.py tests/test_pd.py: 102 passed, 66 deselected.
  • tests/test_split_prefill.py tests/test_input_preparation.py tests/test_mixed_execution.py: 87 passed, 67 deselected.
• Full non-GPU CPU suite: 825 passed, 34 skipped, 328 deselected.
• pre-commit run --all-files: passed, including Ruff lint, Ruff format, and CPU regression tests.

GPU Test Result

Performance matrix: NVIDIA A100 80GB PCIe, Ouro-1.4B BF16, Triton, max_num_seqs=8, max_num_batched_tokens=1024, 100 requests per profile, 64 output tokens, one run per configuration. Tested chunk sizes 32 and 128 with serial (concurrency=1), concurrent (concurrency=8), and arrival (concurrency=4, 2 req/s) profiles, using fixed 512-token inputs and a variable 256/512/768-token workload. Split and WCPB used mixed; Atomic used refill.

WCPB had higher request throughput than Split in all 12 matched profiles, by 0.8% to 22.6%. The runs recorded request and token throughput plus TTFT, TPOT, ITL, and E2E percentiles. Each profile completed 100 requests with zero failures. These are single-run results, not stable performance estimates. For variable-length chunk-128 traffic, WCPB p99 ITL was 191 ms concurrent and 192 ms arrival, versus 102 ms for Split.

Accuracy: on the fixed GSM8K-87 BF16 greedy protocol, Split and Atomic scored 60/87; WCPB scored 63/87. WCPB output token IDs differed from Split on 24/87 prompts, so the higher score does not establish token equivalence or improved model quality. The WCPB path does not meet the token-preservation expectation in this measurement.

Numerical diagnostic: one 303-position prompt trace with chunk size 64, max_num_batched_tokens=128, and four recurrent passes showed maximum absolute output differences of 0.25, 2.703125, 2.75, and 4.21875 by depth. The address audit found no block, depth/position mapping, or write-address mismatch in that trace; it does not establish a general cause.

vllm-rlt Version / Commit: ff05fe15d182af97a30be321c99e20d5d3b17b92

[bjf-frz]

Could you also help clean up --no-atomic-prefill as part of this PR? This is a leftover from my earlier refactoring: whether a batch needs to return to the scheduler between forward stages should be an internal scheduling decision, rather than something users have to configure.

The conditions should be:

• Non-atomic: wavefront needs to regroup chunks between depths.
• Non-atomic: a mixed batch contains decode work that needs exit decisions and batch updates at depth boundaries.
• Non-atomic: the scheduler explicitly needs to yield before the current chunk finishes, either to run waiting decode/speculative work or to perform preemption. These need concrete scheduling rules; simply having waiting requests or enabling preemption should not force every batch to split.
• Atomic: the batch can run through its remaining depths without an intermediate scheduling decision.

Existing capability restrictions still apply: Huginn and shared KV currently require atomic prefill. If a requested feature requires splitting, an unsupported combination should fail clearly.

Please remove the user-facing flag and derive the execution boundary from these requirements, preferably using the existing ExecutionBatch.loop_count. Buffer allocation, progress updates, and resume handling also need to follow the actual execution plan instead of a global atomic/split switch.

Sorry for leaving this cleanup behind in the refactor, and thanks for helping address it here.

[bjf-frz]

Please address the state consolidation and naming comments below.

[bjf-frz]

Please consolidate prefill_reclaimable_tokens and prefill_chunk_start into a single field representing the start of the oldest resident chunk. In the wavefront scheduling path, both track the boundary of fully completed chunks and advance together. Maintaining both adds redundant state and synchronization work across admission, completion, and recovery.

Keep num_prefilled_tokens separate: it tracks partial progress within a chunk, while the chunk boundary must only advance once the entire chunk finishes and its ring slot can be reused. Check the non-wavefront and PD import paths when consolidating initialization and updates.

[bjf-frz]

Please rename buffer_offset to prefill_hidden_row_offset consistently across ScheduledRequest and StateSpan. It represents a row offset into the request's prefill hidden-state buffer, calculated as slot * chunk_size + chunk_offset. Also document that None means the runner derives the offset from token_start - chunk_start.

[bjf-frz]

Please rename the local variable chunk_offset to hidden_row_offset. For wavefront prefill, this value includes the ring-slot offset, so it represents a position in the entire hidden-state buffer rather than an offset within a chunk.

[bjf-frz]

[LyxWxj]

@bjf-frz Thanks for the review and the beautiful pipeline graph for wcpb. As you suggested I removed the user-facing --no-atomic-prefill flag and addressed the three state and offset comments.

• Execution boundaries now come from the selected schedule and ExecutionBatch.loop_count. WCPB returns at recurrent-depth boundaries to regroup ready chunks. Mixed batches return at depth boundaries when decode work needs exit decisions and batch updates. Other batches run through their remaining depths unless the scheduler has a concrete reason to yield. Waiting requests or enabling preemption alone do not force every batch to split; preemption is triggered by actual capacity pressure. Existing atomic-only capability restrictions remain enforced.
• Consolidated prefill_chunk_start into prefill_reclaimable_tokens, which marks the oldest resident chunk boundary. num_prefilled_tokens remains separate for progress within a chunk. Admission, completion, recovery, and PD transfer initialization use the consolidated field.
• Renamed buffer_offset to prefill_hidden_row_offset across ScheduledRequest and StateSpan. None means the runner derives the offset from token_start - chunk_start.
• Renamed the runner’s local offset and InputSegment field to hidden_row_offset, reflecting that the value indexes the full request hidden-state buffer, including the wavefront ring slot.

[bjf-frz]

lgtm