diff --git a/docs/README.md b/docs/README.md index 9f21626f..c3507a8a 100644 --- a/docs/README.md +++ b/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 | diff --git a/docs/assets/Split-prefill-vs-wavefront-chunked-prefill.png b/docs/assets/Split-prefill-vs-wavefront-chunked-prefill.png new file mode 100644 index 00000000..d72fbd3c Binary files /dev/null and b/docs/assets/Split-prefill-vs-wavefront-chunked-prefill.png differ diff --git a/docs/async-speculative.md b/docs/async-speculative.md index 543efc5a..b6fe4557 100644 --- a/docs/async-speculative.md +++ b/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 diff --git a/docs/launching.md b/docs/launching.md index 79ced6bd..6fbc0de5 100644 --- a/docs/launching.md +++ b/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 diff --git a/docs/runtime.md b/docs/runtime.md index 9e3574bc..d98050fc 100644 --- a/docs/runtime.md +++ b/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. diff --git a/docs/scheduler.md b/docs/scheduler.md index 117db13e..b3952188 100644 --- a/docs/scheduler.md +++ b/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 diff --git a/docs/serving.md b/docs/serving.md index e03911ca..99a9121f 100644 --- a/docs/serving.md +++ b/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. diff --git a/docs/split-prefill-wavefront.md b/docs/split-prefill-wavefront.md new file mode 100644 index 00000000..5bbed969 --- /dev/null +++ b/docs/split-prefill-wavefront.md @@ -0,0 +1,126 @@ +# Split prefill and wavefront chunked prefill + +Split prefill limits how much prompt work a request contributes at once. +Wavefront Chunked Prefill Batching (WCPB) uses the same chunk boundary while +batching ready prompt work from different chunks and recurrent depths. Its main +target is better recurrent-core occupancy when prompt batches are sparse or +uneven, without making chunks larger and reducing scheduling opportunities. + +![Split prefill compared with wavefront chunked prefill](assets/Split-prefill-vs-wavefront-chunked-prefill.png) + +The current runner requires one operation per batch: PRELUDE ranges and RECURRENT +ranges are submitted separately. Ordinary prefill can also execute its selected range +through all recurrent depths in one plan. Actual submission counts depend on scheduler +mode, token budget, and ready work; the numbered batches in the figure should not be +read as measured latency or a fixed runtime schedule. + +## Execution model + +The ordinary path selects a prompt range and runs the range through the +remaining recurrent depths before applying progress. In mixed scheduling, the +scheduler may choose a shorter execution boundary when ready decode work needs a +depth decision. `ExecutionBatch.loop_count` records the recurrent work in that +submission. + +With `wavefront_chunked_prefill=True`, prompt work is scheduled in two phases: + +| Phase | Work | Recurrent loops in the submission | +| --- | --- | ---: | +| `PREFILLING_PRELUDE` | Embed staged prompt ranges and prepare inputs | 0 | +| `PREFILLING_RECURRENT` | Run ready prompt ranges at their selected recurrent depths | 1 | + +The scheduler can include several ranges from different chunks, requests, and +depths in one recurrent batch. It obeys the shared token budget, request-owner +limit, KV capacity, and each depth's readiness frontier. In mixed mode, +compatible decode and prompt work can share PRELUDE or RECURRENT batches. The +batch still contains one operation type. + +WCPB is opt-in through `SchedulerConfig.wavefront_chunked_prefill` or the +`--wavefront-chunked-prefill` CLI flag. `prefill_chunk_size` controls the prompt +chunk size. The option works with refill or mixed scheduling; mixed scheduling +does not support speculative work. + +## Progress and dependencies + +WCPB relies on the depth-specific KV planes of `LAST_EXITED` layout. For +example, a later chunk's depth-0 work can use the completed depth-0 prefix while +an earlier chunk advances at a deeper recurrence depth. A deeper task also +needs its preceding depth's hidden state and the causal KV prefix at its own +depth. The scheduler uses per-depth prompt frontiers to admit only ready ranges. + +Progress belongs to the request rather than to a persistent task table: + +- `num_preluded_tokens` tracks the prompt range whose prelude has been staged. +- `prefill_depth_ranges` tracks the ready range at each recurrent depth. +- `num_prefilled_tokens` tracks prompt progress at the final depth, including + progress within a chunk. +- `prefill_reclaimable_tokens` marks the oldest resident chunk boundary. It + advances only after that chunk completes its final depth and its ring slot + can be reused. + +This separation prevents a prelude or partial chunk from advancing the +reclaimable boundary. Prefix publication, prompt completion, and CODA follow +the completed final-depth frontier. + +## Hidden-state ring + +Each request reserves `prefill_chunk_size * total_ut_steps` rows for staged +prompt hidden states. A chunk uses a ring slot derived from its index relative +to the cached-prefix origin. For a chunk-local offset `j` and slot `s`, its row +offset is: + +```text +prefill_hidden_row_offset = s * prefill_chunk_size + j +``` + +`ScheduledRequest` and `StateSpan` carry this row offset. When it is `None`, the +runner derives the row from `token_start - chunk_start`. Intermediate outputs +are scattered to the ring for later depths; after final-depth completion, the +chunk's intermediate state is no longer needed. Asynchronous submissions keep +the backing rows alive until dependent reads complete. + +## Scheduling and recovery + +WCPB advances one recurrent depth per recurrent submission so the scheduler can +regroup the next set of ready ranges. A boundary gives admission and mixed +scheduling a chance to select other work; it does not guarantee that another +request will run. Smaller chunks therefore improve scheduling granularity but +can reduce ordinary prefill occupancy. WCPB tries to recover some occupancy by +batching independent ready chunk/depth ranges. + +The implementation uses the existing request queues, `ScheduledRequest`, and +per-depth frontiers; it does not introduce a generic `PrefillTask` registry. +Preemption rewinds staged work to the request's completed prompt frontier and +recomputes unfinished work after resume. PD transfer progress is reported for +completed chunks, and imported prompt state initializes the same frontiers. +Speculative execution is supported with compatible refill scheduling; mixed +mode rejects speculative configuration. CUDA graph and asynchronous paths use +the existing runner machinery. + +`SHARED` KV is rejected for WCPB. Its position-major dependencies do not permit +the same cross-chunk/depth ordering as `LAST_EXITED` without a different +scheduler dependency graph. + +## Performance and correctness limits + +Controlled A100 measurements show that ordinary small chunks can underfill the +recurrent batch. At chunk size 16, one-run WCPB measurements raised rows per +batch and throughput from 64.00 rows and 275.82 tok/s to 254.49 rows and +412.01 tok/s for four 8,064-token requests. For 32- and 128-request mixed +workloads, measured throughput changed from 337.51 to 628.45 tok/s and from +471.95 to 740.29 tok/s respectively. These are directional results from one +warmup and one measured run per variant, not multi-run medians or universal +speedup claims. + +A separate HTTP matrix used Ouro-1.4B BF16 with Triton on two A100 80GB GPUs, +`max_num_seqs=8`, `max_num_batched_tokens=1024`, chunk sizes 32 and 128, and +fixed 512-token and variable 256/512/768-token prompts. Each serial, concurrent, +and arrival profile had 100 requests and one measured run. WCPB request +throughput was 0.8% to 22.6% higher than Split in the 12 matched profiles, but +some tail-latency metrics were worse. These measurements do not qualify decode +latency under arbitrary live traffic. + +The prior GSM8K-87 BF16 run scored 63/87 for WCPB and 60/87 for Split and +Atomic, but WCPB differed from Split on 24/87 token sequences. The score +difference does not establish output equivalence or improved model quality. +Exact GPU token preservation remains an open validation issue. diff --git a/tests/helpers.py b/tests/helpers.py index daa39d92..cf4a4236 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 """Shared small models, workload-specific engine factories and protocol doubles.""" +from dataclasses import replace from typing import Any import pytest @@ -91,11 +92,20 @@ def make_cpu_cache(config, *, blocks=64): def make_prefill_engine(*, atomic=False, **kwargs): kwargs.setdefault("cache_config", CacheConfig(64, 2)) - kwargs.setdefault("execution_config", ExecutionConfig(atomic_prefill=atomic)) - kwargs.setdefault( - "scheduler_config", - SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2), - ) + scheduler = kwargs.get("scheduler_config") + if scheduler is None: + scheduler = SchedulerConfig( + mode="refill", + wavefront_chunked_prefill=not atomic, + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + ) + elif not atomic and not scheduler.wavefront_chunked_prefill: + if scheduler.mode == "refill": + scheduler = replace(scheduler, wavefront_chunked_prefill=True) + kwargs["scheduler_config"] = scheduler + kwargs.setdefault("execution_config", ExecutionConfig()) return StageEngine(tiny_ouro_model(), **kwargs) @@ -110,9 +120,9 @@ def make_stage_engine(*, role="prefill_decode", atomic=False, asynchronous=False max_num_seqs=3, max_num_batched_tokens=4, prefill_chunk_size=3, + wavefront_chunked_prefill=not atomic and not mixed, ), execution_config=ExecutionConfig( - atomic_prefill=atomic, async_scheduling=asynchronous, static_buffers=True, ), diff --git a/tests/test_backend_selection.py b/tests/test_backend_selection.py index 17439d93..b0c00e4b 100644 --- a/tests/test_backend_selection.py +++ b/tests/test_backend_selection.py @@ -217,8 +217,9 @@ def test_cuda_engine_resolves_auto_before_profiling(packages, monkeypatch, execu class ProfileReached(Exception): pass - def plan_cache(model, cache, scheduler, execution, selected): + def plan_cache(model, cache, scheduler, execution, selected, *, staged_prefill): assert selected == "triton" # FA2 cannot use the configured 16-token pages. + assert staged_prefill is False raise ProfileReached monkeypatch.setattr(llm_engine, "plan_cache", plan_cache) diff --git a/tests/test_engine_core_contract.py b/tests/test_engine_core_contract.py index cf9538a7..ce749ac4 100644 --- a/tests/test_engine_core_contract.py +++ b/tests/test_engine_core_contract.py @@ -20,7 +20,6 @@ def make_engine( preemption=False, token_budget=16, mode="refill", - atomic_prefill=True, prefill_chunk_size=16, ): torch.manual_seed(123) @@ -39,10 +38,7 @@ def make_engine( enable_preemption=preemption, ), exit_config=ExitConfig("ouro_delayed"), - execution_config=ExecutionConfig( - async_scheduling=asynchronous, - atomic_prefill=atomic_prefill, - ), + execution_config=ExecutionConfig(async_scheduling=asynchronous), ) @@ -361,14 +357,16 @@ def test_mixed_recurrent_and_split_prefill_progression_through_scheduler(asynchr engine = make_engine( asynchronous=asynchronous, mode="mixed", - atomic_prefill=False, prefill_chunk_size=2, token_budget=4, ) engine.add_request("decode", [2, 3], params) for _ in range(10): engine.step() - if engine.scheduler.requests["decode"].generated_token_ids: + if ( + engine.scheduler.requests["decode"].generated_token_ids + and engine.scheduler.requests["decode"].phase is RequestPhase.RECURRENT + ): break engine.add_request("prompt", [4, 5, 6, 7], params) saw_mixed_batch = False diff --git a/tests/test_huginn.py b/tests/test_huginn.py index eece70d5..2a606835 100644 --- a/tests/test_huginn.py +++ b/tests/test_huginn.py @@ -239,17 +239,34 @@ def test_prefill_populates_history_without_projecting_unused_logits(model, padde @pytest.mark.parametrize( - "cache,execution,error", + "cache,execution,scheduler,error", [ - (CacheConfig(layout="shared"), ExecutionConfig(), "last_exited"), - (CacheConfig(enable_prefix_caching=True), ExecutionConfig(), "prefix caching"), - (CacheConfig(), ExecutionConfig(atomic_prefill=False), "atomic prefill"), - (CacheConfig(), ExecutionConfig(async_scheduling=True), "synchronous"), + (CacheConfig(layout="shared"), ExecutionConfig(), SchedulerConfig(), "last_exited"), + ( + CacheConfig(enable_prefix_caching=True), + ExecutionConfig(), + SchedulerConfig(), + "prefix caching", + ), + (CacheConfig(), ExecutionConfig(), SchedulerConfig(mode="mixed"), "staged prefill"), + ( + CacheConfig(), + ExecutionConfig(async_scheduling=True), + SchedulerConfig(), + "synchronous", + ), ], ) -def test_unsupported_execution_rejected_before_cache_allocation(model, cache, execution, error): +def test_unsupported_execution_rejected_before_cache_allocation( + model, cache, execution, scheduler, error +): with pytest.raises(ValueError, match=error): - LLMEngine(model, cache_config=cache, execution_config=execution) + LLMEngine( + model, + cache_config=cache, + execution_config=execution, + scheduler_config=scheduler, + ) @pytest.mark.gpu diff --git a/tests/test_input_preparation.py b/tests/test_input_preparation.py index 361e559f..31a9b840 100644 --- a/tests/test_input_preparation.py +++ b/tests/test_input_preparation.py @@ -40,7 +40,7 @@ def test_loop_preparation_reuses_growing_storage_and_only_saves_final_hidden( model, cache_config=CacheConfig(128, 16 if backend == "flash_attn_4" else 2, layout=layout), execution_config=ExecutionConfig( - atomic_prefill=True, static_buffers=True, pad_to_power_of_two=True, prefill_uva=uva + static_buffers=True, pad_to_power_of_two=True, prefill_uva=uva ), exit_config=ExitConfig("ouro"), scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=8), @@ -341,8 +341,8 @@ def test_cpu_preparation_defers_transfers_and_hidden_gather(monkeypatch, device, e = LLMEngine( tiny_ouro_model().to(device), cache_config=CacheConfig(64, 2), - scheduler_config=SchedulerConfig(max_num_seqs=1, max_num_batched_tokens=4), - execution_config=ExecutionConfig(atomic_prefill=False), + scheduler_config=SchedulerConfig(mode="mixed", max_num_seqs=1, max_num_batched_tokens=4), + execution_config=ExecutionConfig(), attention_backend="triton" if device == "cuda" else "torch", ) runner = e.model_runner @@ -534,10 +534,12 @@ def test_forward_consumes_prepared_inputs( cache_config=CacheConfig(128, 2, layout=layout), exit_config=ExitConfig("random_lookahead"), scheduler_config=SchedulerConfig( - max_num_seqs=3, max_num_batched_tokens=8, prefill_chunk_size=3 + mode="refill" if atomic else "mixed", + max_num_seqs=3, + max_num_batched_tokens=8, + prefill_chunk_size=3, ), execution_config=ExecutionConfig( - atomic_prefill=atomic, static_buffers=static_buffers, pad_to_power_of_two=static_buffers, async_scheduling=async_scheduling, diff --git a/tests/test_mixed_execution.py b/tests/test_mixed_execution.py index 1eb343f0..542b8e96 100644 --- a/tests/test_mixed_execution.py +++ b/tests/test_mixed_execution.py @@ -6,7 +6,7 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import SamplingParams, SchedulerConfig from vllm_rlt.core.kv_cache_manager import KVCacheManager -from vllm_rlt.core.scheduler import Scheduler, SchedulerOutput +from vllm_rlt.core.scheduler import ScheduledRequest, Scheduler, SchedulerOutput from vllm_rlt.models import OuroForCausalLM from vllm_rlt.request import ExecutionBatch, Operation, Request, RequestPhase, StateSpan from vllm_rlt.worker.executor import ForwardBatch, ModelRuntime, build_execution_batch @@ -16,6 +16,30 @@ def make_span(name, start, count, depth): return StateSpan(Request(name, [1], SamplingParams()), start, count, depth) +def test_mixed_decode_batch_yields_after_one_loop_without_queued_decode(monkeypatch): + scheduler = Scheduler(SchedulerConfig(mode="mixed"), None, total_ut_steps=4) + scheduled = [ + ScheduledRequest( + Request("decode", [1, 2], SamplingParams()), + token_start=2, + depth=1, + phase=RequestPhase.RECURRENT, + ), + ScheduledRequest( + Request("prompt", [3, 4], SamplingParams()), + token_start=0, + token_count=2, + depth=0, + chunk_start=0, + phase=RequestPhase.PREFILLING_RECURRENT, + ), + ] + monkeypatch.setattr(scheduler, "has_ready", lambda phase: False) + monkeypatch.setattr(scheduler, "_max_contiguous_prefill_loops", lambda item: 4) + + assert scheduler._planned_loop_count(scheduled) == 1 + + def make_cache(config, layout): return KVCacheManager( num_layers=config.num_hidden_layers, @@ -145,7 +169,7 @@ def test_mixed_policy_runs_real_combined_batches_with_matching_outputs( from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig from vllm_rlt.engine.llm_engine import LLMEngine - def run(mode, atomic=False): + def run(mode, wavefront=False): torch.manual_seed(123) engine = LLMEngine( OuroForCausalLM(tiny_ouro_config()).to( @@ -158,9 +182,9 @@ def run(mode, atomic=False): max_num_seqs=3, max_num_batched_tokens=5, prefill_chunk_size=4, + wavefront_chunked_prefill=wavefront, ), execution_config=ExecutionConfig( - atomic_prefill=atomic, cuda_graphs=graphs, async_scheduling=asynchronous, static_buffers=padded, @@ -168,7 +192,26 @@ def run(mode, atomic=False): ), exit_config=ExitConfig(mode="ouro_delayed" if asynchronous else "ouro"), ) - for name, prompt in (("short", [1]), ("long", [2, 3, 4, 5, 6, 7, 8, 9, 10])): + prompts = (("short", [1]), ("long", [2, 3, 4, 5, 6, 7, 8, 9, 10])) + if mode == "mixed": + engine.add_request( + "short", + [1], + SamplingParams( + max_tokens=6, + max_loops=3, + ignore_eos=True, + exit_threshold=0.6, + ), + ) + for _ in range(40): + engine.step() + if engine.scheduler.requests["short"].phase == RequestPhase.RECURRENT: + break + else: + pytest.fail("decode request did not reach recurrent work") + prompts = prompts[1:] + for name, prompt in prompts: engine.add_request( name, prompt, @@ -198,9 +241,19 @@ def record(hidden, kv, cache, **kwargs): if batch is not None and batch.is_mixed: phases = {batch.phase_for(i) for i in batch.requests} if len(phases) == 2: - # Both storage kinds pass through a single model invocation. - if not graphs: + recurrent_mix = phases == { + RequestPhase.RECURRENT, + RequestPhase.PREFILLING_RECURRENT, + } + # Recurrent storage kinds share one model invocation. Prelude + # mixes use the same operation but do not call recurrent_core. + if recurrent_mix and not graphs: assert calls == [batch.num_tokens] + elif not recurrent_mix: + assert phases == { + RequestPhase.PRELUDE, + RequestPhase.PREFILLING_PRELUDE, + } assert batch.num_tokens <= 5 mixed.append(batch) if not engine.has_unfinished_requests(): @@ -210,11 +263,13 @@ def record(hidden, kv, cache, **kwargs): assert engine.cache_manager.num_used_blocks == 0 return outputs, mixed - expected, _ = run("refill", atomic=True) + expected, _ = run("refill") separated, _ = run("refill") - actual, mixed = run("mixed") - assert mixed - assert actual == separated == expected + original_mixed, original_batches = run("mixed") + wavefront_refill, _ = run("refill", wavefront=True) + actual, mixed = run("mixed", wavefront=True) + assert original_batches and mixed + assert actual == wavefront_refill == original_mixed == separated == expected def mixed_ready_engine(): @@ -228,8 +283,14 @@ def mixed_ready_engine(): prefill_chunk_size=4, ) ) - for name, prompt in (("short", [1]), ("long", list(range(2, 11)))): - engine.add_request(name, prompt, SamplingParams(max_tokens=6, ignore_eos=True)) + engine.add_request("short", [1], SamplingParams(max_tokens=6, ignore_eos=True)) + for _ in range(40): + engine.step() + if engine.scheduler.requests["short"].phase == RequestPhase.RECURRENT: + break + else: + pytest.fail("decode request did not reach recurrent work") + engine.add_request("long", list(range(2, 11)), SamplingParams(max_tokens=6, ignore_eos=True)) for _ in range(100): batch = engine.scheduler.schedule() if batch.is_mixed and len({batch.phase_for(i) for i in batch.requests}) == 2: @@ -393,8 +454,8 @@ def run(continuous): engine = LLMEngine( OuroForCausalLM(tiny_ouro_config()), cache_config=CacheConfig(64, 2), + scheduler_config=SchedulerConfig(wavefront_chunked_prefill=chunked), execution_config=ExecutionConfig( - atomic_prefill=False, static_buffers=padded, pad_to_power_of_two=padded, ), diff --git a/tests/test_pd.py b/tests/test_pd.py index 496ed968..22ced40c 100644 --- a/tests/test_pd.py +++ b/tests/test_pd.py @@ -116,7 +116,7 @@ def drain(engine): return outputs -def test_pd_split_transfer_uses_chunk_origin_after_budget_contention(monkeypatch): +def test_pd_wavefront_transfer_uses_chunk_origin_after_budget_contention(monkeypatch): from contextlib import nullcontext from types import SimpleNamespace @@ -138,9 +138,12 @@ def query(self): OuroForCausalLM(config), cache_config=CacheConfig(64, 4), scheduler_config=SchedulerConfig( - max_num_seqs=2, max_num_batched_tokens=5, prefill_chunk_size=4 + max_num_seqs=2, + max_num_batched_tokens=5, + prefill_chunk_size=4, + wavefront_chunked_prefill=True, ), - execution_config=ExecutionConfig(atomic_prefill=False), + execution_config=ExecutionConfig(), stage_config=StageConfig(role="prefill"), ) transport = SimpleNamespace(inflight_bytes=0, busy=lambda tid: False) @@ -174,31 +177,70 @@ def step(): ) assert step() and connector.prefill_tokens == 0 - for _ in range(config.total_ut_steps): - assert step() - assert (engine.scheduler.requests["a"].num_prefilled_tokens, connector.prefill_tokens) == (1, 4) - - for _ in range(config.total_ut_steps): + for _ in range(30): + if len(transfers) == 2: + break assert step() - final_scheduled_request = next( - scheduled_request - for scheduled_request in batches[-1].requests - if scheduled_request.request.request_id == "a" + else: + pytest.fail("wavefront transfer did not complete") + assert any( + (item.token_start, item.token_count, item.chunk_start) == (1, 3, 0) + for batch in batches + for item in batch.requests + if item.request.request_id == "a" ) - assert ( - final_scheduled_request.token_start, - final_scheduled_request.token_count, - final_scheduled_request.chunk_start, - ) == (1, 3, 0) assert (transfers, connector.prefill_tokens) == ([("b", 0, 4, True), ("a", 0, 4, True)], 8) +def test_pd_wavefront_exports_only_completed_chunks(): + from tests.helpers import RecordingConnector, make_prefill_engine + from vllm_rlt.engine.stage_config import StageConfig + + engine = make_prefill_engine( + scheduler_config=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=5, + prefill_chunk_size=3, + wavefront_chunked_prefill=True, + ), + stage_config=StageConfig(role="prefill"), + ) + connector = RecordingConnector() + engine.connector = connector + engine.add_request("pd", list(range(1, 15)), SamplingParams(max_tokens=1)) + request = engine.scheduler.requests["pd"] + transferred = 0 + boundaries = [] + for _ in range(100): + before = len(connector.submissions) + engine.step() + # One callback per owner even when this batch contains several depths. + assert len(connector.submissions) == before + 1 + owner, submission = connector.submissions[-1] + assert owner is request + if submission["progress"].chunk_complete: + assert submission["start"] == transferred + transferred = request.num_prefilled_tokens + boundaries.append(transferred) + assert transferred == len(request.prompt_token_ids) or transferred % 3 == 0 + if submission["progress"].prompt_complete: + break + else: + pytest.fail("wavefront PD prompt did not complete") + + assert boundaries == [3, 6, 9, 12, 14] + assert request.loops_done == engine.model.config.total_ut_steps + assert request.phase.value == "transferring" + engine.close() + + def create_pd( *, graph=False, layout="last_exited", multi=False, - atomic_prefill=True, + mode="refill", + wavefront_chunked_prefill=False, prefill_chunk_size=3, max_num_batched_tokens=5, speculative=False, @@ -224,9 +266,11 @@ def create_pd( prefill_cache_config=CacheConfig(128, 4, layout, enable_prefix_caching=prefix_caching), decode_cache_config=CacheConfig(128, 4, layout, enable_prefix_caching=prefix_caching), prefill_scheduler_config=SchedulerConfig( + mode=mode, max_num_seqs=2, max_num_batched_tokens=max_num_batched_tokens, prefill_chunk_size=prefill_chunk_size, + wavefront_chunked_prefill=wavefront_chunked_prefill, ), decode_scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=2), exit_config=ExitConfig("ouro") @@ -235,7 +279,6 @@ def create_pd( speculative_config=SpeculativeConfig(3) if speculative else None, execution_config=ExecutionConfig( async_scheduling=asynchronous, - atomic_prefill=atomic_prefill, static_buffers=graph, pad_to_power_of_two=graph, cuda_graphs=graph, @@ -287,30 +330,31 @@ def test_pd_generation_refill_cancel_and_reuse_matches_local(graph, layout): @pytest.mark.gpu -def test_pd_split_prefill_matches_atomic_transfer_boundaries_and_overlaps(): +def test_pd_wavefront_matches_atomic_transfer_boundaries_and_overlaps(): prompt = list(range(1, 33)) params = SamplingParams(max_tokens=4, min_loops=1, ignore_eos=True) runs = {} - for atomic_prefill in (True, False): + for wavefront in (False, True): with create_pd( - atomic_prefill=atomic_prefill, + mode="refill", + wavefront_chunked_prefill=wavefront, prefill_chunk_size=16, max_num_batched_tokens=16, ) as e: e.add_request("pd", prompt, params, trace_id="t") output = drain(e)["pd"] timings = next(m["timings"] for m in e.metrics if m["request_id"] == "pd") - runs[atomic_prefill] = ( + runs[wavefront] = ( output.token_ids, output.exit_depths, timings["prefill_transfer_boundaries"], timings["prefill_transfer_overlaps"], ) - atomic, split = runs[True], runs[False] - assert split[:2] == atomic[:2] - assert atomic[2] == split[2] == ((0, 16), (16, 32), 2) - assert split[3] > 0 + atomic, wavefront = runs[False], runs[True] + assert wavefront[:2] == atomic[:2] + assert atomic[2] == wavefront[2] == ((0, 16), (16, 32), 2) + assert wavefront[3] > 0 @pytest.mark.gpu @@ -512,16 +556,16 @@ def test_cancelled_transfer_does_not_recycle_pending_hidden_copy(): @pytest.mark.gpu @pytest.mark.parametrize("asynchronous,graph", [(False, False), (False, True), (True, False)]) -@pytest.mark.parametrize("atomic_prefill", [False, True]) +@pytest.mark.parametrize("wavefront", [False, True]) @pytest.mark.parametrize("prefix_caching", [False, True]) def test_pd_speculative_matches_local_and_retires_cancelled_rounds( - asynchronous, graph, atomic_prefill, prefix_caching + asynchronous, graph, wavefront, prefix_caching ): params = SamplingParams(max_tokens=9, ignore_eos=True) with create_pd( speculative=True, asynchronous=asynchronous, - atomic_prefill=atomic_prefill, + wavefront_chunked_prefill=wavefront, prefix_caching=prefix_caching, graph=graph, ) as engine: diff --git a/tests/test_prefix_growth.py b/tests/test_prefix_growth.py index 1916b10b..c31d4aeb 100644 --- a/tests/test_prefix_growth.py +++ b/tests/test_prefix_growth.py @@ -88,7 +88,7 @@ def test_incremental_growth_matches_reserved_generation(): "stage", [RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT] ) def test_incremental_split_prefill_reserves_prompt_growth(stage): - prompt = list(range(1, 9)) + prompt = list(range(1, 18)) e = LLMEngine( model(), cache_config=CacheConfig(64, 2, incremental_allocation=True), @@ -97,8 +97,9 @@ def test_incremental_split_prefill_reserves_prompt_growth(stage): max_num_batched_tokens=2, prefill_chunk_size=2, enable_preemption=True, + wavefront_chunked_prefill=True, ), - execution_config=ExecutionConfig(atomic_prefill=False), + execution_config=ExecutionConfig(), ) try: e.add_request("split", prompt, SamplingParams(max_tokens=1, ignore_eos=True)) @@ -106,7 +107,10 @@ def test_incremental_split_prefill_reserves_prompt_growth(stage): request = e.scheduler.requests["split"] assert request.phase is RequestPhase.PREFILLING_PRELUDE if stage is RequestPhase.PREFILLING_RECURRENT: - e.step() + for _ in range(8): + if request.phase is RequestPhase.PREFILLING_RECURRENT: + break + e.step() assert request.phase is RequestPhase.PREFILLING_RECURRENT allocated = e.cache_manager.allocated_blocks(request.request_id) diff --git a/tests/test_scheduler_selection.py b/tests/test_scheduler_selection.py index 22809aaf..165159ea 100644 --- a/tests/test_scheduler_selection.py +++ b/tests/test_scheduler_selection.py @@ -21,7 +21,7 @@ def model(): return OuroForCausalLM(tiny_ouro_config()) -def engine(max_num_seqs=2, atomic_prefill=True, **scheduler_options): +def engine(max_num_seqs=2, **scheduler_options): return LLMEngine( model(), cache_config=CacheConfig(64, 2), @@ -30,7 +30,7 @@ def engine(max_num_seqs=2, atomic_prefill=True, **scheduler_options): enable_preemption=True, **scheduler_options, ), - execution_config=ExecutionConfig(atomic_prefill=atomic_prefill), + execution_config=ExecutionConfig(), ) @@ -38,17 +38,41 @@ def params(max_tokens=5): return SamplingParams(max_tokens=max_tokens, ignore_eos=True) -def test_batch_selection_is_never_preempted_by_a_later_item(monkeypatch): +@pytest.mark.parametrize("wavefront", [False, True]) +def test_batch_selection_is_never_preempted_by_a_later_item(monkeypatch, wavefront): """Selecting A must stop B's capacity failure from evicting A.""" - e = engine() + e = engine(wavefront_chunked_prefill=wavefront) for name in ("a", "b"): e.add_request(name, [1, 2], params()) e.scheduler._admit() monkeypatch.setattr(e.cache_manager, "ensure_capacity", lambda rid, frontier: rid != "b") - batch = e.scheduler._take(RequestPhase.PREFILL) + phase = RequestPhase.PREFILLING_PRELUDE if wavefront else RequestPhase.PREFILL + batch = e.scheduler._take(phase) assert [scheduled_request.request.request_id for scheduled_request in batch.requests] == ["a"] assert e.preemption.preemptions == 0 - assert e.scheduler.requests["a"].phase == RequestPhase.PREFILL + assert e.scheduler.requests["a"].phase == phase + + +def test_wavefront_does_not_schedule_a_candidate_preempted_during_selection(monkeypatch): + e = engine(wavefront_chunked_prefill=True, prefill_chunk_size=2) + for name in ("a", "b"): + e.add_request(name, [1, 2, 3, 4], params()) + e.scheduler._admit() + # A needs B's pages. B is already in the local candidate list when suspended. + monkeypatch.setattr( + e.cache_manager, + "ensure_capacity", + lambda rid, frontier: rid == "a" and not e.cache_manager.has_allocation("b"), + ) + + batch = e.scheduler._take(RequestPhase.PREFILLING_PRELUDE) + + assert {item.request.request_id for item in batch.requests} == {"a"} + assert e.preemption.preemptions == 1 + assert list(e.scheduler.queues[RequestPhase.WAITING]) == ["b"] + assert not e.scheduler.queues[RequestPhase.PREFILLING_PRELUDE] + assert e.scheduler.requests["b"].phase == RequestPhase.WAITING + e.close() def test_direct_take_does_not_inherit_an_earlier_batch_selection(monkeypatch): @@ -102,10 +126,9 @@ def test_reused_request_id_is_not_protected_by_the_previous_request(): @pytest.mark.parametrize("boundary", [RequestPhase.PRELUDE, RequestPhase.CODA]) -@pytest.mark.parametrize("atomic_prefill", [True, False], ids=["atomic", "split"]) -def test_refill_runs_boundary_restored_during_admission(boundary, atomic_prefill): +def test_refill_runs_boundary_restored_during_admission(boundary): """A sole resumed boundary must run in the admission step that restores it.""" - e = engine(max_num_seqs=1, atomic_prefill=atomic_prefill) + e = engine(max_num_seqs=1) e.add_request("resumed", [1, 2], params(max_tokens=3)) request = e.scheduler.requests["resumed"] for _ in range(20): diff --git a/tests/test_speculative.py b/tests/test_speculative.py index 5788c6bb..155d9f6f 100644 --- a/tests/test_speculative.py +++ b/tests/test_speculative.py @@ -700,18 +700,22 @@ def test_abort_suspended_speculative_request_allows_id_reuse(): @pytest.mark.parametrize("chunk_size", [1, 3, 8]) @pytest.mark.parametrize("budget", [2, 7]) -def test_split_prefill_and_speculation_share_forward_and_match_atomic(chunk_size, budget): +@pytest.mark.parametrize("wavefront", [False, True]) +def test_split_prefill_and_speculation_share_forward_and_match_atomic( + chunk_size, budget, wavefront +): from vllm_rlt.request import Operation - def run(atomic): + def run(wavefront): e = engine( model(), cache_config=CacheConfig(256, 2), scheduler_config=SchedulerConfig( max_num_batched_tokens=budget, prefill_chunk_size=chunk_size, + wavefront_chunked_prefill=wavefront, ), - execution_config=ExecutionConfig(atomic_prefill=atomic), + execution_config=ExecutionConfig(), ) assert e.speculative_runner.runtime is e.model_runner.runtime operations = set() @@ -732,9 +736,13 @@ def record(batch): assert run(False) == run(True) -@pytest.mark.parametrize("atomic", [False, True]) -def test_common_submit_entry_preserves_prefill_and_speculative_results(atomic, monkeypatch): - e = engine(model(), execution_config=ExecutionConfig(atomic_prefill=atomic)) +@pytest.mark.parametrize("wavefront", [False, True]) +def test_common_submit_entry_preserves_prefill_and_speculative_results(wavefront, monkeypatch): + e = engine( + model(), + scheduler_config=SchedulerConfig(wavefront_chunked_prefill=wavefront), + execution_config=ExecutionConfig(), + ) calls = [] submit = e.executor.submit diff --git a/tests/test_split_prefill.py b/tests/test_split_prefill.py index 9f5a609f..66e5f6bc 100644 --- a/tests/test_split_prefill.py +++ b/tests/test_split_prefill.py @@ -16,6 +16,7 @@ from tests.helpers import make_prefill_engine, tiny_ouro_config, tiny_ouro_model from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig from vllm_rlt.core.scheduling_policy import ( + MixedPolicy, NoRefillPolicy, RefillPolicy, SchedulingPolicy, @@ -112,12 +113,16 @@ def test_recurrent_prepare_uses_absolute_token_start_for_kv_rows(): e.close() -def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(): +@pytest.mark.parametrize("wavefront", [False, True]) +def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(wavefront): from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput e = make_prefill_engine( scheduler_config=SchedulerConfig( - max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2 + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + wavefront_chunked_prefill=wavefront, ) ) e.add_request("mid", [1, 2], SamplingParams(max_tokens=1)) @@ -125,11 +130,20 @@ def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(): for request_id in ("mid", "final"): assert e.cache_manager.allocate(request_id, 2) mid, final = (e.scheduler.requests[rid] for rid in ("mid", "final")) + offset = 4 if wavefront else 0 batch = SchedulerOutput( RequestPhase.PREFILLING_RECURRENT, [ - ScheduledRequest(mid, 0, 2, depth=1, chunk_start=0), - ScheduledRequest(final, 0, 2, depth=3, chunk_start=0, completes_prompt=True), + ScheduledRequest(mid, 0, 2, depth=1, chunk_start=0, prefill_hidden_row_offset=offset), + ScheduledRequest( + final, + 0, + 2, + depth=3, + chunk_start=0, + completes_prompt=True, + prefill_hidden_row_offset=offset, + ), ], ) runner = e.model_runner @@ -139,8 +153,9 @@ def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(): runner.prepare_on_cpu(build_execution_batch(batch, runner.model.config.total_ut_steps)) ) mid_slot, final_slot = runner.state_slots["mid"], runner.state_slots["final"] - runner.prefill_states[mid_slot, :2].fill_(9) - runner.prefill_states[final_slot, :2].fill_(7) + target = slice(offset, offset + 2) + runner.prefill_states[mid_slot, target].fill_(9) + runner.prefill_states[final_slot, target].fill_(7) hidden = torch.arange( 4 * e.model.config.hidden_size, dtype=runner.prefill_states.dtype ).view(4, e.model.config.hidden_size) @@ -149,17 +164,95 @@ def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(): runner._store_split_prefill_outputs(hidden, prepared) assert prepared.batch_inputs.buffers.segments.cpu[:2].tolist() == [ - [mid_slot, 0, 2, 0, 1], - [final_slot, 0, 2, 2, 0], + [mid_slot, offset, 2, 0, 1], + [final_slot, offset, 2, 2, 0], ] - assert torch.equal(runner.prefill_states[mid_slot, :2], hidden[:2]) - assert torch.equal(runner.prefill_states[final_slot, :2], torch.full_like(hidden[2:], 7)) + assert torch.equal(runner.prefill_states[mid_slot, target], hidden[:2]) + assert torch.equal( + runner.prefill_states[final_slot, target], torch.full_like(hidden[2:], 7) + ) assert torch.equal(final.hidden_state, hidden[-1]) assert mid.hidden_state is None finally: e.close() +@pytest.mark.parametrize("mixed", [False, True]) +def test_wavefront_final_depth_does_not_store_any_chunk_hidden(monkeypatch, mixed): + from types import SimpleNamespace + + from vllm_rlt.request import ExecutionBatch, StateSpan + from vllm_rlt.worker.model_runner import BatchInputs, PreparedExecution + + e = make_prefill_engine( + scheduler_config=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, + ) + ) + runner = e.model_runner + e.add_request("r", [1, 2, 3, 4], SamplingParams(max_tokens=1)) + request = e.scheduler.requests["r"] + slot = runner._ensure_request_slot(request) + runner.prefill_states.fill_(7) + spans = ( + StateSpan( + request, + 0, + 2, + 3, + chunk_start=0, + retain_chunk=False, + prefill_hidden_row_offset=4, + ), + StateSpan( + request, + 2, + 2, + 3, + chunk_start=2, + retain_chunk=False, + retain_last=True, + prefill_hidden_row_offset=6, + ), + ) + if not mixed: + # Direct descriptors need the same output policy as the executor. + from dataclasses import replace + + spans = (replace(spans[0], retain_last=False), spans[1]) + batch = ExecutionBatch(Operation.RECURRENT, spans) + buffers = runner.batch_input_buffers[0] + buffers.reset() + prepared = PreparedExecution( + batch, (3, 3), (0, 2), (0, 0), batch_inputs=BatchInputs(buffers, 4, 2, mixed=mixed) + ) + hidden = torch.arange(4 * e.model.config.hidden_size, dtype=runner.prefill_states.dtype).view( + 4, e.model.config.hidden_size + ) + try: + if mixed: + runner._mixed_hidden_sources(batch, buffers) + assert all(chunk is None for _, _, chunk, _ in buffers.hidden_outputs) + runner._scatter_prefill(hidden, prepared) + assert torch.equal(request.hidden_state, hidden[-1]) + else: + # Reaching the CUDA branch would import/launch a kernel. A batch + # with no retained chunks must return before touching device state. + with monkeypatch.context() as patch: + patch.setattr(runner, "device", SimpleNamespace(type="cuda")) + runner._scatter_prefill(hidden, prepared) + runner._store_split_prefill_outputs(hidden, prepared) + assert torch.equal(request.hidden_state, hidden[-1]) + assert torch.equal( + runner.prefill_states[slot], torch.full_like(runner.prefill_states[slot], 7) + ) + finally: + e.close() + + @pytest.mark.gpu def test_cuda_prefill_segments_gather_and_scatter_variable_spans(): from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput @@ -168,9 +261,9 @@ def test_cuda_prefill_segments_gather_and_scatter_variable_spans(): OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16), cache_config=CacheConfig(64, 4), scheduler_config=SchedulerConfig( - max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=4 + mode="mixed", max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=4 ), - execution_config=ExecutionConfig(atomic_prefill=False), + execution_config=ExecutionConfig(), attention_backend="triton", ) try: @@ -213,9 +306,9 @@ def test_cuda_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(): OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16), cache_config=CacheConfig(64, 4), scheduler_config=SchedulerConfig( - max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2 + mode="mixed", max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2 ), - execution_config=ExecutionConfig(atomic_prefill=False), + execution_config=ExecutionConfig(), attention_backend="triton", ) prepared = None @@ -268,9 +361,9 @@ def test_cuda_prefill_submit_keeps_host_and_device_buffers_until_event(use_uva): OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16), cache_config=CacheConfig(64, 4), scheduler_config=SchedulerConfig( - max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=4 + mode="mixed", max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=4 ), - execution_config=ExecutionConfig(atomic_prefill=False, use_uva=use_uva), + execution_config=ExecutionConfig(use_uva=use_uva), attention_backend="triton", ) @@ -366,8 +459,8 @@ def test_submit_transfer_failure_releases_only_prepared_prefill_buffers(): e = LLMEngine( OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16), cache_config=CacheConfig(64, 4), - scheduler_config=SchedulerConfig(max_num_seqs=1, max_num_batched_tokens=4), - execution_config=ExecutionConfig(async_scheduling=True, atomic_prefill=False), + scheduler_config=SchedulerConfig(mode="mixed", max_num_seqs=1, max_num_batched_tokens=4), + execution_config=ExecutionConfig(async_scheduling=True), exit_config=ExitConfig("ouro_delayed"), attention_backend="triton", ) @@ -478,8 +571,8 @@ def test_cuda_submit_readback_failure_releases_batch_input_buffers(monkeypatch): e = LLMEngine( OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16), cache_config=CacheConfig(64, 4), - scheduler_config=SchedulerConfig(max_num_seqs=1, max_num_batched_tokens=4), - execution_config=ExecutionConfig(atomic_prefill=False), + scheduler_config=SchedulerConfig(mode="mixed", max_num_seqs=1, max_num_batched_tokens=4), + execution_config=ExecutionConfig(), attention_backend="triton", ) e.add_request("a", [1, 2, 3, 4], SamplingParams(max_tokens=1)) @@ -543,59 +636,61 @@ def schedule_stage(scheduler, stage): def test_split_prefill_steps_one_depth_per_chunk_and_keeps_the_prompt_ordered(): - e = make_prefill_engine() + e = make_prefill_engine( + scheduler_config=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=2, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, + ) + ) e.add_request("a", [1, 2, 3, 4, 5], SamplingParams(max_tokens=2, ignore_eos=True)) request = e.scheduler.requests["a"] - prompt = request.prompt_token_ids - - # Mirror the pending engine update: each completed stage enqueues the next one. - def apply(batch): - scheduled_request = batch.requests[0] - if batch.phase == RequestPhase.PREFILLING_PRELUDE: - request.num_preluded_tokens = ( - scheduled_request.token_start + scheduled_request.token_count - ) - e.scheduler.enqueue(request, RequestPhase.PREFILLING_RECURRENT) - return (scheduled_request.token_start, scheduled_request.token_count, None) - depth = advance_chunk(request, scheduled_request) - if depth + 1 < TOTAL_DEPTHS: - e.scheduler.enqueue(request, RequestPhase.PREFILLING_RECURRENT) - elif request.num_prefilled_tokens < len(prompt): - e.scheduler.enqueue(request, RequestPhase.PREFILLING_PRELUDE) - return (scheduled_request.token_start, scheduled_request.token_count, depth) - trace = [] + prelude_positions = [] + recurrent_visits = {} for _ in range(40): batch = e.scheduler.schedule() if batch is None: break - assert len(batch.requests) == 1 - trace.append((batch.phase, *apply(batch))) - assert e.scheduler.schedule() is None, "prefill must drain, not schedule forever" + assert all(item.request is request for item in batch.requests) + for item in batch.requests: + if batch.phase == RequestPhase.PREFILLING_PRELUDE: + prelude_positions.extend( + range(item.token_start, item.token_start + item.token_count) + ) + else: + assert batch.phase == RequestPhase.PREFILLING_RECURRENT + assert batch.loop_count == 1 + for position in range(item.token_start, item.token_start + item.token_count): + key = (position, item.depth) + recurrent_visits[key] = recurrent_visits.get(key, 0) + 1 + e.scheduler._apply_wavefront_prefill_batch_progress( + batch.phase, batch.requests, loop_count=batch.loop_count + ) + if request.phase == RequestPhase.CODA: + break - assert trace == [ - (RequestPhase.PREFILLING_PRELUDE, 0, 2, None), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 0), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 1), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 2), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 3), - (RequestPhase.PREFILLING_PRELUDE, 2, 2, None), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 0), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 1), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 2), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 3), - (RequestPhase.PREFILLING_PRELUDE, 4, 1, None), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 0), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 1), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 2), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 3), - ] + assert sorted(prelude_positions) == list(range(len(request.prompt_token_ids))) + assert recurrent_visits == { + (position, depth): 1 + for position in range(len(request.prompt_token_ids)) + for depth in range(TOTAL_DEPTHS) + } assert request.num_prefilled_tokens == 5 assert all(left == right for left, right in request.prefill_depth_ranges) def test_engine_split_prefill_transition_advances_the_same_frontier_contract(): - e = make_prefill_engine() + e = make_prefill_engine( + atomic=True, + scheduler_config=SchedulerConfig( + mode="mixed", + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + ), + ) e.add_request("a", [1, 2, 3, 4, 5], SamplingParams(max_tokens=2, ignore_eos=True)) request = e.scheduler.requests["a"] trace = [] @@ -606,8 +701,10 @@ def test_engine_split_prefill_transition_advances_the_same_frontier_contract(): break assert len(batch.requests) == 1 scheduled_request = batch.requests[0] - progress = e.scheduler._apply_split_prefill_progress( - request, batch.phase, scheduled_request + progress = e.scheduler._apply_prefill_progress( + scheduled_request, + batch.phase, + loop_count=batch.loop_count, ) trace.append( ( @@ -615,27 +712,16 @@ def test_engine_split_prefill_transition_advances_the_same_frontier_contract(): scheduled_request.token_start, scheduled_request.token_count, scheduled_request.depth, + batch.loop_count, ) ) if progress.prompt_complete: break assert trace == [ - (RequestPhase.PREFILLING_PRELUDE, 0, 2, None), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 0), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 1), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 2), - (RequestPhase.PREFILLING_RECURRENT, 0, 2, 3), - (RequestPhase.PREFILLING_PRELUDE, 2, 2, None), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 0), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 1), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 2), - (RequestPhase.PREFILLING_RECURRENT, 2, 2, 3), - (RequestPhase.PREFILLING_PRELUDE, 4, 1, None), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 0), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 1), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 2), - (RequestPhase.PREFILLING_RECURRENT, 4, 1, 3), + (RequestPhase.PREFILLING_PRELUDE, 0, 2, None, TOTAL_DEPTHS), + (RequestPhase.PREFILLING_PRELUDE, 2, 2, None, TOTAL_DEPTHS), + (RequestPhase.PREFILLING_PRELUDE, 4, 1, None, TOTAL_DEPTHS), ] assert request.num_prefilled_tokens == len(request.prompt_token_ids) assert request.phase is RequestPhase.CODA @@ -665,11 +751,185 @@ def test_sync_split_prefill_matches_atomic_tokens_and_exit_depths(chunk_size, to assert split.cache_manager.num_used_blocks == 0 -def test_sync_split_prefill_resumes_after_a_cached_prefix(): +def test_wavefront_prefill_batches_ready_depths_and_reuses_retired_ring_slots(): + scheduler = SchedulerConfig( + max_num_seqs=1, + max_num_batched_tokens=8, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, + ) + e = make_prefill_engine(scheduler_config=scheduler) + e.add_request( + "wave", list(range(1, 11)), SamplingParams(max_tokens=1, min_loops=1, max_loops=1) + ) + request = e.scheduler.requests["wave"] + + prelude = e.scheduler.schedule() + assert prelude.phase is RequestPhase.PREFILLING_PRELUDE + assert len(prelude.requests) == TOTAL_DEPTHS + assert [(item.chunk_start, item.prefill_hidden_row_offset) for item in prelude.requests] == [ + (0, 0), + (2, 2), + (4, 4), + (6, 6), + ] + e._update(prelude, None) + + recurrent_trace = [] + reused = None + for _ in range(20): + batch = e.scheduler.schedule() + assert batch is not None + if batch.phase == RequestPhase.PREFILLING_PRELUDE: + reused = batch.requests[0] + assert (reused.chunk_start, reused.prefill_hidden_row_offset) == (8, 0) + else: + recurrent_trace.append([(item.token_start, item.depth) for item in batch.requests]) + e._update(batch, None) + if request.phase is RequestPhase.CODA: + break + + assert recurrent_trace[:4] == [ + [(0, 0)], + [(0, 1), (2, 0)], + [(0, 2), (2, 1), (4, 0)], + [(0, 3), (2, 2), (4, 1), (6, 0)], + ] + assert reused is not None + assert request.num_prefilled_tokens == len(request.prompt_token_ids) + assert request.prefill_reclaimable_tokens == len(request.prompt_token_ids) + assert request.phase is RequestPhase.CODA + e.close() + + +def test_wavefront_prefill_matches_atomic_tokens_and_keeps_full_depth_first_output(): + wavefront = SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=5, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, + ) + atomic = SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=5, + prefill_chunk_size=2, + ) + prompt = list(range(2, 15)) + reference = make_prefill_engine(atomic=True, scheduler_config=atomic) + actual = make_prefill_engine(scheduler_config=wavefront) + expected_output = run_request(reference, prompt) + actual_output = run_request(actual, prompt) + + assert actual_output.token_ids == expected_output.token_ids + assert actual_output.exit_depths == expected_output.exit_depths + assert actual_output.exit_depths[0] == TOTAL_DEPTHS + assert actual.scheduler.config.wavefront_chunked_prefill + reference.close() + actual.close() + + +@pytest.mark.parametrize("budget", [1, 5, 17]) +def test_wavefront_prefill_preserves_hidden_and_all_depth_kv(budget): + snapshots = [] + for wavefront in (False, True): + engine = make_prefill_engine( + atomic=not wavefront, + scheduler_config=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=budget, + prefill_chunk_size=3, + wavefront_chunked_prefill=wavefront, + ), + ) + for rid, prompt in (("long", list(range(1, 15))), ("short", [2, 3, 4, 5, 6])): + engine.add_request(rid, prompt, SamplingParams(max_tokens=1)) + completed = {} + for _ in range(300): + if len(completed) == 2: + break + engine.step() + for request in engine.scheduler.requests.values(): + if request.phase == RequestPhase.CODA and request.request_id not in completed: + tensors = [request.hidden_state.clone()] + for depth in range(TOTAL_DEPTHS): + for layer in range(engine.model.config.num_hidden_layers): + tensors.extend( + tensor.clone() + for tensor in engine.cache_manager.read( + layer, request.request_id, depth + ) + ) + completed[request.request_id] = tensors + else: + pytest.fail("wavefront prompt snapshots did not complete") + snapshots.append(completed) + engine.close() + + assert snapshots[0].keys() == snapshots[1].keys() + for rid in snapshots[0]: + for expected, actual in zip(snapshots[0][rid], snapshots[1][rid]): + torch.testing.assert_close(actual, expected, atol=4e-5, rtol=4e-5) + + +def test_wavefront_mixed_prelude_prioritizes_decode_before_prompt_window(): + import torch + + scheduler = SchedulerConfig( + mode="mixed", + max_num_seqs=2, + max_num_batched_tokens=6, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, + ) + e = make_prefill_engine(scheduler_config=scheduler) + e.add_request("decode", [1], SamplingParams(max_tokens=2, ignore_eos=True)) + e.add_request("prompt", [2, 3, 4, 5, 6, 7], SamplingParams(max_tokens=2)) + e.scheduler._admit() + decode = e.scheduler.requests["decode"] + e.scheduler.queues[RequestPhase.PREFILLING_PRELUDE].remove("decode") + decode.generated_token_ids.append(9) + decode.hidden_state = torch.zeros(e.model.config.hidden_size) + e.scheduler.enqueue(decode, RequestPhase.PRELUDE) + + batch = e.scheduler.schedule() + assert batch.is_mixed + assert batch.num_tokens == 6 + assert batch.phase_for(batch.requests[0]) is RequestPhase.PRELUDE + prompt_items = [ + item for item in batch.requests if item.phase is RequestPhase.PREFILLING_PRELUDE + ] + assert [item.chunk_start for item in prompt_items] == [0, 2, 4] + assert sum(item.token_count for item in prompt_items) == 5 + e.close() + + +def test_wavefront_prefill_rejects_shared_kv(): + scheduler = SchedulerConfig(wavefront_chunked_prefill=True) + with pytest.raises(ValueError, match="staged prefill requires last_exited"): + make_prefill_engine( + atomic=True, + cache_config=CacheConfig(64, 2, "shared"), + scheduler_config=scheduler, + ) + + +@pytest.mark.parametrize("wavefront", [False, True]) +def test_sync_split_prefill_resumes_after_a_cached_prefix(wavefront): prompt = [1, 2, 3, 4, 5] cache = CacheConfig(64, 2, enable_prefix_caching=True) - scheduler = SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=3, prefill_chunk_size=2) - atomic = make_prefill_engine(atomic=True, cache_config=cache, scheduler_config=scheduler) + scheduler = SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=3, + prefill_chunk_size=2, + wavefront_chunked_prefill=wavefront, + ) + atomic = make_prefill_engine( + atomic=True, + cache_config=cache, + scheduler_config=SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=3, prefill_chunk_size=2 + ), + ) split = make_prefill_engine(cache_config=cache, scheduler_config=scheduler) expected = run_request(atomic, prompt) @@ -684,6 +944,9 @@ def test_sync_split_prefill_resumes_after_a_cached_prefix(): ) split.scheduler._admit() assert split.scheduler.requests["cached"].num_prefilled_tokens == 4 + if wavefront: + assert split.scheduler.requests["cached"].prefill_window_origin == 4 + assert split.scheduler.requests["cached"].prefill_reclaimable_tokens == 4 actual = drain_request(split) assert actual.token_ids == expected.token_ids @@ -691,8 +954,16 @@ def test_sync_split_prefill_resumes_after_a_cached_prefix(): assert split.model_runner.state_slots == {} -def test_sync_split_prefill_abort_releases_slot_before_request_id_reuse(): - split = make_prefill_engine() +@pytest.mark.parametrize("wavefront", [False, True]) +def test_sync_split_prefill_abort_releases_slot_before_request_id_reuse(wavefront): + split = make_prefill_engine( + scheduler_config=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + wavefront_chunked_prefill=wavefront, + ) + ) prompt = [1, 2, 3, 4] split.add_request("reused", prompt, SamplingParams(max_tokens=2, ignore_eos=True)) @@ -713,13 +984,21 @@ def test_sync_split_prefill_abort_releases_slot_before_request_id_reuse(): assert actual.exit_depths == expected.exit_depths -def test_cpu_async_split_prefill_matches_sync_split(): +@pytest.mark.parametrize("wavefront", [False, True]) +def test_cpu_async_split_prefill_matches_sync_split(wavefront): prompt = [1, 2, 3, 4, 5] trace = ExitConfig("trace", depths_by_request={"run": [TOTAL_DEPTHS, 2, 2]}) - sync = make_prefill_engine(exit_config=trace) + scheduler = SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + wavefront_chunked_prefill=wavefront, + ) + sync = make_prefill_engine(exit_config=trace, scheduler_config=scheduler) asynchronous = make_prefill_engine( exit_config=trace, - execution_config=ExecutionConfig(async_scheduling=True, atomic_prefill=False), + scheduler_config=scheduler, + execution_config=ExecutionConfig(async_scheduling=True), ) expected = run_request(sync, prompt) @@ -734,8 +1013,14 @@ def test_cpu_async_split_frontier_advances_once_per_submission(): prompt = [1, 2, 3, 4, 5] trace = ExitConfig("trace", depths_by_request={"run": [TOTAL_DEPTHS, 2, 2]}) e = make_prefill_engine( + scheduler_config=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=2, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, + ), exit_config=trace, - execution_config=ExecutionConfig(async_scheduling=True, atomic_prefill=False), + execution_config=ExecutionConfig(async_scheduling=True), ) e.add_request("run", prompt, SamplingParams(max_tokens=3, max_loops=2, ignore_eos=True)) request = e.scheduler.requests["run"] @@ -751,7 +1036,14 @@ def hold_split_ticket(batch): e.model_runner.events[scheduled_request.request.request_id] = ticket.event return ticket - e.executor.submit = hold_split_ticket + held_batches = [] + + def record_batch(batch): + if batch.phase in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT): + held_batches.append(batch) + return hold_split_ticket(batch) + + e.executor.submit = record_batch for _ in range(32): e.step() if request.phase is RequestPhase.CODA: @@ -760,7 +1052,11 @@ def hold_split_ticket(batch): assert request.phase is RequestPhase.CODA assert request.num_prefilled_tokens == len(prompt) assert all(left == right for left, right in request.prefill_depth_ranges) - assert len(held) == 15 # Three chunks, each with prelude plus four depths. + assert 0 < len(held) < 15 # Wavefront batches some chunk-depth tasks together. + assert all( + batch.loop_count == (0 if batch.phase == RequestPhase.PREFILLING_PRELUDE else 1) + for batch in held_batches + ) assert e.model_runner.split_prefill_submissions == len(held) assert all(not event.complete for event in held) for event in held: @@ -773,7 +1069,7 @@ def hold_split_ticket(batch): def test_split_submission_bound_counts_and_waits_for_oldest_event(): e = make_prefill_engine( exit_config=ExitConfig("ouro_delayed"), - execution_config=ExecutionConfig(async_scheduling=True, atomic_prefill=False), + execution_config=ExecutionConfig(async_scheduling=True), ) e.add_request("wait", [1, 2], SamplingParams(max_tokens=1, ignore_eos=True)) e.scheduler._admit() @@ -843,7 +1139,7 @@ def test_split_preemption_rewinds_and_recomputes_the_resident_chunk(completed_de assert e.scheduler.preempt(requester) assert e.preemption.snapshots["victim"]["phase"] is RequestPhase.PREFILLING_PRELUDE assert request.num_preluded_tokens == request.num_prefilled_tokens - assert request.prefill_chunk_start == request.num_prefilled_tokens + assert request.prefill_reclaimable_tokens == request.num_prefilled_tokens assert all(left == request.num_prefilled_tokens for left, _ in request.prefill_depth_ranges) assert e.cache_manager.num_used_blocks == 0 @@ -857,18 +1153,21 @@ def test_split_preemption_rewinds_and_recomputes_the_resident_chunk(completed_de assert e.cache_manager.num_used_blocks == 0 -def test_priority_admission_preempts_split_prefill_and_matches_reference(): +@pytest.mark.parametrize("wavefront", [False, True]) +def test_priority_admission_preempts_split_prefill_and_matches_reference(wavefront): prompt = [1, 2, 3, 4] params = SamplingParams(max_tokens=3, max_loops=2, ignore_eos=True, priority=10) trace = ExitConfig("trace", depths_by_request={"reference": [4, 2, 2]}) - scheduler = SchedulerConfig( + reference_scheduler = SchedulerConfig( max_num_seqs=1, max_num_batched_tokens=4, prefill_chunk_size=2, policy="priority", ) expected = run_request( - make_prefill_engine(scheduler_config=scheduler, exit_config=trace), prompt, "reference" + make_prefill_engine(scheduler_config=reference_scheduler, exit_config=trace), + prompt, + "reference", ) e = make_prefill_engine( @@ -878,6 +1177,7 @@ def test_priority_admission_preempts_split_prefill_and_matches_reference(): prefill_chunk_size=2, policy="priority", enable_preemption=True, + wavefront_chunked_prefill=wavefront, ), exit_config=ExitConfig("trace", depths_by_request={"low": [4, 2, 2], "high": [4]}), ) @@ -912,14 +1212,18 @@ def test_priority_admission_preempts_split_prefill_and_matches_reference(): @pytest.mark.gpu @pytest.mark.parametrize("async_scheduling", [False, True]) @pytest.mark.parametrize("failure_site", ["prepare_on_gpu", "_execute"]) -def test_cuda_split_failure_drains_writes_before_request_id_reuse( +def test_cuda_wavefront_failure_drains_writes_before_request_id_reuse( monkeypatch, async_scheduling, failure_site ): e = LLMEngine( OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16), cache_config=CacheConfig(64, 4), - scheduler_config=SchedulerConfig(max_num_seqs=1, max_num_batched_tokens=4), - execution_config=ExecutionConfig(atomic_prefill=False, async_scheduling=async_scheduling), + scheduler_config=SchedulerConfig( + wavefront_chunked_prefill=True, + max_num_seqs=1, + max_num_batched_tokens=4, + ), + execution_config=ExecutionConfig(async_scheduling=async_scheduling), exit_config=ExitConfig("trace", depths_by_request={"reused": [4]}), attention_backend="triton", ) @@ -962,8 +1266,9 @@ def fail(batch, *args, **kwargs): @pytest.mark.gpu @pytest.mark.parametrize("async_scheduling", [False, True]) @pytest.mark.parametrize("multi_stream", [False, True]) -def test_cuda_split_prefill_matches_atomic(async_scheduling, multi_stream): - prompt = [1, 2, 3, 4, 5] +@pytest.mark.parametrize("wavefront", [False, True]) +def test_cuda_split_prefill_matches_atomic(async_scheduling, multi_stream, wavefront): + prompt = list(range(1, 10)) params = SamplingParams(max_tokens=3, max_loops=2, ignore_eos=True) execution = dict( async_scheduling=async_scheduling, @@ -976,8 +1281,12 @@ def test_cuda_split_prefill_matches_atomic(async_scheduling, multi_stream): atomic = LLMEngine( atomic_model, cache_config=CacheConfig(128, 2, enable_prefix_caching=True), - scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=4), - execution_config=ExecutionConfig(atomic_prefill=True, **execution), + scheduler_config=SchedulerConfig( + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + ), + execution_config=ExecutionConfig(**execution), exit_config=ExitConfig("ouro_delayed"), attention_backend="triton", ) @@ -994,8 +1303,14 @@ def test_cuda_split_prefill_matches_atomic(async_scheduling, multi_stream): split = LLMEngine( split_model, cache_config=CacheConfig(128, 2, enable_prefix_caching=True), - scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=4), - execution_config=ExecutionConfig(atomic_prefill=False, **execution), + scheduler_config=SchedulerConfig( + mode="refill" if wavefront else "mixed", + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + wavefront_chunked_prefill=wavefront, + ), + execution_config=ExecutionConfig(**execution), exit_config=ExitConfig("ouro_delayed"), attention_backend="triton", ) @@ -1025,12 +1340,16 @@ def run(cuda_graphs): cache = CacheConfig( block_size=16, gpu_memory_utilization=0.5, memory_reserve_bytes=32 * 1024**2 ) - scheduler = SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2) + scheduler = SchedulerConfig( + wavefront_chunked_prefill=True, + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + ) execution = ExecutionConfig( async_scheduling=True, cuda_graphs=cuda_graphs, cuda_graph_memory_reserve_bytes=graph_reserve, - atomic_prefill=False, ) e = LLMEngine( model, @@ -1131,8 +1450,13 @@ def run(cuda_graphs): def test_cuda_async_split_preemption_releases_and_reuses_slot(): prompt = [1, 2, 3, 4] params = SamplingParams(max_tokens=3, max_loops=2, ignore_eos=True, priority=10) - scheduler = SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2) - execution = ExecutionConfig(atomic_prefill=False) + scheduler = SchedulerConfig( + wavefront_chunked_prefill=True, + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + ) + execution = ExecutionConfig() torch.manual_seed(432) reference_model = OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16) reference = LLMEngine( @@ -1154,13 +1478,15 @@ def test_cuda_async_split_preemption_releases_and_reuses_slot(): model, cache_config=CacheConfig(128, 2), scheduler_config=SchedulerConfig( + mode="refill", max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2, policy="priority", enable_preemption=True, + wavefront_chunked_prefill=True, ), - execution_config=ExecutionConfig(async_scheduling=True, atomic_prefill=False), + execution_config=ExecutionConfig(async_scheduling=True), exit_config=ExitConfig( "trace", depths_by_request={ @@ -1227,9 +1553,13 @@ def test_cuda_split_prefill_batches_mixed_depths(attention_backend): model, cache_config=CacheConfig(128, 16, enable_prefix_caching=True), scheduler_config=SchedulerConfig( - max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2 + mode="mixed", + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, ), - execution_config=ExecutionConfig(atomic_prefill=False), + execution_config=ExecutionConfig(), exit_config=ExitConfig("ouro_delayed"), attention_backend=attention_backend, ) @@ -1258,7 +1588,15 @@ def test_cuda_split_prefill_batches_mixed_depths(attention_backend): def test_split_prefill_holds_the_next_prelude_until_the_chunk_is_read(): """One staged chunk per request: its hidden rows are in use until the last depth.""" - e = make_prefill_engine() + e = make_prefill_engine( + atomic=True, + scheduler_config=SchedulerConfig( + mode="mixed", + max_num_seqs=2, + max_num_batched_tokens=4, + prefill_chunk_size=2, + ), + ) e.add_request("a", [1, 2, 3, 4], SamplingParams(max_tokens=2, ignore_eos=True)) e.scheduler._admit() request = e.scheduler.requests["a"] @@ -1395,9 +1733,20 @@ def test_prefill_depth_ranges_track_logical_depths(): ).prefill_depth(5) -def test_split_prefill_policy_is_selected_only_by_the_atomic_prefill_flag(): +def test_staged_prefill_policy_is_selected_by_scheduler_requirements(): assert RefillPolicy is type(make_prefill_engine(atomic=True).scheduler.policy) - assert SplitPrefillPolicy is type(make_prefill_engine().scheduler.policy) + assert SplitPrefillPolicy is type( + make_prefill_engine( + atomic=True, + scheduler_config=SchedulerConfig(wavefront_chunked_prefill=True), + ).scheduler.policy + ) + assert MixedPolicy is type( + make_prefill_engine( + atomic=True, + scheduler_config=SchedulerConfig(mode="mixed"), + ).scheduler.policy + ) assert NoRefillPolicy is type( make_prefill_engine( atomic=True, scheduler_config=SchedulerConfig(mode="no_refill") @@ -1422,12 +1771,17 @@ def test_scheduling_policies_declare_their_prefill_stages(): ( "shared layout", {"cache_config": CacheConfig(64, 2, "shared")}, - "split prefill requires last_exited", + "staged prefill requires last_exited", ), ( - "no_refill mode", - {"scheduler_config": SchedulerConfig(mode="no_refill")}, - "split prefill currently implements refill", + "wavefront no_refill mode", + { + "atomic": True, + "scheduler_config": SchedulerConfig( + mode="no_refill", wavefront_chunked_prefill=True + ), + }, + "staged prefill currently implements refill or mixed", ), ], ) @@ -1444,10 +1798,13 @@ def test_chunk_streams_wait_before_hidden_gather(monkeypatch, async_scheduling, OuroForCausalLM(tiny_ouro_config()).to("cuda", torch.bfloat16), cache_config=CacheConfig(64, 4), scheduler_config=SchedulerConfig( - max_num_seqs=1, max_num_batched_tokens=4, prefill_chunk_size=4 + wavefront_chunked_prefill=True, + max_num_seqs=1, + max_num_batched_tokens=4, + prefill_chunk_size=4, ), execution_config=ExecutionConfig( - atomic_prefill=False, async_scheduling=async_scheduling, multi_stream=multi_stream + async_scheduling=async_scheduling, multi_stream=multi_stream ), exit_config=ExitConfig("trace", depths_by_request={"r": [4]}), attention_backend="triton", @@ -1506,9 +1863,9 @@ def test_chunk_prompt_cache_reuse_override_and_request_id_reuse(device, use_uva) tiny_ouro_model().to(device), cache_config=CacheConfig(64, 2), scheduler_config=SchedulerConfig( - max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2 + mode="mixed", max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2 ), - execution_config=ExecutionConfig(atomic_prefill=False, use_uva=use_uva), + execution_config=ExecutionConfig(use_uva=use_uva), attention_backend="triton" if device == "cuda" else "torch", ) runner = e.model_runner @@ -1560,9 +1917,7 @@ def test_mixed_hidden_buffers_wait_for_strided_decode_and_chunk_producers(multi_ scheduler_config=SchedulerConfig( mode="mixed", max_num_seqs=2, max_num_batched_tokens=4, prefill_chunk_size=2 ), - execution_config=ExecutionConfig( - atomic_prefill=False, multi_stream=multi_stream, use_uva=use_uva - ), + execution_config=ExecutionConfig(multi_stream=multi_stream, use_uva=use_uva), attention_backend="triton", ) runner = e.model_runner diff --git a/tests/test_stage_engine.py b/tests/test_stage_engine.py index f3c82045..71facca1 100644 --- a/tests/test_stage_engine.py +++ b/tests/test_stage_engine.py @@ -124,6 +124,7 @@ def test_independent_stage_resources_preserve_output_after_import(atomic, asynch d.cache_manager.value_cache[dst].copy_(p.cache_manager.value_cache[src]) assert d.accept_import(target, source.hidden_state.clone()) assert d.scheduler.requests["r"] is target + assert target.prefill_reclaimable_tokens == len(prompt) assert p.release_transfer(source) local.add_request("r", prompt, params) assert drain(d) == drain(local) @@ -265,9 +266,7 @@ def test_spawned_mixed_worker_matches_local_without_connector(asynchronous): scheduler_config=SchedulerConfig( mode="mixed", max_num_seqs=3, max_num_batched_tokens=4, prefill_chunk_size=3 ), - execution_config=ExecutionConfig( - atomic_prefill=False, async_scheduling=asynchronous, static_buffers=True - ), + execution_config=ExecutionConfig(async_scheduling=asynchronous, static_buffers=True), exit_config=ExitConfig("ouro_delayed" if asynchronous else "ouro"), ), ) @@ -596,8 +595,9 @@ def test_transfer_reservation_preserves_active_prompt_growth(receiving, preempti max_num_batched_tokens=2, prefill_chunk_size=2, enable_preemption=preemption, + wavefront_chunked_prefill=True, ), - execution_config=ExecutionConfig(atomic_prefill=False), + execution_config=ExecutionConfig(), ) params = SamplingParams(max_tokens=1, ignore_eos=True) try: diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index ae8bd06b..8872d018 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -60,6 +60,7 @@ class SchedulerConfig: max_prefill_batches_before_decode: int = 1 policy: str = "fcfs" enable_preemption: bool = False + wavefront_chunked_prefill: bool = False def __post_init__(self): for name in ( @@ -74,6 +75,8 @@ def __post_init__(self): _positive(name, getattr(self, name)) if type(self.enable_preemption) is not bool: raise ValueError("enable_preemption must be a boolean") + if type(self.wavefront_chunked_prefill) is not bool: + raise ValueError("wavefront_chunked_prefill must be a boolean") if self.policy not in ("fcfs", "priority"): raise ValueError("policy must be fcfs or priority") if self.mode not in ("refill", "no_refill", "mixed"): @@ -116,16 +119,11 @@ class ExecutionConfig: # Select packed prefill metadata staging (and require packed attention). # The historical name is retained; use_uva controls its transfer mechanism. prefill_uva: bool = False - # Split prompt prefill into a staged prelude plus one scheduling boundary per - # recurrence depth. Atomic prefill keeps mainline behavior and remains the - # default until the split path is feature-complete. - atomic_prefill: bool = True def __post_init__(self): for name in ( "prefill_uva", "use_uva", - "atomic_prefill", "async_scheduling", "multi_stream", "static_buffers", diff --git a/vllm_rlt/core/memory.py b/vllm_rlt/core/memory.py index f0a8edc9..220643d9 100644 --- a/vllm_rlt/core/memory.py +++ b/vllm_rlt/core/memory.py @@ -18,7 +18,7 @@ def budget_blocks(available_bytes, bytes_per_block): @torch.inference_mode() -def plan_cache(model, cache, scheduler, execution, backend): +def plan_cache(model, cache, scheduler, execution, backend, *, staged_prefill=None): parameter = next(model.parameters()) config, device = model.config, parameter.device per_block = ( @@ -100,7 +100,14 @@ def plan_cache(model, cache, scheduler, execution, backend): * (8 * config.head_dim + 8) * 4 ) - buffers = execution_buffer_bytes(config, scheduler, cache, execution, parameter.element_size()) + buffers = execution_buffer_bytes( + config, + scheduler, + cache, + execution, + parameter.element_size(), + staged_prefill=staged_prefill, + ) # Dynamic metadata at the largest supported context is small but not zero. metadata_width = math.ceil(config.max_position_embeddings / cache.block_size) metadata_rows = scheduler.max_num_batched_tokens diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 509246b5..f443d742 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -5,7 +5,7 @@ from enum import Enum, auto from vllm_rlt.config import SchedulerConfig -from vllm_rlt.core.scheduling_policy import create_scheduling_policy +from vllm_rlt.core.scheduling_policy import create_scheduling_policy, requires_staged_prefill from vllm_rlt.request import FinishReason, Request, RequestPhase from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput @@ -33,6 +33,9 @@ class ScheduledRequest: chunk_start: int | None = None completes_prompt: bool = False phase: RequestPhase | None = None + # Row offset in the request's prefill hidden-state buffer. None lets the runner + # derive token_start - chunk_start. + prefill_hidden_row_offset: int | None = None @dataclass(frozen=True) @@ -40,6 +43,15 @@ class SchedulerOutput: # None means a mixed batch; each scheduled request carries its own phase. phase: RequestPhase | None requests: list[ScheduledRequest] + # Actual number of recurrent stages in this submission. None is reserved for + # direct callers that construct a batch without passing through Scheduler. + loop_count: int | None = None + + def __post_init__(self): + if self.loop_count is not None and ( + type(self.loop_count) is not int or self.loop_count < 0 + ): + raise ValueError("loop_count must be a nonnegative integer") @property def is_speculative(self) -> bool: @@ -71,6 +83,7 @@ class PrefillProgress: chunk_complete: bool prompt_complete: bool + transfer_start: int | None = None @dataclass @@ -122,7 +135,6 @@ def __init__( config: SchedulerConfig, cache_manager, speculative_config=None, - atomic_prefill: bool = True, *, produces_tokens: bool = True, total_ut_steps: int | None = None, @@ -151,11 +163,14 @@ def __init__( # Optional connector readiness and export predicates. Never mutate queues or blocks. self.execution_ready = None self.exports_request = None - # One scheduling boundary per chunk (atomic) or per recurrence depth (split). - # The policy is derived from the same flag so the two cannot disagree. - self.atomic_prefill = atomic_prefill + # Staged prompt state is required by wavefront regrouping and mixed + # recurrent batches. Other prompt chunks execute all depths atomically. + self.wavefront_chunked_prefill = config.wavefront_chunked_prefill + self.staged_prefill = requires_staged_prefill(config, produces_tokens=produces_tokens) self.policy = create_scheduling_policy( - config, speculative_config, atomic_prefill, produces_tokens=produces_tokens + config, + speculative_config, + produces_tokens=produces_tokens, ) def add_request(self, request: Request): @@ -438,7 +453,8 @@ def _commit_admission(self, request: Request, plan: AdmissionPlan) -> bool: # A reused prompt prefix is already prefilled at every depth, so the # staged window starts closed at that frontier and opens for the rest. request.num_preluded_tokens = plan.cached_tokens - request.prefill_chunk_start = plan.cached_tokens + request.prefill_window_origin = plan.cached_tokens + request.prefill_reclaimable_tokens = plan.cached_tokens if request.prefill_depth_ranges: request.prefill_depth_ranges[:] = Request.prefill_boundaries( len(request.prompt_token_ids), @@ -449,10 +465,10 @@ def _commit_admission(self, request: Request, plan: AdmissionPlan) -> bool: return True def _enter_prefill(self, request): - if self.atomic_prefill: - self.enqueue(request, RequestPhase.PREFILL) - else: + if self.staged_prefill: self.enqueue(request, RequestPhase.PREFILLING_PRELUDE) + else: + self.enqueue(request, RequestPhase.PREFILL) def _admit(self): """Admit waiting requests: order -> slot -> restore/allocate -> defer/enter. @@ -532,9 +548,10 @@ def _admit(self): def _make_scheduled_request( self, request: Request, phase: RequestPhase, token_budget: int ) -> ScheduledRequest: - """Snapshot one request's execution range without advancing its progress. + """Snapshot one execution range without advancing request progress. Prompt length 10, prefilled 4, chunk 4, budget 6 -> range [4, 8). + Wavefront expands a queued request into multiple ranges in take_phases(). Decode phases contribute one position; its loop may differ by request. """ if phase == RequestPhase.PREFILL: @@ -556,7 +573,7 @@ def _make_scheduled_request( start = request.num_prefilled_tokens if start >= request.num_preluded_tokens: return ScheduledRequest( - request, start, 0, depth=0, chunk_start=request.prefill_chunk_start + request, start, 0, depth=0, chunk_start=request.prefill_reclaimable_tokens ) first_token_depth = request.prefill_depth(start) # One run per depth frontier: every token of the run already completed @@ -576,7 +593,7 @@ def _make_scheduled_request( start, count, depth=first_token_depth, - chunk_start=request.prefill_chunk_start, + chunk_start=request.prefill_reclaimable_tokens, completes_prompt=( first_token_depth == len(request.prefill_depth_ranges) - 1 and start + count == len(request.prompt_token_ids) @@ -592,6 +609,52 @@ def _make_scheduled_request( return ScheduledRequest(request, request.position, depth=request.loops_done) return ScheduledRequest(request) + def _make_wavefront_task( + self, request, phase, token_budget, *, start, depth=None + ) -> ScheduledRequest | None: + """Snapshot one chunk within the bounded per-request wavefront ring.""" + chunk_size = self.config.prefill_chunk_size + origin = request.prefill_window_origin + if start < origin: + raise RuntimeError("wavefront work precedes its cached-prefix origin") + chunk_start = origin + (start - origin) // chunk_size * chunk_size + chunk_offset = start - chunk_start + chunk_end = min(len(request.prompt_token_ids), chunk_start + chunk_size) + total_depths = len(request.prefill_depth_ranges) + slot = ((chunk_start - origin) // chunk_size) % total_depths + prefill_hidden_row_offset = slot * chunk_size + chunk_offset + + if phase == RequestPhase.PREFILLING_PRELUDE: + window_end = min( + len(request.prompt_token_ids), + request.prefill_reclaimable_tokens + total_depths * chunk_size, + ) + end = min(chunk_end, window_end) + else: + if depth is None: + raise ValueError("wavefront recurrent work requires an explicit depth") + left, right = request.prefill_depth_ranges[depth] + if left != start: + raise RuntimeError("wavefront recurrent task is not at its depth frontier") + end = min(chunk_end, right, request.num_preluded_tokens) + + count = min(token_budget, end - start) + if count <= 0: + return None + return ScheduledRequest( + request, + start, + count, + depth=depth, + chunk_start=chunk_start, + completes_prompt=( + phase == RequestPhase.PREFILLING_RECURRENT + and depth == total_depths - 1 + and start + count == len(request.prompt_token_ids) + ), + prefill_hidden_row_offset=prefill_hidden_row_offset, + ) + def _ensure_execution_capacity( self, request: Request, frontier: int, *, selected: set[str] ) -> bool: @@ -612,11 +675,16 @@ def _ensure_execution_capacity( def _prelude_window_open(self, request: Request) -> bool: """Whether the next prompt chunk may be staged for this request. - Split prefill keeps one staged chunk: deeper depths of the current chunk - still read its hidden rows, so the next chunk is staged only after the - current one reached the last depth. Wavefront chunked prefill widens the - window to several chunks, and then this bound becomes a chunk count. + Split prefill keeps one staged chunk. WCPB opens a bounded ring window, + and retires one ring slot after each contiguous final-depth chunk. """ + if self.wavefront_chunked_prefill: + window_end = min( + len(request.prompt_token_ids), + request.prefill_reclaimable_tokens + + self.config.prefill_chunk_size * len(request.prefill_depth_ranges), + ) + return request.num_preluded_tokens < window_end return request.num_preluded_tokens == request.num_prefilled_tokens def has_ready(self, phase): @@ -637,13 +705,21 @@ def take_phases(self, phases, *, eligible=None): selected: set[str] = set() for phase in phases: queue = self.queues[phase] - remaining = len(queue) - while ( - queue - and token_budget - and len(scheduled_requests) < self.config.max_num_seqs - and remaining + if self.wavefront_chunked_prefill and phase in ( + RequestPhase.PREFILLING_PRELUDE, + RequestPhase.PREFILLING_RECURRENT, ): + token_budget = self._take_wavefront_phase( + phase, + queue, + token_budget, + scheduled_requests, + selected, + eligible, + ) + continue + remaining = len(queue) + while queue and token_budget and len(selected) < self.config.max_num_seqs and remaining: remaining -= 1 request = self.requests[queue.popleft()] if eligible is not None and not eligible(request): @@ -687,16 +763,189 @@ def take_phases(self, phases, *, eligible=None): # Fairness follows the requests actually selected, including a mixed attempt # whose other queue could not reserve capacity. selected_phases = {scheduled_request.phase for scheduled_request in scheduled_requests} + prompt_phases = selected_phases.intersection(self.policy.PREFILL_PHASES) recorded = ( RequestPhase.RECURRENT if RequestPhase.RECURRENT in selected_phases + else next(iter(prompt_phases)) + if prompt_phases else scheduled_requests[0].phase ) self.policy.record_batch(recorded) + loop_count = self._planned_loop_count(scheduled_requests) + scheduled_requests = [ + replace( + item, + completes_prompt=( + self._scheduled_prefill_completes(item, loop_count) + if item.phase + in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + else item.completes_prompt + ), + ) + for item in scheduled_requests + ] return SchedulerOutput( - next(iter(selected_phases)) if len(selected_phases) == 1 else None, scheduled_requests + next(iter(selected_phases)) if len(selected_phases) == 1 else None, + scheduled_requests, + loop_count, ) + def _planned_loop_count(self, scheduled_requests): + phases = {item.phase for item in scheduled_requests} + if phases == {RequestPhase.PREFILL}: + return self.total_ut_steps + if RequestPhase.SPECULATIVE in phases or RequestPhase.CODA in phases: + return 0 + if RequestPhase.PRELUDE in phases: + return 0 + if self.wavefront_chunked_prefill: + return 1 if RequestPhase.PREFILLING_RECURRENT in phases else 0 + if RequestPhase.PREFILLING_PRELUDE in phases: + # In mixed mode, stage just the embeddings when recurrent decode work + # is ready so the next depth can join that batch. Without a reason to + # yield, PRELUDE can execute the full prompt depth in this submission. + if self.config.mode == "mixed" and self.has_ready(RequestPhase.RECURRENT): + return 0 + return self.total_ut_steps + if RequestPhase.PREFILLING_RECURRENT in phases: + if self.config.mode == "mixed" and ( + RequestPhase.RECURRENT in phases or self.has_ready(RequestPhase.RECURRENT) + ): + return 1 + return min(self._max_contiguous_prefill_loops(item) for item in scheduled_requests) + if RequestPhase.RECURRENT in phases: + return 1 + raise ValueError(f"cannot derive recurrent boundary for phases: {phases}") + + def _batch_loop_count(self, batch): + if batch.loop_count is not None: + return batch.loop_count + phases = {batch.phase_for(item) for item in batch.requests} + if RequestPhase.PREFILL in phases: + return self.total_ut_steps + if phases.intersection((RequestPhase.PRELUDE, RequestPhase.CODA, RequestPhase.SPECULATIVE)): + return 0 + if phases.intersection((RequestPhase.PREFILLING_RECURRENT, RequestPhase.RECURRENT)): + return 1 + return 0 + + def _max_contiguous_prefill_loops(self, scheduled_request): + """Count consecutive depth frontiers that already cover this exact span.""" + start = scheduled_request.token_start + end = start + scheduled_request.token_count + depth = scheduled_request.depth + loops = 0 + for left, right in scheduled_request.request.prefill_depth_ranges[depth:]: + if left != start or right < end: + break + loops += 1 + return max(1, loops) + + def _scheduled_prefill_completes(self, scheduled_request, loop_count): + phase = scheduled_request.phase + depth = scheduled_request.depth if phase == RequestPhase.PREFILLING_RECURRENT else 0 + return ( + loop_count > 0 + and depth + loop_count == self.total_ut_steps + and scheduled_request.token_start + scheduled_request.token_count + == len(scheduled_request.request.prompt_token_ids) + ) + + def _take_wavefront_phase( + self, phase, queue, token_budget, scheduled_requests, selected, eligible + ): + """Round-robin each request's chunks/depths into the shared token budget.""" + remaining = len(queue) + candidates = [] + owner_limit = self.config.max_num_seqs - len(selected) + while queue and remaining and len(candidates) < owner_limit and token_budget: + remaining -= 1 + request = self.requests[queue.popleft()] + if eligible is not None and not eligible(request): + queue.append(request.request_id) + continue + if phase == RequestPhase.PREFILLING_PRELUDE and not self._prelude_window_open(request): + queue.append(request.request_id) + continue + if phase == RequestPhase.PREFILLING_RECURRENT and not any( + start < min(end, request.num_preluded_tokens) + for start, end in request.prefill_depth_ranges + ): + queue.append(request.request_id) + continue + candidates.append(request) + + if not candidates: + return token_budget + + added = set() + blocked = set() + if phase == RequestPhase.PREFILLING_PRELUDE: + cursors = {request.request_id: request.num_preluded_tokens for request in candidates} + for _ in range(len(candidates[0].prefill_depth_ranges)): + for request in candidates: + request_id = request.request_id + if token_budget <= 0: + break + if request_id in blocked or request.phase != phase: + continue + start = cursors[request_id] + if start >= len(request.prompt_token_ids): + continue + scheduled_request = self._make_wavefront_task( + request, phase, token_budget, start=start + ) + if scheduled_request is None: + continue + frontier = scheduled_request.token_start + scheduled_request.token_count + if not self._ensure_execution_capacity(request, frontier, selected=selected): + blocked.add(request_id) + continue + scheduled_requests.append(replace(scheduled_request, phase=phase)) + selected.add(request_id) + added.add(request_id) + cursors[request_id] += scheduled_request.token_count + token_budget -= scheduled_request.token_count + else: + tasks = [] + for depth in range(len(candidates[0].prefill_depth_ranges)): + for owner_index, request in enumerate(candidates): + request_id = request.request_id + if token_budget <= 0: + break + if request_id in blocked or request.phase != phase: + continue + start, end = request.prefill_depth_ranges[depth] + end = min(end, request.num_preluded_tokens) + if start >= end: + continue + scheduled_request = self._make_wavefront_task( + request, phase, token_budget, depth=depth, start=start + ) + if scheduled_request is None: + continue + frontier = scheduled_request.token_start + scheduled_request.token_count + if not self._ensure_execution_capacity(request, frontier, selected=selected): + blocked.add(request_id) + continue + tasks.append((owner_index, scheduled_request)) + selected.add(request_id) + added.add(request_id) + token_budget -= scheduled_request.token_count + + tasks.sort(key=lambda item: (item[0], item[1].token_start)) + scheduled_requests.extend( + replace(scheduled_request, phase=phase) for _, scheduled_request in tasks + ) + + for request in candidates: + # Growing another candidate may have suspended this request and put + # it in WAITING. It must not also return to its old prompt queue. + if request.request_id not in added and request.phase == phase: + queue.append(request.request_id) + return token_budget + def schedule(self, *, prefer_recurrent=False) -> SchedulerOutput | None: """Choose queued requests or an admission opportunity, then build a batch. @@ -719,6 +968,7 @@ def _apply_split_prefill_progress( phase: RequestPhase, scheduled_request: ScheduledRequest, *, + loop_count: int | None = None, event=None, ) -> PrefillProgress: """Advance host-side logical progress for a submitted prefill request. @@ -729,24 +979,39 @@ def _apply_split_prefill_progress( total_depths = self.total_ut_steps start = scheduled_request.token_start end = start + scheduled_request.token_count + if loop_count is None: + loop_count = 0 if phase == RequestPhase.PREFILLING_PRELUDE else 1 if phase == RequestPhase.PREFILLING_PRELUDE: - request.prefill_chunk_start = start + request.prefill_reclaimable_tokens = start request.num_preluded_tokens = end + depth = 0 + else: + depth = scheduled_request.depth + if loop_count <= 0: + if phase != RequestPhase.PREFILLING_PRELUDE: + raise ValueError("recurrent prefill progress requires at least one loop") self.enqueue(request, RequestPhase.PREFILLING_RECURRENT) return PrefillProgress(False, False) - # Assertion: phase == RequestPhase.PREFILLING_RECURRENT - depth = scheduled_request.depth boundaries = request.prefill_depth_ranges - boundaries[depth][0] = end - if depth + 1 < total_depths: - boundaries[depth + 1][1] = end - else: - request.num_prefilled_tokens = end + if depth + loop_count > total_depths: + raise ValueError("prefill execution range exceeds model loop depth") + for offset in range(loop_count): + current_depth = depth + offset + if phase == RequestPhase.PREFILLING_RECURRENT and ( + boundaries[current_depth][0] != start or boundaries[current_depth][1] < end + ): + raise RuntimeError("prefill execution range crosses an incomplete depth frontier") + boundaries[current_depth][0] = end + if current_depth + 1 < total_depths: + boundaries[current_depth + 1][1] = end + else: + request.num_prefilled_tokens = end - chunk_complete = end == request.num_preluded_tokens - if depth + 1 < total_depths or not chunk_complete: + final_depth_reached = depth + loop_count == total_depths + chunk_complete = final_depth_reached and end == request.num_preluded_tokens + if not final_depth_reached or not chunk_complete: self.enqueue(request, RequestPhase.PREFILLING_RECURRENT) return PrefillProgress(False, False) @@ -756,7 +1021,7 @@ def _apply_split_prefill_progress( request.num_prefilled_tokens, event, ) - request.prefill_chunk_start = end + request.prefill_reclaimable_tokens = end prompt_complete = end == len(request.prompt_token_ids) if prompt_complete: request.loops_done = total_depths @@ -767,7 +1032,12 @@ def _apply_split_prefill_progress( return PrefillProgress(True, prompt_complete) def _apply_prefill_progress( - self, scheduled_request: ScheduledRequest, phase: RequestPhase, *, event=None + self, + scheduled_request: ScheduledRequest, + phase: RequestPhase, + *, + loop_count: int | None = None, + event=None, ) -> PrefillProgress: """One prompt frontier update for local and remotely consumed outputs.""" request = scheduled_request.request @@ -783,7 +1053,109 @@ def _apply_prefill_progress( request, self._prompt_destination(request) if final else RequestPhase.PREFILL ) return PrefillProgress(True, final) - return self._apply_split_prefill_progress(request, phase, scheduled_request, event=event) + return self._apply_split_prefill_progress( + request, phase, scheduled_request, loop_count=loop_count, event=event + ) + + def _apply_wavefront_prefill_batch_progress( + self, + phase: RequestPhase, + scheduled_requests: list[ScheduledRequest], + *, + loop_count: int = 1, + event=None, + ) -> tuple[ScheduledRequest, PrefillProgress]: + """Apply all spans for one WCPB request before advancing its lifecycle.""" + if not scheduled_requests: + raise ValueError("wavefront progress requires at least one span") + scheduled_requests = sorted(scheduled_requests, key=lambda item: item.token_start) + request = scheduled_requests[0].request + if any(item.request is not request for item in scheduled_requests): + raise ValueError("wavefront progress spans must share one request") + if phase not in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT): + raise ValueError(f"invalid wavefront prefill phase: {phase}") + expected_loops = 0 if phase == RequestPhase.PREFILLING_PRELUDE else 1 + if loop_count != expected_loops: + raise ValueError("wavefront prefill requires one recurrent depth per scheduling step") + + transfer_start = request.prefill_reclaimable_tokens + if phase == RequestPhase.PREFILLING_PRELUDE: + frontier = request.num_preluded_tokens + for item in scheduled_requests: + if item.token_start != frontier: + raise RuntimeError("wavefront prelude contains a prompt-range gap") + frontier += item.token_count + request.num_preluded_tokens = frontier + progress = PrefillProgress(False, False, transfer_start) + representative = scheduled_requests[0] + next_phase = ( + RequestPhase.PREFILLING_PRELUDE + if self._prelude_window_open(request) + else RequestPhase.PREFILLING_RECURRENT + ) + else: + chunk_complete = False + representative = scheduled_requests[0] + for item in scheduled_requests: + depth = item.depth + start = item.token_start + end = start + item.token_count + boundaries = request.prefill_depth_ranges + if boundaries[depth][0] != start: + raise RuntimeError("stale wavefront recurrent depth frontier") + boundaries[depth][0] = end + if depth + 1 < self.total_ut_steps: + boundaries[depth + 1][1] = end + continue + + request.num_prefilled_tokens = end + chunk_end = min( + len(request.prompt_token_ids), + item.chunk_start + self.config.prefill_chunk_size, + ) + if end == chunk_end: + if item.chunk_start != request.prefill_reclaimable_tokens: + raise RuntimeError("wavefront completed prompt chunks out of order") + request.prefill_reclaimable_tokens = chunk_end + chunk_complete = True + representative = item + + prompt_complete = request.num_prefilled_tokens == len(request.prompt_token_ids) + progress = PrefillProgress(chunk_complete, prompt_complete, transfer_start) + if chunk_complete: + self.cache_manager.publish_prefix( + request.request_id, + request.prompt_token_ids, + request.num_prefilled_tokens, + event, + ) + if prompt_complete: + request.loops_done = self.total_ut_steps + next_phase = self._prompt_destination(request) + else: + next_phase = ( + RequestPhase.PREFILLING_PRELUDE + if self._prelude_window_open(request) + else RequestPhase.PREFILLING_RECURRENT + ) + + self.enqueue(request, next_phase) + return representative, progress + + def _wavefront_prefill_groups(self, batch: SchedulerOutput): + groups = {} + if not self.wavefront_chunked_prefill: + return groups + for item in batch.requests: + phase = batch.phase_for(item) + request = item.request + if ( + phase in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + and self.requests.get(request.request_id) is request + and request.phase != RequestPhase.FINISHED + ): + groups.setdefault((request.request_id, phase), []).append(item) + return groups def _should_exit(self, request: Request) -> bool: """Evaluate the policy after recording the completed loop's actual hazard.""" @@ -840,6 +1212,7 @@ def update_from_output( if result is None: result = ModelRunnerOutput() update = SchedulerUpdate() + wavefront_groups = self._wavefront_prefill_groups(batch) for row, scheduled_request in enumerate(batch.requests): request = scheduled_request.request # Ignore results for a request that is no longer registered: it may @@ -851,6 +1224,22 @@ def update_from_output( ): continue phase = batch.phase_for(scheduled_request) + if self.wavefront_chunked_prefill and phase in ( + RequestPhase.PREFILLING_PRELUDE, + RequestPhase.PREFILLING_RECURRENT, + ): + grouped = wavefront_groups.pop((request.request_id, phase), None) + if grouped is not None: + event = ( + request_events.get(request.request_id) + if request_events is not None + else result.completion + ) + representative, progress = self._apply_wavefront_prefill_batch_progress( + phase, grouped, loop_count=self._batch_loop_count(batch), event=event + ) + update.prefill.append((representative, phase, progress)) + continue if phase in ( RequestPhase.PREFILL, RequestPhase.PREFILLING_PRELUDE, @@ -861,7 +1250,12 @@ def update_from_output( if request_events is not None else result.completion ) - progress = self._apply_prefill_progress(scheduled_request, phase, event=event) + progress = self._apply_prefill_progress( + scheduled_request, + phase, + loop_count=self._batch_loop_count(batch), + event=event, + ) update.prefill.append((scheduled_request, phase, progress)) elif phase == RequestPhase.PRELUDE: request.loops_done = 0 @@ -953,6 +1347,7 @@ def update_after_submission( ) -> SchedulerUpdate: """Advance logical submissions without requiring CPU token/score delivery.""" update = SchedulerUpdate() + wavefront_groups = self._wavefront_prefill_groups(batch) for row, scheduled_request in enumerate(batch.requests): request = scheduled_request.request if ( @@ -961,6 +1356,22 @@ def update_after_submission( ): continue phase = batch.phase_for(scheduled_request) + if self.wavefront_chunked_prefill and phase in ( + RequestPhase.PREFILLING_PRELUDE, + RequestPhase.PREFILLING_RECURRENT, + ): + grouped = wavefront_groups.pop((request.request_id, phase), None) + if grouped is not None: + event = ( + request_events.get(request.request_id) + if request_events is not None + else completion + ) + representative, progress = self._apply_wavefront_prefill_batch_progress( + phase, grouped, loop_count=self._batch_loop_count(batch), event=event + ) + update.prefill.append((representative, phase, progress)) + continue if phase in ( RequestPhase.PREFILL, RequestPhase.PREFILLING_PRELUDE, @@ -971,7 +1382,12 @@ def update_after_submission( if request_events is not None else completion ) - progress = self._apply_prefill_progress(scheduled_request, phase, event=event) + progress = self._apply_prefill_progress( + scheduled_request, + phase, + loop_count=self._batch_loop_count(batch), + event=event, + ) update.prefill.append((scheduled_request, phase, progress)) elif phase == RequestPhase.PRELUDE: request.loops_done = 0 diff --git a/vllm_rlt/core/scheduling_policy.py b/vllm_rlt/core/scheduling_policy.py index feba1174..39b405ec 100644 --- a/vllm_rlt/core/scheduling_policy.py +++ b/vllm_rlt/core/scheduling_policy.py @@ -39,12 +39,12 @@ def record_batch(self, phase: RequestPhase) -> None: class PromptOnlyPolicy(SchedulingPolicy): """Select prompt work only; resource ownership and execution stay shared.""" - def __init__(self, config, *, atomic_prefill): + def __init__(self, config, *, staged_prefill): super().__init__(config) self.PREFILL_PHASES = ( - (RequestPhase.PREFILL,) - if atomic_prefill - else (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + if staged_prefill + else (RequestPhase.PREFILL,) ) def schedule(self, scheduler, *, prefer_recurrent=False): @@ -128,12 +128,12 @@ def schedule( class SpeculativePolicy(SchedulingPolicy): """Schedule draft/verify rounds with independently chunked prompt work.""" - def __init__(self, config, *, atomic_prefill=True): + def __init__(self, config): super().__init__(config) self.PREFILL_PHASES = ( - (RequestPhase.PREFILL,) - if atomic_prefill - else (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + if requires_staged_prefill(config, produces_tokens=True) + else (RequestPhase.PREFILL,) ) def schedule(self, scheduler, *, prefer_recurrent=False): @@ -154,13 +154,13 @@ def schedule(self, scheduler, *, prefer_recurrent=False): class SplitPrefillPolicy(RefillPolicy): - """Refill ordering where each recurrence depth is its own batch boundary.""" + """Refill ordering for staged prompt work and recurrent depth boundaries.""" PREFILL_PHASES = (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) class MixedPolicy(SplitPrefillPolicy): - """Choose compatible prompt/generation loops; execution is policy-agnostic.""" + """Choose compatible prompt/generation work; execution is policy-agnostic.""" def schedule(self, scheduler, *, prefer_recurrent=False): if scheduler.has_ready(RequestPhase.RECURRENT) and scheduler.has_ready( @@ -169,13 +169,21 @@ def schedule(self, scheduler, *, prefer_recurrent=False): return scheduler.take_phases( (RequestPhase.RECURRENT, RequestPhase.PREFILLING_RECURRENT) ) - return super().schedule(scheduler, prefer_recurrent=prefer_recurrent) + if prefer_recurrent and scheduler.has_ready(RequestPhase.RECURRENT): + return scheduler._take(RequestPhase.RECURRENT) + if ( + scheduler.wavefront_chunked_prefill + and scheduler.has_ready(RequestPhase.PRELUDE) + and scheduler.has_ready(RequestPhase.PREFILLING_PRELUDE) + ): + return scheduler.take_phases((RequestPhase.PRELUDE, RequestPhase.PREFILLING_PRELUDE)) + batch = super().schedule(scheduler, prefer_recurrent=prefer_recurrent) + return batch def create_scheduling_policy( scheduler_config: SchedulerConfig, speculative_config: SpeculativeConfig | None, - atomic_prefill: bool = True, *, produces_tokens: bool = True, ) -> SchedulingPolicy: @@ -184,22 +192,31 @@ def create_scheduling_policy( Failing here keeps an unimplemented mode from silently changing semantics or returning an empty schedule forever. """ + staged_prefill = requires_staged_prefill(scheduler_config, produces_tokens=produces_tokens) + # Prompt-only phases do not use prompt/generation interleaving rules. if not produces_tokens: - return PromptOnlyPolicy(scheduler_config, atomic_prefill=atomic_prefill) + return PromptOnlyPolicy( + scheduler_config, + staged_prefill=requires_staged_prefill(scheduler_config, produces_tokens=False), + ) # The remaining compatibility checks apply to generation scheduling. if scheduler_config.mode == "mixed": - if atomic_prefill or speculative_config is not None: - raise ValueError("mixed scheduling requires non-atomic, non-speculative work") - elif not atomic_prefill and speculative_config is None and scheduler_config.mode != "refill": - raise ValueError("split prefill currently implements refill scheduling only") + if speculative_config is not None: + raise ValueError("mixed scheduling does not support speculative work") + elif staged_prefill and speculative_config is None and scheduler_config.mode != "refill": + raise ValueError("staged prefill currently implements refill or mixed scheduling") if scheduler_config.mode == "mixed": return MixedPolicy(scheduler_config) if speculative_config is not None: - return SpeculativePolicy(scheduler_config, atomic_prefill=atomic_prefill) - if atomic_prefill: - policy_cls = RefillPolicy if scheduler_config.mode == "refill" else NoRefillPolicy - return policy_cls(scheduler_config) - return SplitPrefillPolicy(scheduler_config) + return SpeculativePolicy(scheduler_config) + if staged_prefill: + return SplitPrefillPolicy(scheduler_config) + policy_cls = RefillPolicy if scheduler_config.mode == "refill" else NoRefillPolicy + return policy_cls(scheduler_config) + + +def requires_staged_prefill(config: SchedulerConfig, *, produces_tokens: bool = True) -> bool: + return config.wavefront_chunked_prefill or (produces_tokens and config.mode == "mixed") diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 2db2639e..40a6c9a8 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -6,6 +6,7 @@ from vllm_rlt.core.kv_cache_manager import KVCacheManager from vllm_rlt.core.memory import plan_cache from vllm_rlt.core.scheduler import Scheduler +from vllm_rlt.core.scheduling_policy import requires_staged_prefill from vllm_rlt.engine.connector import StageConnector from vllm_rlt.engine.preemption import PreemptionManager from vllm_rlt.engine.stage_config import StageConfig @@ -58,9 +59,14 @@ def __init__( raise ValueError("Huginn prefix caching is not yet supported") self.exit_config = exit_config or ExitConfig() self.execution_config = execution_config or ExecutionConfig() + staged_prefill = requires_staged_prefill( + scheduler_config, produces_tokens=self.stage_config.produces_tokens + ) if getattr(model, "requires_boundary_kv", False): - if not self.execution_config.atomic_prefill: - raise ValueError("Huginn requires atomic prefill") + if staged_prefill: + raise ValueError( + "Huginn does not support staged prefill required by this scheduler" + ) if self.execution_config.async_scheduling or self.execution_config.prefill_uva: raise ValueError("Huginn requires synchronous scheduling without prefill UVA") if speculative_config is not None: @@ -114,9 +120,8 @@ def __init__( parameter.device.type != "cuda" or cache_config.layout != "last_exited" ): raise ValueError("prefill_uva requires CUDA packed prefill with last_exited KV") - if not self.execution_config.atomic_prefill: - if cache_config.layout != "last_exited": - raise ValueError("split prefill requires last_exited KV planes") + if staged_prefill and cache_config.layout != "last_exited": + raise ValueError("staged prefill requires last_exited KV planes") attention = create_backend( attention_backend, parameter.device, @@ -138,6 +143,7 @@ def __init__( scheduler_config, self.execution_config, attention.info["selected_backend"], + staged_prefill=staged_prefill, ) self.cache_manager = KVCacheManager( num_layers=config.num_hidden_layers, @@ -160,7 +166,6 @@ def __init__( scheduler_config, self.cache_manager, speculative_config, - atomic_prefill=self.execution_config.atomic_prefill, produces_tokens=self.stage_config.produces_tokens, total_ut_steps=config.total_ut_steps, eos_token_id=config.eos_token_id, @@ -177,6 +182,7 @@ def __init__( exit_config=self.exit_config, execution_config=self.execution_config, scheduler_config=scheduler_config, + staged_prefill=self.scheduler.staged_prefill, rng=self.rng, ) self.speculative_runner = ( @@ -482,6 +488,8 @@ def _notify_connector_prefill(self, scheduled_request, phase, progress): start=( scheduled_request.token_start if phase == RequestPhase.PREFILL + else progress.transfer_start + if progress.transfer_start is not None else scheduled_request.chunk_start ), progress=progress, @@ -510,7 +518,8 @@ def reserve_transfer( count = len(prompt_token_ids) hit = plan.cached_tokens request.num_prefilled_tokens = request.num_preluded_tokens = hit - request.prefill_chunk_start = hit + request.prefill_window_origin = hit + request.prefill_reclaimable_tokens = hit request.prefill_depth_ranges[:] = Request.prefill_boundaries( count, self.model.config.total_ut_steps, start=hit ) @@ -577,6 +586,8 @@ def accept_import(self, request, hidden, *, cached_tokens=0): self.cache_manager.mark_imported_prefix(request.request_id, count, start=cached_tokens) self.cache_manager.publish_prefix(request.request_id, request.prompt_token_ids, count) request.num_prefilled_tokens = request.num_preluded_tokens = count + request.prefill_window_origin = count + request.prefill_reclaimable_tokens = count request.loops_done = self.model.config.total_ut_steps request.hidden_state = hidden # Keep RECEIVING until the scheduler grants a compute slot. The lease @@ -586,14 +597,15 @@ def accept_import(self, request, hidden, *, cached_tokens=0): def _rewind_split_prefill(self, request): """Rewind an in-flight chunk; the caller rebuilds its queue membership.""" - if self.execution_config.atomic_prefill or request.phase not in ( + if request.phase not in ( RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT, ): return frontier = request.num_prefilled_tokens request.num_preluded_tokens = frontier - request.prefill_chunk_start = frontier + request.prefill_window_origin = frontier + request.prefill_reclaimable_tokens = frontier request.prefill_depth_ranges[:] = Request.prefill_boundaries( len(request.prompt_token_ids), self.model.config.total_ut_steps, diff --git a/vllm_rlt/entrypoints/cli.py b/vllm_rlt/entrypoints/cli.py index 5d0d1e12..a9b8025c 100644 --- a/vllm_rlt/entrypoints/cli.py +++ b/vllm_rlt/entrypoints/cli.py @@ -125,6 +125,7 @@ def main(): max_num_batched_tokens=args.max_num_batched_tokens, mode=args.mode, prefill_chunk_size=getattr(args, "prefill_chunk_size", 128), + wavefront_chunked_prefill=getattr(args, "wavefront_chunked_prefill", False), max_prefill_batches_before_decode=getattr(args, "max_prefill_batches_before_decode", 1), admission_scan_limit=getattr(args, "admission_scan_limit", 64), max_admission_bypasses=getattr(args, "max_admission_bypasses", 8), diff --git a/vllm_rlt/entrypoints/pd_serve.py b/vllm_rlt/entrypoints/pd_serve.py index 5730c17b..7a63621b 100644 --- a/vllm_rlt/entrypoints/pd_serve.py +++ b/vllm_rlt/entrypoints/pd_serve.py @@ -61,6 +61,7 @@ def load_engine(args): max_num_seqs=args.prefill_max_num_seqs, max_num_batched_tokens=args.prefill_max_num_batched_tokens, prefill_chunk_size=args.prefill_chunk_size, + wavefront_chunked_prefill=getattr(args, "wavefront_chunked_prefill", False), ), decode_scheduler_config=SchedulerConfig( policy=getattr(args, "scheduling_policy", "fcfs"), diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index f369ed06..a166705b 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -34,12 +34,6 @@ def add_runtime_args(parser): action="store_false", help="Copy metadata to the GPU instead of reading mapped host memory", ) - parser.add_argument( - "--no-atomic-prefill", - dest="atomic_prefill", - action="store_false", - help="Split prompt prefill into a staged prelude plus one boundary per recurrence depth", - ) parser.add_argument("--kv-layout", choices=["last_exited", "shared"], default="last_exited") parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) parser.add_argument("--kv-cache-memory-bytes", type=int) @@ -56,7 +50,7 @@ def add_runtime_args(parser): parser.add_argument( "--cuda-graphs", action="store_true", - help="Capture decode recurrent cores and speculative coda; prefill/sampling stay eager", + help="Capture decode/wavefront recurrent cores and speculative coda; sampling stays eager", ) parser.add_argument("--cuda-graph-max-batch-size", type=int, default=128) parser.add_argument("--cuda-graph-max-graphs", type=int, default=16) @@ -69,6 +63,11 @@ def add_runtime_args(parser): parser.add_argument("--static-buffers", action="store_true") parser.add_argument("--pad-to-power-of-two", action="store_true") parser.add_argument("--prefill-chunk-size", type=int, default=128) + parser.add_argument( + "--wavefront-chunked-prefill", + action="store_true", + help="Batch ready prompt chunks across recurrence depths", + ) parser.add_argument("--max-prefill-batches-before-decode", type=int, default=1) parser.add_argument("--admission-scan-limit", type=int, default=64) parser.add_argument("--max-admission-bypasses", type=int, default=8) @@ -109,7 +108,6 @@ def runtime_configs(args): execution_config=ExecutionConfig( prefill_uva=args.prefill_uva, use_uva=args.use_uva, - atomic_prefill=args.atomic_prefill, cuda_graphs=args.cuda_graphs, cuda_graph_max_batch_size=args.cuda_graph_max_batch_size, cuda_graph_max_graphs=args.cuda_graph_max_graphs, diff --git a/vllm_rlt/entrypoints/serve.py b/vllm_rlt/entrypoints/serve.py index 0c011228..ff6b8f18 100644 --- a/vllm_rlt/entrypoints/serve.py +++ b/vllm_rlt/entrypoints/serve.py @@ -53,6 +53,7 @@ def load_engine(args): max_num_batched_tokens=args.max_num_batched_tokens, mode=args.mode, prefill_chunk_size=getattr(args, "prefill_chunk_size", 128), + wavefront_chunked_prefill=getattr(args, "wavefront_chunked_prefill", False), max_prefill_batches_before_decode=getattr(args, "max_prefill_batches_before_decode", 1), admission_scan_limit=getattr(args, "admission_scan_limit", 64), max_admission_bypasses=getattr(args, "max_admission_bypasses", 8), diff --git a/vllm_rlt/pd/connector.py b/vllm_rlt/pd/connector.py index c36865b0..121a6538 100644 --- a/vllm_rlt/pd/connector.py +++ b/vllm_rlt/pd/connector.py @@ -187,7 +187,10 @@ def prefill_submitted(self, request, *, phase, start, progress, event, stream): transfer = self.transfers[request.request_id] if transfer.request is not request: raise RuntimeError("stale transfer generation") - if phase == RequestPhase.PREFILLING_PRELUDE: + if phase == RequestPhase.PREFILLING_PRELUDE or ( + phase == RequestPhase.PREFILLING_RECURRENT + and self.engine.scheduler.config.wavefront_chunked_prefill + ): transfer.transfer_overlaps += self.transport.busy(transfer.tid) if not progress.chunk_complete: transfer.last_event = event diff --git a/vllm_rlt/request.py b/vllm_rlt/request.py index 537e2bf1..ef784437 100644 --- a/vllm_rlt/request.py +++ b/vllm_rlt/request.py @@ -47,14 +47,16 @@ class Request: num_speculative_rounds: int = 0 input_token_tensor: torch.Tensor | None = field(default=None, repr=False) num_prefilled_tokens: int = 0 - # Split prefill. [num_prefilled_tokens, num_preluded_tokens) is the staged - # chunk: its prelude output is resident and it advances one depth per batch. + # Split prefill: [num_prefilled_tokens, num_preluded_tokens) has resident + # hidden states; wavefront keeps multiple chunks at distinct depths here. num_preluded_tokens: int = 0 - # Start of the resident chunk. It stays fixed while the final depth advances - # a partial chunk and is used for runner pool row addressing. - prefill_chunk_start: int = 0 # Split prefill only: per-depth logical prompt frontiers. prefill_depth_ranges: list[list[int]] = field(default_factory=list, repr=False) + # Wavefront chunk slots are assigned relative to this cached-prefix origin. + prefill_window_origin: int = 0 + # End of the completed prefix and start of the oldest resident chunk. Advance + # only after final-depth completion releases that chunk's hidden-state slot. + prefill_reclaimable_tokens: int = 0 # Host-side decode progress; ModelRunner events track pending device work. loops_done: int = 0 pending_exit_depth: int | None = None @@ -118,8 +120,8 @@ class StateSpan: # for boundary-layer KV addressing when required. CODA consumes the final # recurrent hidden state; its actual exit depth remains in request.loops_done. depth: int - # None addresses the request's single resident hidden row. Otherwise the - # span addresses [token_start-chunk_start, ...] in its resident chunk. + # None lets the runner derive token_start - chunk_start. Otherwise this is a + # row offset into the request's prefill hidden-state buffer, including its ring slot. chunk_start: int | None = None retain_chunk: bool = False retain_last: bool = True @@ -127,6 +129,9 @@ class StateSpan: # For chunk PRELUDE, None reads the immutable request prompt range directly. # Explicit tokens override that source (also used by direct execution callers). input_tokens: tuple[int, ...] | None = None + # Row offset in the request's prefill hidden-state buffer. None lets the runner + # derive token_start - chunk_start. + prefill_hidden_row_offset: int | None = None @dataclass(frozen=True) diff --git a/vllm_rlt/worker/buffers.py b/vllm_rlt/worker/buffers.py index 7a2de728..2829b0d3 100644 --- a/vllm_rlt/worker/buffers.py +++ b/vllm_rlt/worker/buffers.py @@ -10,21 +10,24 @@ from vllm_rlt.worker.tensors import UvaBackedTensor, copy_tensor, use_uva_on -def prefill_stage_rows(config, scheduler, execution) -> int: +def prefill_stage_rows(config, scheduler, *, staged_prefill=None) -> int: """Hidden rows per request that staged prompt prefill keeps resident. - Split prefill stages one chunk at a time, so a chunk is the unit that must - survive between two scheduling boundaries. Wavefront chunked prefill advances - several chunks of one prompt at different depths concurrently; it then keeps a - chunk per recurrence depth in flight, and this becomes - ``scheduler.prefill_chunk_size * config.total_ut_steps``. + Mixed scheduling can stage one chunk per request when decode joins at a + depth boundary. Wavefront keeps at most one chunk per model depth in flight. """ - if execution.atomic_prefill: + if staged_prefill is None: + staged_prefill = scheduler.wavefront_chunked_prefill or scheduler.mode == "mixed" + if not staged_prefill: return 0 + if scheduler.wavefront_chunked_prefill: + return scheduler.prefill_chunk_size * config.total_ut_steps return scheduler.prefill_chunk_size -def execution_buffer_bytes(config, scheduler, cache, execution, element_size): +def execution_buffer_bytes( + config, scheduler, cache, execution, element_size, *, staged_prefill=None +): rows = scheduler.max_num_batched_tokens if execution.pad_to_power_of_two: rows = 1 << (rows - 1).bit_length() @@ -45,12 +48,14 @@ def execution_buffer_bytes(config, scheduler, cache, execution, element_size): (4 * min(rows, scheduler.max_num_seqs) + 2 * scheduler.max_num_seqs) * planes * 8 * 8 ) # Split prefill holds the staged chunk's hidden rows across batch boundaries. - total += scheduler.max_num_seqs * prefill_stage_rows(config, scheduler, execution) * hidden + stage_rows = prefill_stage_rows(config, scheduler, staged_prefill=staged_prefill) + total += scheduler.max_num_seqs * stage_rows * hidden # Split recurrent batches gather staged rows into a contiguous model input. - if not execution.atomic_prefill: + if stage_rows: total += rows * hidden + span_capacity = rows if scheduler.wavefront_chunked_prefill else scheduler.max_num_seqs total += MAX_CONCURRENT_SUBMISSIONS * ( - scheduler.max_num_seqs * (5 + 4 * 8) * 8 + max(rows, scheduler.prefill_chunk_size) * 8 + span_capacity * (5 + 4 * 8) * 8 + max(rows, scheduler.prefill_chunk_size) * 8 ) return total diff --git a/vllm_rlt/worker/executor.py b/vllm_rlt/worker/executor.py index 87ba27df..bad146d4 100644 --- a/vllm_rlt/worker/executor.py +++ b/vllm_rlt/worker/executor.py @@ -22,6 +22,7 @@ def build_execution_batch(scheduled_batch: SchedulerOutput, total_depths: int) - spans = [] operations = set() loop_counts = set() + planned_loops = scheduled_batch.loop_count for scheduled_request in scheduled_batch.requests: phase = scheduled_batch.phase_for(scheduled_request) request = scheduled_request.request @@ -49,6 +50,13 @@ def build_execution_batch(scheduled_batch: SchedulerOutput, total_depths: int) - if chunk or phase == RequestPhase.PREFILL or scheduled_request.depth is not None else request.position ) + execution_loops = ( + planned_loops + if planned_loops is not None + else total_depths + if phase == RequestPhase.PREFILL + else int(operation == Operation.RECURRENT) + ) spans.append( StateSpan( request=request, @@ -56,20 +64,21 @@ def build_execution_batch(scheduled_batch: SchedulerOutput, total_depths: int) - token_count=scheduled_request.token_count, depth=depth, chunk_start=scheduled_request.chunk_start if chunk else None, - retain_chunk=chunk and (operation == Operation.PRELUDE or depth + 1 < total_depths), + retain_chunk=chunk and depth + execution_loops < total_depths, retain_last=scheduled_request.completes_prompt if chunk else True, emit_gate=phase == RequestPhase.RECURRENT, input_tokens=tuple( request.prompt_token_ids[position : position + scheduled_request.token_count] ) - if prompt_input and not chunk + if prompt_input and (not chunk or execution_loops > 0) else None, + prefill_hidden_row_offset=( + scheduled_request.prefill_hidden_row_offset if chunk else None + ), ) ) operations.add(operation) - loop_counts.add( - total_depths if phase == RequestPhase.PREFILL else int(operation == Operation.RECURRENT) - ) + loop_counts.add(execution_loops) if len(operations) != 1 or len(loop_counts) != 1: raise ValueError("batch requests must share an operation and execution boundary") return ExecutionBatch(operations.pop(), tuple(spans), loop_counts.pop()) diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index 58e3414f..ada23d20 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -53,7 +53,7 @@ def reset(self): class InputSegment(NamedTuple): state_slot: int - chunk_offset: int + hidden_row_offset: int token_count: int batch_row_start: int keep_intermediate: bool @@ -158,6 +158,7 @@ def __init__( exit_config=None, execution_config=None, scheduler_config=None, + staged_prefill=None, rng: RngRegistry | None = None, ): self.model = model.eval() @@ -210,7 +211,6 @@ def __init__( torch.nn.init.zeros_(head.bias) self.lookahead_head = head.to(device=self.device, dtype=parameter.dtype).eval() self.lookahead_head.requires_grad_(False) - self.split_prefill = not self.execution_config.atomic_prefill self.core_stream = self.boundary_stream = self.copy_stream = self.readback_stream = None if self.device.type == "cuda": self.core_stream = torch.cuda.Stream(device=self.device) @@ -260,14 +260,15 @@ def __init__( self.states = None # Prompt hidden states staged between split-prefill scheduling boundaries. # Sized by the same helper the memory plan reserves with. - stage_rows = prefill_stage_rows(model.config, scheduler, self.execution_config) + stage_rows = prefill_stage_rows(model.config, scheduler, staged_prefill=staged_prefill) + self.staged_prefill = stage_rows > 0 self.prefill_states = ( torch.empty( (scheduler.max_num_seqs, stage_rows, model.config.hidden_size), device=self.device, dtype=parameter.dtype, ) - if self.split_prefill + if self.staged_prefill else None ) prefill_rows = scheduler.max_num_batched_tokens @@ -279,13 +280,15 @@ def __init__( device=self.device, dtype=parameter.dtype, ) - if self.split_prefill + if self.staged_prefill else None ) self.batch_input_buffers = ( [ BatchInputBuffers( - scheduler.max_num_seqs, + scheduler.max_num_batched_tokens + if scheduler.wavefront_chunked_prefill + else scheduler.max_num_seqs, max(scheduler.max_num_batched_tokens, scheduler.prefill_chunk_size), self.device, parameter.dtype, @@ -293,7 +296,7 @@ def __init__( ) for _ in range(MAX_CONCURRENT_SUBMISSIONS) ] - if self.split_prefill + if self.staged_prefill else [] ) # Immutable prompt rows are filled once per request/runner-slot lifetime. @@ -304,7 +307,7 @@ def __init__( torch.int64, self.device, ) - if self.split_prefill + if self.staged_prefill else None ) self.prompt_token_owners = {} @@ -395,18 +398,22 @@ def _prefill_segments(self, batch, buffers): max_rows = 0 for index, span in enumerate(batch.spans): state_slot = self._ensure_request_slot(span.request) - chunk_offset = span.token_start - span.chunk_start + hidden_row_offset = ( + span.prefill_hidden_row_offset + if span.prefill_hidden_row_offset is not None + else span.token_start - span.chunk_start + ) keep_intermediate = span.retain_chunk segment = InputSegment( state_slot, - chunk_offset, + hidden_row_offset, span.token_count, row_start, keep_intermediate, ) buffers.segments.np[index] = ( segment.state_slot, - segment.chunk_offset, + segment.hidden_row_offset, segment.token_count, segment.batch_row_start, int(segment.keep_intermediate), @@ -451,8 +458,14 @@ def _mixed_hidden_sources(self, batch, buffers): slot = self._ensure_request_slot(request) chunk = None if span.chunk_start is not None: - start = span.token_start - span.chunk_start - chunk = self.prefill_states[slot, start : start + span.token_count] + hidden_row_offset = ( + span.prefill_hidden_row_offset + if span.prefill_hidden_row_offset is not None + else span.token_start - span.chunk_start + ) + chunk = self.prefill_states[ + slot, hidden_row_offset : hidden_row_offset + span.token_count + ] source = chunk else: source = request.hidden_state @@ -530,7 +543,7 @@ def _gather_prefill(self, prepared): hidden[start:end].copy_( self.prefill_states[ segment.state_slot, - segment.chunk_offset : segment.chunk_offset + segment.token_count, + segment.hidden_row_offset : segment.hidden_row_offset + segment.token_count, ] ) return hidden @@ -540,6 +553,8 @@ def _scatter_prefill(self, hidden, prepared): if metadata.mixed: self._mixed_hidden(hidden, prepared, scatter=True) return + if not any(span.retain_chunk for span in prepared.batch.spans): + return if self.device.type == "cuda": from vllm_rlt.kernels.routing import scatter_segments_kernel @@ -568,7 +583,7 @@ def _scatter_prefill(self, hidden, prepared): continue self.prefill_states[ segment.state_slot, - segment.chunk_offset : segment.chunk_offset + segment.token_count, + segment.hidden_row_offset : segment.hidden_row_offset + segment.token_count, ].copy_( hidden[segment.batch_row_start : segment.batch_row_start + segment.token_count] )