From 4d9782b63bd5a108dba552137cc8bd5bc749f136 Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Thu, 22 Oct 2026 15:10:22 +0800 Subject: [PATCH 1/5] refactor(scheduler): separate request status from scheduled work introduce RequestStatus for request lifecycle and WorkKey for pending work while keeping Operation as the model execution boundary. Move scheduling to WorkKey queues with explicit admission slots, and carry work identity through execution, preemption, async speculaive and PD paths. Update regression tests, scheduler documentation, and benchmark work traces. Signed-off-by: LyxWxj --- benchmarks/cdb_runtime.py | 2 +- docs/scheduler.md | 123 +++++---- tests/test_async_pipeline.py | 29 +- tests/test_async_speculative.py | 16 +- tests/test_async_speculative_lifecycle.py | 8 +- tests/test_async_state.py | 4 +- tests/test_cdb_runtime.py | 30 +- tests/test_engine.py | 41 +-- tests/test_input_preparation.py | 14 +- tests/test_mixed_execution.py | 50 ++-- tests/test_pd.py | 8 +- tests/test_prefix_growth.py | 26 +- tests/test_scheduler_selection.py | 60 ++-- tests/test_serving.py | 14 +- tests/test_speculative.py | 35 ++- tests/test_split_prefill.py | 271 ++++++++++-------- tests/test_stage_engine.py | 36 +-- tests/test_transfer_connector.py | 54 ++-- vllm_rlt/core/execution.py | 111 ++++++++ vllm_rlt/core/kv_cache_manager.py | 2 +- vllm_rlt/core/scheduler.py | 317 +++++++++++++--------- vllm_rlt/core/scheduling_policy.py | 184 +++++++------ vllm_rlt/engine/connector.py | 5 +- vllm_rlt/engine/preemption.py | 11 +- vllm_rlt/engine/stage_engine.py | 152 +++++------ vllm_rlt/pd/connector.py | 13 +- vllm_rlt/pd/engine.py | 4 +- vllm_rlt/request.py | 108 +------- vllm_rlt/worker/async_speculative.py | 3 +- vllm_rlt/worker/executor.py | 59 ++-- vllm_rlt/worker/model_runner.py | 9 +- vllm_rlt/worker/speculative.py | 2 +- 32 files changed, 997 insertions(+), 804 deletions(-) create mode 100644 vllm_rlt/core/execution.py diff --git a/benchmarks/cdb_runtime.py b/benchmarks/cdb_runtime.py index 0b75e849..2250b1e4 100644 --- a/benchmarks/cdb_runtime.py +++ b/benchmarks/cdb_runtime.py @@ -28,7 +28,7 @@ def drive(engine): if batch is not None: steps.append( dict( - stage=batch.phase.value, + works=sorted(work.name.lower() for work in batch.works), effective_rows=batch.num_tokens, submitted_rows=engine.model_runner.last_submitted_size, ) diff --git a/docs/scheduler.md b/docs/scheduler.md index 156555b8..29f4ac00 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -23,14 +23,11 @@ active-request limit and the KV budget, including growth reserved for existing requests. Priority/FCFS and bounded bypass decide which waiting requests get a chance; they do not replace the resource checks. -If resources fit, the request enters the appropriate phase queue in B. A new -prompt starts prefill; a fully imported prompt enters CODA after admission; -restoration returns to the saved phase. No fit means waiting. Where preemption -is enabled, a safe victim can be suspended and later restored through this same -admission path. These decisions are inside A, not a separate post-processing -step. The request sent back to WAITING by preemption is the **victim**, not -necessarily the candidate whose admission was blocked. Restoration is retried -through resource checks and returns the victim to its saved phase. +If resources fit, the request becomes ACTIVE and enters the appropriate work +queue in B. A new prompt starts prompt work; a fully imported prompt enters +CODA after admission; restoration returns to the saved `WorkKey`. No fit means +waiting. Where preemption is enabled, a safe victim can be suspended and later +restored through this same admission path. These decisions are inside A, not a separate post-processing step. The request sent back to WAITING by preemption is the **victim**, not necessarily the candidate whose admission was blocked. Restoration is retried through resource checks and returns the victim to its saved work. The amber **PD import → A** line starts only after the connector has observed commit and all expected chunks. Its destination Request and prompt KV mapping @@ -38,19 +35,22 @@ already exist; receiving data alone does not grant an execution slot. ### B — Ready queues: what does each request need next? -B is a collection of phase queues, not a model computation. Each Request retains +B is a collection of work queues, not a model computation. Each Request retains its own position, prompt frontier and recurrent depth. The diagram groups queues by the work they represent: -- **Atomic prompt:** PRELUDE and all recurrent depths of a prompt chunk run in one - submission. Chunking and yielding between recurrent depths are separate choices. -- **Split prompt:** PRELUDE prepares a chunk, then RECURRENT work can be selected - at individual depths. Prompt tokens still traverse the full model depth. -- **Decode:** PRELUDE prepares the current token, RECURRENT advances its hidden - state, and CODA produces logits after recurrence exits. -- **Speculative round:** the local controller drafts candidate tokens and verifies - them. A round expands into model calls rather than becoming another model - operation alongside PRELUDE, RECURRENT and CODA. +- **Atomic prompt (`PROMPT_ATOMIC`):** PRELUDE and all recurrent depths of a + prompt chunk run in one submission. Chunking and yielding between recurrent + depths are separate choices. +- **Split prompt (`PROMPT_PRELUDE`, `PROMPT_RECURRENT`):** PRELUDE prepares a + chunk, then RECURRENT work can be selected at individual depths. Prompt tokens + still traverse the full model depth. +- **Decode (`DECODE_PRELUDE`, `DECODE_RECURRENT`, `DECODE_CODA`):** PRELUDE + prepares the current token, RECURRENT advances its hidden state, and CODA + produces logits after recurrence exits. +- **Speculative round (`SPECULATIVE`):** the local controller drafts candidate + tokens and verifies them. A round expands into model calls rather than becoming + another model operation alongside PRELUDE, RECURRENT and CODA. Queue membership is not proof that GPU or transport work has completed. C checks scheduling eligibility; execution events order the actual device consumers. @@ -70,7 +70,7 @@ peer conceptual choices: 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; +`SplitPrefillPolicy` uses refill ordering with split prompt work; `MixedPolicy` extends that behavior with recurrent co-batching. Prompt-only is a stage-role constraint. Speculation selects a local draft/verify @@ -119,8 +119,7 @@ more general than “embedding” and “LM head.” ### E — Track completion: what is still using device state or buffers? Pending submissions retain result/readback storage and completion evidence. -They are execution tickets, not additional Request phases. The diagram does not -prescribe a universal queue count shared by every controller. +They are execution tickets, not additional Request work items. The diagram does not prescribe a universal queue count shared by every controller. Events prevent buffer reuse while earlier GPU consumers still read them. A synchronous path collects host-dependent results when needed. An asynchronous @@ -135,7 +134,7 @@ GPU dependencies or authorize early resource reuse. ### F — Apply progress: what should each request do next? StageEngine applies progress and result-dependent decisions per Request, then -calls `Scheduler.enqueue()` for its next phase. The four labels show common +calls `Scheduler.enqueue_work()` for its next `WorkKey`. The four labels show common continuations: another loop, another prompt chunk, the next generated token, or CODA after recurrence exits. The blue return arrow leads back to B, not A: continuing requests keep their admitted resources. @@ -209,7 +208,7 @@ arcs on the right indicate separate paths, not joins. | 3. Ranking | Lower service priority first, then larger estimated footprint; pick one victim. | | 4. Suspend | Synchronize, rewind unfinished split prefill if applicable, snapshot, release device resources and retain RNG. | | 5. Requester | A retries once, then selects or requeues. B continues shared admission, including restore or fresh allocation/import handling, and defers if admission fails. | -| 6. Victim | Shared WAITING, slot check, allocation, tensor restoration, synchronization, snapshot deletion and saved-phase enqueue. Slot shortage returns to B; failed KV allocation retains the snapshot and waits. | +| 6. Victim | Shared WAITING, slot check, allocation, tensor restoration, synchronization, snapshot deletion and saved-work enqueue. Slot shortage returns to B; failed KV allocation retains the snapshot and waits. | | Red exits | Cancellation discards suspended state and aborts; snapshot/restore exceptions propagate without transactional rollback in this manager. | The detailed rules below describe configuration restrictions, scan/bypass @@ -220,22 +219,18 @@ preemption policy; it does not expand the separate connector import protocol. | Data | While suspended | On resumption | Current owner / implementation | | --- | --- | --- | --- | -| Request identity and logical progress: token history, position, counters | The original Request stays in the registry; its ID moves to shared WAITING. Its phase becomes WAITING; the resume phase is saved separately. | Reuse the same Request and enqueue the saved phase. | `Scheduler.requests`, `Scheduler.queues`, `PreemptionManager.snapshots` | +| Request identity and logical progress: token history, position, counters | The original Request stays in the registry; its `status` becomes WAITING; the resume work is saved separately. | Reuse the same Request, activate it and enqueue the saved `WorkKey`. | `Scheduler.requests`, `Scheduler.waiting`, `Scheduler.work_queues`, `PreemptionManager.snapshots` | | Sampling RNG | Retained by the runner's RNG manager; not serialized into the CPU snapshot. | Continue the same random stream. | `ModelRunner.suspend()` preserves RNG but frees the runner execution slot. | | KV tensors | Copy allocated KV pages to CPU, then release the request's GPU allocation. Shared/cached pages need not all become free. | Allocate fresh GPU blocks and copy the saved pages into them. | `KVCacheManager.snapshot()` / `restore()` | | KV logical metadata: capacity, page count, written state per plane | Save capacity and page count; deep-copy written-state metadata into `KVSnapshot`. | Restore written-state metadata onto the new allocation. | `KVSnapshot` and cache allocation | | Physical KV mapping and batch execution metadata: block tables, write locations, row indices | Do not preserve the old batch as a resumable execution object. The old allocation is released. | Allocation establishes physical block mappings; subsequent batch preparation derives execution metadata from current mappings and progress. | Cache allocation and `ModelRunner.prepare_on_cpu()` / `prepare_on_gpu()` | | Intermediate hidden state and input-token tensor | When present, copy to CPU and clear their references on the Request. | Copy back to the cache device and attach to the same Request. | `PreemptionManager` | -There is no dedicated queue containing snapshots. WAITING contains request IDs; -`snapshots[request_id]` contains `phase`, `kv`, `hidden` and `token`. Persistent +There is no dedicated queue containing snapshots. `Scheduler.waiting` contains request IDs; `snapshots[request_id]` contains `work`, `kv`, `hidden` and `token`. Persistent logical state must not be confused with physical-address metadata that becomes 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. +For split prefill, `_rewind_split_prefill()` first resets an unfinished chunk to `num_prefilled_tokens` and returns `WorkKey.PROMPT_PRELUDE`. The snapshot therefore records the rewound work, not necessarily the loop at which the request was interrupted. The unfinished chunk is replayed after restoration. #### Current policy: triggers and eligibility @@ -273,7 +268,7 @@ Memory-pressure selection uses this ranking even under FCFS; only priority admission requires the victim to be strictly lower priority than the requester. If no safe victim exists, or the one allocation retry still fails, batch -construction returns the requesting ID to its phase queue and considers other +construction returns the requesting ID to its work queue and considers other candidates within the scan/token/sequence limits. It does not force an unsafe eviction or keep evicting until that individual request fits. A later scheduling attempt can try again; other candidates in the same call can also trigger their @@ -283,11 +278,11 @@ own capacity checks. With no selected work, no batch is returned. | Transition | Ordered actions | | --- | --- | -| Suspend device state | Synchronize the runner's configured streams; rewind unfinished split prefill; snapshot KV, phase and optional hidden/token tensors to CPU. | +| Suspend device state | Synchronize the runner's configured streams; rewind unfinished split prefill; snapshot KV, pending `WorkKey` and optional hidden/token tensors to CPU. | | Relinquish resources | Suspend runner resources while retaining RNG; poll cache prefixes; free the victim's cache allocation; clear hidden/token references; store the snapshot. | -| Move the request | Remove the victim's old queue entries and enqueue its ID into shared WAITING. | +| Move the request | Remove the victim's pending work and enqueue its ID into shared WAITING. | | Reconsider admission | FCFS preserves waiting order; priority mode sorts by numeric priority. Admission is bounded by `admission_scan_limit` and checks active slots before attempting restoration. | -| Restore | If allocation succeeds, copy KV and optional hidden/token tensors to device, synchronize the current CUDA stream, delete the CPU snapshot and enqueue the saved phase. If allocation fails, retain the snapshot and defer. | +| Restore | If allocation succeeds, copy KV and optional hidden/token tensors to device, synchronize the current CUDA stream, delete the CPU snapshot and enqueue the saved `WorkKey`. If allocation fails, retain the snapshot and defer. | | Abort a suspended request | Discard its snapshot, release runner resources/RNG and abort it through the scheduler. | Deferred IDs are put back at the front in their existing order. Fresh admission @@ -427,8 +422,9 @@ forward chains. Panel B shows one mixed recurrent selection: | B | One decode token | 1 | 3 | | C | A prefill span | 4 | 0 | -`MixedPolicy` considers decode and prefill recurrent queues. -`Scheduler.take_phases()` checks eligibility, shared token/sequence budgets and KV +`MixedPolicy` considers decode and prompt recurrent queues. +`Scheduler.take_works((WorkKey.DECODE_RECURRENT, WorkKey.PROMPT_RECURRENT))` +checks eligibility, shared token/sequence budgets and KV capacity, retaining each selected request's token range and depth. The example assumes all three fit; a blocked or over-budget request need not be selected. @@ -554,10 +550,10 @@ The four numbered steps are actions with explicit owners: | Step | Owner and action | Handoff | | --- | --- | --- | -| 01 | Scheduling policy chooses ready, compatible phases | Prompt and decode recurrent queues in this example | +| 01 | Scheduling policy chooses ready, compatible work | Prompt and decode recurrent queues in this example | | 02 | Scheduler selects token ranges/depths under readiness and resource limits | `SchedulerOutput` with per-request snapshots | | 03 | Executor/runner prepares and executes the selected batch | Result storage and completion events | -| 04 | StageEngine applies progress and calls `Scheduler.enqueue()` | Updated requests in their next phase queues; return to 01 | +| 04 | StageEngine applies progress and calls `Scheduler.enqueue_work()` | Updated requests in their next work queues; return to 01 | Steps 01 and 02 are parts of one scheduling call, not separate runtime stages. Policy chooses which queues to consider; batch construction checks individual @@ -588,7 +584,7 @@ Suppose three requests have recurrent work ready: With a six-token budget and room for three requests, mixed scheduling selects B and C from the generation recurrent queue, then fills the remaining four rows with A's prompt span. These six token rows share one RECURRENT execution, even -though their request phases and depths differ. In code, recurrent `depth` is +though their `WorkKey` values and depths differ. In code, recurrent `depth` is zero-based: A, B and C carry 1, 2 and 0 respectively. The batch retains each request's identity, token range and depth. A prompt span @@ -603,22 +599,23 @@ 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 +`DECODE_RECURRENT` and `PROMPT_RECURRENT`, not arbitrary operations or speculative verification rows. ## Who owns each decision? | Component | Responsibility | Code | | --- | --- | --- | -| `Scheduler` | Request registry, phase queues, admission, capacity checks, batch selection and queue transitions | [scheduler.py](../vllm_rlt/core/scheduler.py) | -| `SchedulingPolicy` implementations | Choose when to admit and which ready phase or compatible phases to select; track fairness/cohort state | [scheduling_policy.py](../vllm_rlt/core/scheduling_policy.py) | +| `Scheduler` | Request registry, waiting/work queues, admission slots, capacity checks, batch selection and queue transitions | [scheduler.py](../vllm_rlt/core/scheduler.py) | +| `SchedulingPolicy` implementations | Choose when to admit and which ready `WorkKey` or compatible work groups to select; track fairness/cohort state | [scheduling_policy.py](../vllm_rlt/core/scheduling_policy.py) | | `StageEngine` | Drive scheduling/submission, apply logical progress and output decisions, coordinate retirement | [stage_engine.py](../vllm_rlt/engine/stage_engine.py) | | `Executor` and runners | Bind selected work, prepare tensors, submit model computation and expose results/events | [executor.py](../vllm_rlt/worker/executor.py) | | `KVCacheManager` | Allocate/grow pages, track validity and protect borrowed allocations | [kv_cache_manager.py](../vllm_rlt/core/kv_cache_manager.py) | | Connector | Report transport readiness and manage transfers; never choose loop progress or mutate scheduler queues | [connector.py](../vllm_rlt/engine/connector.py) | -A **stage** is a deployment/resource boundary. A **request phase** describes -what that request needs next. PRELUDE, RECURRENT and CODA are model operations. +A **stage** is a deployment/resource boundary. `RequestStatus` describes the +request lifecycle, while `WorkKey` describes what it needs next. PRELUDE, +RECURRENT and CODA are model operations. A **scheduling mode** chooses batching policy. These are separate concepts: `prefill_decode` is a stage role; `mixed` is a scheduling mode. @@ -631,11 +628,11 @@ StageEngine.step → Scheduler.schedule → policy.schedule → _admit, when the policy allows new or restored work - → _take(phase) or take_phases(phases) + → take_works(compatible_work_keys) → readiness check → range snapshot → KV capacity check → SchedulerOutput → Executor.submit - → StageEngine applies progress/results and enqueues the next phase + → StageEngine applies progress/results and enqueues the next WorkKey ``` This is a logical chain, not a requirement to block after every submission. @@ -643,10 +640,11 @@ The asynchronous path can submit dependent work before collecting an earlier result, subject to its event and exit-policy contracts. `ScheduledRequest` holds a reference to the existing local `Request` plus a -snapshot of the selected token range, depth, chunk start and phase. It does not -copy the request or construct GPU metadata. `SchedulerOutput.phase` is the -common phase for homogeneous work, or `None` when selected phases differ; -`phase_for()` retrieves the phase of an individual selection. +snapshot of the selected `WorkKey`, token range, depth and chunk start. It does +not copy the request or construct GPU metadata. `SchedulerOutput.works` reports +the set of work types selected in the batch, and each `ScheduledRequest.work` +preserves the immutable per-request work identity. `is_mixed` is derived from +that set; batch inspection should use `works` and each item's `work`. `Executor` binds ordinary selections into `ExecutionBatch` spans. All spans must share a model operation and execution boundary. Different recurrent depths are @@ -692,16 +690,16 @@ update cache bookkeeping; `_allocate()` performs the actual allocation. ## Constructing the batch -`take_phases()` visits the supplied phase queues in order under one token budget, +`take_works()` visits the supplied work queues in order under one token budget, request-count limit and selected-request protection set. For each candidate it: -1. Checks connector readiness and phase-specific eligibility. +1. Checks connector readiness and work-specific eligibility. 2. Snapshots a contiguous token range and recurrent depth. 3. Ensures KV capacity covers the exclusive token frontier. 4. Adds the request to the selected set and deducts its token count. -Blocked requests return to the queue tail. Each phase scans at most the queue -length observed when that phase begins, so a blocked request cannot cycle +Blocked requests return to the queue tail. Each work queue scans at most the queue +length observed when that queue begins, so a blocked request cannot cycle forever in one call. An empty selection returns `None` without changing fairness counters. Selecting a request does not advance its completed progress. @@ -711,8 +709,8 @@ waits. For split recurrence, a span is additionally limited by the ready depth frontier and the prompt tokens whose prelude is already resident. `execution_ready` is the connector's nonblocking readiness predicate. Both -`has_ready()` and batch construction consult it, so a blocked transfer does not -hide another runnable phase. This is backpressure: the connector can defer work +`has_ready_work()` and batch construction consult it, so a blocked transfer does +not hide another runnable work type. This is backpressure: the connector can defer work when it cannot safely accept more transfer state. It does not wait for transport completion inside the scheduler. @@ -723,10 +721,10 @@ 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, using atomic or split prompt work | | `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 | +| `SplitPrefillPolicy` | Refill ordering with separate `PROMPT_PRELUDE` and `PROMPT_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 | | `SpeculativePolicy` | Select CODA bootstrap, prompt work and local draft/verify rounds, with a prefill fairness bound | @@ -771,8 +769,8 @@ transitions; receiving data does not mean decode has started. ## Progress and resource safety -`StageEngine` applies per-request progress, then calls `Scheduler.enqueue()` for -the next phase. In the async path, host progress can describe submitted work; +`StageEngine` applies per-request progress, then calls `Scheduler.enqueue_work()` +for the next `WorkKey`. In the async path, host progress can describe submitted work; CUDA events establish when the corresponding device state is ready. Queue membership alone is not completion evidence. @@ -784,7 +782,8 @@ requests; memory-pressure selection has its own priority/footprint ranking. The preemption manager copies/restores device state. Its resume callback returns `None` for no snapshot, `False` for insufficient capacity, or the saved -`RequestPhase` after restoration. Scheduler enqueues that phase. Preemption may +`WorkKey` after restoration. Scheduler activates the request and enqueues that +work. Preemption may synchronize and copy state; ordinary scheduling does not add a global barrier. Finishing removes queue entries, records the finish reason and requests KV @@ -795,7 +794,7 @@ their own lifetimes beyond one CPU scheduling call. ## Code and regression entry points Start with `Scheduler.schedule()` and `create_scheduling_policy()`, then follow -`take_phases()` and `_make_scheduled_request()`. Read `_admit()` and +`take_works()` and `_make_scheduled_work()`. Read `_admit()` and `_plan_admission()` for capacity behavior. Progress application lives in `StageEngine._update()` and `_apply_split_prefill_progress()`, with separate asynchronous result handling in the same engine. diff --git a/tests/test_async_pipeline.py b/tests/test_async_pipeline.py index 76c5dc7c..eeb312a1 100644 --- a/tests/test_async_pipeline.py +++ b/tests/test_async_pipeline.py @@ -5,9 +5,9 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import Operation, WorkKey from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import Operation, RequestPhase from vllm_rlt.worker.model_runner import Submission @@ -37,7 +37,8 @@ def drain(engine): for _ in range(300): if not engine.has_unfinished_requests(): assert engine.cache_manager.num_used_blocks == 0 - assert all(not q for q in engine.scheduler.queues.values()) + assert not engine.scheduler.waiting + assert all(not q for q in engine.scheduler.work_queues.values()) return results results.extend(engine.step()) pytest.fail("async pipeline failed to drain") @@ -71,7 +72,7 @@ def test_delayed_delivery_uses_snapshot_depths_and_bounds_placeholders( def ready(ticket): # Force delivery to lag until the one-output speculation bound drains it. - return False if ticket.batch.phase == RequestPhase.CODA else original_ready(ticket) + return False if WorkKey.DECODE_CODA in ticket.batch.works else original_ready(ticket) monkeypatch.setattr(Submission, "ready", ready) actual = drain(engine) @@ -107,11 +108,11 @@ def prelude(tokens): monkeypatch.setattr(engine.model, "prelude", prelude) engine.step() - assert engine.last_schedule.phase == RequestPhase.PRELUDE + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_PRELUDE,)) assert not held.complete and not request.generated_token_ids assert request.loops_done == 0 engine.step() - assert engine.last_schedule.phase == RequestPhase.RECURRENT + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_RECURRENT,)) assert not held.complete and not request.generated_token_ids engine.model_runner.synchronize() assert torch.equal(observed[0], ticket.device_values) @@ -180,7 +181,7 @@ def execute_with_delay(batch, prepared=None): injected and not delayed.query() and multi_stream - and batch.phase == RequestPhase.RECURRENT + and WorkKey.DECODE_RECURRENT in batch.works ): # H2D completes on the copy stream while the preceding core is # still busy. It must not depend on core or boundary completion. @@ -188,7 +189,7 @@ def execute_with_delay(batch, prepared=None): runner.copy_stream.synchronize() assert not delayed.query() result = execute(batch, prepared) - if batch.phase == RequestPhase.RECURRENT and not injected: + if WorkKey.DECODE_RECURRENT in batch.works and not injected: # Extend a forward to make the ordering test deterministic. This is # a topology test, not performance evidence. torch.cuda._sleep(150_000_000) @@ -233,7 +234,7 @@ def test_bounded_pipeline_survives_slow_gpu_and_dma_bank_reuse(monkeypatch, mult def execute(batch, prepared=None): nonlocal maximum maximum = max(maximum, len(runner.submissions.events)) - if batch.phase == RequestPhase.RECURRENT: + if WorkKey.DECODE_RECURRENT in batch.works: torch.cuda._sleep(10_000_000) return original(batch, prepared) @@ -258,7 +259,7 @@ def test_abort_speculative_core_then_reuse_id(monkeypatch, multi_stream): ticket.event = HeldEvent() engine.step() engine.step() - assert engine.last_schedule.phase == RequestPhase.RECURRENT + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_RECURRENT,)) output = engine.abort_request("reuse") assert output.finish_reason == "abort" and output.token_ids == [] assert engine.cache_manager.num_used_blocks == 0 @@ -281,7 +282,7 @@ def test_completed_core_refills_even_when_coda_cpu_delivery_is_pending(monkeypat def submit(batch): ticket = original(batch) if ( - batch.phase == RequestPhase.CODA + WorkKey.DECODE_CODA in batch.works and len(batch.requests) == 1 and batch.requests[0].request.request_id == "fast" ): @@ -297,10 +298,10 @@ def submit(batch): assert target engine.model_runner.synchronize() # GPU is done; ONLY CPU delivery lags. engine.step() - assert engine.last_schedule.phase == RequestPhase.PRELUDE + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_PRELUDE,)) assert not held.complete engine.step() - assert engine.last_schedule.phase == RequestPhase.RECURRENT + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_RECURRENT,)) assert {i.request.request_id for i in engine.last_schedule.requests} == {"fast", "slow"} assert not held.complete target[0][0].event = target[0][1] @@ -324,7 +325,7 @@ def test_device_progress_does_not_wait_for_readback_stream(): compute_done.synchronize() assert not ticket.ready() engine.step() # Prelude consumes the device token, not the host result. - assert engine.last_schedule.phase == RequestPhase.PRELUDE + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_PRELUDE,)) runner.events["a"].synchronize() assert not ticket.ready() assert not engine.scheduler.requests["a"].generated_token_ids @@ -342,7 +343,7 @@ def test_fixed_depth_async_does_not_read_back_unused_gates(monkeypatch, device): def record(batch): ticket = submit(batch) - if batch.phase == RequestPhase.RECURRENT: + if WorkKey.DECODE_RECURRENT in batch.works: recurrent.append(ticket) assert not ticket.has_result assert ticket.slot is None diff --git a/tests/test_async_speculative.py b/tests/test_async_speculative.py index 5fd08117..0d1d529b 100644 --- a/tests/test_async_speculative.py +++ b/tests/test_async_speculative.py @@ -14,9 +14,10 @@ SchedulerConfig, SpeculativeConfig, ) +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.scheduler import Scheduler from vllm_rlt.engine.llm_engine import LLMEngine -from vllm_rlt.request import Request, RequestPhase +from vllm_rlt.request import Request def drain(engine): @@ -33,12 +34,12 @@ def test_scheduler_uses_upper_bound_and_releases_rejected_budget(): scheduler = Scheduler(SchedulerConfig(), None, SpeculativeConfig(3)) r = Request("r", [2, 3], SamplingParams(max_tokens=8), generated_token_ids=[4]) r.num_output_placeholders = 4 - scheduled_request = scheduler._make_scheduled_request(r, RequestPhase.SPECULATIVE, 16) + scheduled_request = scheduler._make_scheduled_work(r, WorkKey.SPECULATIVE, 16) assert (scheduled_request.token_start, scheduled_request.token_count) == (6, 3) # One committed token instead of the four reserved: recover three slots. r.num_output_placeholders = 0 r.generated_token_ids.append(5) - scheduled_request = scheduler._make_scheduled_request(r, RequestPhase.SPECULATIVE, 16) + scheduled_request = scheduler._make_scheduled_work(r, WorkKey.SPECULATIVE, 16) assert (scheduled_request.token_start, scheduled_request.token_count) == (3, 4) @@ -295,10 +296,9 @@ def test_active_page_tables_are_contiguous_and_do_not_touch_inactive_storage(): runner.buffers = [buffer] requests = [Request(rid, [2], SamplingParams()) for rid in ("short", "long")] batch = SchedulerOutput( - RequestPhase.SPECULATIVE, [ - ScheduledRequest(requests[0], 1, 2), - ScheduledRequest(requests[1], 7, 2), + ScheduledRequest(requests[0], WorkKey.SPECULATIVE, 1, 2), + ScheduledRequest(requests[1], WorkKey.SPECULATIVE, 7, 2), ], ) _, _, host, device = runner._prepare_round(batch) @@ -318,7 +318,9 @@ def test_active_page_tables_are_contiguous_and_do_not_touch_inactive_storage(): # Simulate retirement, then shrink both dimensions using the same storage. buffer.leased = False before = buffer.host_tables.clone() - smaller = SchedulerOutput(RequestPhase.SPECULATIVE, [ScheduledRequest(requests[0], 0, 1)]) + smaller = SchedulerOutput( + [ScheduledRequest(requests[0], WorkKey.SPECULATIVE, 0, 1)] + ) _, _, host, device = runner._prepare_round(smaller) assert host.shape == device.shape == (1, 4, 1) assert host.is_contiguous() and device.is_contiguous() diff --git a/tests/test_async_speculative_lifecycle.py b/tests/test_async_speculative_lifecycle.py index f0a58b58..523fe58d 100644 --- a/tests/test_async_speculative_lifecycle.py +++ b/tests/test_async_speculative_lifecycle.py @@ -13,10 +13,10 @@ tiny_ouro_model, ) from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig, SpeculativeConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import RequestPhase from vllm_rlt.serving.worker import EngineWorker from vllm_rlt.worker.async_speculative import ( AsyncSpeculativeRunner, @@ -94,7 +94,9 @@ def prepared_engine(*, bad_second=False, first_frontier=4, retirement_ready=True ) tickets = [] for round_id, (rows, frontier) in enumerate([(3, first_frontier), (2, 6)]): - scheduled_request = ScheduledRequest(request, token_start=1 + round_id, token_count=rows) + scheduled_request = ScheduledRequest( + request, WorkKey.SPECULATIVE, token_start=1 + round_id, token_count=rows + ) host = torch.tensor( list(range(10 + 3 * round_id, 10 + 3 * round_id + rows)) + [rows - 1, frontier] ) @@ -108,7 +110,7 @@ def prepared_engine(*, bad_second=False, first_frontier=4, retirement_ready=True ) tickets.append( SpeculativeRoundTicket( - SchedulerOutput(RequestPhase.SPECULATIVE, [scheduled_request]), + SchedulerOutput([scheduled_request]), buffer, (state,), (round_id,), diff --git a/tests/test_async_state.py b/tests/test_async_state.py index 83ba405c..79e060fd 100644 --- a/tests/test_async_state.py +++ b/tests/test_async_state.py @@ -5,9 +5,9 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import RequestPhase def engine(*, asynchronous=True, static=False, graph=False, shared=False, use_uva=True): @@ -125,7 +125,7 @@ def test_execution_error_retires_routing_readers_before_request_reuse(monkeypatc def fail(batch, prepared=None): result = execute(batch, prepared) - if batch.phase == RequestPhase.RECURRENT: + if WorkKey.DECODE_RECURRENT in batch.works: torch.cuda._sleep(10_000_000) raise RuntimeError("injected failure after GPU submission") return result diff --git a/tests/test_cdb_runtime.py b/tests/test_cdb_runtime.py index 8b345dcb..bdd363cc 100644 --- a/tests/test_cdb_runtime.py +++ b/tests/test_cdb_runtime.py @@ -8,11 +8,12 @@ from tests.helpers import LocalLLM as LLM from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.kv_cache_manager import KVCacheManager from vllm_rlt.core.memory import budget_blocks from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import RequestPhase +from vllm_rlt.request import RequestStatus from vllm_rlt.worker.model_runner import Submission @@ -191,14 +192,17 @@ def test_short_requests_bypass_then_protected_long_request_gets_memory(): engine.add_request("short", [2], params) # 4 blocks engine.add_request("later", [3], params) engine.scheduler._admit() - assert engine.scheduler.requests["short"].phase == RequestPhase.PREFILL - assert list(engine.scheduler.queues[RequestPhase.WAITING]) == ["long", "later"] + assert ( + engine.scheduler.pending_work(engine.scheduler.requests["short"]) + == WorkKey.PROMPT_ATOMIC + ) + assert list(engine.scheduler.waiting) == ["long", "later"] engine.abort_request("short") engine.scheduler._admit() - assert engine.scheduler.requests["later"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["later"].status == RequestStatus.WAITING engine.abort_request("running") engine.scheduler._admit() - assert engine.scheduler.requests["long"].phase == RequestPhase.PREFILL + assert engine.scheduler.pending_work(engine.scheduler.requests["long"]) == WorkKey.PROMPT_ATOMIC assert set(drain(engine)) == {"long", "later"} @@ -228,12 +232,12 @@ def test_async_submits_next_core_before_collecting_previous_signal(monkeypatch): events = [] def submit(batch): - if batch.phase == RequestPhase.RECURRENT: + if WorkKey.DECODE_RECURRENT in batch.works: events.append("submit") return original_submit(batch) def collect(ticket): - if ticket.batch.phase == RequestPhase.RECURRENT: + if WorkKey.DECODE_RECURRENT in ticket.batch.works: events.append("collect") return original_collect(ticket) @@ -375,7 +379,7 @@ def synchronize(self): def submit(batch): ticket = original(batch) if ( - batch.phase == RequestPhase.CODA + WorkKey.DECODE_CODA in batch.works and len(batch.requests) == 1 and batch.requests[0].request.request_id == "fast" ): @@ -392,11 +396,11 @@ def submit(batch): engine.step() # The pending sample can already re-enter prelude on device. Its next # recurrent batch includes BOTH requests, without waiting for CPU delivery. - assert engine.last_schedule.phase == RequestPhase.PRELUDE + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_PRELUDE,)) assert engine.last_schedule.requests[0].request.request_id == "fast" assert not held.complete engine.step() - assert engine.last_schedule.phase == RequestPhase.RECURRENT + assert engine.last_schedule.works == frozenset((WorkKey.DECODE_RECURRENT,)) assert {i.request.request_id for i in engine.last_schedule.requests} == {"fast", "slow"} assert not held.complete held.complete = True @@ -495,11 +499,11 @@ def test_cuda_boundary_and_core_can_overlap(monkeypatch): def execute(batch, prepared=None): label = None - if batch.phase == RequestPhase.CODA and len(batch.requests) == 1: + if WorkKey.DECODE_CODA in batch.works and len(batch.requests) == 1: request = batch.requests[0].request if request.request_id == "fast" and len(request.generated_token_ids) == 1: label = "coda" - if batch.phase == RequestPhase.RECURRENT: + if WorkKey.DECODE_RECURRENT in batch.works: if any( i.request.request_id == "slow" and i.request.loops_done == 2 @@ -513,7 +517,7 @@ def execute(batch, prepared=None): # Preparation is outside _execute now. Keep both measured operations # alive long enough to test stream independence, regardless of which # one the host submits first; a tiny coda can finish before core launch. - if batch.phase == RequestPhase.RECURRENT or label == "coda": + if WorkKey.DECODE_RECURRENT in batch.works or label == "coda": torch.cuda._sleep(30_000_000) result = original(batch, prepared) end.record() diff --git a/tests/test_engine.py b/tests/test_engine.py index 1ad00979..45aa1511 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -5,9 +5,10 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import LLM, CacheConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase +from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestStatus def tiny_model(): @@ -22,7 +23,7 @@ def drain(engine, limit=1000): return outputs, trace emitted = engine.step() batch = engine.last_schedule - trace.append((batch.phase, tuple(i.request.request_id for i in batch.requests))) + trace.append((next(iter(batch.works)), tuple(i.request.request_id for i in batch.requests))) for output in emitted: if output.finished: outputs[output.request_id] = output @@ -80,7 +81,7 @@ def test_refill_recycles_exited_token_before_slow_token_coda(mode): engine.add_request("fast", [1], SamplingParams(max_tokens=3, exit_threshold=0, ignore_eos=True)) engine.add_request("slow", [2], SamplingParams(max_tokens=3, exit_threshold=1, ignore_eos=True)) outputs, trace = drain(engine) - recurrent = [ids for stage, ids in trace if stage == RequestPhase.RECURRENT] + recurrent = [ids for stage, ids in trace if stage == WorkKey.DECODE_RECURRENT] # In refill mode fast returns with its next token while slow is on its first. assert recurrent[:4] == ( [("fast", "slow")] * 2 + [("slow", "fast")] * 2 @@ -96,7 +97,7 @@ def test_memory_pressure_waits_then_reuses_pages_and_abort_reclaims(): for name in ["a", "b", "c"]: engine.add_request(name, [2, 3], params) engine.step() - assert engine.scheduler.requests["b"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["b"].status == RequestStatus.WAITING assert engine.abort_request("b").finish_reason == "abort" assert engine.cache_manager.num_used_blocks == 8 assert engine.abort_request("a").finished @@ -141,7 +142,7 @@ def test_eos_finishes_without_extra_decode_and_invalid_requests_do_not_poison_qu outputs, trace = drain(engine) assert outputs["eos"].token_ids == [0] assert outputs["eos"].finish_reason == "stop" - assert [stage for stage, _ in trace] == [RequestPhase.PREFILL, RequestPhase.CODA] + assert [stage for stage, _ in trace] == [WorkKey.PROMPT_ATOMIC, WorkKey.DECODE_CODA] assert engine.cache_manager.num_used_blocks == 0 @@ -241,7 +242,7 @@ def forbidden(*args, **kwargs): def test_coda_sampling_is_prepared_before_execution(monkeypatch): - from vllm_rlt.request import ExecutionBatch, Operation, StateSpan + from vllm_rlt.core.execution import ExecutionBatch, Operation, StateSpan engine = LLMEngine(tiny_model(), cache_config=CacheConfig(64, 2)) runner = engine.model_runner @@ -340,7 +341,7 @@ def test_repeated_short_arrivals_cannot_starve_existing_decode(mode): for _ in range(100): engine.step() batch = engine.last_schedule - if batch.phase == RequestPhase.RECURRENT and any( + if WorkKey.DECODE_RECURRENT in batch.works and any( scheduled_request.request.request_id == "long" for scheduled_request in batch.requests ): @@ -361,16 +362,16 @@ def test_failed_prefill_batch_does_not_consume_decode_fairness_budget(monkeypatc engine.add_request(rid, [1, 2], SamplingParams(max_tokens=2, ignore_eos=True)) scheduler = engine.scheduler scheduler._admit() - scheduler.queues[RequestPhase.PREFILL].remove("decode") - scheduler.enqueue(scheduler.requests["decode"], RequestPhase.RECURRENT) + scheduler.work_queues[WorkKey.PROMPT_ATOMIC].remove("decode") + scheduler.enqueue_work(scheduler.requests["decode"], WorkKey.DECODE_RECURRENT) ensure_capacity = engine.cache_manager.ensure_capacity monkeypatch.setattr( engine.cache_manager, "ensure_capacity", lambda rid, frontier: rid != "prefill" ) assert scheduler.schedule() is None monkeypatch.setattr(engine.cache_manager, "ensure_capacity", ensure_capacity) - assert scheduler.schedule().phase == RequestPhase.PREFILL - assert scheduler.schedule().phase == RequestPhase.RECURRENT + assert scheduler.schedule().works == frozenset((WorkKey.PROMPT_ATOMIC,)) + assert scheduler.schedule().works == frozenset((WorkKey.DECODE_RECURRENT,)) def test_failed_recurrent_batch_keeps_decode_due(monkeypatch): @@ -379,15 +380,15 @@ def test_failed_recurrent_batch_keeps_decode_due(monkeypatch): engine.add_request(rid, [1, 2], SamplingParams(max_tokens=2, ignore_eos=True)) scheduler = engine.scheduler scheduler._admit() - scheduler.queues[RequestPhase.PREFILL].remove("decode") - scheduler.enqueue(scheduler.requests["decode"], RequestPhase.RECURRENT) - assert scheduler.schedule().phase == RequestPhase.PREFILL + scheduler.work_queues[WorkKey.PROMPT_ATOMIC].remove("decode") + scheduler.enqueue_work(scheduler.requests["decode"], WorkKey.DECODE_RECURRENT) + assert scheduler.schedule().works == frozenset((WorkKey.PROMPT_ATOMIC,)) ensure_capacity = engine.cache_manager.ensure_capacity monkeypatch.setattr(engine.cache_manager, "ensure_capacity", lambda rid, frontier: False) assert scheduler.schedule() is None - scheduler.enqueue(scheduler.requests["prefill"], RequestPhase.PREFILL) + scheduler.enqueue_work(scheduler.requests["prefill"], WorkKey.PROMPT_ATOMIC) monkeypatch.setattr(engine.cache_manager, "ensure_capacity", ensure_capacity) - assert scheduler.schedule().phase == RequestPhase.RECURRENT + assert scheduler.schedule().works == frozenset((WorkKey.DECODE_RECURRENT,)) def test_abort_removes_all_queued_work_and_serializes_reason(): @@ -396,12 +397,14 @@ def test_abort_removes_all_queued_work_and_serializes_reason(): engine.scheduler._admit() request = engine.scheduler.requests["r"] # A delayed termination must remove every queued occurrence. - engine.scheduler.queues[RequestPhase.PREFILL].append("r") - engine.scheduler.enqueue(request, RequestPhase.RECURRENT) + engine.scheduler.work_queues[WorkKey.PROMPT_ATOMIC].append("r") + with pytest.raises(ValueError, match="already has pending work"): + engine.scheduler.enqueue_work(request, WorkKey.DECODE_RECURRENT) output = engine.abort_request("r") assert request.finish_reason is FinishReason.ABORT assert type(output.finish_reason) is str and output.finish_reason == "abort" - assert all("r" not in queue for queue in engine.scheduler.queues.values()) + assert "r" not in engine.scheduler.waiting + assert all("r" not in queue for queue in engine.scheduler.work_queues.values()) assert engine.cache_manager.num_used_blocks == 0 assert not engine.has_unfinished_requests() diff --git a/tests/test_input_preparation.py b/tests/test_input_preparation.py index 361e559f..3c1aa000 100644 --- a/tests/test_input_preparation.py +++ b/tests/test_input_preparation.py @@ -5,10 +5,11 @@ 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.execution import ExecutionBatch, Operation, StateSpan from vllm_rlt.core.kv_cache_manager import KVCacheManager from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import ExecutionBatch, Operation, Request, StateSpan +from vllm_rlt.request import Request from vllm_rlt.worker.buffers import PackedMetadataBuffer, Workspace, stage_metadata from vllm_rlt.worker.metadata import KVMetadata @@ -188,7 +189,8 @@ def checked_execute(scheduled, prepared): @pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) @pytest.mark.parametrize("phase", ["prepare_on_cpu", "prepare_on_gpu", "_execute"]) def test_loop_input_failure_releases_ownership_before_pool_wrap(monkeypatch, device, phase): - from vllm_rlt.request import ExecutionBatch, Operation, Request, StateSpan + from vllm_rlt.core.execution import ExecutionBatch, Operation, StateSpan + from vllm_rlt.request import Request model = OuroForCausalLM(tiny_ouro_config()).to(device) engine = LLMEngine( @@ -335,7 +337,7 @@ def test_packed_metadata_supports_query_lengths_and_per_request_depths(use_uva, @pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) @pytest.mark.parametrize("operation", [Operation.PRELUDE, Operation.RECURRENT]) def test_cpu_preparation_defers_transfers_and_hidden_gather(monkeypatch, device, operation): - from vllm_rlt.request import ExecutionBatch, StateSpan + from vllm_rlt.core.execution import ExecutionBatch, StateSpan from vllm_rlt.worker.tensors import TensorTransfer, UvaBuffer e = LLMEngine( @@ -381,7 +383,7 @@ def unexpected(*args, **kwargs): def test_cpu_preparation_failure_returns_acquired_buffer(monkeypatch): - from vllm_rlt.request import ExecutionBatch, StateSpan + from vllm_rlt.core.execution import ExecutionBatch, StateSpan e = make_prefill_engine() runner = e.model_runner @@ -409,7 +411,7 @@ def fail(*args, **kwargs): @pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) def test_prelude_gpu_preparation_uses_cpu_snapshot_and_preserves_new_token_reference(device): - from vllm_rlt.request import ExecutionBatch, StateSpan + from vllm_rlt.core.execution import ExecutionBatch, StateSpan e = LLMEngine( tiny_ouro_model().to(device), @@ -448,7 +450,7 @@ def test_prelude_gpu_preparation_uses_cpu_snapshot_and_preserves_new_token_refer def test_prelude_workspace_owned_once_across_gpu_prepare_and_forward( monkeypatch, device, fail_preparation ): - from vllm_rlt.request import ExecutionBatch, StateSpan + from vllm_rlt.core.execution import ExecutionBatch, StateSpan e = LLMEngine( tiny_ouro_model().to(device), diff --git a/tests/test_mixed_execution.py b/tests/test_mixed_execution.py index 12305176..61f76313 100644 --- a/tests/test_mixed_execution.py +++ b/tests/test_mixed_execution.py @@ -5,10 +5,11 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import ExecutionBatch, Operation, StateSpan, WorkKey from vllm_rlt.core.kv_cache_manager import KVCacheManager from vllm_rlt.core.scheduler import Scheduler, SchedulerOutput from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import ExecutionBatch, Operation, Request, RequestPhase, StateSpan +from vllm_rlt.request import Request, RequestStatus from vllm_rlt.worker.executor import ForwardBatch, ModelRuntime, build_execution_batch @@ -102,8 +103,8 @@ def test_scheduled_recurrence_does_not_follow_mutated_request_progress(): request = Request("r", [1, 2, 3], SamplingParams()) request.loops_done = 2 scheduler = Scheduler(SchedulerConfig(), None) - scheduled_request = scheduler._make_scheduled_request(request, RequestPhase.RECURRENT, 1) - batch = SchedulerOutput(RequestPhase.RECURRENT, [scheduled_request]) + scheduled_request = scheduler._make_scheduled_work(request, WorkKey.DECODE_RECURRENT, 1) + batch = SchedulerOutput([scheduled_request]) request.loops_done = 3 request.generated_token_ids.append(4) assert build_execution_batch(batch, 4).addresses() == (["r"], [2], [2]) @@ -182,12 +183,18 @@ def run(mode, atomic=False): outputs = {} mixed = [] calls = [] + prelude = engine.model.prelude recurrent = engine.model.recurrent_prepared + def record_prelude(token_ids): + calls.append(len(token_ids)) + return prelude(token_ids) + def record(hidden, kv, cache, **kwargs): calls.append(len(hidden)) return recurrent(hidden, kv, cache, **kwargs) + engine.model.prelude = record_prelude engine.model.recurrent_prepared = record for _ in range(300): calls.clear() @@ -196,8 +203,8 @@ def record(hidden, kv, cache, **kwargs): outputs[output.request_id] = (output.token_ids, output.exit_depths) batch = engine.last_schedule if batch is not None and batch.is_mixed: - phases = {batch.phase_for(i) for i in batch.requests} - if len(phases) == 2: + works = {item.work for item in batch.requests} + if len(works) == 2: # Both storage kinds pass through a single model invocation. if not graphs: assert calls == [batch.num_tokens] @@ -232,9 +239,14 @@ def mixed_ready_engine(): engine.add_request(name, prompt, 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: + works = batch.works if batch is not None else frozenset() + if { + WorkKey.DECODE_RECURRENT, + WorkKey.PROMPT_RECURRENT, + } <= works: return engine, batch - engine._update(batch, engine.executor.submit(batch).collect()) + if batch is not None: + engine._update(batch, engine.executor.submit(batch).collect()) pytest.fail("no mixed batch became ready") @@ -273,14 +285,14 @@ def test_late_mixed_result_cannot_advance_a_reused_request_id(): old = next( scheduled_request.request for scheduled_request in batch.requests - if batch.phase_for(scheduled_request) == RequestPhase.RECURRENT + if scheduled_request.work == WorkKey.DECODE_RECURRENT ) engine.abort_request(old.request_id) engine.add_request(old.request_id, [8, 9], SamplingParams(max_tokens=2)) replacement = engine.scheduler.requests[old.request_id] engine._update(batch, result) assert replacement is not old - assert replacement.phase == RequestPhase.WAITING + assert replacement.status == RequestStatus.WAITING assert replacement.loops_done == 0 assert replacement.generated_token_ids == [] other = next( @@ -294,12 +306,12 @@ def test_late_mixed_result_cannot_advance_a_reused_request_id(): def test_mixed_selection_protects_requests_selected_from_another_queue(monkeypatch): engine, batch = mixed_ready_engine() for scheduled_request in batch.requests: - engine.scheduler.enqueue(scheduled_request.request, batch.phase_for(scheduled_request)) + engine.scheduler.enqueue_work(scheduled_request.request, scheduled_request.work) selected = next( - i.request.request_id for i in batch.requests if batch.phase_for(i) == RequestPhase.RECURRENT + i.request.request_id for i in batch.requests if i.work == WorkKey.DECODE_RECURRENT ) blocked = next( - i.request.request_id for i in batch.requests if batch.phase_for(i) != RequestPhase.RECURRENT + i.request.request_id for i in batch.requests if i.work != WorkKey.DECODE_RECURRENT ) exclusions = [] monkeypatch.setattr( @@ -314,7 +326,7 @@ def preempt(request, *, excluded): actual = engine.scheduler.schedule() assert [i.request.request_id for i in actual.requests] == [selected] assert exclusions == [frozenset({selected})] - assert list(engine.scheduler.queues[RequestPhase.PREFILLING_RECURRENT]) == [blocked] + assert list(engine.scheduler.work_queues[WorkKey.PROMPT_RECURRENT]) == [blocked] def test_runtime_closes_metadata_iterator_when_loop_execution_fails(): @@ -452,12 +464,16 @@ def test_mixed_attempt_with_one_selected_phase_is_not_a_mixed_batch(): from dataclasses import replace engine, batch = mixed_ready_engine() - for span in batch.requests: - engine.scheduler.enqueue(span.request, span.phase) + for item in batch.requests: + engine.scheduler.enqueue_work(item.request, item.work) + compatible = {WorkKey.DECODE_RECURRENT, WorkKey.PROMPT_RECURRENT} + for work, queue in engine.scheduler.work_queues.items(): + if work not in compatible: + queue.clear() engine.scheduler.config = replace(engine.scheduler.config, max_num_batched_tokens=1) selected = engine.scheduler.schedule() assert selected is not None - assert {span.phase for span in selected.requests} == {RequestPhase.RECURRENT} - assert selected.phase == RequestPhase.RECURRENT + assert {item.work for item in selected.requests} == {WorkKey.DECODE_RECURRENT} + assert selected.works == frozenset((WorkKey.DECODE_RECURRENT,)) assert not selected.is_mixed engine.close() diff --git a/tests/test_pd.py b/tests/test_pd.py index 84f2c89e..9aa4e6f3 100644 --- a/tests/test_pd.py +++ b/tests/test_pd.py @@ -364,7 +364,7 @@ def test_engine_yields_while_waiting_for_remote_kv(async_scheduling): from vllm_rlt.core.scheduler import Scheduler from vllm_rlt.engine.preemption import PreemptionManager - from vllm_rlt.request import Request, RequestPhase + from vllm_rlt.request import Request, RequestStatus engine = object.__new__(LLMEngine) engine.async_speculative = False @@ -373,7 +373,7 @@ def test_engine_yields_while_waiting_for_remote_kv(async_scheduling): engine.execution_config = ExecutionConfig(async_scheduling=async_scheduling) engine.scheduler = Scheduler(SchedulerConfig(), Mock()) engine.cache_manager = engine.scheduler.cache_manager - receiving = Request("remote", [1], SamplingParams(max_tokens=1), phase=RequestPhase.RECEIVING) + receiving = Request("remote", [1], SamplingParams(max_tokens=1), status=RequestStatus.RECEIVING) engine.scheduler.requests[receiving.request_id] = receiving engine._inflight = [] engine._pending_coda = [] @@ -391,7 +391,7 @@ def test_engine_yields_while_waiting_for_remote_kv(async_scheduling): engine.model_runner.synchronize.assert_not_called() if async_scheduling: # Reap the last local output while another request still awaits KV. - active = Request("local", [1], SamplingParams(max_tokens=1), phase=RequestPhase.CODA) + active = Request("local", [1], SamplingParams(max_tokens=1), status=RequestStatus.ACTIVE) engine.scheduler.requests[active.request_id] = active active.num_output_placeholders = 1 engine.model = SimpleNamespace(config=SimpleNamespace(eos_token_id=2)) @@ -407,7 +407,7 @@ def test_engine_yields_while_waiting_for_remote_kv(async_scheduling): assert outputs[0].token_ids == [3] assert list(engine.scheduler.requests) == ["remote"] # A lost runnable request must still trigger the original invariant. - receiving.phase = RequestPhase.RECURRENT + receiving.status = RequestStatus.ACTIVE with pytest.raises(RuntimeError, match="scheduler made no progress"): engine.step() diff --git a/tests/test_prefix_growth.py b/tests/test_prefix_growth.py index 1916b10b..83945e5a 100644 --- a/tests/test_prefix_growth.py +++ b/tests/test_prefix_growth.py @@ -3,10 +3,11 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, ExecutionConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.kv_cache_manager import KVCacheManager from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import RequestPhase +from vllm_rlt.request import RequestStatus def model(): @@ -85,7 +86,7 @@ def test_incremental_growth_matches_reserved_generation(): @pytest.mark.parametrize( - "stage", [RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT] + "stage", [WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT] ) def test_incremental_split_prefill_reserves_prompt_growth(stage): prompt = list(range(1, 9)) @@ -104,10 +105,10 @@ def test_incremental_split_prefill_reserves_prompt_growth(stage): e.add_request("split", prompt, SamplingParams(max_tokens=1, ignore_eos=True)) e.scheduler._admit() request = e.scheduler.requests["split"] - assert request.phase is RequestPhase.PREFILLING_PRELUDE - if stage is RequestPhase.PREFILLING_RECURRENT: + assert e.scheduler.pending_work(request) is WorkKey.PROMPT_PRELUDE + if stage is WorkKey.PROMPT_RECURRENT: e.step() - assert request.phase is RequestPhase.PREFILLING_RECURRENT + assert e.scheduler.pending_work(request) is WorkKey.PROMPT_RECURRENT allocated = e.cache_manager.allocated_blocks(request.request_id) expected = max(0, e.cache_manager.required_blocks(len(prompt)) - allocated) @@ -325,8 +326,6 @@ def test_async_preemption_restores_prelude_token_in_a_new_slot(): # (LIFO), so the victim resumes in another slot whose device token belongs # to a different request. Distinct prompts make a stale token visible. from vllm_rlt import ExecutionConfig, ExitConfig - from vllm_rlt.request import RequestPhase - torch.manual_seed(72) m = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to(device="cuda", dtype=torch.bfloat16) results, victim_stages = [], [] @@ -349,7 +348,7 @@ def test_async_preemption_restores_prelude_token_in_a_new_slot(): def record_victim(*args, **kwargs): victim = select(*args, **kwargs) if victim is not None: - victim_stages.append(victim.phase) + victim_stages.append(e.scheduler.pending_work(victim)) return victim e.scheduler._select_preemption_victim = record_victim @@ -362,7 +361,7 @@ def record_victim(*args, **kwargs): ) results.append({rid: out.token_ids for rid, out in finish(e).items()}) assert not e.preemption.snapshots and e.cache_manager.num_used_blocks == 0 - assert RequestPhase.PRELUDE in victim_stages + assert WorkKey.DECODE_PRELUDE in victim_stages assert results[0] == results[1] @@ -389,10 +388,11 @@ def test_watermark_only_limits_incremental_admission(incremental): scheduler.add_request(Request(request_id, [1, 2, 3, 4], SamplingParams(max_tokens=1))) scheduler._admit() assert cache.watermark_blocks == (8 if incremental else 0) - assert scheduler.requests["a"].phase == RequestPhase.PREFILL + assert scheduler.pending_work(scheduler.requests["a"]) == WorkKey.PROMPT_ATOMIC # Full allocation admits both requests into the 16-block pool despite the # supplied ratio; incremental allocation retains the configured headroom. - assert scheduler.requests["b"].phase == ( - RequestPhase.WAITING if incremental else RequestPhase.PREFILL - ) + if incremental: + assert scheduler.requests["b"].status == RequestStatus.WAITING + else: + assert scheduler.pending_work(scheduler.requests["b"]) == WorkKey.PROMPT_ATOMIC assert cache.num_free_blocks == (8 if incremental else 0) diff --git a/tests/test_scheduler_selection.py b/tests/test_scheduler_selection.py index 22809aaf..04e8015b 100644 --- a/tests/test_scheduler_selection.py +++ b/tests/test_scheduler_selection.py @@ -11,9 +11,10 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, ExecutionConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import RequestPhase +from vllm_rlt.request import Request, RequestStatus def model(): @@ -38,6 +39,21 @@ def params(max_tokens=5): return SamplingParams(max_tokens=max_tokens, ignore_eos=True) +def test_enqueue_work_requires_registered_active_request_and_unique_pending_work(): + e = engine() + e.add_request("r", [1, 2], params()) + request = e.scheduler.requests["r"] + with pytest.raises(ValueError, match="registered request"): + e.scheduler.enqueue_work(Request("r", [1, 2], params()), WorkKey.PROMPT_ATOMIC) + with pytest.raises(ValueError, match="only active requests"): + e.scheduler.enqueue_work(request, WorkKey.PROMPT_ATOMIC) + + e.scheduler._admit() + with pytest.raises(ValueError, match="already has pending work"): + e.scheduler.enqueue_work(request, WorkKey.DECODE_RECURRENT) + e.close() + + def test_batch_selection_is_never_preempted_by_a_later_item(monkeypatch): """Selecting A must stop B's capacity failure from evicting A.""" e = engine() @@ -45,10 +61,10 @@ def test_batch_selection_is_never_preempted_by_a_later_item(monkeypatch): 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) + batch = e.scheduler.take_works((WorkKey.PROMPT_ATOMIC,)) 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 batch.requests[0].work == WorkKey.PROMPT_ATOMIC def test_direct_take_does_not_inherit_an_earlier_batch_selection(monkeypatch): @@ -58,7 +74,7 @@ def test_direct_take_does_not_inherit_an_earlier_batch_selection(monkeypatch): e.add_request(name, [1, 2], params()) e.scheduler._admit() # A first batch selects both requests. - assert e.scheduler._take(RequestPhase.PREFILL) is not None + assert e.scheduler.take_works((WorkKey.PROMPT_ATOMIC,)) is not None seen = [] @@ -69,9 +85,9 @@ def spy(requester, *, priority_only=False, excluded=frozenset()): # The Scheduler holds the callback it was given, so replace that reference. monkeypatch.setattr(e.scheduler, "preempt", spy) for name in ("a", "b"): - e.scheduler.enqueue(e.scheduler.requests[name], RequestPhase.PREFILL) + e.scheduler.enqueue_work(e.scheduler.requests[name], WorkKey.PROMPT_ATOMIC) monkeypatch.setattr(e.cache_manager, "ensure_capacity", lambda rid, frontier: rid != "b") - assert e.scheduler._take(RequestPhase.PREFILL) is not None + assert e.scheduler.take_works((WorkKey.PROMPT_ATOMIC,)) is not None # Only "a", selected in this batch, is excluded. The first batch's # selection is not still in force, and the set is handed over frozen. assert seen == [("b", frozenset({"a"}))] @@ -82,7 +98,7 @@ def test_request_selected_in_an_earlier_batch_is_still_preemptable(): e = engine(max_num_seqs=1) e.add_request("victim", [1, 2], params()) e.step() - assert e.scheduler.requests["victim"].phase != RequestPhase.WAITING + assert e.scheduler.requests["victim"].status == RequestStatus.ACTIVE e.add_request("requester", [3, 4], params()) assert e.scheduler.preempt(e.scheduler.requests["requester"]) assert "victim" in e.preemption.snapshots @@ -101,7 +117,7 @@ def test_reused_request_id_is_not_protected_by_the_previous_request(): assert "r" in e.preemption.snapshots -@pytest.mark.parametrize("boundary", [RequestPhase.PRELUDE, RequestPhase.CODA]) +@pytest.mark.parametrize("boundary", [WorkKey.DECODE_PRELUDE, WorkKey.DECODE_CODA]) @pytest.mark.parametrize("atomic_prefill", [True, False], ids=["atomic", "split"]) def test_refill_runs_boundary_restored_during_admission(boundary, atomic_prefill): """A sole resumed boundary must run in the admission step that restores it.""" @@ -109,7 +125,7 @@ def test_refill_runs_boundary_restored_during_admission(boundary, atomic_prefill e.add_request("resumed", [1, 2], params(max_tokens=3)) request = e.scheduler.requests["resumed"] for _ in range(20): - if request.phase == boundary: + if e.scheduler.pending_work(request) == boundary: break e.step() else: @@ -118,21 +134,21 @@ def test_refill_runs_boundary_restored_during_admission(boundary, atomic_prefill generated_before = len(request.generated_token_ids) e.add_request("pressure", [3, 4], params(max_tokens=1)) assert e.scheduler.preempt(e.scheduler.requests["pressure"]) - assert e.preemption.snapshots["resumed"]["phase"] == boundary - assert request.phase == RequestPhase.WAITING + assert e.preemption.snapshots["resumed"]["work"] == boundary + assert request.status == RequestStatus.WAITING e.abort_request("pressure") outputs = e.step() assert e.preemption.resumptions == 1 - assert e.last_schedule.phase == boundary - if boundary == RequestPhase.PRELUDE: + assert e.last_schedule.works == frozenset((boundary,)) + if boundary == WorkKey.DECODE_PRELUDE: assert not outputs - assert request.phase == RequestPhase.RECURRENT + assert e.scheduler.pending_work(request) == WorkKey.DECODE_RECURRENT assert len(request.generated_token_ids) == generated_before else: assert [output.request_id for output in outputs] == ["resumed"] - assert request.phase == RequestPhase.PRELUDE + assert e.scheduler.pending_work(request) == WorkKey.DECODE_PRELUDE assert len(request.generated_token_ids) == generated_before + 1 @@ -146,7 +162,7 @@ def test_refill_restored_coda_waits_for_recurrent_batchmate(): e.add_request("core", [1, 2], SamplingParams(max_tokens=2, ignore_eos=True)) core = e.scheduler.requests["core"] for _ in range(20): - if core.phase == RequestPhase.RECURRENT: + if e.scheduler.pending_work(core) == WorkKey.DECODE_RECURRENT: break e.step() else: @@ -159,23 +175,23 @@ def test_refill_restored_coda_waits_for_recurrent_batchmate(): ) e.step() restored = e.scheduler.requests["restored-coda"] - assert core.phase == RequestPhase.RECURRENT - assert restored.phase == RequestPhase.CODA + assert e.scheduler.pending_work(core) == WorkKey.DECODE_RECURRENT + assert e.scheduler.pending_work(restored) == WorkKey.DECODE_CODA e.add_request("pressure", [5, 6], params(max_tokens=1)) assert e.scheduler.preempt(e.scheduler.requests["pressure"]) - assert e.preemption.snapshots["restored-coda"]["phase"] == RequestPhase.CODA + assert e.preemption.snapshots["restored-coda"]["work"] == WorkKey.DECODE_CODA e.abort_request("pressure") e.step() assert e.preemption.resumptions == 1 - assert e.last_schedule.phase == RequestPhase.RECURRENT - assert restored.phase == RequestPhase.CODA + assert e.last_schedule.works == frozenset((WorkKey.DECODE_RECURRENT,)) + assert e.scheduler.pending_work(restored) == WorkKey.DECODE_CODA for _ in range(20): e.step() - if e.last_schedule.phase == RequestPhase.CODA: + if e.last_schedule.works == frozenset((WorkKey.DECODE_CODA,)): break else: pytest.fail("CODA requests did not form a batch") diff --git a/tests/test_serving.py b/tests/test_serving.py index 223ab911..2f34f5aa 100644 --- a/tests/test_serving.py +++ b/tests/test_serving.py @@ -15,9 +15,9 @@ from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import RequestPhase from vllm_rlt.serving.protocol import ( CompletionRequest, IncrementalText, @@ -46,7 +46,7 @@ class PauseAt: """Hold one engine step while a test exercises HTTP or worker lifecycle.""" def __init__(self, stage, *, fail=False): - self.phase, self.fail = stage, fail + self.work, self.fail = stage, fail self.entered, self.release = threading.Event(), threading.Event() self.engine = None @@ -55,7 +55,7 @@ def __call__(self, engine): submit = engine.executor.submit def pause(batch): - if batch.phase == self.phase and not self.entered.is_set(): + if self.work in batch.works and not self.entered.is_set(): self.entered.set() assert self.release.wait(5) if self.fail: @@ -369,7 +369,7 @@ async def run(): @pytest.mark.parametrize("expired", [False, True]) def test_shutdown_cleans_http_and_engine_state(caplog, expired): async def run(): - pause, cleaned = PauseAt(RequestPhase.PREFILL), asyncio.Event() + pause, cleaned = PauseAt(WorkKey.PROMPT_ATOMIC), asyncio.Event() app = create_app( lambda: factory(hook=pause), model="test-model", @@ -498,14 +498,14 @@ def record(batch): owner_threads.append(threading.get_ident()) trace.append( ( - batch.phase, + batch.works, [ scheduled_request.request.request_id for scheduled_request in batch.requests ], ) ) - if batch.phase == RequestPhase.RECURRENT and not blocked.is_set(): + if WorkKey.DECODE_RECURRENT in batch.works and not blocked.is_set(): blocked.set() assert release.wait(5) return submit(batch) @@ -583,7 +583,7 @@ def load(): def test_request_failure_and_disconnect_cleanup(monkeypatch, failure, stream, max_tokens): from vllm_rlt.serving import worker as worker_module - pause = PauseAt(RequestPhase.RECURRENT, fail=failure == "engine") + pause = PauseAt(WorkKey.DECODE_RECURRENT, fail=failure == "engine") class BadTokenizer(TinyTokenizer): calls = 0 diff --git a/tests/test_speculative.py b/tests/test_speculative.py index eb9b2a69..b215d2f4 100644 --- a/tests/test_speculative.py +++ b/tests/test_speculative.py @@ -16,10 +16,11 @@ SchedulerConfig, SpeculativeConfig, ) +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import Request, RequestPhase +from vllm_rlt.request import Request from vllm_rlt.worker.sampler import Sampler from vllm_rlt.worker.speculative import greedy_accept, rejection_sample @@ -126,7 +127,7 @@ def record(hidden, batch, cache, **kwargs): ) try: runner.submit( - SchedulerOutput(RequestPhase.SPECULATIVE, [ScheduledRequest(r, 3, 4)]) + SchedulerOutput([ScheduledRequest(r, WorkKey.SPECULATIVE, 3, 4)]) ).collect() finally: hook.remove() @@ -343,7 +344,7 @@ def spy(request_id, params, device): def test_speculative_greedy_execution_never_acquires_a_generator(monkeypatch): # The CODA path has the same check in test_engine.py; this one proves the - # contract on a request that really executes through RequestPhase.SPECULATIVE, so a + # contract on a request that really executes through WorkKey.SPECULATIVE, so a # generator that is created but unused would still be caught. e = engine(model()) @@ -357,10 +358,10 @@ def forbidden(*args, **kwargs): if not e.has_unfinished_requests(): break e.step() - stages.append(e.last_schedule.phase) + stages.append(e.last_schedule.works) else: pytest.fail("speculative engine did not drain") - assert RequestPhase.SPECULATIVE in stages + assert any(WorkKey.SPECULATIVE in works for works in stages) assert e.rng.peek("r") is None @@ -377,7 +378,7 @@ def test_speculative_path_advances_the_registry_generator(): def step_until_speculative(): for _ in range(10): e.step() - if e.last_schedule.phase == RequestPhase.SPECULATIVE: + if WorkKey.SPECULATIVE in e.last_schedule.works: return True return False @@ -524,7 +525,7 @@ def test_greedy_round_reads_back_once_with_ragged_requests(): e.step() finally: torch.cuda.set_sync_debug_mode("default") - if e.last_schedule.phase == RequestPhase.SPECULATIVE: + if WorkKey.SPECULATIVE in e.last_schedule.works: rounds += 1 assert sum("synchroniz" in str(w.message) for w in caught) == 1 assert rounds > 1 @@ -568,7 +569,7 @@ def run(graphs): @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): - from vllm_rlt.request import Operation + from vllm_rlt.core.execution import Operation def run(atomic): e = engine( @@ -606,7 +607,7 @@ def test_common_submit_entry_preserves_prefill_and_speculative_results(atomic, m submit = e.executor.submit def record(batch): - calls.append((batch.phase, batch.is_speculative)) + calls.append((batch.works, batch.is_speculative)) return submit(batch) monkeypatch.setattr(e.executor, "submit", record) @@ -618,13 +619,14 @@ def record(batch): assert actual.exit_depths == expected.exit_depths assert any(flag for _, flag in calls) assert any(not flag for _, flag in calls) - assert all(flag == (stage == RequestPhase.SPECULATIVE) for stage, flag in calls) + assert all(flag == (works == frozenset((WorkKey.SPECULATIVE,))) for works, flag in calls) assert e.cache_manager.num_used_blocks == 0 def test_common_submit_rejects_speculation_without_controller(): e = LLMEngine(model()) - batch = SchedulerOutput(RequestPhase.SPECULATIVE, []) + request = Request("r", [2], SamplingParams()) + batch = SchedulerOutput([ScheduledRequest(request, WorkKey.SPECULATIVE, 0, 1)]) with pytest.raises(ValueError, match="configured round controller"): e.executor.submit(batch).collect() @@ -644,7 +646,7 @@ def collect(submission): "r", [2, 3, 4], SamplingParams(max_tokens=9, ignore_eos=True, temperature=temperature) ) try: - while e.scheduler.requests["r"].phase != RequestPhase.SPECULATIVE: + while e.scheduler.pending_work(e.scheduler.requests["r"]) != WorkKey.SPECULATIVE: e.step() batch = e._schedule() ticket = e.executor.submit(batch) @@ -668,7 +670,7 @@ def test_speculative_inputs_reuse_owned_buffers_without_private_cache_staging( params = SamplingParams(max_tokens=30, ignore_eos=True) e.add_request("r", [2, 3, 4], params) try: - while e.scheduler.requests["r"].phase != RequestPhase.SPECULATIVE: + while e.scheduler.pending_work(e.scheduler.requests["r"]) != WorkKey.SPECULATIVE: e.step() batch = e._schedule() runner = e.speculative_runner @@ -704,11 +706,14 @@ def test_speculative_input_slots_survive_three_submits_before_collection(use_uva try: for i in range(3): e.add_request(str(i), [i + 2, i + 3], SamplingParams(max_tokens=12, ignore_eos=True)) - while any(r.phase != RequestPhase.SPECULATIVE for r in e.scheduler.requests.values()): + while any( + e.scheduler.pending_work(request) != WorkKey.SPECULATIVE + for request in e.scheduler.requests.values() + ): e.step() batch = e._schedule() assert len(batch.requests) == 3 - batches = [SchedulerOutput(RequestPhase.SPECULATIVE, [span]) for span in batch.requests] + batches = [SchedulerOutput([span]) for span in batch.requests] runner = e.speculative_runner expected = [runner.submit(part).collect() for part in batches] # The third submission wraps the two input slots before any host collection. diff --git a/tests/test_split_prefill.py b/tests/test_split_prefill.py index 9a771d80..ded859e8 100644 --- a/tests/test_split_prefill.py +++ b/tests/test_split_prefill.py @@ -15,6 +15,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.execution import Operation, WorkKey from vllm_rlt.core.scheduling_policy import ( NoRefillPolicy, RefillPolicy, @@ -24,7 +25,7 @@ ) from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import Operation, Request, RequestPhase +from vllm_rlt.request import Request, RequestStatus from vllm_rlt.worker.executor import build_execution_batch TOTAL_DEPTHS = 4 @@ -59,10 +60,9 @@ def test_cpu_prefill_segments_gather_and_scatter_contiguous_spans(): e.add_request("b", [5, 6, 7, 8], SamplingParams(max_tokens=1)) a, b = (e.scheduler.requests[request_id] for request_id in ("a", "b")) batch = SchedulerOutput( - RequestPhase.PREFILLING_PRELUDE, [ - ScheduledRequest(a, 1, 2, chunk_start=0), - ScheduledRequest(b, 2, 1, chunk_start=1), + ScheduledRequest(a, WorkKey.PROMPT_PRELUDE, 1, 2, chunk_start=0), + ScheduledRequest(b, WorkKey.PROMPT_PRELUDE, 2, 1, chunk_start=1), ], ) prepared = e.model_runner.prepare_on_gpu( @@ -95,8 +95,16 @@ def test_recurrent_prepare_uses_absolute_token_start_for_kv_rows(): request = e.scheduler.requests["a"] assert e.cache_manager.allocate("a", 4) batch = SchedulerOutput( - RequestPhase.PREFILLING_RECURRENT, - [ScheduledRequest(request, token_start=2, token_count=1, depth=1, chunk_start=1)], + [ + ScheduledRequest( + request, + WorkKey.PROMPT_RECURRENT, + token_start=2, + token_count=1, + depth=1, + chunk_start=1, + ) + ], ) try: prepared = e.model_runner.prepare_on_gpu( @@ -126,10 +134,17 @@ def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(): assert e.cache_manager.allocate(request_id, 2) mid, final = (e.scheduler.requests[rid] for rid in ("mid", "final")) 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, WorkKey.PROMPT_RECURRENT, 0, 2, depth=1, chunk_start=0), + ScheduledRequest( + final, + WorkKey.PROMPT_RECURRENT, + 0, + 2, + depth=3, + chunk_start=0, + completes_prompt=True, + ), ], ) runner = e.model_runner @@ -178,10 +193,9 @@ def test_cuda_prefill_segments_gather_and_scatter_variable_spans(): e.add_request("b", [5, 6, 7, 8], SamplingParams(max_tokens=1)) a, b = (e.scheduler.requests[request_id] for request_id in ("a", "b")) batch = SchedulerOutput( - RequestPhase.PREFILLING_PRELUDE, [ - ScheduledRequest(a, 1, 2, chunk_start=0), - ScheduledRequest(b, 2, 1, chunk_start=1), + ScheduledRequest(a, WorkKey.PROMPT_PRELUDE, 1, 2, chunk_start=0), + ScheduledRequest(b, WorkKey.PROMPT_PRELUDE, 2, 1, chunk_start=1), ], ) prepared = e.model_runner.prepare_on_gpu( @@ -226,10 +240,17 @@ def test_cuda_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(): assert e.cache_manager.allocate(request_id, 2) mid, final = (e.scheduler.requests[rid] for rid in ("mid", "final")) 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, WorkKey.PROMPT_RECURRENT, 0, 2, depth=1, chunk_start=0), + ScheduledRequest( + final, + WorkKey.PROMPT_RECURRENT, + 0, + 2, + depth=3, + chunk_start=0, + completes_prompt=True, + ), ], ) runner = e.model_runner @@ -292,8 +313,7 @@ def synchronize(self): a, b = (e.scheduler.requests[rid] for rid in ("a", "b")) batches = [ SchedulerOutput( - RequestPhase.PREFILLING_PRELUDE, - [ScheduledRequest(request, 0, 2, chunk_start=0)], + [ScheduledRequest(request, WorkKey.PROMPT_PRELUDE, 0, 2, chunk_start=0)], ) for request in (a, b) ] @@ -375,8 +395,16 @@ def test_submit_transfer_failure_releases_only_prepared_prefill_buffers(): assert e.cache_manager.allocate("a", 4) request = e.scheduler.requests["a"] batch = SchedulerOutput( - RequestPhase.PREFILLING_RECURRENT, - [ScheduledRequest(request, token_start=0, token_count=2, depth=0, chunk_start=0)], + [ + ScheduledRequest( + request, + WorkKey.PROMPT_RECURRENT, + token_start=0, + token_count=2, + depth=0, + chunk_start=0, + ) + ], ) runner = e.model_runner original_prepare = runner.prepare_on_cpu @@ -421,8 +449,11 @@ def test_submit_does_not_retry_failed_routing_cleanup(monkeypatch): e.add_request("a", [1, 2], SamplingParams(max_tokens=1)) request = e.scheduler.requests["a"] batch = SchedulerOutput( - RequestPhase.PREFILLING_PRELUDE, - [ScheduledRequest(request, token_start=0, token_count=2, chunk_start=0)], + [ + ScheduledRequest( + request, WorkKey.PROMPT_PRELUDE, token_start=0, token_count=2, chunk_start=0 + ) + ], ) runner = e.model_runner @@ -485,8 +516,11 @@ def test_cuda_submit_readback_failure_releases_batch_input_buffers(monkeypatch): e.add_request("a", [1, 2, 3, 4], SamplingParams(max_tokens=1)) request = e.scheduler.requests["a"] batch = SchedulerOutput( - RequestPhase.PREFILLING_PRELUDE, - [ScheduledRequest(request, token_start=0, token_count=2, chunk_start=0)], + [ + ScheduledRequest( + request, WorkKey.PROMPT_PRELUDE, token_start=0, token_count=2, chunk_start=0 + ) + ], ) runner = e.model_runner original_readback = runner._readback_slot @@ -537,7 +571,7 @@ def advance_chunk(request, scheduled_request): def schedule_stage(scheduler, stage): - batch = scheduler._take(stage) + batch = scheduler.take_works((stage,)) assert batch is not None, f"no {stage.value} batch was selected" return batch @@ -551,17 +585,17 @@ def test_split_prefill_steps_one_depth_per_chunk_and_keeps_the_prompt_ordered(): # 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: + if scheduled_request.work == WorkKey.PROMPT_PRELUDE: request.num_preluded_tokens = ( scheduled_request.token_start + scheduled_request.token_count ) - e.scheduler.enqueue(request, RequestPhase.PREFILLING_RECURRENT) + e.scheduler.enqueue_work(request, WorkKey.PROMPT_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) + e.scheduler.enqueue_work(request, WorkKey.PROMPT_RECURRENT) elif request.num_prefilled_tokens < len(prompt): - e.scheduler.enqueue(request, RequestPhase.PREFILLING_PRELUDE) + e.scheduler.enqueue_work(request, WorkKey.PROMPT_PRELUDE) return (scheduled_request.token_start, scheduled_request.token_count, depth) trace = [] @@ -570,25 +604,25 @@ def apply(batch): if batch is None: break assert len(batch.requests) == 1 - trace.append((batch.phase, *apply(batch))) + trace.append((batch.requests[0].work, *apply(batch))) assert e.scheduler.schedule() is None, "prefill must drain, not schedule forever" 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), + (WorkKey.PROMPT_PRELUDE, 0, 2, None), + (WorkKey.PROMPT_RECURRENT, 0, 2, 0), + (WorkKey.PROMPT_RECURRENT, 0, 2, 1), + (WorkKey.PROMPT_RECURRENT, 0, 2, 2), + (WorkKey.PROMPT_RECURRENT, 0, 2, 3), + (WorkKey.PROMPT_PRELUDE, 2, 2, None), + (WorkKey.PROMPT_RECURRENT, 2, 2, 0), + (WorkKey.PROMPT_RECURRENT, 2, 2, 1), + (WorkKey.PROMPT_RECURRENT, 2, 2, 2), + (WorkKey.PROMPT_RECURRENT, 2, 2, 3), + (WorkKey.PROMPT_PRELUDE, 4, 1, None), + (WorkKey.PROMPT_RECURRENT, 4, 1, 0), + (WorkKey.PROMPT_RECURRENT, 4, 1, 1), + (WorkKey.PROMPT_RECURRENT, 4, 1, 2), + (WorkKey.PROMPT_RECURRENT, 4, 1, 3), ] assert request.num_prefilled_tokens == 5 assert all(left == right for left, right in request.prefill_depth_ranges) @@ -606,10 +640,12 @@ def test_engine_split_prefill_transition_advances_the_same_frontier_contract(): break assert len(batch.requests) == 1 scheduled_request = batch.requests[0] - progress = e._apply_split_prefill_progress(request, batch.phase, scheduled_request) + progress = e._apply_split_prefill_progress( + request, scheduled_request.work, scheduled_request + ) trace.append( ( - batch.phase, + scheduled_request.work, scheduled_request.token_start, scheduled_request.token_count, scheduled_request.depth, @@ -619,24 +655,24 @@ def test_engine_split_prefill_transition_advances_the_same_frontier_contract(): 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), + (WorkKey.PROMPT_PRELUDE, 0, 2, None), + (WorkKey.PROMPT_RECURRENT, 0, 2, 0), + (WorkKey.PROMPT_RECURRENT, 0, 2, 1), + (WorkKey.PROMPT_RECURRENT, 0, 2, 2), + (WorkKey.PROMPT_RECURRENT, 0, 2, 3), + (WorkKey.PROMPT_PRELUDE, 2, 2, None), + (WorkKey.PROMPT_RECURRENT, 2, 2, 0), + (WorkKey.PROMPT_RECURRENT, 2, 2, 1), + (WorkKey.PROMPT_RECURRENT, 2, 2, 2), + (WorkKey.PROMPT_RECURRENT, 2, 2, 3), + (WorkKey.PROMPT_PRELUDE, 4, 1, None), + (WorkKey.PROMPT_RECURRENT, 4, 1, 0), + (WorkKey.PROMPT_RECURRENT, 4, 1, 1), + (WorkKey.PROMPT_RECURRENT, 4, 1, 2), + (WorkKey.PROMPT_RECURRENT, 4, 1, 3), ] assert request.num_prefilled_tokens == len(request.prompt_token_ids) - assert request.phase is RequestPhase.CODA + assert e.scheduler.pending_work(request) is WorkKey.DECODE_CODA @pytest.mark.parametrize( @@ -742,7 +778,7 @@ def test_cpu_async_split_frontier_advances_once_per_submission(): def hold_split_ticket(batch): ticket = submit(batch) - if batch.phase in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT): + if batch.has_split_prefill_work: ticket.event = HeldEvent() held.append(ticket.event) for scheduled_request in batch.requests: @@ -752,10 +788,10 @@ def hold_split_ticket(batch): e.executor.submit = hold_split_ticket for _ in range(32): e.step() - if request.phase is RequestPhase.CODA: + if e.scheduler.pending_work(request) is WorkKey.DECODE_CODA: break - assert request.phase is RequestPhase.CODA + assert e.scheduler.pending_work(request) is WorkKey.DECODE_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. @@ -835,18 +871,18 @@ def test_split_preemption_rewinds_and_recomputes_the_resident_chunk(completed_de for _ in range(completed_depths + 1): e.step() - assert request.phase is RequestPhase.PREFILLING_RECURRENT + assert e.scheduler.pending_work(request) is WorkKey.PROMPT_RECURRENT requester = Request("pressure", [9], SamplingParams(max_tokens=1)) assert e.scheduler.preempt(requester) - assert e.preemption.snapshots["victim"]["phase"] is RequestPhase.PREFILLING_PRELUDE + assert e.preemption.snapshots["victim"]["work"] is WorkKey.PROMPT_PRELUDE assert request.num_preluded_tokens == request.num_prefilled_tokens assert request.prefill_chunk_start == 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 e.scheduler._admit() - assert request.phase is RequestPhase.PREFILLING_PRELUDE + assert e.scheduler.pending_work(request) is WorkKey.PROMPT_PRELUDE actual = drain_request(e) assert actual.token_ids == expected.token_ids assert actual.exit_depths == expected.exit_depths @@ -882,7 +918,7 @@ def test_priority_admission_preempts_split_prefill_and_matches_reference(): e.add_request("low", prompt, params) e.step() e.step() - assert e.scheduler.requests["low"].phase is RequestPhase.PREFILLING_RECURRENT + assert e.scheduler.pending_work(e.scheduler.requests["low"]) is WorkKey.PROMPT_RECURRENT e.add_request( "high", [5, 6], @@ -930,7 +966,7 @@ def fail(batch, *args, **kwargs): is_loop = ( (batch.batch.operation == Operation.RECURRENT) if failure_site == "prepare_on_gpu" - else (batch.phase == RequestPhase.PREFILLING_RECURRENT) + else (batch.works == frozenset((WorkKey.PROMPT_RECURRENT,))) ) if is_loop: torch.cuda._sleep(10_000_000) @@ -1043,10 +1079,10 @@ def run(cuda_graphs): e.step() e.step() split_request = e.scheduler.requests["split"] - assert split_request.phase is RequestPhase.PREFILLING_RECURRENT + assert e.scheduler.pending_work(split_request) is WorkKey.PROMPT_RECURRENT assert split_request.prefill_depth_ranges[0] == [2, 4] # Keep the staged chunk resident while decode triggers graph capture. - e.scheduler.queues[RequestPhase.PREFILLING_RECURRENT].remove("split") + e.scheduler.work_queues[WorkKey.PROMPT_RECURRENT].remove("split") e.add_request( "decode", @@ -1057,7 +1093,7 @@ def run(cuda_graphs): for _ in range(100): e.step() batch = e.last_schedule - if batch is not None and batch.phase is RequestPhase.RECURRENT: + if batch is not None and batch.works == frozenset((WorkKey.DECODE_RECURRENT,)): decode_recurrent_runs += sum( scheduled_request.request.request_id == "decode" for scheduled_request in batch.requests @@ -1066,7 +1102,7 @@ def run(cuda_graphs): break else: pytest.fail("decode request did not reach its second recurrent step") - assert split_request.phase is RequestPhase.PREFILLING_RECURRENT + assert e.scheduler.pending_work(split_request) is WorkKey.PROMPT_RECURRENT assert split_request.num_prefilled_tokens == 0 assert split_request.num_preluded_tokens == 2 @@ -1077,7 +1113,7 @@ def run(cuda_graphs): e.model_runner.graphs.replays, e.model_runner.graphs.fallbacks, ) - e.scheduler.enqueue(split_request, RequestPhase.PREFILLING_RECURRENT) + e.scheduler.enqueue_work(split_request, WorkKey.PROMPT_RECURRENT) finished = {} for _ in range(200): for output in e.step(): @@ -1173,7 +1209,7 @@ def test_cuda_async_split_preemption_releases_and_reuses_slot(): e.add_request("victim", prompt, params) e.step() e.step() - assert e.scheduler.requests["victim"].phase is RequestPhase.PREFILLING_RECURRENT + assert e.scheduler.pending_work(e.scheduler.requests["victim"]) is WorkKey.PROMPT_RECURRENT state = e.model_runner.async_state slot = state.slots["victim"] old_event = e.model_runner.events["victim"] @@ -1189,7 +1225,7 @@ def test_cuda_async_split_preemption_releases_and_reuses_slot(): e.add_request("pressure-b", [7, 8], pressure_params) e.scheduler._admit() e.step() - assert e.scheduler.requests["victim"].phase is RequestPhase.WAITING + assert e.scheduler.requests["victim"].status is RequestStatus.WAITING assert slot in state.slots.values() assert len(state.slots) == 2 occupant = next(rid for rid, assigned in state.slots.items() if assigned == slot) @@ -1241,14 +1277,14 @@ def test_cuda_split_prefill_batches_mixed_depths(attention_backend): e.add_request("b", [3, 4], params) e.step() # Prelude for both requests. - recurrent = e.scheduler.queues[RequestPhase.PREFILLING_RECURRENT] + recurrent = e.scheduler.work_queues[WorkKey.PROMPT_RECURRENT] recurrent.clear() recurrent.append("a") e.step() # Advance only request a to depth one. recurrent.append("b") e.step() # Request a is depth one; request b is depth zero. - assert e.last_schedule.phase is RequestPhase.PREFILLING_RECURRENT + assert e.last_schedule.works == frozenset((WorkKey.PROMPT_RECURRENT,)) assert {scheduled_request.depth for scheduled_request in e.last_schedule.requests} == {0, 1} finally: e.close() @@ -1261,19 +1297,19 @@ def test_split_prefill_holds_the_next_prelude_until_the_chunk_is_read(): e.scheduler._admit() request = e.scheduler.requests["a"] - prelude = schedule_stage(e.scheduler, RequestPhase.PREFILLING_PRELUDE) + prelude = schedule_stage(e.scheduler, WorkKey.PROMPT_PRELUDE) scheduled_request = prelude.requests[0] request.num_preluded_tokens = scheduled_request.token_start + scheduled_request.token_count - e.scheduler.enqueue(request, RequestPhase.PREFILLING_RECURRENT) - batch = schedule_stage(e.scheduler, RequestPhase.PREFILLING_RECURRENT) + e.scheduler.enqueue_work(request, WorkKey.PROMPT_RECURRENT) + batch = schedule_stage(e.scheduler, WorkKey.PROMPT_RECURRENT) advance_chunk(request, batch.requests[0]) # Depth 0 of the first chunk is done, but its hidden rows are still the input # of depth 1, so staging the next chunk must wait. - e.scheduler.enqueue(request, RequestPhase.PREFILLING_PRELUDE) - assert e.scheduler._take(RequestPhase.PREFILLING_PRELUDE) is None - assert list(e.scheduler.queues[RequestPhase.PREFILLING_PRELUDE]) == ["a"] - e.scheduler.queues[RequestPhase.PREFILLING_PRELUDE].clear() + e.scheduler.enqueue_work(request, WorkKey.PROMPT_PRELUDE) + assert e.scheduler.take_works((WorkKey.PROMPT_PRELUDE,)) is None + assert list(e.scheduler.work_queues[WorkKey.PROMPT_PRELUDE]) == ["a"] + e.scheduler.work_queues[WorkKey.PROMPT_PRELUDE].clear() assert e.scheduler._prelude_window_open(request) is False @@ -1289,8 +1325,8 @@ def test_split_prefill_reports_no_depth_work_before_the_prelude(): e.scheduler._admit() request = e.scheduler.requests["a"] - scheduled_request = e.scheduler._make_scheduled_request( - request, RequestPhase.PREFILLING_RECURRENT, 4 + scheduled_request = e.scheduler._make_scheduled_work( + request, WorkKey.PROMPT_RECURRENT, 4 ) assert scheduled_request.token_count == 0 assert request.num_prefilled_tokens == 0 @@ -1305,7 +1341,7 @@ def test_split_prefill_batches_stay_within_the_token_budget(): for name in ("a", "b", "c"): e.add_request(name, [1, 2, 3, 4], SamplingParams(max_tokens=1, ignore_eos=True)) e.scheduler._admit() - batch = schedule_stage(e.scheduler, RequestPhase.PREFILLING_PRELUDE) + batch = schedule_stage(e.scheduler, WorkKey.PROMPT_PRELUDE) assert sum(i.token_count for i in batch.requests) <= 3 assert len(batch.requests) == 2 @@ -1318,32 +1354,36 @@ def test_split_policy_honors_decode_due_before_split_prefill(): policy = e.scheduler.policy policy.prefill_batches_since_recurrent = 1 scheduler = SimpleNamespace( - queues={ - RequestPhase.RECURRENT: deque(["decode"]), - RequestPhase.PRELUDE: deque(), - RequestPhase.CODA: deque(), - RequestPhase.PREFILLING_PRELUDE: deque(["prefill"]), + work_queues={ + WorkKey.DECODE_RECURRENT: deque(["decode"]), + WorkKey.DECODE_PRELUDE: deque(), + WorkKey.DECODE_CODA: deque(), + WorkKey.PROMPT_PRELUDE: deque(["prefill"]), + WorkKey.PROMPT_ATOMIC: deque(), + WorkKey.PROMPT_RECURRENT: deque(), + WorkKey.SPECULATIVE: deque(), }, - _take=lambda stage: stage, + take_works=lambda works: works[0], ) - scheduler.has_ready = lambda stage: bool(scheduler.queues[stage]) - assert policy.schedule(scheduler) is RequestPhase.RECURRENT + scheduler.has_ready_work = lambda work: bool(scheduler.work_queues[work]) + scheduler.has_ready_works = lambda *works: any(scheduler.has_ready_work(work) for work in works) + assert policy.schedule(scheduler) is WorkKey.DECODE_RECURRENT e.close() def test_split_prefill_stages_count_toward_decode_fairness(): e = make_prefill_engine(scheduler_config=SchedulerConfig(max_prefill_batches_before_decode=2)) policy = e.scheduler.policy - assert policy.PREFILL_PHASES == ( - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, + assert policy.PREFILL_WORKS == ( + WorkKey.PROMPT_PRELUDE, + WorkKey.PROMPT_RECURRENT, ) - policy.record_batch(RequestPhase.PREFILLING_PRELUDE) + policy.record_batch((WorkKey.PROMPT_PRELUDE,)) assert not policy.decode_due - policy.record_batch(RequestPhase.PREFILLING_RECURRENT) + policy.record_batch((WorkKey.PROMPT_RECURRENT,)) assert policy.decode_due - policy.record_batch(RequestPhase.RECURRENT) + policy.record_batch((WorkKey.DECODE_RECURRENT,)) assert not policy.decode_due e.close() @@ -1361,8 +1401,8 @@ def test_split_prefill_stages_after_a_reused_prompt_prefix(): request.num_prefilled_tokens = 2 request.num_preluded_tokens = 2 assert e.scheduler._prelude_window_open(request) - scheduled_request = e.scheduler._make_scheduled_request( - request, RequestPhase.PREFILLING_PRELUDE, 4 + scheduled_request = e.scheduler._make_scheduled_work( + request, WorkKey.PROMPT_PRELUDE, 4 ) assert (scheduled_request.token_start, scheduled_request.token_count) == (2, 2) assert request.prefill_depth(scheduled_request.token_start) == 0 @@ -1403,14 +1443,14 @@ def test_split_prefill_policy_is_selected_only_by_the_atomic_prefill_flag(): ) -def test_scheduling_policies_declare_their_prefill_stages(): - assert SchedulingPolicy.PREFILL_PHASES == (RequestPhase.PREFILL,) - assert RefillPolicy.PREFILL_PHASES == (RequestPhase.PREFILL,) - assert NoRefillPolicy.PREFILL_PHASES == (RequestPhase.PREFILL,) - assert SpeculativePolicy.PREFILL_PHASES == (RequestPhase.PREFILL,) - assert SplitPrefillPolicy.PREFILL_PHASES == ( - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, +def test_scheduling_policies_declare_their_prompt_work(): + assert SchedulingPolicy.PREFILL_WORKS == (WorkKey.PROMPT_ATOMIC,) + assert RefillPolicy.PREFILL_WORKS == (WorkKey.PROMPT_ATOMIC,) + assert NoRefillPolicy.PREFILL_WORKS == (WorkKey.PROMPT_ATOMIC,) + assert SpeculativePolicy.PREFILL_WORKS == (WorkKey.PROMPT_ATOMIC,) + assert SplitPrefillPolicy.PREFILL_WORKS == ( + WorkKey.PROMPT_PRELUDE, + WorkKey.PROMPT_RECURRENT, ) @@ -1513,8 +1553,7 @@ def test_chunk_prompt_cache_reuse_override_and_request_id_reuse(device, use_uva) def prepare(request, start, override=None): scheduled = SchedulerOutput( - RequestPhase.PREFILLING_PRELUDE, - [ScheduledRequest(request, start, 2, chunk_start=start)], + [ScheduledRequest(request, WorkKey.PROMPT_PRELUDE, start, 2, chunk_start=start)], ) batch = build_execution_batch(scheduled, e.model.config.total_ut_steps) assert batch.spans[0].input_tokens is None @@ -1550,7 +1589,7 @@ def prepare(request, start, override=None): @pytest.mark.gpu @pytest.mark.parametrize("use_uva", [False, True]) def test_mixed_hidden_buffers_wait_for_strided_decode_and_chunk_producers(multi_stream, use_uva): - from vllm_rlt.request import ExecutionBatch, StateSpan + from vllm_rlt.core.execution import ExecutionBatch, StateSpan e = LLMEngine( tiny_ouro_model().to("cuda"), diff --git a/tests/test_stage_engine.py b/tests/test_stage_engine.py index f3c82045..662f70cf 100644 --- a/tests/test_stage_engine.py +++ b/tests/test_stage_engine.py @@ -5,10 +5,11 @@ from tests.helpers import RecordingConnector, make_stage_engine, tiny_ouro_config from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.engine.stage_config import StageConfig from vllm_rlt.engine.stage_engine import StageEngine from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import RequestPhase +from vllm_rlt.request import RequestStatus def drain(engine): @@ -49,8 +50,8 @@ def test_prefill_handoff_matches_local_execution_without_sampling(atomic, asynch break else: pytest.fail("prompt did not complete") - assert request.phase == RequestPhase.TRANSFERRING - assert original.phase == RequestPhase.CODA + assert request.status == RequestStatus.TRANSFERRING + assert local.scheduler.pending_work(original) == WorkKey.DECODE_CODA assert request.generated_token_ids == [] torch.testing.assert_close(request.hidden_state, original.hidden_state) assert all(r is request for r, _ in connector.submissions) @@ -71,11 +72,11 @@ def test_connector_backpressure_preserves_other_requests_and_queue_membership(): for _ in range(20): engine.step() assert engine.scheduler.requests["a"].num_preluded_tokens == 0 - assert engine.scheduler.requests["b"].phase == RequestPhase.TRANSFERRING + assert engine.scheduler.requests["b"].status == RequestStatus.TRANSFERRING connector.blocked.clear() for _ in range(20): engine.step() - assert engine.scheduler.requests["a"].phase == RequestPhase.TRANSFERRING + assert engine.scheduler.requests["a"].status == RequestStatus.TRANSFERRING engine.close() @@ -87,7 +88,7 @@ def test_cancelled_import_cannot_activate_reused_request_id(): engine.reserve_transfer("r", [3], SamplingParams(max_tokens=2), receiving=True) current = engine.scheduler.requests["r"] assert not engine.accept_import(old, torch.zeros(1)) - assert current.phase == RequestPhase.RECEIVING + assert current.status == RequestStatus.RECEIVING assert current.hidden_state is None engine.close() @@ -110,7 +111,7 @@ def test_independent_stage_resources_preserve_output_after_import(atomic, asynch assert d.step() == [] for _ in range(100): p.step() - if source.phase == RequestPhase.TRANSFERRING: + if source.status == RequestStatus.TRANSFERRING: break else: pytest.fail("source prompt did not complete") @@ -188,7 +189,7 @@ def test_prefill_decode_stage_routes_per_request_and_keeps_local_decode_running( connector.blocked.clear() for _ in range(100): engine.step() - if engine.scheduler.requests["remote"].phase == RequestPhase.TRANSFERRING: + if engine.scheduler.requests["remote"].status == RequestStatus.TRANSFERRING: break else: pytest.fail("remote request did not reach its handoff boundary") @@ -221,7 +222,7 @@ def test_prefill_does_not_hide_a_lost_runnable_request(): engine = make_stage_engine(role="prefill") engine.connector = RecordingConnector() engine.add_request("r", [1, 2], SamplingParams(max_tokens=1)) - engine.scheduler.queues[RequestPhase.WAITING].clear() + engine.scheduler.waiting.clear() with pytest.raises(RuntimeError, match="no progress"): engine.step() engine.close() @@ -239,13 +240,13 @@ def test_prefill_admission_waits_for_transfer_to_release_kv(): engine.add_request("first", [1, 2], SamplingParams(max_tokens=1)) engine.step() first = engine.scheduler.requests["first"] - assert first.phase == RequestPhase.TRANSFERRING + assert first.status == RequestStatus.TRANSFERRING engine.add_request("second", [1, 2], SamplingParams(max_tokens=1)) assert engine.step() == [] - assert engine.scheduler.requests["second"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["second"].status == RequestStatus.WAITING engine.release_transfer(first) engine.step() - assert engine.scheduler.requests["second"].phase == RequestPhase.TRANSFERRING + assert engine.scheduler.requests["second"].status == RequestStatus.TRANSFERRING engine.close() @@ -388,7 +389,7 @@ def test_speculative_policy_schedules_newly_imported_coda(): assert engine.step() == [] engine.accept_import(request, torch.zeros(1, engine.model.config.hidden_size)) batch = engine.scheduler.schedule() - assert batch.phase == RequestPhase.CODA + assert batch.works == frozenset((WorkKey.DECODE_CODA,)) assert batch.requests[0].request is request finally: engine.close() @@ -578,7 +579,7 @@ def test_backpressured_export_allows_queued_arrival_to_wait(asynchronous): connector.blocked.add("active") engine.reserve_transfer("waiting", [3, 4], SamplingParams(max_tokens=1), receiving=False) assert engine.step() == [] - assert engine.scheduler.requests["waiting"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["waiting"].status == RequestStatus.WAITING connector.blocked.clear() engine.step() assert engine.last_schedule is not None @@ -607,7 +608,8 @@ def test_transfer_reservation_preserves_active_prompt_growth(receiving, preempti assert engine.reserve_transfer("remote", [1, 2], params, receiving=receiving) is None assert not engine.has_request("remote") assert not engine.cache_manager.has_allocation("remote") - assert all("remote" not in queue for queue in engine.scheduler.queues.values()) + assert "remote" not in engine.scheduler.waiting + assert all("remote" not in queue for queue in engine.scheduler.work_queues.values()) assert len(drain(engine)["local"][0]) == 1 # A rejected generation can retry after the local request releases its KV. remote = engine.reserve_transfer("remote", [1, 2], params, receiving=receiving) @@ -646,7 +648,7 @@ def test_decode_import_preemption_restores_outputs_and_rng(device): source_engine.connector = RecordingConnector() params = SamplingParams(max_tokens=5, temperature=0.7, seed=41, ignore_eos=True) source = source_engine.reserve_transfer("source", [1, 2, 3, 4], params, receiving=False) - while source.phase != RequestPhase.TRANSFERRING: + while source.status != RequestStatus.TRANSFERRING: source_engine.step() results = [] try: @@ -685,7 +687,7 @@ def import_request(name): waiting = import_request("waiting") if preempt: assert engine.scheduler.preempt(waiting) - assert victim.phase == RequestPhase.WAITING + assert victim.status == RequestStatus.WAITING results.append(drain(engine)) assert engine.preemption.resumptions == int(preempt) assert not engine.preemption.snapshots diff --git a/tests/test_transfer_connector.py b/tests/test_transfer_connector.py index 7c285891..c0a83d94 100644 --- a/tests/test_transfer_connector.py +++ b/tests/test_transfer_connector.py @@ -7,9 +7,10 @@ from tests.helpers import RecordingConnector, make_stage_engine from vllm_rlt import SamplingParams +from vllm_rlt.core.execution import WorkKey from vllm_rlt.pd.config import PDConfig from vllm_rlt.pd.connector import PDStageConnector -from vllm_rlt.request import RequestPhase +from vllm_rlt.request import RequestStatus def make_connector(engine, slots=4): @@ -58,7 +59,7 @@ def test_import_requires_commit_all_sequences_and_capacity(commit_first): assert sum(w.active for w in transfers) == 3 blocked = next(w for w in transfers if not w.active) active = next(w for w in transfers if w.active) - assert blocked.request.phase == RequestPhase.RECEIVING + assert blocked.request.status == RequestStatus.RECEIVING assert engine.cache_manager.has_transfer_lease(blocked.tid) assert engine.scheduler.requests[active.tid] is active.request connector.cancel(active.tid) @@ -76,8 +77,9 @@ def test_cancelled_export_holds_slot_and_kv_until_local_and_remote_access_end(): event = transfer.last_event = Mock() event.query.return_value = False connector.cancel(transfer.tid) - assert transfer.request.phase == RequestPhase.TRANSFERRING - assert all(transfer.tid not in q for q in engine.scheduler.queues.values()) + assert transfer.request.status == RequestStatus.TRANSFERRING + assert transfer.tid not in engine.scheduler.waiting + assert all(transfer.tid not in q for q in engine.scheduler.work_queues.values()) connector.progress() assert not connector.free_slots assert engine.cache_manager.has_transfer_lease(transfer.tid) @@ -108,7 +110,7 @@ def test_cancelled_import_waits_for_safe_notification_and_old_generation_is_igno current = reserve(connector, "generation-2") connector.progress() # Old completion cannot activate the reused staging slot. connector.commit(old.tid, 1) - assert current.request.phase == RequestPhase.RECEIVING + assert current.request.status == RequestStatus.RECEIVING assert current.received == set() assert engine.cache_manager.has_transfer_lease(current.tid) connector.cancel(current.tid) @@ -120,7 +122,7 @@ def test_compute_release_checks_event_and_identity_without_releasing_kv(): engine = make_stage_engine(role="prefill") engine.connector = RecordingConnector() request = engine.reserve_transfer("r", [1, 2], SamplingParams(max_tokens=1), receiving=False) - while request.phase != RequestPhase.TRANSFERRING: + while request.status != RequestStatus.TRANSFERRING: engine.step() event = Mock() event.query.return_value = False @@ -153,7 +155,7 @@ def test_mixed_connector_handles_both_transfer_directions_by_request(): engine.scheduler._admit() connector.progress() assert incoming.active - assert outgoing.request.phase == RequestPhase.PREFILLING_PRELUDE + assert engine.scheduler.pending_work(outgoing.request) == WorkKey.PROMPT_PRELUDE connector.cancel(outgoing.tid) connector.cancel(incoming.tid) connector.progress() @@ -188,6 +190,7 @@ def test_ack_alone_cannot_release_unsubmitted_or_locally_pending_export(): engine = make_stage_engine() connector = make_connector(engine, slots=1) transfer = reserve(connector, "out", receiving=False) + engine.scheduler._admit() connector.acknowledge(transfer.tid) connector.progress() assert transfer.tid in connector.transfers @@ -199,13 +202,13 @@ def test_ack_alone_cannot_release_unsubmitted_or_locally_pending_export(): transfer.chunks.append((event, 0, [], [], 1)) connector.progress() assert not transfer.committed - assert not transfer.request.compute_released + assert engine.scheduler.has_admission_slot(transfer.request) assert engine.cache_manager.has_transfer_lease(transfer.tid) event.query.return_value = True connector.transport.submit.return_value = True connector.transport.busy.return_value = True connector.progress() - assert transfer.committed and transfer.request.compute_released + assert transfer.committed and not engine.scheduler.has_admission_slot(transfer.request) assert not connector.free_slots connector.transport.busy.return_value = False connector.progress() @@ -256,15 +259,15 @@ def test_import_and_local_prompt_share_priority_admission(local_priority, remote connector.commit("remote", 1) connector.transport.poll.return_value = ([], [("peer", "remote", 0)]) connector.progress() - assert engine.scheduler.requests["local"].phase == RequestPhase.WAITING - assert engine.scheduler.requests["remote"].phase == RequestPhase.RECEIVING + assert engine.scheduler.requests["local"].status == RequestStatus.WAITING + assert engine.scheduler.requests["remote"].status == RequestStatus.RECEIVING engine.scheduler._admit() - assert engine.scheduler.requests[winner].phase not in ( - RequestPhase.WAITING, - RequestPhase.RECEIVING, - ) + assert engine.scheduler.requests[winner].status == RequestStatus.ACTIVE loser = "remote" if winner == "local" else "local" - assert engine.scheduler.requests[loser].phase in (RequestPhase.WAITING, RequestPhase.RECEIVING) + assert engine.scheduler.requests[loser].status in ( + RequestStatus.WAITING, + RequestStatus.RECEIVING, + ) connector.cancel("remote") engine.close() @@ -276,17 +279,21 @@ def test_export_holds_compute_admission_until_final_staging_read_finishes(): engine.scheduler.config = replace(engine.scheduler.config, max_num_seqs=1) connector = make_connector(engine) transfer = reserve(connector, "export", receiving=False) + engine.scheduler._admit() engine.suspend_transfer(transfer.request) engine.add_request("local", [2], SamplingParams(max_tokens=1)) event = Mock() event.query.return_value = False assert not engine.release_transfer_compute(transfer.request, ready_event=event) engine.scheduler._admit() - assert engine.scheduler.requests["local"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["local"].status == RequestStatus.WAITING event.query.return_value = True assert engine.release_transfer_compute(transfer.request, ready_event=event) engine.scheduler._admit() - assert engine.scheduler.requests["local"].phase == RequestPhase.PREFILLING_PRELUDE + assert ( + engine.scheduler.pending_work(engine.scheduler.requests["local"]) + == WorkKey.PROMPT_PRELUDE + ) assert engine.cache_manager.has_transfer_lease("export") connector.cancel("export") connector.progress() @@ -311,16 +318,19 @@ def test_export_reservation_obeys_shared_compute_admission(local_priority, remot peer="peer", ) remote = engine.scheduler.requests["remote"] - assert remote.phase == RequestPhase.WAITING + assert remote.status == RequestStatus.WAITING assert engine.cache_manager.has_transfer_lease("remote") engine.scheduler._admit() - assert engine.scheduler.requests[winner].phase == RequestPhase.PREFILLING_PRELUDE + assert ( + engine.scheduler.pending_work(engine.scheduler.requests[winner]) + == WorkKey.PROMPT_PRELUDE + ) loser = "remote" if winner == "local" else "local" - assert engine.scheduler.requests[loser].phase == RequestPhase.WAITING + assert engine.scheduler.requests[loser].status == RequestStatus.WAITING # Reserving another export must not bypass an occupied compute slot either. reserve(connector, "queued", receiving=False) engine.scheduler._admit() - assert engine.scheduler.requests["queued"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["queued"].status == RequestStatus.WAITING connector.cancel("queued") connector.cancel("remote") engine.close() diff --git a/vllm_rlt/core/execution.py b/vllm_rlt/core/execution.py new file mode 100644 index 00000000..3450ae23 --- /dev/null +++ b/vllm_rlt/core/execution.py @@ -0,0 +1,111 @@ +"""Model operations and immutable execution bindings for scheduled work.""" + +from dataclasses import dataclass, replace +from enum import Enum + +from vllm_rlt.request import Request + + +class Operation(Enum): + PRELUDE = "prelude" + RECURRENT = "recurrent" + CODA = "coda" + + +class WorkKey(Enum): + PROMPT_ATOMIC = ("prompt_atomic", Operation.PRELUDE) + PROMPT_PRELUDE = ("prompt_prelude", Operation.PRELUDE) + PROMPT_RECURRENT = ("prompt_recurrent", Operation.RECURRENT) + DECODE_PRELUDE = ("decode_prelude", Operation.PRELUDE) + DECODE_RECURRENT = ("decode_recurrent", Operation.RECURRENT) + DECODE_CODA = ("decode_coda", Operation.CODA) + SPECULATIVE = ("speculative", None) + + @property + def operation(self) -> Operation | None: + return self.value[1] + + +@dataclass(frozen=True) +class StateSpan: + """Immutable execution range borrowing the original Request for state ownership.""" + + request: Request + token_start: int + token_count: int + # Zero-based starting recurrent depth, not the request's completed-loop count. + # CODA uses zero for boundary-layer KV; its exit depth remains on Request. + depth: int + # None addresses the request's single resident hidden row. Otherwise the + # span addresses [token_start - chunk_start, ...] in its resident chunk. + chunk_start: int | None = None + retain_chunk: bool = False + retain_last: bool = True + emit_gate: bool = False + # None reads the immutable prompt range for chunk PRELUDE; explicit tokens override it. + input_tokens: tuple[int, ...] | None = None + + +@dataclass(frozen=True) +class ExecutionBatch: + """Selected request ranges bound to one model operation and loop boundary.""" + + operation: Operation + spans: tuple[StateSpan, ...] + # Recurrent calls following PRELUDE, or starting at RECURRENT. + loop_count: int = 0 + has_split_prefill_work: bool | None = None + + def __post_init__(self): + if self.loop_count < 0: + raise ValueError("loop_count must be nonnegative") + if self.has_split_prefill_work is None: + object.__setattr__(self, "has_split_prefill_work", self.has_chunks) + elif self.has_chunks and not self.has_split_prefill_work: + raise ValueError("chunk spans require split-prefill work") + if self.operation == Operation.RECURRENT and self.loop_count == 0: + object.__setattr__(self, "loop_count", 1) + if self.operation == Operation.CODA and self.loop_count: + raise ValueError("CODA cannot include recurrent work") + for span in self.spans: + if span.token_start < 0 or span.token_count <= 0 or span.depth < 0: + raise ValueError("state spans require nonnegative addresses and positive rows") + + def at_loop(self, offset: int): + """Bind a loop view without changing request progress or storage ownership.""" + if not 0 <= offset < self.loop_count: + raise ValueError("loop offset outside execution range") + return replace( + self, + operation=Operation.RECURRENT, + loop_count=1, + spans=tuple(replace(span, depth=span.depth + offset) for span in self.spans), + ) + + def validate_depths(self, total_depths: int): + if any(span.depth + self.loop_count > total_depths for span in self.spans): + raise ValueError("execution range exceeds model loop depth") + + @property + def num_rows(self): + return sum(span.token_count for span in self.spans) + + @property + def has_chunks(self): + return any(span.chunk_start is not None for span in self.spans) + + @property + def all_chunks(self): + return bool(self.spans) and all(span.chunk_start is not None for span in self.spans) + + @property + def emit_gate(self): + return any(span.emit_gate for span in self.spans) + + def addresses(self): + ids, depths, positions = [], [], [] + for span in self.spans: + ids.extend([span.request.request_id] * span.token_count) + depths.extend([span.depth] * span.token_count) + positions.extend(range(span.token_start, span.token_start + span.token_count)) + return ids, depths, positions diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 3bb6bf02..244015fb 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -23,7 +23,7 @@ create_backend, validate_backend_name, ) -from vllm_rlt.request import ExecutionBatch +from vllm_rlt.core.execution import ExecutionBatch from vllm_rlt.worker.buffers import stage_metadata from vllm_rlt.worker.metadata import DeviceKVMetadata, KVMetadata, MetadataBuilder diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 1200f218..5bfa185f 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -5,24 +5,27 @@ from enum import Enum, auto from vllm_rlt.config import SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.scheduling_policy import create_scheduling_policy -from vllm_rlt.request import FinishReason, Request, RequestPhase +from vllm_rlt.request import FinishReason, Request, RequestStatus # Items of these phases cover a contiguous span of prompt tokens, so their KV # frontier is the span end rather than the request's decode position. -_SPAN_PHASES = ( - RequestPhase.PREFILL, - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, - RequestPhase.SPECULATIVE, +_SPAN_WORKS = frozenset( + ( + WorkKey.PROMPT_ATOMIC, + WorkKey.PROMPT_PRELUDE, + WorkKey.PROMPT_RECURRENT, + WorkKey.SPECULATIVE, + ) ) -# Stages that must reserve KV capacity for that frontier before selection. -_CAPACITY_PHASES = _SPAN_PHASES + (RequestPhase.PRELUDE, RequestPhase.RECURRENT) +_CAPACITY_WORKS = _SPAN_WORKS | {WorkKey.DECODE_PRELUDE, WorkKey.DECODE_RECURRENT} @dataclass(frozen=True) class ScheduledRequest: request: Request + work: WorkKey # PREFILL and SPECULATIVE use a contiguous span; ordinary decode uses one row. token_start: int = 0 token_count: int = 1 @@ -31,29 +34,33 @@ class ScheduledRequest: depth: int | None = None chunk_start: int | None = None completes_prompt: bool = False - phase: RequestPhase | None = None @dataclass(frozen=True) class SchedulerOutput: - # None means a mixed batch; each scheduled request carries its own phase. - phase: RequestPhase | None requests: list[ScheduledRequest] @property - def is_speculative(self) -> bool: - """Execution-mode flag for this batch, not the engine-wide configuration.""" - return self.phase == RequestPhase.SPECULATIVE + def works(self) -> frozenset[WorkKey]: + return frozenset(item.work for item in self.requests) @property def is_mixed(self) -> bool: - return self.phase is None + return len(self.works) > 1 + + @property + def has_split_prefill_work(self) -> bool: + return any( + item.work in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) + for item in self.requests + ) - def phase_for(self, scheduled_request) -> RequestPhase: - phase = scheduled_request.phase or self.phase - if phase is None: - raise ValueError("mixed batches require a phase on every scheduled request") - return phase + @property + def is_speculative(self) -> bool: + """Execution-mode flag for this batch, not the engine-wide configuration.""" + return bool(self.requests) and all( + item.work == WorkKey.SPECULATIVE for item in self.requests + ) @property def num_tokens(self) -> int: @@ -115,7 +122,9 @@ def __init__( self.cache_manager = cache_manager self.speculative_config = speculative_config self.requests: dict[str, Request] = {} - self.queues: dict[RequestPhase, deque[str]] = {s: deque() for s in RequestPhase} + self.waiting: deque[str] = deque() + self.work_queues: dict[WorkKey, deque[str]] = {work: deque() for work in WorkKey} + self.admission_slots: set[str] = set() # Bound by Engine when preemption is enabled. These callbacks can copy # device state; queue transitions remain owned by this scheduler. self.suspend_callback = None @@ -132,36 +141,110 @@ def __init__( def add_request(self, request: Request): if request.request_id in self.requests: raise ValueError(f"duplicate request ID: {request.request_id}") + request.status = RequestStatus.WAITING self.requests[request.request_id] = request - self.queues[RequestPhase.WAITING].append(request.request_id) + self.waiting.append(request.request_id) + + def mark_receiving(self, request: Request): + if self.requests.get(request.request_id) is not request: + raise ValueError("only the registered request can enter receiving") + if request.status != RequestStatus.WAITING: + raise ValueError("only a waiting request can enter receiving") + self._remove_queued_request(request) + request.status = RequestStatus.RECEIVING + + def queue_for_admission(self, request: Request): + if self.requests.get(request.request_id) is not request: + raise ValueError("only the registered request can wait for admission") + if request.status not in (RequestStatus.WAITING, RequestStatus.RECEIVING): + raise ValueError("only waiting or receiving requests can await admission") + if request.request_id not in self.waiting: + self.waiting.append(request.request_id) + + def activate_request(self, request: Request): + """Mark a resource-ready request active and reserve its compute slot.""" + if self.requests.get(request.request_id) is not request: + raise ValueError("only the registered request can be activated") + if request.status not in (RequestStatus.WAITING, RequestStatus.RECEIVING): + raise ValueError("only waiting or receiving requests can be activated") + if not self.has_admission_capacity(): + raise RuntimeError("no admission slot is available") + request.status = RequestStatus.ACTIVE + self.admission_slots.add(request.request_id) + + def has_admission_capacity(self) -> bool: + return len(self.admission_slots) < self.config.max_num_seqs + + def has_admission_slot(self, request: Request) -> bool: + return ( + self.requests.get(request.request_id) is request + and request.request_id in self.admission_slots + ) - def enqueue(self, request: Request, phase: RequestPhase): - request.phase = phase - self.queues[phase].append(request.request_id) + def release_admission_slot(self, request: Request): + if self.requests.get(request.request_id) is not request: + return False + self.admission_slots.discard(request.request_id) + return True + + def enqueue_work(self, request: Request, work: WorkKey): + if self.requests.get(request.request_id) is not request: + raise ValueError("only the registered request can enter a work queue") + if request.status != RequestStatus.ACTIVE: + raise ValueError("only active requests can enter a work queue") + if self.pending_work(request) is not None: + raise ValueError(f"request {request.request_id!r} already has pending work") + self.work_queues[work].append(request.request_id) + + def pending_work(self, request: Request) -> WorkKey | None: + for work, queue in self.work_queues.items(): + if request.request_id in queue: + return work + return None + + def remove_pending_work(self, request: Request) -> WorkKey | None: + pending = None + for work, queue in self.work_queues.items(): + if request.request_id in queue: + pending = work if pending is None else pending + while request.request_id in queue: + queue.remove(request.request_id) + return pending + + def _remove_queued_request(self, request: Request): + while request.request_id in self.waiting: + self.waiting.remove(request.request_id) + self.remove_pending_work(request) + + def park_waiting(self, request: Request): + self._remove_queued_request(request) + self.release_admission_slot(request) + request.status = RequestStatus.WAITING + self.waiting.append(request.request_id) def suspend_for_transfer(self, request): """Stop scheduling without freeing resources that a transfer still borrows.""" if self.requests.get(request.request_id) is not request: return False - for queue in self.queues.values(): - while request.request_id in queue: - queue.remove(request.request_id) - request.phase = RequestPhase.TRANSFERRING + self._remove_queued_request(request) + request.status = RequestStatus.TRANSFERRING return True def finish(self, request: Request, reason: FinishReason): - # A delayed EOS can stop a request already queued for its next phase. - for queue in self.queues.values(): - while request.request_id in queue: - queue.remove(request.request_id) + self.mark_finished(request, reason) + self.release_admission_slot(request) self.cache_manager.free(request.request_id) - request.phase = RequestPhase.FINISHED - request.finish_reason = reason request.hidden_state = None request.input_token_tensor = None request.num_output_placeholders = 0 request.num_speculative_rounds = 0 - self.requests.pop(request.request_id) + self.requests.pop(request.request_id, None) + + def mark_finished(self, request: Request, reason: FinishReason): + """Mark logical completion while a runner may still own device resources.""" + self._remove_queued_request(request) + request.status = RequestStatus.FINISHED + request.finish_reason = reason def abort(self, request_id: str) -> Request: request = self.requests[request_id] @@ -179,7 +262,7 @@ def _order_waiting_requests(self) -> deque[str]: Lower values go first: -5, 0, 10. Equal values retain current queue order. FCFS preserves queue order, but admission may still bypass blocked requests. """ - waiting = self.queues[RequestPhase.WAITING] + waiting = self.waiting if self.config.policy == "priority": ordered = sorted(waiting, key=lambda rid: self.requests[rid].sampling_params.priority) waiting.clear() @@ -195,7 +278,9 @@ def _is_preemption_candidate(self, request, requesting_request, *, priority_only """ if ( request is requesting_request - or request.phase in (RequestPhase.WAITING, RequestPhase.RECEIVING) + or request.status != RequestStatus.ACTIVE + or request.request_id not in self.admission_slots + or self.pending_work(request) is None or request.num_output_placeholders or request.request_id in excluded ): @@ -240,36 +325,25 @@ def preempt(self, requesting_request, *, priority_only=False, excluded=frozenset if victim is None: return False self.suspend_callback(victim) - for queue in self.queues.values(): - while victim.request_id in queue: - queue.remove(victim.request_id) - self.enqueue(victim, RequestPhase.WAITING) + self.park_waiting(victim) return True - def _ensure_active_slot(self, request: Request, active_count: int) -> int | None: - """Return the updated active count, or None when admission must stop. - - A full engine may make room for a higher-priority arrival. Scheduler - suspends a selected victim and moves it to WAITING before admitting the - arrival. This operation can synchronize the runner and copy KV to CPU. - Admission runs before any batch is selected, so no request is excluded - from preemption here. - """ - if active_count < self.config.max_num_seqs: - return active_count - if ( + def _ensure_admission_slot(self, request: Request) -> bool: + """Make capacity for one admission, preempting only when policy allows.""" + if self.has_admission_capacity(): + return True + preempted = ( self.config.policy == "priority" and self.suspend_callback is not None and self.preempt(request, priority_only=True, excluded=frozenset()) - ): - return active_count - 1 - return None + ) + return bool(preempted and self.has_admission_capacity()) def _try_resume(self, request: Request) -> ResumeResult: """Translate the existing restoration callback's three-way result. None: no snapshot, continue fresh admission. - RequestPhase: resources/state restored; enqueue the saved phase. + WorkKey: resources/state restored; enqueue the saved work. False: snapshot exists, but restoration cannot obtain capacity yet. """ restored = self.resume_callback(request) if self.resume_callback is not None else None @@ -277,7 +351,8 @@ def _try_resume(self, request: Request) -> ResumeResult: return ResumeResult.NOT_PREEMPTED if restored is False: return ResumeResult.BLOCKED - self.enqueue(request, restored) + self.activate_request(request) + self.enqueue_work(request, restored) return ResumeResult.RESTORED def _reserved_growth_blocks(self, *, reserve_outputs: bool) -> int: @@ -291,14 +366,13 @@ def _reserved_growth_blocks(self, *, reserve_outputs: bool) -> int: cache = self.cache_manager reserved = 0 for request in self.requests.values(): - # Imported/reserved requests already own pages before compute admission. - # Without preemption their future decode growth must remain affordable too. + # Progress, not queue state, remains valid while a prompt ticket is in flight. + active_prompt = ( + request.status == RequestStatus.ACTIVE + and request.num_prefilled_tokens < len(request.prompt_token_ids) + ) has_allocation = cache.has_allocation(request.request_id) - if request.phase not in ( - RequestPhase.PREFILL, - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, - ) and not (reserve_outputs and has_allocation): + if not active_prompt and not (reserve_outputs and has_allocation): continue tokens = len(request.prompt_token_ids) if reserve_outputs: @@ -307,7 +381,9 @@ def _reserved_growth_blocks(self, *, reserve_outputs: bool) -> int: reserved += max(0, cache.required_blocks(tokens) - allocated) return reserved - def _plan_admission(self, request: Request, active_count: int) -> AdmissionPlan: + def _plan_admission( + self, request: Request, *, has_allocated_requests: bool | None = None + ) -> AdmissionPlan: """Describe prefix reuse, physical allocation and logical admission cost. With 8 prompt tokens and 5 outputs, capacity is 12: the last sampled @@ -336,7 +412,12 @@ def _plan_admission(self, request: Request, active_count: int) -> AdmissionPlan: required = cache.required_blocks(budget_tokens) - len(prefix) * cache.storage_depths reserved = self._reserved_growth_blocks(reserve_outputs=reserve_outputs) # Allow an otherwise idle engine to start a large request. - admission_headroom_blocks = cache.watermark_blocks if active_count else 0 + has_occupants = ( + bool(self.admission_slots) + if has_allocated_requests is None + else has_allocated_requests + ) + admission_headroom_blocks = cache.watermark_blocks if has_occupants else 0 # num_free_blocks includes evictable prefix-only blocks. Reusing those # blocks makes them non-evictable: do not count the same capacity twice. # Example: free=8 fresh + 4 cached, demand=12 with a 4-block prefix. @@ -361,7 +442,7 @@ def reserve_transfer(self, request: Request) -> AdmissionPlan | None: This operation grants no compute slot and leaves queue ownership unchanged. """ occupied = any(self.cache_manager.has_allocation(rid) for rid in self.requests) - plan = self._plan_admission(request, int(occupied)) + plan = self._plan_admission(request, has_allocated_requests=occupied) if self.cache_manager.incremental_allocation: plan = replace(plan, initial_tokens=len(request.prompt_token_ids)) if not self._allocate(request, plan): @@ -399,10 +480,11 @@ def _commit_admission(self, request: Request, plan: AdmissionPlan) -> bool: return True def _enter_prefill(self, request): + self.activate_request(request) if self.atomic_prefill: - self.enqueue(request, RequestPhase.PREFILL) + self.enqueue_work(request, WorkKey.PROMPT_ATOMIC) else: - self.enqueue(request, RequestPhase.PREFILLING_PRELUDE) + self.enqueue_work(request, WorkKey.PROMPT_PRELUDE) def _admit(self): """Admit waiting requests: order -> slot -> restore/allocate -> defer/enter. @@ -411,10 +493,6 @@ def _admit(self): admission increments earlier blocked requests' bypass counts; reaching the limit stops further bypasses so active requests can drain. """ - active_count = sum( - r.phase not in (RequestPhase.WAITING, RequestPhase.RECEIVING) and not r.compute_released - for r in self.requests.values() - ) waiting = self._order_waiting_requests() deferred = [] scan_count = min(len(waiting), self.config.admission_scan_limit) @@ -423,18 +501,17 @@ def _admit(self): break # Peek before popping: preemption may append its victim to WAITING. request = self.requests[waiting[0]] - available_count = self._ensure_active_slot(request, active_count) - if available_count is None: + if not self._ensure_admission_slot(request): break - active_count = available_count request_id = waiting.popleft() request = self.requests[request_id] - if request.phase == RequestPhase.RECEIVING: + if request.status == RequestStatus.RECEIVING: # Transfer completion has already made KV and hidden valid. # Imported and local requests compete in this same admission queue. self.cache_manager.unpin_transfer(request_id, request_id) - self.enqueue(request, RequestPhase.CODA) + self.activate_request(request) + self.enqueue_work(request, WorkKey.DECODE_CODA) elif self.cache_manager.has_allocation( request_id ) and self.cache_manager.has_transfer_lease(request_id): @@ -443,20 +520,18 @@ def _admit(self): else: resume_result = self._try_resume(request) if resume_result == ResumeResult.RESTORED: - active_count += 1 continue if resume_result == ResumeResult.BLOCKED: deferred.append(request_id) continue - plan = self._plan_admission(request, active_count) + plan = self._plan_admission(request) if not self._commit_admission(request, plan): deferred.append(request_id) if request.admission_bypasses >= self.config.max_admission_bypasses: break continue - active_count += 1 for blocked_id in deferred: self.requests[blocked_id].admission_bypasses += 1 if any( @@ -467,34 +542,33 @@ def _admit(self): # [A, B] deferred must precede the untouched tail in the same order. waiting.extendleft(reversed(deferred)) - def _make_scheduled_request( - self, request: Request, phase: RequestPhase, token_budget: int + def _make_scheduled_work( + self, request: Request, work: WorkKey, token_budget: int ) -> ScheduledRequest: - """Snapshot one request's execution range without advancing its progress. + """Snapshot one work item's execution range without advancing progress. Prompt length 10, prefilled 4, chunk 4, budget 6 -> range [4, 8). - Decode phases contribute one position; its loop may differ by request. + Decode work contributes one position; recurrence depth varies per request. """ - if phase == RequestPhase.PREFILL: + if work == WorkKey.PROMPT_ATOMIC: start = request.num_prefilled_tokens count = min( token_budget, self.config.prefill_chunk_size, len(request.prompt_token_ids) - start ) - return ScheduledRequest(request, start, count) - # The two split phases are unreachable while atomic prefill is enabled. - if phase == RequestPhase.PREFILLING_PRELUDE: + return ScheduledRequest(request, work, start, count) + if work == WorkKey.PROMPT_PRELUDE: start = request.num_preluded_tokens count = min( token_budget, self.config.prefill_chunk_size, len(request.prompt_token_ids) - start, ) - return ScheduledRequest(request, start, count, chunk_start=start) - if phase == RequestPhase.PREFILLING_RECURRENT: + return ScheduledRequest(request, work, start, count, chunk_start=start) + if work == WorkKey.PROMPT_RECURRENT: 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, work, start, 0, depth=0, chunk_start=request.prefill_chunk_start ) first_token_depth = request.prefill_depth(start) # One run per depth frontier: every token of the run already completed @@ -511,6 +585,7 @@ def _make_scheduled_request( ) return ScheduledRequest( request, + work, start, count, depth=first_token_depth, @@ -520,15 +595,15 @@ def _make_scheduled_request( and start + count == len(request.prompt_token_ids) ), ) - if phase == RequestPhase.SPECULATIVE: + if work == WorkKey.SPECULATIVE: # K candidates plus one bonus distribution. At the output limit, # K=0 is an ordinary fixed-depth step and needs no extra KV slot. remaining = request.sampling_params.max_tokens - request.num_scheduled_outputs count = min(self.speculative_config.num_speculative_tokens + 1, remaining, token_budget) - return ScheduledRequest(request, request.position, count) - if phase == RequestPhase.RECURRENT: - return ScheduledRequest(request, request.position, depth=request.loops_done) - return ScheduledRequest(request) + return ScheduledRequest(request, work, request.position, count) + if work == WorkKey.DECODE_RECURRENT: + return ScheduledRequest(request, work, request.position, depth=request.loops_done) + return ScheduledRequest(request, work) def _ensure_execution_capacity( self, request: Request, frontier: int, *, selected: set[str] @@ -558,24 +633,24 @@ def _prelude_window_open(self, request: Request) -> bool: """ return request.num_preluded_tokens == request.num_prefilled_tokens - def has_ready(self, phase): - """A blocked transport queue must not hide another runnable phase.""" - queue = self.queues[phase] + def has_ready_work(self, work: WorkKey): + """A blocked transport queue must not hide another runnable work type.""" + queue = self.work_queues[work] if self.execution_ready is None: return bool(queue) return any(self.execution_ready(self.requests[rid]) for rid in queue) - def _take(self, phase: RequestPhase) -> SchedulerOutput | None: - return self.take_phases((phase,)) + def has_ready_works(self, *works: WorkKey) -> bool: + return any(self.has_ready_work(work) for work in works) - def take_phases(self, phases, *, eligible=None): + def take_works(self, works, *, eligible=None): """Select compatible requests across queues under one budget/reservation set.""" eligible = eligible if eligible is not None else self.execution_ready token_budget = self.config.max_num_batched_tokens scheduled_requests = [] selected: set[str] = set() - for phase in phases: - queue = self.queues[phase] + for work in works: + queue = self.work_queues[work] remaining = len(queue) while ( queue @@ -585,33 +660,35 @@ def take_phases(self, phases, *, eligible=None): ): remaining -= 1 request = self.requests[queue.popleft()] + if request.request_id in selected: + continue 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( + if work == WorkKey.PROMPT_PRELUDE and not self._prelude_window_open( request ): queue.append(request.request_id) continue - if phase == RequestPhase.SPECULATIVE and ( + if work == WorkKey.SPECULATIVE and ( request.num_speculative_rounds >= 2 or request.num_scheduled_outputs >= request.sampling_params.max_tokens ): queue.append(request.request_id) continue - scheduled_request = self._make_scheduled_request(request, phase, token_budget) + scheduled_request = self._make_scheduled_work(request, work, token_budget) if ( - phase in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + work in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) and scheduled_request.token_count <= 0 ): queue.append(request.request_id) continue - if phase in _CAPACITY_PHASES: + if work in _CAPACITY_WORKS: # A scheduled span declares a prompt frontier; decode phases advance the # request's own position by one. frontier = ( scheduled_request.token_start + scheduled_request.token_count - if phase in _SPAN_PHASES + if work in _SPAN_WORKS else request.position + 1 ) if not self._ensure_execution_capacity(request, frontier, selected=selected): @@ -619,22 +696,14 @@ def take_phases(self, phases, *, eligible=None): continue # Selecting A must prevent B's later capacity check from evicting A. selected.add(request.request_id) - scheduled_requests.append(replace(scheduled_request, phase=phase)) + scheduled_requests.append(scheduled_request) token_budget -= scheduled_request.token_count if not scheduled_requests: return 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} - recorded = ( - RequestPhase.RECURRENT - if RequestPhase.RECURRENT in selected_phases - else scheduled_requests[0].phase - ) - self.policy.record_batch(recorded) - return SchedulerOutput( - next(iter(selected_phases)) if len(selected_phases) == 1 else None, scheduled_requests - ) + self.policy.record_batch(item.work for item in scheduled_requests) + return SchedulerOutput(scheduled_requests) def schedule(self, *, prefer_recurrent=False) -> SchedulerOutput | None: """Choose queued requests or an admission opportunity, then build a batch. diff --git a/vllm_rlt/core/scheduling_policy.py b/vllm_rlt/core/scheduling_policy.py index feba1174..8e2cd86f 100644 --- a/vllm_rlt/core/scheduling_policy.py +++ b/vllm_rlt/core/scheduling_policy.py @@ -1,10 +1,10 @@ -"""Request-phase selection and fairness state for loop-level scheduling.""" +"""Work selection and fairness state for loop-level scheduling.""" from enum import Enum, auto from typing import TYPE_CHECKING from vllm_rlt.config import SchedulerConfig, SpeculativeConfig -from vllm_rlt.request import RequestPhase +from vllm_rlt.core.execution import WorkKey if TYPE_CHECKING: from vllm_rlt.core.scheduler import Scheduler, SchedulerOutput @@ -16,9 +16,16 @@ class NoRefillPhase(Enum): CODA = auto() +PROMPT_WORKS = frozenset( + (WorkKey.PROMPT_ATOMIC, WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) +) +DECODE_PROGRESS_WORKS = frozenset((WorkKey.DECODE_RECURRENT, WorkKey.SPECULATIVE)) +PRELUDE_WORKS = (WorkKey.DECODE_PRELUDE, WorkKey.PROMPT_PRELUDE) +RECURRENT_WORKS = (WorkKey.DECODE_RECURRENT, WorkKey.PROMPT_RECURRENT) + + class SchedulingPolicy: - # Prompt phases supported by this policy also count against the decode bound. - PREFILL_PHASES = (RequestPhase.PREFILL,) + PREFILL_WORKS = (WorkKey.PROMPT_ATOMIC,) def __init__(self, config: SchedulerConfig): self.config = config @@ -28,12 +35,24 @@ def __init__(self, config: SchedulerConfig): def decode_due(self) -> bool: return self.prefill_batches_since_recurrent >= self.config.max_prefill_batches_before_decode - def record_batch(self, phase: RequestPhase) -> None: - """Count nonempty scheduled batches, not attempts or GPU completions.""" - if phase in self.PREFILL_PHASES: - self.prefill_batches_since_recurrent += 1 - elif phase in (RequestPhase.RECURRENT, RequestPhase.SPECULATIVE): + def work_group(self, work: WorkKey) -> tuple[WorkKey, ...]: + return (work,) + + def has_ready_exact_work(self, scheduler, work: WorkKey) -> bool: + return scheduler.has_ready_work(work) + + def has_ready_work(self, scheduler, work: WorkKey) -> bool: + return scheduler.has_ready_works(*self.work_group(work)) + + def take_work(self, scheduler, work: WorkKey): + return scheduler.take_works(self.work_group(work)) + + def record_batch(self, works) -> None: + works = frozenset(works) + if works & DECODE_PROGRESS_WORKS: self.prefill_batches_since_recurrent = 0 + elif works & PROMPT_WORKS: + self.prefill_batches_since_recurrent += 1 class PromptOnlyPolicy(SchedulingPolicy): @@ -41,16 +60,16 @@ class PromptOnlyPolicy(SchedulingPolicy): def __init__(self, config, *, atomic_prefill): super().__init__(config) - self.PREFILL_PHASES = ( - (RequestPhase.PREFILL,) + self.PREFILL_WORKS = ( + (WorkKey.PROMPT_ATOMIC,) if atomic_prefill - else (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + else (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) ) def schedule(self, scheduler, *, prefer_recurrent=False): scheduler._admit() - for phase in self.PREFILL_PHASES: - batch = scheduler.take_phases((phase,)) + for work in self.PREFILL_WORKS: + batch = self.take_work(scheduler, work) if batch is not None: return batch return None @@ -60,32 +79,32 @@ class RefillPolicy(SchedulingPolicy): def schedule( self, scheduler: "Scheduler", *, prefer_recurrent: bool = False ) -> "SchedulerOutput | None": - # Give independent core work an overlap opportunity before refilling. - if prefer_recurrent and scheduler.has_ready(RequestPhase.RECURRENT): - return scheduler._take(RequestPhase.RECURRENT) - if scheduler.has_ready(RequestPhase.PRELUDE): - return scheduler._take(RequestPhase.PRELUDE) - if scheduler.has_ready(RequestPhase.CODA) and ( - len(scheduler.queues[RequestPhase.CODA]) >= self.config.min_coda_batch_size - or not scheduler.has_ready(RequestPhase.RECURRENT) + if prefer_recurrent and self.has_ready_exact_work(scheduler, WorkKey.DECODE_RECURRENT): + return self.take_work(scheduler, WorkKey.DECODE_RECURRENT) + if self.has_ready_exact_work(scheduler, WorkKey.DECODE_PRELUDE): + return self.take_work(scheduler, WorkKey.DECODE_PRELUDE) + if self.has_ready_exact_work(scheduler, WorkKey.DECODE_CODA) and ( + len(scheduler.work_queues[WorkKey.DECODE_CODA]) >= self.config.min_coda_batch_size + or not self.has_ready_exact_work(scheduler, WorkKey.DECODE_RECURRENT) ): - return scheduler._take(RequestPhase.CODA) - if self.decode_due and scheduler.has_ready(RequestPhase.RECURRENT): - return scheduler._take(RequestPhase.RECURRENT) + return self.take_work(scheduler, WorkKey.DECODE_CODA) + if self.decode_due and self.has_ready_exact_work(scheduler, WorkKey.DECODE_RECURRENT): + return self.take_work(scheduler, WorkKey.DECODE_RECURRENT) + scheduler._admit() - for phase in self.PREFILL_PHASES: - if scheduler.has_ready(phase): - return scheduler._take(phase) - # Admission can restore a snapshot directly to a boundary phase. - if scheduler.has_ready(RequestPhase.PRELUDE): - return scheduler._take(RequestPhase.PRELUDE) - if scheduler.has_ready(RequestPhase.CODA) and ( - len(scheduler.queues[RequestPhase.CODA]) >= self.config.min_coda_batch_size - or not scheduler.has_ready(RequestPhase.RECURRENT) + for work in self.PREFILL_WORKS: + if self.has_ready_work(scheduler, work): + return self.take_work(scheduler, work) + # Admission can restore a request directly to a boundary work queue. + if self.has_ready_exact_work(scheduler, WorkKey.DECODE_PRELUDE): + return self.take_work(scheduler, WorkKey.DECODE_PRELUDE) + if self.has_ready_exact_work(scheduler, WorkKey.DECODE_CODA) and ( + len(scheduler.work_queues[WorkKey.DECODE_CODA]) >= self.config.min_coda_batch_size + or not self.has_ready_exact_work(scheduler, WorkKey.DECODE_RECURRENT) ): - return scheduler._take(RequestPhase.CODA) - if scheduler.has_ready(RequestPhase.RECURRENT): - return scheduler._take(RequestPhase.RECURRENT) + return self.take_work(scheduler, WorkKey.DECODE_CODA) + if self.has_ready_exact_work(scheduler, WorkKey.DECODE_RECURRENT): + return self.take_work(scheduler, WorkKey.DECODE_RECURRENT) return None @@ -98,30 +117,29 @@ def schedule( self, scheduler: "Scheduler", *, prefer_recurrent: bool = False ) -> "SchedulerOutput | None": if self.phase == NoRefillPhase.CORE: - if scheduler.has_ready(RequestPhase.RECURRENT): - return scheduler._take(RequestPhase.RECURRENT) + if self.has_ready_work(scheduler, WorkKey.DECODE_RECURRENT): + return self.take_work(scheduler, WorkKey.DECODE_RECURRENT) self.phase = NoRefillPhase.CODA if self.phase == NoRefillPhase.CODA: - if scheduler.has_ready(RequestPhase.CODA): - return scheduler._take(RequestPhase.CODA) + if self.has_ready_work(scheduler, WorkKey.DECODE_CODA): + return self.take_work(scheduler, WorkKey.DECODE_CODA) self.phase = NoRefillPhase.FILL - # Bound prefill runs so continuous arrivals cannot starve decode. if self.decode_due: - if scheduler.has_ready(RequestPhase.PRELUDE): - return scheduler._take(RequestPhase.PRELUDE) - if scheduler.has_ready(RequestPhase.RECURRENT): + if self.has_ready_work(scheduler, WorkKey.DECODE_PRELUDE): + return self.take_work(scheduler, WorkKey.DECODE_PRELUDE) + if self.has_ready_work(scheduler, WorkKey.DECODE_RECURRENT): self.phase = NoRefillPhase.CORE - return scheduler._take(RequestPhase.RECURRENT) + return self.take_work(scheduler, WorkKey.DECODE_RECURRENT) scheduler._admit() - if scheduler.has_ready(RequestPhase.PREFILL): - return scheduler._take(RequestPhase.PREFILL) - if scheduler.has_ready(RequestPhase.CODA): # First output after full-depth prompt prefill. - return scheduler._take(RequestPhase.CODA) - if scheduler.has_ready(RequestPhase.PRELUDE): - return scheduler._take(RequestPhase.PRELUDE) - if scheduler.has_ready(RequestPhase.RECURRENT): + if self.has_ready_work(scheduler, WorkKey.PROMPT_ATOMIC): + return self.take_work(scheduler, WorkKey.PROMPT_ATOMIC) + if self.has_ready_work(scheduler, WorkKey.DECODE_CODA): + return self.take_work(scheduler, WorkKey.DECODE_CODA) + if self.has_ready_work(scheduler, WorkKey.DECODE_PRELUDE): + return self.take_work(scheduler, WorkKey.DECODE_PRELUDE) + if self.has_ready_work(scheduler, WorkKey.DECODE_RECURRENT): self.phase = NoRefillPhase.CORE - return scheduler._take(RequestPhase.RECURRENT) + return self.take_work(scheduler, WorkKey.DECODE_RECURRENT) return None @@ -130,46 +148,43 @@ class SpeculativePolicy(SchedulingPolicy): def __init__(self, config, *, atomic_prefill=True): super().__init__(config) - self.PREFILL_PHASES = ( - (RequestPhase.PREFILL,) + self.PREFILL_WORKS = ( + (WorkKey.PROMPT_ATOMIC,) if atomic_prefill - else (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + else (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) ) def schedule(self, scheduler, *, prefer_recurrent=False): - if scheduler.has_ready(RequestPhase.CODA): - return scheduler._take(RequestPhase.CODA) - if self.decode_due and scheduler.has_ready(RequestPhase.SPECULATIVE): - return scheduler._take(RequestPhase.SPECULATIVE) + if self.has_ready_exact_work(scheduler, WorkKey.DECODE_CODA): + return self.take_work(scheduler, WorkKey.DECODE_CODA) + if self.decode_due and self.has_ready_exact_work(scheduler, WorkKey.SPECULATIVE): + return self.take_work(scheduler, WorkKey.SPECULATIVE) scheduler._admit() - # Admission can import a complete prompt directly into CODA. - if scheduler.has_ready(RequestPhase.CODA): - return scheduler._take(RequestPhase.CODA) - for phase in self.PREFILL_PHASES: - if scheduler.has_ready(phase): - return scheduler._take(phase) - if scheduler.has_ready(RequestPhase.SPECULATIVE): - return scheduler._take(RequestPhase.SPECULATIVE) + if self.has_ready_exact_work(scheduler, WorkKey.DECODE_CODA): + return self.take_work(scheduler, WorkKey.DECODE_CODA) + for work in self.PREFILL_WORKS: + if self.has_ready_work(scheduler, work): + return self.take_work(scheduler, work) + if self.has_ready_exact_work(scheduler, WorkKey.SPECULATIVE): + return self.take_work(scheduler, WorkKey.SPECULATIVE) return None class SplitPrefillPolicy(RefillPolicy): - """Refill ordering where each recurrence depth is its own batch boundary.""" + """Refill ordering where each prompt recurrence depth is its own batch.""" - PREFILL_PHASES = (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) + PREFILL_WORKS = (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) class MixedPolicy(SplitPrefillPolicy): - """Choose compatible prompt/generation loops; execution is policy-agnostic.""" + """Combine compatible prompt and decode work while sharing refill ordering.""" - def schedule(self, scheduler, *, prefer_recurrent=False): - if scheduler.has_ready(RequestPhase.RECURRENT) and scheduler.has_ready( - RequestPhase.PREFILLING_RECURRENT - ): - return scheduler.take_phases( - (RequestPhase.RECURRENT, RequestPhase.PREFILLING_RECURRENT) - ) - return super().schedule(scheduler, prefer_recurrent=prefer_recurrent) + def work_group(self, work: WorkKey) -> tuple[WorkKey, ...]: + if work in PRELUDE_WORKS: + return PRELUDE_WORKS + if work in RECURRENT_WORKS: + return RECURRENT_WORKS + return (work,) def create_scheduling_policy( @@ -179,16 +194,9 @@ def create_scheduling_policy( *, produces_tokens: bool = True, ) -> SchedulingPolicy: - """Pick one policy; reject combinations no policy implements yet. - - Failing here keeps an unimplemented mode from silently changing - semantics or returning an empty schedule forever. - """ - # Prompt-only phases do not use prompt/generation interleaving rules. + """Pick one policy and reject combinations no policy implements.""" if not produces_tokens: return PromptOnlyPolicy(scheduler_config, atomic_prefill=atomic_prefill) - - # 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") diff --git a/vllm_rlt/engine/connector.py b/vllm_rlt/engine/connector.py index 9c56346b..67ec01c4 100644 --- a/vllm_rlt/engine/connector.py +++ b/vllm_rlt/engine/connector.py @@ -2,8 +2,9 @@ from typing import Protocol +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.scheduler import PrefillProgress -from vllm_rlt.request import Request, RequestPhase +from vllm_rlt.request import Request class StageConnector(Protocol): @@ -19,7 +20,7 @@ def prefill_submitted( self, request: Request, *, - phase: RequestPhase, + work: WorkKey, start: int | None, progress: PrefillProgress, event, diff --git a/vllm_rlt/engine/preemption.py b/vllm_rlt/engine/preemption.py index 427a2aa6..f425ece5 100644 --- a/vllm_rlt/engine/preemption.py +++ b/vllm_rlt/engine/preemption.py @@ -23,11 +23,14 @@ def discard_snapshot(self, request_id: str) -> None: def suspend(self, victim): """Save the selected request's resources; Scheduler owns its queue transition.""" e = self.engine + work = e.scheduler.pending_work(victim) + if work is None: + raise ValueError("a preemption victim must have pending work") e.model_runner.synchronize() - e._rewind_split_prefill(victim) + work = e._rewind_split_prefill(victim, work) cache = e.cache_manager snapshot = dict( - phase=victim.phase, + work=work, kv=cache.snapshot(victim.request_id), hidden=None if victim.hidden_state is None else victim.hidden_state.cpu().clone(), token=None @@ -44,7 +47,7 @@ def suspend(self, victim): self.preemptions += 1 def resume(self, request): - """Return the saved phase, None without a snapshot, or False when blocked.""" + """Return saved work, None without a snapshot, or False when blocked.""" state = self.snapshots.get(request.request_id) if state is None: return None @@ -63,4 +66,4 @@ def resume(self, request): torch.cuda.current_stream(cache.device).synchronize() del self.snapshots[request.request_id] self.resumptions += 1 - return state["phase"] + return state["work"] diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 30258e13..9f313dc6 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -3,6 +3,7 @@ from vllm_rlt.attention import BackendCapabilities, create_backend from vllm_rlt.config import CacheConfig, ExecutionConfig, ExitConfig, SchedulerConfig +from vllm_rlt.core.execution import WorkKey from vllm_rlt.core.kv_cache_manager import KVCacheManager from vllm_rlt.core.memory import plan_cache from vllm_rlt.core.scheduler import PrefillProgress, Scheduler @@ -10,7 +11,7 @@ from vllm_rlt.engine.preemption import PreemptionManager from vllm_rlt.engine.stage_config import StageConfig from vllm_rlt.profiling import Profiler -from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase +from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestStatus from vllm_rlt.sampling_params import SamplingParams from vllm_rlt.worker.executor import Executor from vllm_rlt.worker.model_runner import ModelRunner @@ -321,7 +322,7 @@ def request_ids(self): def has_request(self, request_id): request = self.scheduler.requests.get(request_id) - return request is not None and request.phase != RequestPhase.FINISHED + return request is not None and request.status != RequestStatus.FINISHED def stats(self): return dict( @@ -352,7 +353,7 @@ def _discard_speculative(self): def abort_request(self, request_id: str) -> RequestOutput: request = self.scheduler.requests[request_id] - if request.phase == RequestPhase.FINISHED: + if request.status == RequestStatus.FINISHED: return RequestOutput.from_request(request) if self.async_speculative: request = self.scheduler.requests[request_id] @@ -461,7 +462,7 @@ def _update_speculative(self, batch, results): request.request_id, scheduled_request.token_start + emitted ) request.loops_done = 0 - self.scheduler.enqueue(request, RequestPhase.SPECULATIVE) + self.scheduler.enqueue_work(request, WorkKey.SPECULATIVE) outputs.append(RequestOutput.from_request(request)) return outputs @@ -475,31 +476,30 @@ def _waiting_for_connector(self): transfer_waiters = { r.request_id for r in self.scheduler.requests.values() - if r.phase in (RequestPhase.RECEIVING, RequestPhase.TRANSFERRING) - or ( - r.request_id in self._connector_blocked - and r.request_id in self.scheduler.queues[r.phase] - ) + if r.status in (RequestStatus.RECEIVING, RequestStatus.TRANSFERRING) + or (r.request_id in self._connector_blocked and self.scheduler.pending_work(r)) } return all( r.request_id in transfer_waiters # An admitted export may retain a compute slot while backpressured, # even before it reaches TRANSFERRING. Queued arrivals must yield too. # A WAITING request absent from its queue remains an invariant error. - or (transfer_waiters and r.request_id in self.scheduler.queues[RequestPhase.WAITING]) + or (transfer_waiters and r.request_id in self.scheduler.waiting) for r in self.scheduler.requests.values() ) - def _prompt_destination(self, request): - exporting = self.connector is not None and self.connector.exports(request) - return RequestPhase.TRANSFERRING if exporting else RequestPhase.CODA + def _schedule_prompt_destination(self, request): + if self.connector is not None and self.connector.exports(request): + self.scheduler.suspend_for_transfer(request) + else: + self.scheduler.enqueue_work(request, WorkKey.DECODE_CODA) - def _apply_prefill_progress(self, batch, scheduled_request): + def _apply_prefill_progress(self, scheduled_request): """One prompt frontier update for local and remotely consumed outputs.""" request = scheduled_request.request - phase = batch.phase_for(scheduled_request) + work = scheduled_request.work event = self.model_runner.events.get(request.request_id) - if phase == RequestPhase.PREFILL: + if work == WorkKey.PROMPT_ATOMIC: request.num_prefilled_tokens += scheduled_request.token_count final = request.num_prefilled_tokens == len(request.prompt_token_ids) self.cache_manager.publish_prefix( @@ -507,20 +507,20 @@ def _apply_prefill_progress(self, batch, scheduled_request): ) if final: request.loops_done = self.model.config.total_ut_steps - self.scheduler.enqueue( - request, self._prompt_destination(request) if final else RequestPhase.PREFILL - ) + self._schedule_prompt_destination(request) + else: + self.scheduler.enqueue_work(request, WorkKey.PROMPT_ATOMIC) progress = PrefillProgress(True, final) start = scheduled_request.token_start else: progress = self._apply_split_prefill_progress( - request, phase, scheduled_request, event=event + request, work, scheduled_request, event=event ) start = scheduled_request.chunk_start if self.connector is not None and self.connector.exports(request): self.connector.prefill_submitted( request, - phase=phase, + work=work, start=start, progress=progress, event=event, @@ -556,8 +556,7 @@ def reserve_transfer( # Reservation owns transport resources, not a compute slot. Outgoing # requests compete with local arrivals in the scheduler's WAITING queue. if receiving: - self.scheduler.queues[RequestPhase.WAITING].remove(request_id) - request.phase = RequestPhase.RECEIVING + self.scheduler.mark_receiving(request) return request def suspend_transfer(self, request): @@ -572,13 +571,13 @@ def release_transfer_compute(self, request, *, ready_event): """ if self.scheduler.requests.get(request.request_id) is not request: return False - if request.phase != RequestPhase.TRANSFERRING: + if request.status != RequestStatus.TRANSFERRING: raise ValueError("request is not awaiting transfer completion") if not ready_event.query(): return False self.model_runner.release(request.request_id) request.hidden_state = None - request.compute_released = True + self.scheduler.release_admission_slot(request) return True def release_transfer(self, request): @@ -587,7 +586,7 @@ def release_transfer(self, request): return False rid = request.request_id if self.async_speculative: - if request.phase != RequestPhase.FINISHED: + if request.status != RequestStatus.FINISHED: self._finish(request, FinishReason.ABORT) self._retire_speculative() if self.scheduler.requests.get(rid) is request: @@ -608,8 +607,8 @@ def accept_import(self, request, hidden, *, cached_tokens=0): if self.scheduler.requests.get(request.request_id) is not request: return False if ( - request.phase != RequestPhase.RECEIVING - or request.request_id in self.scheduler.queues[RequestPhase.WAITING] + request.status != RequestStatus.RECEIVING + or request.request_id in self.scheduler.waiting ): raise ValueError("request is not awaiting an import") count = len(request.prompt_token_ids) @@ -620,11 +619,11 @@ def accept_import(self, request, hidden, *, cached_tokens=0): request.hidden_state = hidden # Keep RECEIVING until the scheduler grants a compute slot. The lease # protects the imported allocation while it waits alongside local work. - self.scheduler.queues[RequestPhase.WAITING].append(request.request_id) + self.scheduler.queue_for_admission(request) return True def _apply_split_prefill_progress( - self, request, phase, scheduled_request, *, event=None + self, request, work, scheduled_request, *, event=None ) -> PrefillProgress: """Advance host-side logical progress for a submitted prefill request. @@ -635,13 +634,12 @@ def _apply_split_prefill_progress( start = scheduled_request.token_start end = start + scheduled_request.token_count - if phase == RequestPhase.PREFILLING_PRELUDE: + if work == WorkKey.PROMPT_PRELUDE: request.prefill_chunk_start = start request.num_preluded_tokens = end - self.scheduler.enqueue(request, RequestPhase.PREFILLING_RECURRENT) + self.scheduler.enqueue_work(request, WorkKey.PROMPT_RECURRENT) return PrefillProgress(False, False) - # Assertion: stage == RequestPhase.PREFILLING_RECURRENT depth = scheduled_request.depth boundaries = request.prefill_depth_ranges boundaries[depth][0] = end @@ -652,7 +650,7 @@ def _apply_split_prefill_progress( chunk_complete = end == request.num_preluded_tokens if depth + 1 < total_depths or not chunk_complete: - self.scheduler.enqueue(request, RequestPhase.PREFILLING_RECURRENT) + self.scheduler.enqueue_work(request, WorkKey.PROMPT_RECURRENT) return PrefillProgress(False, False) self.cache_manager.publish_prefix( @@ -665,19 +663,15 @@ def _apply_split_prefill_progress( prompt_complete = end == len(request.prompt_token_ids) if prompt_complete: request.loops_done = total_depths - next_stage = self._prompt_destination(request) + self._schedule_prompt_destination(request) else: - next_stage = RequestPhase.PREFILLING_PRELUDE - self.scheduler.enqueue(request, next_stage) + self.scheduler.enqueue_work(request, WorkKey.PROMPT_PRELUDE) return PrefillProgress(True, prompt_complete) - 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 ( - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, - ): - return + def _rewind_split_prefill(self, request, work): + """Discard the resident partial chunk and return the work to resume.""" + if work not in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT): + return work frontier = request.num_prefilled_tokens request.num_preluded_tokens = frontier request.prefill_chunk_start = frontier @@ -686,7 +680,7 @@ def _rewind_split_prefill(self, request): self.model.config.total_ut_steps, start=frontier, ) - request.phase = RequestPhase.PREFILLING_PRELUDE + return WorkKey.PROMPT_PRELUDE def _update(self, batch, result) -> list[RequestOutput]: if batch.is_speculative: @@ -699,21 +693,21 @@ def _update(self, batch, result) -> list[RequestOutput]: # Compare object identity so an old result cannot update the new request. if self.scheduler.requests.get(request.request_id) is not request: continue - phase = batch.phase_for(scheduled_request) + work = scheduled_request.work params = request.sampling_params - if phase in ( - RequestPhase.PREFILL, - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, + if work in ( + WorkKey.PROMPT_ATOMIC, + WorkKey.PROMPT_PRELUDE, + WorkKey.PROMPT_RECURRENT, ): - self._apply_prefill_progress(batch, scheduled_request) - elif phase == RequestPhase.PRELUDE: + self._apply_prefill_progress(scheduled_request) + elif work == WorkKey.DECODE_PRELUDE: request.loops_done = 0 request.remaining_probability = 1.0 request.pending_exit_depth = None self._pending_exit_signals.pop(request.request_id, None) - self.scheduler.enqueue(request, RequestPhase.RECURRENT) - elif phase == RequestPhase.RECURRENT: + self.scheduler.enqueue_work(request, WorkKey.DECODE_RECURRENT) + elif work == WorkKey.DECODE_RECURRENT: request.loops_done += 1 if self.exit_config.mode == "trace": should_exit = self._trace_exit(request) @@ -724,10 +718,10 @@ def _update(self, batch, result) -> list[RequestOutput]: should_exit = self._delayed_exit(request, result[index]) if should_exit: self.model_runner.finalize(request) - self.scheduler.enqueue(request, RequestPhase.CODA) + self.scheduler.enqueue_work(request, WorkKey.DECODE_CODA) else: - self.scheduler.enqueue(request, RequestPhase.RECURRENT) - elif phase == RequestPhase.CODA: + self.scheduler.enqueue_work(request, WorkKey.DECODE_RECURRENT) + elif work == WorkKey.DECODE_CODA: token_id = result[index] request.generated_token_ids.append(token_id) request.exit_depths.append(request.loops_done) @@ -738,12 +732,12 @@ def _update(self, batch, result) -> list[RequestOutput]: elif len(request.generated_token_ids) >= params.max_tokens: self._finish(request, FinishReason.LENGTH) else: - self.scheduler.enqueue( - request, - RequestPhase.SPECULATIVE + next_work = ( + WorkKey.SPECULATIVE if self.speculative_config - else RequestPhase.PRELUDE, + else WorkKey.DECODE_PRELUDE ) + self.scheduler.enqueue_work(request, next_work) outputs.append(RequestOutput.from_request(request)) return outputs @@ -762,11 +756,7 @@ def _finish(self, request, reason: FinishReason): if self.async_speculative: # Logical finish is immediate; allocation ownership survives until # the final submitted round has stopped accessing the request. - request.phase = RequestPhase.FINISHED - request.finish_reason = reason - for queue in self.scheduler.queues.values(): - while request.request_id in queue: - queue.remove(request.request_id) + self.scheduler.mark_finished(request, reason) self._retiring_speculative[request.request_id] = request return self.model_runner.release(request.request_id) @@ -805,7 +795,7 @@ def _collect_speculative(self): request = scheduled_request.request if ( self.scheduler.requests.get(request.request_id) is not request - or request.phase == RequestPhase.FINISHED + or request.status == RequestStatus.FINISHED ): runner.discarded_rounds += 1 continue @@ -837,10 +827,9 @@ def _collect_speculative(self): runner.stats.accepted_tokens += min(result.accepted_count, emitted) # In particular, do NOT truncate KV here: the GPU may already be # writing the following round at a later actual position. - if request.phase != RequestPhase.FINISHED: - queue = self.scheduler.queues[RequestPhase.SPECULATIVE] - if request.request_id not in queue: - self.scheduler.enqueue(request, RequestPhase.SPECULATIVE) + if request.status != RequestStatus.FINISHED: + if self.scheduler.pending_work(request) is None: + self.scheduler.enqueue_work(request, WorkKey.SPECULATIVE) outputs.append(RequestOutput.from_request(request)) ticket.retire() self._pending_speculative.pop(0) @@ -857,14 +846,14 @@ def _step_async_speculative(self): batch = self._schedule() self.last_schedule = batch if batch is not None: - if batch.phase == RequestPhase.SPECULATIVE: + if batch.is_speculative: ticket = self.executor.submit(batch) self._pending_speculative.append(ticket) for scheduled_request in batch.requests: request = scheduled_request.request request.num_output_placeholders += scheduled_request.token_count request.num_speculative_rounds += 1 - self.scheduler.enqueue(request, RequestPhase.SPECULATIVE) + self.scheduler.enqueue_work(request, WorkKey.SPECULATIVE) else: # Bootstrap uses the same submit contract; only actual host # results (the first sampled token) require collection here. @@ -956,7 +945,8 @@ def _step_async(self): # splitting on a briefly pending coda readback fragments the batch. prefer_recurrent=self._overlap_boundary and any( - t.batch.phase == RequestPhase.RECURRENT and not t.ready() for t in self._inflight + WorkKey.DECODE_RECURRENT in t.batch.works and not t.ready() + for t in self._inflight ) ) self._overlap_boundary = False @@ -968,7 +958,7 @@ def _step_async(self): elif not self._waiting_for_connector(): raise RuntimeError("scheduler made no progress") return outputs - if batch.phase == RequestPhase.CODA: + if WorkKey.DECODE_CODA in batch.works: # Bound speculation to one output per request. Its next prelude and # core may run before delivery (including a possible EOS), but never # sample another token until that output's stop decision is known. @@ -987,7 +977,7 @@ def _step_async(self): return outputs ticket = self.executor.submit(batch) self._inflight.append(ticket) - if batch.phase == RequestPhase.CODA: + if WorkKey.DECODE_CODA in batch.works: self._pending_coda.append(ticket) runner = self.model_runner self._overlap_boundary = ( @@ -999,15 +989,15 @@ def _step_async(self): request.num_output_placeholders += 1 if request.num_scheduled_outputs < request.sampling_params.max_tokens: request.input_token_tensor = ticket.device_values[index] - self.scheduler.enqueue(request, RequestPhase.PRELUDE) - elif not any(batch.phase_for(i) == RequestPhase.RECURRENT for i in batch.requests): + self.scheduler.enqueue_work(request, WorkKey.DECODE_PRELUDE) + elif WorkKey.DECODE_RECURRENT not in batch.works: self._update(batch, None) else: # Submit r FIRST. While the GPU runs r, consume r-1's signal to # determine whether this token may enter r+1. No speculative extra loop. exited = [] for index, scheduled_request in enumerate(batch.requests): - if batch.phase_for(scheduled_request) != RequestPhase.RECURRENT: + if scheduled_request.work != WorkKey.DECODE_RECURRENT: self._update(replace(batch, requests=[scheduled_request]), None) continue request = scheduled_request.request @@ -1029,7 +1019,7 @@ def _step_async(self): if should_exit: request.pending_exit_depth = request.loops_done exited.append(request) - self.scheduler.enqueue(request, RequestPhase.CODA) + self.scheduler.enqueue_work(request, WorkKey.DECODE_CODA) else: if ( self.exit_config.mode in ("ouro_delayed", "random_lookahead") @@ -1041,6 +1031,6 @@ def _step_async(self): request.position, request.loops_done, ) - self.scheduler.enqueue(request, RequestPhase.RECURRENT) + self.scheduler.enqueue_work(request, WorkKey.DECODE_RECURRENT) self.model_runner.finalize_many(exited) return outputs diff --git a/vllm_rlt/pd/connector.py b/vllm_rlt/pd/connector.py index c36865b0..1431f6a3 100644 --- a/vllm_rlt/pd/connector.py +++ b/vllm_rlt/pd/connector.py @@ -8,7 +8,8 @@ import torch -from vllm_rlt.request import Request, RequestPhase +from vllm_rlt.core.execution import WorkKey +from vllm_rlt.request import Request, RequestStatus from .transport import NixlConnector, kv_segments, partition_segments @@ -183,11 +184,11 @@ def can_schedule(self, request): and self.transport.inflight_bytes < self.config.max_inflight_bytes ) - def prefill_submitted(self, request, *, phase, start, progress, event, stream): + def prefill_submitted(self, request, *, work, 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 work == WorkKey.PROMPT_PRELUDE: transfer.transfer_overlaps += self.transport.busy(transfer.tid) if not progress.chunk_complete: transfer.last_event = event @@ -370,7 +371,7 @@ def progress(self): continue if ( transfer.compute_done - and not transfer.request.compute_released + and self.engine.scheduler.has_admission_slot(transfer.request) and self.engine.release_transfer_compute( transfer.request, ready_event=transfer.last_event ) @@ -398,7 +399,7 @@ def progress(self): if ( transfer.acknowledged and transfer.committed - and transfer.request.compute_released + and not self.engine.scheduler.has_admission_slot(transfer.request) and not self.transport.busy(transfer.tid) ): self._remove(transfer) @@ -408,7 +409,7 @@ def progress(self): elif transfer.ready: if transfer.request.request_id not in self.engine.scheduler.requests: self._forget(transfer) - elif not transfer.active and transfer.request.phase != RequestPhase.RECEIVING: + elif not transfer.active and transfer.request.status != RequestStatus.RECEIVING: transfer.active = True self.notify( "activated", tid=transfer.tid, seconds=time.monotonic() - transfer.started diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index ce2b3b7d..dff18ead 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -17,7 +17,7 @@ from vllm_rlt.models import OuroConfig, resolve_model_config from vllm_rlt.profiling import ProfileConfig from vllm_rlt.profiling_artifacts import _write_json_atomic, profile_timestamp -from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase +from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestStatus from vllm_rlt.sampling_params import SamplingParams from .config import PDConfig @@ -392,7 +392,7 @@ def abort_request(self, request_id): self._send(w.d, "cancel_safe", tid=w.tid) else: self._send(w.p, "cancel", tid=w.tid) - request.phase = RequestPhase.FINISHED + request.status = RequestStatus.FINISHED request.finish_reason = FinishReason.ABORT return RequestOutput.from_request(request) diff --git a/vllm_rlt/request.py b/vllm_rlt/request.py index 537e2bf1..308358ee 100644 --- a/vllm_rlt/request.py +++ b/vllm_rlt/request.py @@ -1,6 +1,6 @@ """Request state, immutable execution descriptions and public request outputs.""" -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from enum import Enum import torch @@ -8,26 +8,14 @@ from vllm_rlt.sampling_params import SamplingParams -class RequestPhase(str, Enum): +class RequestStatus(str, Enum): WAITING = "waiting" RECEIVING = "receiving" TRANSFERRING = "transferring" - PREFILL = "prefill" - PREFILLING_PRELUDE = "prefilling_prelude" - PREFILLING_RECURRENT = "prefilling_recurrent" - PRELUDE = "prelude" - RECURRENT = "recurrent" - SPECULATIVE = "speculative" - CODA = "coda" + ACTIVE = "active" FINISHED = "finished" -class Operation(Enum): - PRELUDE = "prelude" - RECURRENT = "recurrent" - CODA = "coda" - - class FinishReason(str, Enum): STOP = "stop" LENGTH = "length" @@ -39,7 +27,7 @@ class Request: request_id: str prompt_token_ids: list[int] sampling_params: SamplingParams - phase: RequestPhase = RequestPhase.WAITING + status: RequestStatus = RequestStatus.WAITING generated_token_ids: list[int] = field(default_factory=list) exit_depths: list[int] = field(default_factory=list) # Scheduling advances before coda's CPU output delivery. @@ -61,8 +49,6 @@ class Request: admission_bypasses: int = 0 remaining_probability: float = 1.0 hidden_state: torch.Tensor | None = field(default=None, repr=False) - # An export can retain KV/transport ownership after releasing compute slots. - compute_released: bool = False finish_reason: FinishReason | None = None exit_trace: tuple[int, ...] = field(default=(), repr=False) @@ -106,90 +92,6 @@ def input_token_id(self) -> int: return self.generated_token_ids[-1] -@dataclass(frozen=True) -class StateSpan: - """Immutable execution range borrowing the original Request for state ownership.""" - - request: Request - token_start: int - token_count: int - # Zero-based starting recurrent depth, not the request's completed-loop count. - # Scheduled PRELUDE starts at 0; CODA executes no recurrent loop and uses 0 - # 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. - chunk_start: int | None = None - retain_chunk: bool = False - retain_last: bool = True - emit_gate: bool = False - # 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 - - -@dataclass(frozen=True) -class ExecutionBatch: - """Selected request ranges bound to one model operation and loop boundary.""" - - operation: Operation - spans: tuple[StateSpan, ...] - # Number of recurrent operations following PRELUDE, or starting at RECURRENT. - # PRELUDE + N loops retains activations locally until the submission boundary. - loop_count: int = 0 - - def __post_init__(self): - if self.loop_count < 0: - raise ValueError("loop_count must be nonnegative") - if self.operation == Operation.RECURRENT and self.loop_count == 0: - object.__setattr__(self, "loop_count", 1) - if self.operation == Operation.CODA and self.loop_count: - raise ValueError("CODA cannot include recurrent work") - for span in self.spans: - if span.token_start < 0 or span.token_count <= 0 or span.depth < 0: - raise ValueError("state spans require nonnegative addresses and positive rows") - - def at_loop(self, offset: int): - """Bind a loop view without changing request progress or storage ownership.""" - if not 0 <= offset < self.loop_count: - raise ValueError("loop offset outside execution range") - return replace( - self, - operation=Operation.RECURRENT, - loop_count=1, - spans=tuple(replace(span, depth=span.depth + offset) for span in self.spans), - ) - - def validate_depths(self, total_depths: int): - if any(span.depth + self.loop_count > total_depths for span in self.spans): - raise ValueError("execution range exceeds model loop depth") - - @property - def num_rows(self): - return sum(span.token_count for span in self.spans) - - @property - def has_chunks(self): - return any(span.chunk_start is not None for span in self.spans) - - @property - def all_chunks(self): - return bool(self.spans) and all(span.chunk_start is not None for span in self.spans) - - @property - def emit_gate(self): - return any(span.emit_gate for span in self.spans) - - def addresses(self): - ids, depths, positions = [], [], [] - for span in self.spans: - ids.extend([span.request.request_id] * span.token_count) - depths.extend([span.depth] * span.token_count) - positions.extend(range(span.token_start, span.token_start + span.token_count)) - return ids, depths, positions - - @dataclass(frozen=True) class RequestOutput: request_id: str @@ -207,7 +109,7 @@ def from_request(cls, request: Request): prompt_token_ids=list(request.prompt_token_ids), token_ids=list(request.generated_token_ids), exit_depths=list(request.exit_depths), - finished=request.phase == RequestPhase.FINISHED, + finished=request.status == RequestStatus.FINISHED, finish_reason=request.finish_reason.value if request.finish_reason is not None else None, diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index b5e5a017..49439f0f 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -5,9 +5,10 @@ import torch +from vllm_rlt.core.execution import Operation from vllm_rlt.core.kv_cache_manager import _Allocation from vllm_rlt.core.scheduler import SchedulerOutput -from vllm_rlt.request import Operation, Request +from vllm_rlt.request import Request from vllm_rlt.worker.executor import ForwardBatch from vllm_rlt.worker.sampler import RngRegistry from vllm_rlt.worker.speculative import SpeculativeResult, SpeculativeRunner diff --git a/vllm_rlt/worker/executor.py b/vllm_rlt/worker/executor.py index 87ba27df..13d54e42 100644 --- a/vllm_rlt/worker/executor.py +++ b/vllm_rlt/worker/executor.py @@ -9,8 +9,8 @@ import torch +from vllm_rlt.core.execution import ExecutionBatch, Operation, StateSpan, WorkKey from vllm_rlt.core.scheduler import SchedulerOutput -from vllm_rlt.request import ExecutionBatch, Operation, RequestPhase, StateSpan def build_execution_batch(scheduled_batch: SchedulerOutput, total_depths: int) -> ExecutionBatch: @@ -20,33 +20,36 @@ def build_execution_batch(scheduled_batch: SchedulerOutput, total_depths: int) - The runner prepares its backend-specific metadata once from these bindings. """ spans = [] - operations = set() - loop_counts = set() + batch_operation = None + batch_loop_count = None for scheduled_request in scheduled_batch.requests: - phase = scheduled_batch.phase_for(scheduled_request) + work = scheduled_request.work request = scheduled_request.request - chunk = phase in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT) - prompt_input = phase in (RequestPhase.PREFILL, RequestPhase.PREFILLING_PRELUDE) - operation = { - RequestPhase.PREFILL: Operation.PRELUDE, - RequestPhase.PREFILLING_PRELUDE: Operation.PRELUDE, - RequestPhase.PRELUDE: Operation.PRELUDE, - RequestPhase.PREFILLING_RECURRENT: Operation.RECURRENT, - RequestPhase.RECURRENT: Operation.RECURRENT, - RequestPhase.CODA: Operation.CODA, - }.get(phase) + operation = work.operation if operation is None: - raise ValueError(f"phase {phase} does not describe a model operation") + raise ValueError(f"work {work.name} is handled by the speculative runner") + chunk = work in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) + prompt_input = work in (WorkKey.PROMPT_ATOMIC, WorkKey.PROMPT_PRELUDE) + loop_count = ( + total_depths if work == WorkKey.PROMPT_ATOMIC else int(operation == Operation.RECURRENT) + ) + if batch_operation is None: + batch_operation, batch_loop_count = operation, loop_count + elif operation != batch_operation or loop_count != batch_loop_count: + raise ValueError( + f"request {request.request_id} work {work.name} has " + f"operation={operation.value}, loop_count={loop_count}; " + f"batch expects operation={batch_operation.value}, " + f"loop_count={batch_loop_count}" + ) depth = ( scheduled_request.depth if scheduled_request.depth is not None else request.loops_done ) if operation != Operation.RECURRENT: - # PRELUDE starts from 0. CODA has no recurrent work; 0 addresses - # boundary-layer KV, not the request's actual completed-loop count. depth = 0 position = ( scheduled_request.token_start - if chunk or phase == RequestPhase.PREFILL or scheduled_request.depth is not None + if chunk or work == WorkKey.PROMPT_ATOMIC or scheduled_request.depth is not None else request.position ) spans.append( @@ -56,9 +59,10 @@ 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 (work == WorkKey.PROMPT_PRELUDE or depth + 1 < total_depths), retain_last=scheduled_request.completes_prompt if chunk else True, - emit_gate=phase == RequestPhase.RECURRENT, + emit_gate=work == WorkKey.DECODE_RECURRENT, input_tokens=tuple( request.prompt_token_ids[position : position + scheduled_request.token_count] ) @@ -66,13 +70,14 @@ def build_execution_batch(scheduled_batch: SchedulerOutput, total_depths: int) - else None, ) ) - operations.add(operation) - loop_counts.add( - total_depths if phase == RequestPhase.PREFILL else int(operation == Operation.RECURRENT) - ) - 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()) + if batch_operation is None: + raise ValueError("cannot build an execution batch without scheduled work") + return ExecutionBatch( + batch_operation, + tuple(spans), + batch_loop_count, + has_split_prefill_work=scheduled_batch.has_split_prefill_work, + ) class Executor: diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index 6c0f077e..c53baf88 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -8,8 +8,9 @@ import torch from vllm_rlt.config import ExecutionConfig, ExitConfig, SchedulerConfig +from vllm_rlt.core.execution import ExecutionBatch, Operation from vllm_rlt.core.scheduler import SchedulerOutput -from vllm_rlt.request import ExecutionBatch, Operation, Request +from vllm_rlt.request import Request from vllm_rlt.worker.buffers import ( PackedMetadataBuffer, Workspace, @@ -1154,7 +1155,7 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu execution needs only ExecutionBatch; neither path waits for readback here unless bounded buffers must be recycled. """ - split_prefill = batch.has_chunks + has_split_prefill_work = batch.has_split_prefill_work # A fused prelude + loops stays resident on core; chunking does not # change the stream of an otherwise standalone prelude or coda. stream = self.core_stream if batch.loop_count else self.boundary_stream @@ -1162,7 +1163,7 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu routing_finish_attempted = False slot = None try: - if self.submissions.begin() and split_prefill: + if self.submissions.begin() and has_split_prefill_work: self.split_prefill_submission_waits += 1 prepared = self.prepare_on_cpu(batch) source_stream = ( @@ -1228,7 +1229,7 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu for span in batch.spans: self.events[span.request.request_id] = compute_done self.submissions.complete(compute_done) - if split_prefill: + if has_split_prefill_work: self.split_prefill_submissions += 1 return Submission( scheduled_batch, diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 55d12777..05d6d2f9 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -9,8 +9,8 @@ import torch +from vllm_rlt.core.execution import Operation from vllm_rlt.core.scheduler import SchedulerOutput -from vllm_rlt.request import Operation from vllm_rlt.worker.buffers import stage_metadata from vllm_rlt.worker.cuda_graph import CodaGraphs, RecurrentGraphs from vllm_rlt.worker.executor import ForwardBatch, ModelRuntime From 8dd4058a03ed9fe417a5e3709e9df158f8581e37 Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Thu, 22 Oct 2026 16:03:31 +0800 Subject: [PATCH 2/5] fix tests: align split prefill graph test with WorkKey queue ownership Signed-off-by: LyxWxj --- tests/test_split_prefill.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_split_prefill.py b/tests/test_split_prefill.py index ded859e8..26e467ab 100644 --- a/tests/test_split_prefill.py +++ b/tests/test_split_prefill.py @@ -1082,7 +1082,8 @@ def run(cuda_graphs): assert e.scheduler.pending_work(split_request) is WorkKey.PROMPT_RECURRENT assert split_request.prefill_depth_ranges[0] == [2, 4] # Keep the staged chunk resident while decode triggers graph capture. - e.scheduler.work_queues[WorkKey.PROMPT_RECURRENT].remove("split") + parked_work = e.scheduler.remove_pending_work(split_request) + assert parked_work is WorkKey.PROMPT_RECURRENT e.add_request( "decode", @@ -1102,7 +1103,7 @@ def run(cuda_graphs): break else: pytest.fail("decode request did not reach its second recurrent step") - assert e.scheduler.pending_work(split_request) is WorkKey.PROMPT_RECURRENT + assert e.scheduler.pending_work(split_request) is None assert split_request.num_prefilled_tokens == 0 assert split_request.num_preluded_tokens == 2 @@ -1114,6 +1115,7 @@ def run(cuda_graphs): e.model_runner.graphs.fallbacks, ) e.scheduler.enqueue_work(split_request, WorkKey.PROMPT_RECURRENT) + assert e.scheduler.pending_work(split_request) is WorkKey.PROMPT_RECURRENT finished = {} for _ in range(200): for output in e.step(): From 17c13dbe1802fa7eeda9bb9c8151e83e441627c3 Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Thu, 22 Oct 2026 16:55:07 +0800 Subject: [PATCH 3/5] format files: tests/test_cdb_runtime,test_prefix_growth.py, test_transfer_connector.py, vllm_rlt/core/scheduling_policy.py, vllm_rlt/worker/executor.py Signed-off-by: LyxWxj --- tests/test_cdb_runtime.py | 3 +-- tests/test_prefix_growth.py | 5 ++--- tests/test_transfer_connector.py | 6 ++---- vllm_rlt/core/scheduling_policy.py | 4 +--- vllm_rlt/worker/executor.py | 3 +-- 5 files changed, 7 insertions(+), 14 deletions(-) diff --git a/tests/test_cdb_runtime.py b/tests/test_cdb_runtime.py index bdd363cc..58a45465 100644 --- a/tests/test_cdb_runtime.py +++ b/tests/test_cdb_runtime.py @@ -193,8 +193,7 @@ def test_short_requests_bypass_then_protected_long_request_gets_memory(): engine.add_request("later", [3], params) engine.scheduler._admit() assert ( - engine.scheduler.pending_work(engine.scheduler.requests["short"]) - == WorkKey.PROMPT_ATOMIC + engine.scheduler.pending_work(engine.scheduler.requests["short"]) == WorkKey.PROMPT_ATOMIC ) assert list(engine.scheduler.waiting) == ["long", "later"] engine.abort_request("short") diff --git a/tests/test_prefix_growth.py b/tests/test_prefix_growth.py index 83945e5a..5ba1607b 100644 --- a/tests/test_prefix_growth.py +++ b/tests/test_prefix_growth.py @@ -85,9 +85,7 @@ def test_incremental_growth_matches_reserved_generation(): assert outputs[0].exit_depths == outputs[1].exit_depths -@pytest.mark.parametrize( - "stage", [WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT] -) +@pytest.mark.parametrize("stage", [WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT]) def test_incremental_split_prefill_reserves_prompt_growth(stage): prompt = list(range(1, 9)) e = LLMEngine( @@ -326,6 +324,7 @@ def test_async_preemption_restores_prelude_token_in_a_new_slot(): # (LIFO), so the victim resumes in another slot whose device token belongs # to a different request. Distinct prompts make a stale token visible. from vllm_rlt import ExecutionConfig, ExitConfig + torch.manual_seed(72) m = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to(device="cuda", dtype=torch.bfloat16) results, victim_stages = [], [] diff --git a/tests/test_transfer_connector.py b/tests/test_transfer_connector.py index c0a83d94..e84792e8 100644 --- a/tests/test_transfer_connector.py +++ b/tests/test_transfer_connector.py @@ -291,8 +291,7 @@ def test_export_holds_compute_admission_until_final_staging_read_finishes(): assert engine.release_transfer_compute(transfer.request, ready_event=event) engine.scheduler._admit() assert ( - engine.scheduler.pending_work(engine.scheduler.requests["local"]) - == WorkKey.PROMPT_PRELUDE + engine.scheduler.pending_work(engine.scheduler.requests["local"]) == WorkKey.PROMPT_PRELUDE ) assert engine.cache_manager.has_transfer_lease("export") connector.cancel("export") @@ -322,8 +321,7 @@ def test_export_reservation_obeys_shared_compute_admission(local_priority, remot assert engine.cache_manager.has_transfer_lease("remote") engine.scheduler._admit() assert ( - engine.scheduler.pending_work(engine.scheduler.requests[winner]) - == WorkKey.PROMPT_PRELUDE + engine.scheduler.pending_work(engine.scheduler.requests[winner]) == WorkKey.PROMPT_PRELUDE ) loser = "remote" if winner == "local" else "local" assert engine.scheduler.requests[loser].status == RequestStatus.WAITING diff --git a/vllm_rlt/core/scheduling_policy.py b/vllm_rlt/core/scheduling_policy.py index 8e2cd86f..26d48d33 100644 --- a/vllm_rlt/core/scheduling_policy.py +++ b/vllm_rlt/core/scheduling_policy.py @@ -16,9 +16,7 @@ class NoRefillPhase(Enum): CODA = auto() -PROMPT_WORKS = frozenset( - (WorkKey.PROMPT_ATOMIC, WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT) -) +PROMPT_WORKS = frozenset((WorkKey.PROMPT_ATOMIC, WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)) DECODE_PROGRESS_WORKS = frozenset((WorkKey.DECODE_RECURRENT, WorkKey.SPECULATIVE)) PRELUDE_WORKS = (WorkKey.DECODE_PRELUDE, WorkKey.PROMPT_PRELUDE) RECURRENT_WORKS = (WorkKey.DECODE_RECURRENT, WorkKey.PROMPT_RECURRENT) diff --git a/vllm_rlt/worker/executor.py b/vllm_rlt/worker/executor.py index 13d54e42..a6f419b3 100644 --- a/vllm_rlt/worker/executor.py +++ b/vllm_rlt/worker/executor.py @@ -59,8 +59,7 @@ 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 (work == WorkKey.PROMPT_PRELUDE or depth + 1 < total_depths), + retain_chunk=chunk and (work == WorkKey.PROMPT_PRELUDE or depth + 1 < total_depths), retain_last=scheduled_request.completes_prompt if chunk else True, emit_gate=work == WorkKey.DECODE_RECURRENT, input_tokens=tuple( From 74d1f5767b0c3bc68c56afda916e29a6bc4f1d55 Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Mon, 26 Oct 2026 16:09:17 +0800 Subject: [PATCH 4/5] refactor: complete RequestStatus and WorkKey scheduler migration Signed-off-by: LyxWxj --- benchmarks/preemption.py | 4 +- docs/diagrams/README.md | 2 +- docs/diagrams/index.html | 2 +- docs/diagrams/loop-scheduler.png | Bin 599131 -> 692528 bytes docs/diagrams/loop-scheduler.svg | 8 +- docs/diagrams/scheduler-overview.png | Bin 966436 -> 1163003 bytes docs/diagrams/scheduler-overview.svg | 10 +- docs/diagrams/scheduler-preemption-detail.png | Bin 393464 -> 462809 bytes docs/diagrams/scheduler-preemption-detail.svg | 4 +- docs/diagrams/scheduler-victim-selection.png | Bin 1221497 -> 1500485 bytes docs/diagrams/scheduler-victim-selection.svg | 10 +- docs/diagrams/stage-architecture.png | Bin 551507 -> 664606 bytes docs/diagrams/stage-architecture.svg | 2 +- docs/diagrams/stage-components.png | Bin 824983 -> 1099981 bytes docs/diagrams/stage-components.svg | 2 +- docs/scheduler.md | 49 +++++--- docs/unified_execution.md | 4 +- tests/test_engine.py | 25 ++-- tests/test_mixed_execution.py | 112 ++++++++++++++---- tests/test_scheduler_selection.py | 62 +++++++++- tests/test_split_prefill.py | 17 ++- tests/test_transfer_connector.py | 23 +++- vllm_rlt/core/scheduler.py | 62 ++++++---- vllm_rlt/core/scheduling_policy.py | 9 ++ vllm_rlt/engine/preemption.py | 3 + vllm_rlt/engine/stage_engine.py | 5 +- vllm_rlt/pd/connector.py | 2 +- 27 files changed, 317 insertions(+), 100 deletions(-) diff --git a/benchmarks/preemption.py b/benchmarks/preemption.py index 9dc56360..51a8ee9f 100644 --- a/benchmarks/preemption.py +++ b/benchmarks/preemption.py @@ -98,7 +98,9 @@ def drive(engine, requests, state): next_request = next(pending, None) state["max_waiting"] = max( state["max_waiting"], - sum(request.phase.value == "waiting" for request in engine.scheduler.requests.values()), + sum( + request.status.value == "waiting" for request in engine.scheduler.requests.values() + ), ) if engine.has_unfinished_requests(): for output in engine.step(): diff --git a/docs/diagrams/README.md b/docs/diagrams/README.md index 56061270..2aa50fe0 100644 --- a/docs/diagrams/README.md +++ b/docs/diagrams/README.md @@ -34,7 +34,7 @@ do not imply that every submission executes all three operations. `scheduler-overview.svg` is the editable 1800 × 2600 source and `scheduler-overview.png` its 2× export. The vertical A–F spine shows admission, -phase queues, policy/batch selection, execution, completion tracking and progress. +WorkKey queues, policy/batch selection, execution, completion tracking and progress. Read top to bottom; one return arrow on the left leads from F to B. A includes waiting, safe preemption and restoration; KV growth pressure in C returns to the same preemption logic. C separates prefill granularity from batching and shows supported combinations. diff --git a/docs/diagrams/index.html b/docs/diagrams/index.html index 3859ae3d..62cd0f69 100644 --- a/docs/diagrams/index.html +++ b/docs/diagrams/index.html @@ -28,7 +28,7 @@ +
Complete scheduler paths: admission, WorkKey queues, policy selection, ordinary/speculative execution, pending results, PD transfer and retirement.