[FEAT] Implement wavefront chunked prefill batching

docs/README.md

@@ -14,6 +14,7 @@ an environment to receiving the first generated response. The
 | [KV layout examples](kv_layout_computation.md) | SHARED and LAST_EXITED semantics and worked attention examples |
 | [KV cache walkthrough](kv_cache_walkthrough.md) | Capacity, validity, shared metadata, transfer pins and snapshot/restore |
 | [Runtime execution and configuration](runtime.md) | Batch contracts, preparation, transfers, execution options and compatibility |
+| [Split prefill and WCPB](split-prefill-wavefront.md) | Chunk/depth scheduling, hidden-state ring, KV dependencies and validation limits |
 | [Asynchronous self-speculation](async-speculative.md) | Shared execution, device-owned positions and round lifetimes |
 | [Asynchronous scheduling](https://github.com/InferMatrixLab/vllm-rlt/pull/30) | CPU/GPU pipelining and single-stream or multi-stream execution |
 | [FlashAttention](https://github.com/InferMatrixLab/vllm-rlt/pull/30) | FA2/FA3/FA4 installation, hardware selection, and constraints |

docs/async-speculative.md

@@ -77,10 +77,10 @@ the controller uses the same `ModelRuntime` for PRELUDE, RECURRENT and CODA.
 It retains its own draft/verify ordering, candidate state and acceptance logic.
 There is no cross-stage handoff for each speculative step.
 
-Prompt bootstrap uses the ordinary runner. Refill scheduling can use atomic
-prefill or `atomic_prefill=False` to yield between prompt loops. PD can import
-the completed prompt into the decode stage. These combinations do not enable
-mixed scheduling of prompt rows with speculative verification rows.
+Prompt bootstrap uses the ordinary runner. Refill scheduling returns to the
+scheduler after each prompt chunk, where ready speculative work can be selected.
+PD can import the completed prompt into the decode stage. These combinations do
+not enable mixed scheduling of prompt rows with speculative verification rows.
 
 
 The synchronous runner reads the acceptance result before constructing the next

docs/launching.md

@@ -280,10 +280,11 @@ PyTorch/CUDA versions from step 4, and your Git revision (`git rev-parse HEAD`).
 
 ## Choosing scheduling and deployment separately
 
-The default atomic prefill path and loop-yielding prefill use the same execution
-architecture. On supported models with LAST_EXITED KV, the CLI/server flags
-`--mode mixed --no-atomic-prefill` allow compatible prompt and decode recurrence
-to share a batch. This is a batching choice, not PD disaggregation.
+Prompt chunks normally execute their full recurrence depth in one submission.
+On supported models with LAST_EXITED KV, `--mode mixed` lets the scheduler
+insert a depth boundary when decode work can join a prompt recurrent batch.
+`--wavefront-chunked-prefill` enables per-depth chunk regrouping. These are
+batching choices, not PD disaggregation.
 
 Stage roles are `prefill`, `decode` and `prefill_decode`; `mixed` is a scheduling
 mode, not a stage role. Speculation requires refill scheduling and cannot join

docs/runtime.md

@@ -106,30 +106,61 @@ Ordinary and speculative runners share the model runtime without sharing their
 progress-control algorithms.
 
 `ExecutionBatch.loop_count` describes the recurrence included in a submission.
-PRELUDE can include subsequent recurrent loops; RECURRENT can advance several
-depths while loading and storing hidden state once. Only the final gate is
-returned from such a multi-loop call. A caller needing an exit decision after
-each depth must yield at those boundaries. Atomic prefill runs every depth of a
-prompt chunk locally through the same submission path.
-
-Speculation with refill scheduling can use loop-yielding chunked prefill through
-`ExecutionConfig(atomic_prefill=False)`. This does not allow speculative verification
-rows in mixed batches; see [speculative constraints](#fixed-loop-self-speculative-decoding).
+The scheduler sets it for each batch: ordinary prompt chunks execute their full
+depth locally, WCPB returns after each depth to regroup ready chunks, and a mixed
+prompt/decode recurrent batch returns after one depth for exit decisions and
+progress updates. When mixed mode has no ready decode work to join, prompt
+PRELUDE can include the full recurrence in the same submission. RECURRENT can
+advance several depths when the selected prompt frontier permits it.
+
+Speculation with refill scheduling is scheduled at prompt-chunk boundaries and
+does not require a prefill switch. It does not allow speculative verification
+rows in mixed batches; see
+[speculative constraints](#fixed-loop-self-speculative-decoding).
 
 ```python
 scheduler_config = SchedulerConfig(
     mode="mixed", max_num_seqs=8, max_num_batched_tokens=128,
     prefill_chunk_size=64,
 )
-execution_config = ExecutionConfig(atomic_prefill=False)
+execution_config = ExecutionConfig()
 ```
 
-CLI and serving accept `--mode mixed --no-atomic-prefill`. Mixed selects compatible recurrent work; PRELUDE and CODA remain separate operations.
+`ExecutionConfig` controls execution resources; the scheduler derives the
+recurrent boundary from the selected batch. Mixed selects compatible recurrent
+work; WCPB can also combine prompt and decode PRELUDE rows, while CODA remains
+separate.
 See [batch selection and feedback](scheduler.md#detail-mixed-loop-selection-and-feedback)
 for the selection order, shared budget and per-request progress.
 
-This policy currently requires non-atomic prompt scheduling and LAST_EXITED KV;
-it does not mix speculative rounds. Existing model restrictions still apply.
+### Wavefront chunked prefill
+
+`SchedulerConfig(wavefront_chunked_prefill=True)` opts into WCPB and requires
+LAST_EXITED KV. It can run with
+refill or mixed scheduling. When combined with speculation, use refill mode;
+speculation cannot use preemption.
+
+The scheduler stages up to `total_ut_steps` chunks per request, then batches one
+ready chunk for each prompt depth. For four depths and chunk size two, recurrent
+batches advance `(chunk0, depth0)`, then `(chunk0, depth1)` with
+`(chunk1, depth0)`, then add `(chunk2, depth0)`, and so on. Each request is
+limited to a `total_ut_steps * prefill_chunk_size` resident hidden-state window;
+the global token budget and `max_num_seqs` still apply. Prompt state uses a ring
+window, while KV writes retain their absolute prompt positions and depth planes.
+
+Prompt work always reaches the model's full depth before CODA produces the first
+token. A request's `SamplingParams.max_loops` limits later decode work and does
+not shorten prompt prefill. Prefix reuse, preemption, PD export, and async
+execution follow the same full-depth frontier. CUDA Graph handling remains on
+the existing recurrent paths; WCPB does not add a new graph mode.
+
+Enable it from the CLI with `--wavefront-chunked-prefill`.
+Combine it with `--mode mixed` when prompt and decode recurrent work should share
+compatible batches; the wavefront option also works without mixed scheduling.
+
+This policy requires staged prompt state and LAST_EXITED KV. Speculative
+rounds are scheduled separately from mixed recurrent batches. Existing model
+restrictions still apply.
 Mixed graphs and asynchronous device execution have regression cases marked
 `gpu`; CPU tests validate the corresponding state machine, not GPU overlap or
 performance. Mixed gather/scatter uses the shared transfer machinery; this
@@ -201,8 +232,9 @@ for the design.
 
 SHARED retains one physical KV plane and completes every loop for a prompt
 position before advancing that request. It is not interchangeable with packed,
-depth-major LAST_EXITED prefill. SHARED requires atomic prompt execution; split
-prefill and mixed scheduling require LAST_EXITED. See the
+depth-major LAST_EXITED prefill. SHARED supports full-depth prompt batches;
+mixed and wavefront scheduling require LAST_EXITED because they can create staged
+prompt work. See the
 [prefill ordering example](kv_layout_computation.md#5-prefill-execution-order)
 for the causal-history difference and its capacity implications.
 

docs/scheduler.md

@@ -60,22 +60,23 @@ scheduling eligibility; execution events order the actual device consumers.
 Selection combines different decisions; the implementation class names are not
 peer conceptual choices:
 
-- **Prefill granularity — atomic or split:** does a prompt chunk finish all its
-  recurrent depths in one submission, or yield at loop boundaries?
+- **Prefill boundary — full depth or scheduler yield:** does this batch finish
+  every remaining recurrent depth locally, or return at a loop boundary for a
+  concrete scheduling decision?
 - **Batching — refill or no-refill:** can prompt work replenish ongoing
   generation, or must the current cohort drain first? **Mixed** additionally
   combines compatible prompt and decode recurrent spans in the same batch;
   refilling alone does not imply a mixed batch.
 
 The ordinary-generation rows in C show currently supported combinations:
-atomic with refill/no-refill, and split with refill/mixed. These are compatibility
-rules, not a claim that the two dimensions mean the same thing.
-`SplitPrefillPolicy` uses refill ordering with split prompt phases;
-`MixedPolicy` extends that behavior with recurrent co-batching.
+full-depth prompt chunks with refill/no-refill, and scheduler-selected depth
+boundaries with refill/mixed. These are compatibility rules, not a claim that
+the two dimensions mean the same thing. `SplitPrefillPolicy` is selected when
+WCPB requires staged prompt phases; `MixedPolicy` adds recurrent co-batching.
 
 Prompt-only is a stage-role constraint. Speculation selects a local draft/verify
-controller; it currently requires refill, supports atomic/split prefill subject
-to model constraints, and excludes mixed. See
+controller; it currently requires refill, uses full-depth prompt chunks unless
+WCPB is enabled, and excludes mixed. See
 [policies and execution granularity](#policies-and-execution-granularity).
 
 The colored tiles illustrate a mixed recurrent batch: R1 is a prompt span at
@@ -237,7 +238,8 @@ invalid when GPU resources are reassigned.
 For split prefill, `_rewind_split_prefill()` first resets an unfinished chunk to
 `num_prefilled_tokens` and sets the resume phase to `PREFILLING_PRELUDE`. The
 snapshot therefore records the rewound phase, not necessarily the loop at which
-the request was interrupted. The unfinished chunk is replayed after restoration.
+the request was interrupted. Unfinished chunks are replayed after restoration;
+WCPB also resets its ring origin to this full-depth frontier.
 
 #### Current policy: triggers and eligibility
 
@@ -447,7 +449,8 @@ The runner submits one recurrent batch and waits on the preceding device events
 of its members, deduplicating shared events. It does not execute three separate
 request forwards, combine their KV histories, or require their depths to match.
 PRELUDE and CODA are different operations and are not packed into this recurrent
-submission. Current mixed scheduling requires non-atomic, non-speculative work.
+submission. Mixed scheduling excludes speculative work and uses a depth boundary
+when ready decode work can join prompt recurrent work.
 
 Logical progress then branches per request: A can continue at depth 2; B can enter
 CODA if its exit is established; C can continue its prompt span at depth 1.
@@ -480,6 +483,16 @@ Atomic prefill can execute PRELUDE plus all recurrent loops in one submission,
 retaining the same preparation, dependency and completion contract. Therefore
 one arrow between model operations does not always mean a new scheduler submission.
 
+With WCPB enabled, one request can stage up to `total_ut_steps` chunks and retain
+one ready chunk per recurrence depth. For four depths, after chunk 0 completes
+depth 0, the next recurrent batch can contain `chunk0@depth1` and
+`chunk1@depth0`; the following batch can add `chunk2@depth0`. A task is ready
+only when its hidden input exists and all earlier KV rows at that same
+LAST_EXITED depth are valid. The request's full-depth frontier advances only
+through final-depth work, and a ring slot is reusable only after its chunk
+finishes that depth. The implementation flattens selected token rows for the
+model call and keeps each chunk/depth range in its metadata.
+
 Speculative draft/verify execution uses its own round controller. Its host decision
 points and buffer ownership must be read from that controller; the ordinary
 runner timeline is not a complete speculative-round diagram.
@@ -604,9 +617,11 @@ exit decisions. The scheduler can then form a different batch. Prompt tokens
 still traverse the full model depth; mixing does not apply decode early exit to
 prompt rows.
 
-PRELUDE and CODA are selected separately. Mixed scheduling currently combines
-`RECURRENT` and `PREFILLING_RECURRENT`, not arbitrary operations or speculative
-verification rows.
+Without WCPB, PRELUDE and CODA are selected separately. With WCPB, MixedPolicy
+may combine decode PRELUDE rows with prompt-chunk PRELUDE rows because both use
+the same operation. Mixed recurrent batches combine `RECURRENT` and
+`PREFILLING_RECURRENT` spans, including multiple depths from one prompt request;
+CODA and speculative verification remain separate operations.
 
 ## Who owns each decision?
 
@@ -726,26 +741,29 @@ independent, peer configuration dimensions.
 
 | Policy | Selection behavior |
 | --- | --- |
-| `PromptOnlyPolicy` | Admit and select prompt work only, using atomic or split prompt phases |
+| `PromptOnlyPolicy` | Admit and select prompt work only; WCPB uses staged prompt phases |
 | `RefillPolicy` | Prefer ready PRELUDE, eligible CODA and due recurrent work; admit and interleave prompt work |
 | `NoRefillPolicy` | Prepare a cohort, drain its recurrence, then process CODA before filling the next cohort |
-| `SplitPrefillPolicy` | Refill ordering with separate prompt PRELUDE and RECURRENT scheduling boundaries |
-| `MixedPolicy` | If both recurrent queues have ready work, select generation rows first and fill with prompt spans; otherwise use split-prefill ordering |
+| `SplitPrefillPolicy` | Refill ordering with staged prompt phases for WCPB |
+| `MixedPolicy` | If both recurrent queues have ready work, select generation rows first and fill with prompt spans; otherwise use the batch's full-depth plan when no depth boundary is needed |
 | `SpeculativePolicy` | Select CODA bootstrap, prompt work and local draft/verify rounds, with a prefill fairness bound |
 
 `create_scheduling_policy()` first handles prompt-only stages, then checks
 compatible generation configurations and selects a policy. Mixed generation
-requires `atomic_prefill=False` and no speculative configuration. Split prefill
-requires LAST_EXITED KV at engine construction; model-specific restrictions
-still apply. Ordinary split generation currently uses refill or mixed, not
-no-refill. Speculative generation requires refill.
-
-Atomic prefill changes the scheduling boundary, not the execution architecture.
-An atomic prompt chunk runs its prelude and all recurrent depths locally in one
-submission. Split prefill yields between operations/depths. Both bind into the
-same execution path. Chunk size and loop submission size are distinct choices.
-The current split path retains one staged prompt chunk per request; it does not
-stage the next chunk while the current chunk still needs its hidden rows.
+selects staged prompt work automatically and excludes speculative configuration.
+WCPB also requires staged prompt state. Both require LAST_EXITED KV; model
+restrictions still apply. Speculative generation requires refill.
+
+`SchedulerOutput.loop_count` is the actual execution boundary. A regular prompt
+chunk runs its prelude and all recurrent depths locally in one submission.
+Mixed mode inserts a boundary when ready decode work can join prompt recurrent
+work, and it returns after every depth in such a mixed batch. With no ready
+decode work, the prompt batch can run through its remaining depths. WCPB always
+returns after one depth to regroup ready chunks. Chunk size and loop submission
+size are distinct choices; WCPB retains up to one staged chunk per model depth.
+Ordinary prompt chunks already end at the scheduler's configured chunk boundary,
+where decode/speculative work and admission or preemption are reconsidered.
+Merely having waiters or enabling preemption does not split a chunk's recurrence.
 
 Fairness counts nonempty scheduled batches, not attempts or GPU completions.
 Prompt batches increment the configured prefill counter; generation RECURRENT

docs/serving.md

@@ -42,9 +42,12 @@ remain separate from the stage-local model execution contract.
 
 Mixed scheduling means batching compatible prompt/decode recurrent work within
 one stage. It does not mean a worker role, and PD does not require it. For
-supported models, select it with `--mode mixed --no-atomic-prefill` and
-LAST_EXITED KV. PD can configure speculative decoding on the decode stage, but
-mixed scheduling and speculative rounds are not a supported combination.
+supported models, select it with `--mode mixed` and LAST_EXITED KV; the
+scheduler chooses depth boundaries when decode work can join a recurrent batch.
+Add `--wavefront-chunked-prefill` to batch ready prompt chunks from different
+recurrent depths; PD can enable this on its prefill scheduler.
+Speculative decoding requires refill scheduling and cannot join mixed recurrent
+batches. PD can configure speculation on its decode stage.
 See [runtime configuration](runtime.md#execution-batches-and-mixed-scheduling)
 for constraints and [the architecture](unified_execution.md) for ownership.