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/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 @@
Scheduler overview Follow A–F from top to bottom, then return to B along the left edge. Admission includes waiting, preemption and restoration; C separates prefill granularity from batching.
-
+
Loop-level scheduler Select individual loops, mix compatible requests, then regroup from independent progress.
Zoom 100% Fit width
diff --git a/docs/diagrams/loop-scheduler.png b/docs/diagrams/loop-scheduler.png
index 6d1888de..f13fad1c 100644
Binary files a/docs/diagrams/loop-scheduler.png and b/docs/diagrams/loop-scheduler.png differ
diff --git a/docs/diagrams/loop-scheduler.svg b/docs/diagrams/loop-scheduler.svg
index b3576137..e69da98e 100644
--- a/docs/diagrams/loop-scheduler.svg
+++ b/docs/diagrams/loop-scheduler.svg
@@ -14,7 +14,7 @@
Select operation · token range · recurrent depth
6 token rows / 3 requests
Mixed policy · split prefill
-01 CHOOSE READY PHASES
+01 CHOOSE READY WORK
02 BUILD THE BATCH
@@ -73,8 +73,8 @@
Policy selects ready prompt + decode recurrent queues.
Ranges · depths · KV / token budget
-policy.schedule() → compatible phases
-take_phases() → SchedulerOutput
+policy.schedule() → compatible WorkKeys
+take_works() → SchedulerOutput
03 Execute batch
Executor → runner → runtime
@@ -86,7 +86,7 @@
04 Apply progress & requeue
StageEngine → Scheduler queues
-StageEngine updates each request and calls Scheduler.enqueue() for its next phase.
+StageEngine updates each request and queues its next WorkKey.
A Next: loop 3
diff --git a/docs/diagrams/scheduler-overview.png b/docs/diagrams/scheduler-overview.png
index 367dc59b..63d46a84 100644
Binary files a/docs/diagrams/scheduler-overview.png and b/docs/diagrams/scheduler-overview.png differ
diff --git a/docs/diagrams/scheduler-overview.svg b/docs/diagrams/scheduler-overview.svg
index df5c0025..3ff4737b 100644
--- a/docs/diagrams/scheduler-overview.svg
+++ b/docs/diagrams/scheduler-overview.svg
@@ -16,20 +16,20 @@
fits
-Admit / restore phase
+Admit / restore work
safe victim
Preempt → snapshot
-Victim returns to WAITING; restores its saved phase.
+Victim returns to WAITING; restores its saved WorkKey.
no fit → wait
Priority / fairness · free KV · reserved growth · protected users
B
-Ready queues — next phase per request
+Work queues — one pending WorkKey per request
SCHEDULER
Atomic prompt
@@ -123,7 +123,7 @@
Next token
CODA
-Update Request → enqueue(next phase), or stop scheduling when finished.
+Update Request → enqueue the next WorkKey, or finish.
@@ -132,7 +132,7 @@
finished
-CONTINUE · enqueue the next phase
+CONTINUE · enqueue the next WorkKey
KV growth blocked → safe victim / otherwise defer
diff --git a/docs/diagrams/scheduler-preemption-detail.png b/docs/diagrams/scheduler-preemption-detail.png
index 19c04e76..c6bfc0c8 100644
Binary files a/docs/diagrams/scheduler-preemption-detail.png and b/docs/diagrams/scheduler-preemption-detail.png differ
diff --git a/docs/diagrams/scheduler-preemption-detail.svg b/docs/diagrams/scheduler-preemption-detail.svg
index e3114e4b..3799d8d5 100644
--- a/docs/diagrams/scheduler-preemption-detail.svg
+++ b/docs/diagrams/scheduler-preemption-detail.svg
@@ -27,11 +27,11 @@
physical addresses / rows
-R2 · active phase
+R2 · active request
R2 · WAITING
-R2 · saved phase
+R2 · saved WorkKey
diff --git a/docs/diagrams/scheduler-victim-selection.png b/docs/diagrams/scheduler-victim-selection.png
index 78660ac9..ce9f9c1a 100644
Binary files a/docs/diagrams/scheduler-victim-selection.png and b/docs/diagrams/scheduler-victim-selection.png differ
diff --git a/docs/diagrams/scheduler-victim-selection.svg b/docs/diagrams/scheduler-victim-selection.svg
index 93c50f9b..2d311297 100644
--- a/docs/diagrams/scheduler-victim-selection.svg
+++ b/docs/diagrams/scheduler-victim-selection.svg
@@ -134,7 +134,7 @@
-KV + h / token + phase
+KV + hidden / token + WorkKey
Free GPU resources; retain RNG
@@ -155,7 +155,7 @@
no
-Requeue in original phase
+Keep WAITING position
Other requests can still be selected.
B · Slot freed → shared admission
@@ -171,7 +171,7 @@
yes
-Queue next phase
+Enqueue saved WorkKey
no
@@ -215,7 +215,7 @@
CPU → new GPU allocation → synchronize
-Delete snapshot → enqueue saved phase
+Delete snapshot → enqueue saved WorkKey
@@ -225,7 +225,7 @@
RECURRENT
CODA
-One saved phase; queue labels above are examples.
+One saved WorkKey; queue labels above are examples.
Split-prefill rewind resumes at chunk PRELUDE.
Failed restore defers without nested KV preemption.
Shared admission: bounded scan; FCFS / priority.
diff --git a/docs/diagrams/stage-architecture.png b/docs/diagrams/stage-architecture.png
index 27a01719..aca0da74 100644
Binary files a/docs/diagrams/stage-architecture.png and b/docs/diagrams/stage-architecture.png differ
diff --git a/docs/diagrams/stage-architecture.svg b/docs/diagrams/stage-architecture.svg
index b1366f17..eab21eef 100644
--- a/docs/diagrams/stage-architecture.svg
+++ b/docs/diagrams/stage-architecture.svg
@@ -118,5 +118,5 @@
Owned resources: model · KV cache · request state · reusable buffers
Blue: scheduler-driven execution and feedback. Gray inside Execution: model operation order. Dashed gray: stage transfer.
-Mixed = compatible recurrent work across requests, prefill/decode phases and loop depths within one stage. Speculative rounds are not included.
+Mixed = compatible recurrent work across requests, prompt/decode work and loop depths within one stage. Speculative rounds are not included.
diff --git a/docs/diagrams/stage-components.png b/docs/diagrams/stage-components.png
index 8393bddc..70295dea 100644
Binary files a/docs/diagrams/stage-components.png and b/docs/diagrams/stage-components.png differ
diff --git a/docs/diagrams/stage-components.svg b/docs/diagrams/stage-components.svg
index d886e7fd..43b63a45 100644
--- a/docs/diagrams/stage-components.svg
+++ b/docs/diagrams/stage-components.svg
@@ -56,7 +56,7 @@
Apply progress · exit / stop · coordinate resource retirement
Scheduler
-Policy · admission · phase queues
+Policy · admission · WorkKey queues
Request registry · loop selection
operation / token range / depth
diff --git a/docs/scheduler.md b/docs/scheduler.md
index 6f1c99e1..23571ed4 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.
@@ -72,7 +72,7 @@ The ordinary-generation rows in C show currently supported combinations:
full-depth prompt chunks with refill/no-refill, and scheduler-selected depth
boundaries with refill/mixed. These are compatibility rules, not a claim that
the two dimensions mean the same thing. `SplitPrefillPolicy` is selected when
-WCPB requires staged prompt phases; `MixedPolicy` adds recurrent co-batching.
+WCPB requires staged prompt work; `MixedPolicy` adds recurrent co-batching.
Prompt-only is a stage-role constraint. Speculation selects a local draft/verify
controller; it currently requires refill, uses full-depth prompt chunks unless
@@ -120,8 +120,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,10 +134,10 @@ GPU dependencies or authorize early resource reuse.
### F — Apply progress: what should each request do next?
-`Scheduler` applies progress and result-dependent decisions per `Request` (via
-`update_from_output()` and `update_after_submission()`) and enqueues its next
-phase, while `StageEngine` applies side effects (connector notifications, runner
-finalization, and retirement). The four labels show common
+`Scheduler` applies progress and result-dependent decisions per `Request` using
+the selected `WorkKey` (`update_from_output()` or `update_after_submission()`),
+then enqueues its next work. `StageEngine` applies side effects such as connector
+notifications, runner finalization, and retirement. 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.
@@ -212,7 +211,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
@@ -223,23 +222,21 @@ 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. Unfinished chunks are replayed after restoration;
-WCPB also resets its ring origin to this full-depth frontier.
+`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. WCPB also resets its ring origin to this full-depth frontier.
#### Current policy: triggers and eligibility
@@ -277,7 +274,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
@@ -287,11 +284,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
@@ -431,8 +428,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` first considers decode and prompt recurrent queues together.
+`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.
@@ -450,7 +448,9 @@ of its members, deduplicating shared events. It does not execute three separate
request forwards, combine their KV histories, or require their depths to match.
PRELUDE and CODA are different operations and are not packed into this recurrent
submission. Mixed scheduling excludes speculative work and uses a depth boundary
-when ready decode work can join prompt recurrent work.
+when ready decode work can join prompt recurrent work. Compatible decode and
+prompt PRELUDE rows may also share a boundary-zero batch; async input routing
+keeps their request-state and chunk-state sources distinct.
Logical progress then branches per request: A can continue at depth 2; B can enter
CODA if its exit is established; C can continue its prompt span at depth 1.
@@ -572,10 +572,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 | Scheduler applies progress (`update_from_output` / `update_after_submission`) and enqueues next phases; StageEngine handles pending tickets, connector hooks and retirement | Updated requests in their next phase queues; return to 01 |
+| 04 | Scheduler applies progress (`update_from_output` / `update_after_submission`) and enqueues next work; StageEngine handles tickets, connector hooks and retirement | 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
@@ -606,7 +606,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
@@ -620,25 +620,29 @@ exit decisions. The scheduler can then form a different batch. Prompt tokens
still traverse the full model depth; mixing does not apply decode early exit to
prompt rows.
-Without WCPB, PRELUDE and CODA are selected separately. With WCPB, MixedPolicy
-may combine decode PRELUDE rows with prompt-chunk PRELUDE rows because both use
-the same operation. Mixed recurrent batches combine `RECURRENT` and
-`PREFILLING_RECURRENT` spans, including multiple depths from one prompt request;
-CODA and speculative verification remain separate operations.
+MixedPolicy prioritizes a ready `DECODE_RECURRENT` + `PROMPT_RECURRENT` batch
+before refill's PRELUDE, CODA, admission and fairness decisions. Its compatible
+PRELUDE group can also combine decode and prompt rows because both use the same
+operation and recurrent boundary zero; CODA remains separate. Async mixed-PRELUDE
+execution stores prompt rows in chunk state and decode rows in request state, then
+the following recurrent batch routes them through their usual state paths.
+Compatible recurrent spans can have different depths, and WCPB may contribute
+multiple spans from one prompt request.
## Who owns each decision?
| Component | Responsibility | Code |
| --- | --- | --- |
-| `Scheduler` | Request registry, phase queues, admission, capacity checks, batch selection, logical progress, exit/termination decisions 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, logical progress, exit decisions 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, retain pending tickets, apply connector/runner side effects and coordinate retirement | [stage_engine.py](../vllm_rlt/engine/stage_engine.py) |
| `Executor` and runners | Bind selected work, prepare tensors, submit model computation and expose `ModelRunnerOutput` / 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.
@@ -651,11 +655,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
- → Scheduler.update_from_output / update_after_submission
+ → Scheduler.update_from_output / update_after_submission applies per-WorkKey progress
→ StageEngine applies runner/connector side effects and retirement
```
@@ -664,10 +668,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
@@ -713,16 +718,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.
@@ -732,8 +737,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.
@@ -744,11 +749,11 @@ independent, peer configuration dimensions.
| Policy | Selection behavior |
| --- | --- |
-| `PromptOnlyPolicy` | Admit and select prompt work only; WCPB uses staged 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 staged prompt phases for WCPB |
-| `MixedPolicy` | If both recurrent queues have ready work, select generation rows first and fill with prompt spans; otherwise use the batch's full-depth plan when no depth boundary is needed |
+| `SplitPrefillPolicy` | Refill ordering with separate `PROMPT_PRELUDE` and `PROMPT_RECURRENT` scheduling boundaries |
+| `MixedPolicy` | First combine ready decode and prompt recurrent work; otherwise use refill ordering, with compatible decode/prompt PRELUDE grouping and separate CODA |
| `SpeculativePolicy` | Select CODA bootstrap, prompt work and local draft/verify rounds, with a prefill fairness bound |
`create_scheduling_policy()` first handles prompt-only stages, then checks
@@ -759,9 +764,11 @@ restrictions still apply. Speculative generation requires refill.
`SchedulerOutput.loop_count` is the actual execution boundary. A regular prompt
chunk runs its prelude and all recurrent depths locally in one submission.
-Mixed mode inserts a boundary when ready decode work can join prompt recurrent
-work, and it returns after every depth in such a mixed batch. With no ready
-decode work, the prompt batch can run through its remaining depths. WCPB always
+Mixed mode first selects compatible recurrent work when both decode and prompt
+queues are ready; it returns after every depth in such a mixed batch. Otherwise
+it follows refill ordering. Decode and prompt PRELUDE rows may share one
+boundary-zero batch. With no ready decode work, the prompt batch can run through
+its remaining depths. WCPB always
returns after one depth to regroup ready chunks. Chunk size and loop submission
size are distinct choices; WCPB retains up to one staged chunk per model depth.
Ordinary prompt chunks already end at the scheduler's configured chunk boundary,
@@ -769,19 +776,31 @@ where decode/speculative work and admission or preemption are reconsidered.
Merely having waiters or enabling preemption does not split a chunk's recurrence.
Fairness counts nonempty scheduled batches, not attempts or GPU completions.
-Prompt batches increment the configured prefill counter; generation RECURRENT
-or SPECULATIVE resets it. In a mixed selection, a selected generation recurrent
-row causes the recurrent accounting rule to apply. If only prompt rows actually
-fit, prompt accounting applies. PRELUDE/CODA do not change the counter.
+Prompt work increments the configured prefill counter; generation RECURRENT or
+SPECULATIVE resets it. A mixed PRELUDE batch containing prompt rows counts as
+prompt work, while decode PRELUDE-only and CODA-only batches do not change the
+counter. In mixed recurrent selection, a selected generation recurrent row
+causes the recurrent accounting rule to apply; if only prompt rows fit, prompt
+accounting applies.
`prefer_recurrent` is an overlap hint, not a CUDA-event query. No-refill instead
uses its own FILL/CORE/CODA cohort state.
+**Python API migration.** `RequestPhase` and `Request.stage` were removed.
+Lifecycle callers use `RequestStatus`; scheduling callers use `WorkKey`, and
+model calls use `Operation` from `vllm_rlt.core.execution`. `ExecutionBatch` and
+`StateSpan` also live in `vllm_rlt.core.execution`. Downstream code importing
+these types from `vllm_rlt.request` must update its imports; no compatibility
+aliases are provided. `benchmarks.cdb_runtime` trace entries also changed from
+the `stage` key to `works`; trace consumers must update their schema handling.
+
## Imported requests use the same admission path
A destination stage creates its own local Request during transfer reservation;
subsequent chunk notifications reuse that object. `reserve_transfer()` reserves
the whole prompt mapping under the same growth budget used for local requests,
then pins the allocation. It does not grant a compute slot.
+An unadmitted outgoing reservation therefore does not count against
+`max_num_seqs`; admission acquires its slot when the scheduler selects it.
After the connector verifies commit and all expected transfer completions,
`accept_import()` places the existing request in the shared admission queue.
@@ -795,7 +814,7 @@ transitions; receiving data does not mean decode has started.
## Progress and resource safety
-`Scheduler` applies per-request progress and enqueues the next phase through
+`Scheduler` applies per-request progress and enqueues the next `WorkKey` through
`update_from_output()` / `update_after_submission()`; `StageEngine` applies the
side effects (connector notifications, runner finalization and retirement). In
the async path, host progress can describe submitted work; CUDA events establish
@@ -810,7 +829,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
@@ -821,7 +841,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
`Scheduler.update_from_output()`, `Scheduler.update_after_submission()`, and
`Scheduler._apply_split_prefill_progress()`, coordinated by `StageEngine._update()`
diff --git a/docs/unified_execution.md b/docs/unified_execution.md
index e86f15a0..55da7b54 100644
--- a/docs/unified_execution.md
+++ b/docs/unified_execution.md
@@ -121,7 +121,7 @@ handle commands, advance transport, call `engine.step()`, send available outputs
Registering a request and completing its generation are separate events.
**Engine and scheduler.** StageEngine calls Scheduler to select work. Scheduler
-owns its local Request registry, phase queues and policy, checks admission/KV
+owns its local Request registry, work queues and policy, checks admission/KV
capacity, returns selected operations, token ranges and depths, and applies
per-request progress, exit and termination decisions through
`update_from_output()` and `update_after_submission()`. It does not invoke
@@ -158,7 +158,7 @@ and optional stream overlap remain execution choices behind this interface.
| --- | --- |
| Worker process, channel and host-side ID mappings | StageClient; connected workers are coordinated by the deployment host |
| Loaded model and local service loop | StageWorker; the engine and runners reference the same loaded model |
-| Request registry, phase queues and selection policy | Scheduler; StageEngine coordinates lifecycle, pending tickets and runner/connector side effects |
+| Request registry, WorkKey queues and selection policy | Scheduler; StageEngine coordinates lifecycle, pending tickets and runner/connector side effects |
| KV pages, mappings, validity and references/pins | KVCacheManager, shared by local scheduling and execution |
| RNG registry and pending results | StageEngine; runners use the shared RNG, and tickets retain their result storage |
| Input buffers, GPU events and readback storage | Runner/controller owners, retained until consumers finish |
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 4d6d1871..1b175cbf 100644
--- a/tests/test_async_speculative.py
+++ b/tests/test_async_speculative.py
@@ -15,9 +15,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
tiny = tiny_ouro_model
bootstrap = bootstrap_speculation
@@ -79,12 +80,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)
@@ -353,10 +354,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)
@@ -376,7 +376,7 @@ 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 b3d469f7..e8249ba7 100644
--- a/tests/test_async_speculative_lifecycle.py
+++ b/tests/test_async_speculative_lifecycle.py
@@ -11,10 +11,10 @@
tiny_ouro_config,
)
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,
@@ -92,7 +92,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]
)
@@ -106,7 +108,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,),
@@ -375,7 +377,7 @@ def test_ticket_result_format_does_not_depend_on_statistics(mode, with_stats):
request = Request("r", [1], SamplingParams(max_tokens=8))
batch = SchedulerOutput(
- RequestPhase.SPECULATIVE, [ScheduledRequest(request, token_start=1, token_count=2)]
+ [ScheduledRequest(request, WorkKey.SPECULATIVE, token_start=1, token_count=2)]
)
values = [11, 0, 1, 2, 1, 1] if mode == "pipeline" else [11, 12, 0, 2]
buffer = SimpleNamespace(
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..58a45465 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,16 @@ 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 +231,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 +378,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 +395,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 +498,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 +516,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..bc51f252 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,12 @@ 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(
+ (
+ tuple(item.work for item in batch.requests),
+ tuple(item.request.request_id for item in batch.requests),
+ )
+ )
for output in emitted:
if output.finished:
outputs[output.request_id] = output
@@ -80,7 +86,11 @@ 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 works, ids in trace
+ if works and all(work == WorkKey.DECODE_RECURRENT for work in works)
+ ]
# 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 +106,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 +151,10 @@ 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 [works for works, _ in trace] == [
+ (WorkKey.PROMPT_ATOMIC,),
+ (WorkKey.DECODE_CODA,),
+ ]
assert engine.cache_manager.num_used_blocks == 0
@@ -241,7 +254,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 +353,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 +374,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.remove_pending_work(scheduler.requests["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 +392,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.remove_pending_work(scheduler.requests["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():
@@ -395,13 +408,14 @@ def test_abort_removes_all_queued_work_and_serializes_reason():
engine.add_request("r", [1, 2], SamplingParams(max_tokens=2))
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)
+ # Work identity remains indexed until the request is retired.
+ 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_engine_core_contract.py b/tests/test_engine_core_contract.py
index ce749ac4..d1a1f75a 100644
--- a/tests/test_engine_core_contract.py
+++ b/tests/test_engine_core_contract.py
@@ -8,9 +8,10 @@
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
+from vllm_rlt.request import RequestStatus
from vllm_rlt.worker.model_runner import Submission
@@ -56,7 +57,10 @@ def wait_for_signal(engine, request_id):
for _ in range(100):
engine.step()
request = engine.scheduler.requests[request_id]
- if request.generated_token_ids and request.phase == RequestPhase.RECURRENT:
+ if (
+ request.generated_token_ids
+ and engine.scheduler.pending_work(request) == WorkKey.DECODE_RECURRENT
+ ):
if request_id in engine._pending_exit_signals:
return request
pytest.fail("recurrent signal was not retained")
@@ -81,19 +85,19 @@ def test_delayed_signal_consumption_at_max_and_threshold(
engine.add_request("a", [2, 3], params)
probabilities = []
consumed = {}
- enqueue = engine.scheduler.enqueue
+ enqueue = engine.scheduler.enqueue_work
delayed_signal = engine.scheduler._delayed_signal
- def record_enqueue(request, phase):
- if phase == RequestPhase.CODA and request.generated_token_ids:
+ def record_enqueue(request, work):
+ if work == WorkKey.DECODE_CODA and request.generated_token_ids:
probabilities.append(request.remaining_probability)
- return enqueue(request, phase)
+ return enqueue(request, work)
def record_signal(request, score, depth):
consumed.setdefault(len(request.generated_token_ids), []).append(depth)
return delayed_signal(request, score, depth)
- monkeypatch.setattr(engine.scheduler, "enqueue", record_enqueue)
+ monkeypatch.setattr(engine.scheduler, "enqueue_work", record_enqueue)
monkeypatch.setattr(engine.scheduler, "_delayed_signal", record_signal)
result = drain(engine)[-1]
assert result.exit_depths == expected_depths
@@ -113,7 +117,7 @@ def test_async_pending_signal_survives_preemption_and_resume():
retained = engine._pending_exit_signals["a"]
engine.add_request("new", [4], SamplingParams(max_tokens=1, ignore_eos=True))
assert engine.scheduler.preempt(engine.scheduler.requests["new"])
- assert request.phase == RequestPhase.WAITING
+ assert request.status == RequestStatus.WAITING
assert engine._pending_exit_signals["a"] is retained
actual = drain(engine)
assert [out for out in actual if out.request_id == "a"][-1] == expected
@@ -300,13 +304,13 @@ def record_submit(scheduled_batch=None, *, batch=None):
ticket = (
submit(scheduled_batch, batch=batch) if batch is not None else submit(scheduled_batch)
)
- if ticket.batch.phase == RequestPhase.RECURRENT:
+ if WorkKey.DECODE_RECURRENT in ticket.batch.works:
recurrent_submissions += 1
submitted_at[id(ticket)] = recurrent_submissions
return ticket
def record_collect(ticket):
- if ticket.batch.phase == RequestPhase.RECURRENT:
+ if WorkKey.DECODE_RECURRENT in ticket.batch.works:
observed.append((submitted_at[id(ticket)], recurrent_submissions))
return collect(ticket)
@@ -365,7 +369,8 @@ def test_mixed_recurrent_and_split_prefill_progression_through_scheduler(asynchr
engine.step()
if (
engine.scheduler.requests["decode"].generated_token_ids
- and engine.scheduler.requests["decode"].phase is RequestPhase.RECURRENT
+ and engine.scheduler.pending_work(engine.scheduler.requests["decode"])
+ is WorkKey.DECODE_RECURRENT
):
break
engine.add_request("prompt", [4, 5, 6, 7], params)
@@ -378,10 +383,9 @@ def test_mixed_recurrent_and_split_prefill_progression_through_scheduler(asynchr
if out.finished:
actual[out.request_id] = out
batch = engine.last_schedule
- if batch is not None and {batch.phase_for(item) for item in batch.requests} == {
- RequestPhase.RECURRENT,
- RequestPhase.PREFILLING_RECURRENT,
- }:
+ if batch is not None and batch.works == frozenset(
+ (WorkKey.DECODE_RECURRENT, WorkKey.PROMPT_RECURRENT)
+ ):
saw_mixed_batch = True
assert saw_mixed_batch
assert actual["decode"].token_ids == expected["decode"].token_ids
diff --git a/tests/test_input_preparation.py b/tests/test_input_preparation.py
index 6f580135..5506ef65 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
@@ -189,7 +190,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(
@@ -336,7 +338,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(
@@ -382,7 +384,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
@@ -410,7 +412,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),
@@ -449,7 +451,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 542b8e96..f4a9f96a 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 ScheduledRequest, 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
@@ -21,20 +22,20 @@ def test_mixed_decode_batch_yields_after_one_loop_without_queued_decode(monkeypa
scheduled = [
ScheduledRequest(
Request("decode", [1, 2], SamplingParams()),
+ WorkKey.DECODE_RECURRENT,
token_start=2,
depth=1,
- phase=RequestPhase.RECURRENT,
),
ScheduledRequest(
Request("prompt", [3, 4], SamplingParams()),
+ WorkKey.PROMPT_RECURRENT,
token_start=0,
token_count=2,
depth=0,
chunk_start=0,
- phase=RequestPhase.PREFILLING_RECURRENT,
),
]
- monkeypatch.setattr(scheduler, "has_ready", lambda phase: False)
+ monkeypatch.setattr(scheduler, "has_ready_work", lambda work: False)
monkeypatch.setattr(scheduler, "_max_contiguous_prefill_loops", lambda item: 4)
assert scheduler._planned_loop_count(scheduled) == 1
@@ -126,8 +127,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])
@@ -206,7 +207,10 @@ def run(mode, wavefront=False):
)
for _ in range(40):
engine.step()
- if engine.scheduler.requests["short"].phase == RequestPhase.RECURRENT:
+ if (
+ engine.scheduler.pending_work(engine.scheduler.requests["short"])
+ == WorkKey.DECODE_RECURRENT
+ ):
break
else:
pytest.fail("decode request did not reach recurrent work")
@@ -224,13 +228,20 @@ def run(mode, wavefront=False):
)
outputs = {}
mixed = []
+ recurrent_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()
@@ -239,27 +250,24 @@ 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:
- recurrent_mix = phases == {
- RequestPhase.RECURRENT,
- RequestPhase.PREFILLING_RECURRENT,
- }
- # Recurrent storage kinds share one model invocation. Prelude
- # mixes use the same operation but do not call recurrent_core.
- if recurrent_mix and not graphs:
+ works = {item.work for item in batch.requests}
+ if len(works) == 2:
+ if works == {WorkKey.DECODE_RECURRENT, WorkKey.PROMPT_RECURRENT}:
+ recurrent_mixed.append(batch)
+ elif works == {WorkKey.DECODE_PRELUDE, WorkKey.PROMPT_PRELUDE}:
+ assert batch.loop_count == 0
+ else:
+ pytest.fail(f"unexpected mixed work batch: {works}")
+ if not graphs:
assert calls == [batch.num_tokens]
- elif not recurrent_mix:
- assert phases == {
- RequestPhase.PRELUDE,
- RequestPhase.PREFILLING_PRELUDE,
- }
assert batch.num_tokens <= 5
mixed.append(batch)
if not engine.has_unfinished_requests():
break
else:
pytest.fail("mixed scheduling stalled")
+ if mode == "mixed":
+ assert recurrent_mixed
assert engine.cache_manager.num_used_blocks == 0
return outputs, mixed
@@ -272,6 +280,81 @@ def record(hidden, kv, cache, **kwargs):
assert actual == wavefront_refill == original_mixed == separated == expected
+@pytest.mark.parametrize(
+ "device,graphs",
+ [
+ ("cpu", False),
+ pytest.param("cuda", False, marks=pytest.mark.gpu),
+ pytest.param("cuda", True, marks=pytest.mark.gpu),
+ ],
+)
+def test_mixed_async_prelude_batches_decode_and_prompt_rows(device, graphs):
+ from tests.helpers import LocalLLM as LLM
+ from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig
+ from vllm_rlt.engine.llm_engine import LLMEngine
+
+ dtype = torch.bfloat16 if device == "cuda" else torch.float32
+ torch.manual_seed(123)
+ e = LLMEngine(
+ OuroForCausalLM(tiny_ouro_config()).to(device=device, dtype=dtype),
+ attention_backend="triton" if device == "cuda" else "torch",
+ cache_config=CacheConfig(128, 2),
+ scheduler_config=SchedulerConfig(
+ mode="mixed",
+ max_num_seqs=2,
+ max_num_batched_tokens=4,
+ prefill_chunk_size=2,
+ wavefront_chunked_prefill=True,
+ ),
+ execution_config=ExecutionConfig(
+ async_scheduling=True,
+ cuda_graphs=graphs,
+ cuda_graph_max_batch_size=4,
+ ),
+ exit_config=ExitConfig(mode="ouro_delayed"),
+ )
+ outputs = {}
+ decode_params = SamplingParams(max_tokens=2, min_loops=2, max_loops=2, ignore_eos=True)
+ prompt_params = SamplingParams(max_tokens=1, ignore_eos=True)
+ try:
+ e.add_request("decode", [2, 3], decode_params)
+ decode = e.scheduler.requests["decode"]
+ for _ in range(100):
+ e.step()
+ if e.scheduler.pending_work(decode) is WorkKey.DECODE_PRELUDE:
+ break
+ else:
+ pytest.fail("decode request did not reach its next PRELUDE")
+
+ e.add_request("prompt", [4, 5, 6, 7], prompt_params)
+ e.scheduler._admit()
+ e.step()
+ assert e.last_schedule.works == frozenset((WorkKey.DECODE_PRELUDE, WorkKey.PROMPT_PRELUDE))
+ assert e.last_schedule.loop_count == 0
+
+ for _ in range(100):
+ for output in e.step():
+ if output.finished:
+ outputs[output.request_id] = output
+ if len(outputs) == 2:
+ break
+ assert set(outputs) == {"decode", "prompt"}
+ assert e.cache_manager.num_used_blocks == 0
+ torch.manual_seed(123)
+ reference = LLM(
+ OuroForCausalLM(tiny_ouro_config()).to(device=device, dtype=dtype),
+ attention_backend="triton" if device == "cuda" else "torch",
+ ).generate([[2, 3], [4, 5, 6, 7]], [decode_params, prompt_params])
+ assert [outputs[rid].token_ids for rid in ("decode", "prompt")] == [
+ result.token_ids for result in reference
+ ]
+ assert [outputs[rid].exit_depths for rid in ("decode", "prompt")] == [
+ result.exit_depths for result in reference
+ ]
+ finally:
+ e.close()
+
+
def mixed_ready_engine():
from tests.helpers import make_prefill_engine
@@ -286,16 +369,24 @@ def mixed_ready_engine():
engine.add_request("short", [1], SamplingParams(max_tokens=6, ignore_eos=True))
for _ in range(40):
engine.step()
- if engine.scheduler.requests["short"].phase == RequestPhase.RECURRENT:
+ if (
+ engine.scheduler.pending_work(engine.scheduler.requests["short"])
+ == WorkKey.DECODE_RECURRENT
+ ):
break
else:
pytest.fail("decode request did not reach recurrent work")
engine.add_request("long", list(range(2, 11)), SamplingParams(max_tokens=6, ignore_eos=True))
for _ in range(100):
batch = engine.scheduler.schedule()
- if batch.is_mixed and len({batch.phase_for(i) for i in batch.requests}) == 2:
+ 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")
@@ -334,14 +425,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(
@@ -355,12 +446,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(
@@ -375,7 +466,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():
@@ -511,16 +602,16 @@ def test_invalid_loop_range_fails_before_writing_kv():
torch.testing.assert_close(cache.key_cache, before)
-def test_mixed_attempt_with_one_selected_phase_is_not_a_mixed_batch():
- from dataclasses import replace
-
+def test_mixed_recurrent_work_keeps_priority_over_coda():
engine, batch = mixed_ready_engine()
- for span in batch.requests:
- engine.scheduler.enqueue(span.request, span.phase)
- engine.scheduler.config = replace(engine.scheduler.config, max_num_batched_tokens=1)
+ for item in batch.requests:
+ engine.scheduler.enqueue_work(item.request, item.work)
+ engine.add_request("other", [5, 6], SamplingParams(max_tokens=2, ignore_eos=True))
+ engine.scheduler._admit()
+ coda = engine.scheduler.requests["other"]
+ engine.scheduler.remove_pending_work(coda)
+ engine.scheduler.enqueue_work(coda, WorkKey.DECODE_CODA)
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 not selected.is_mixed
+ assert selected.works == frozenset((WorkKey.DECODE_RECURRENT, WorkKey.PROMPT_RECURRENT))
engine.close()
diff --git a/tests/test_pd.py b/tests/test_pd.py
index 90a4b38f..5aa96081 100644
--- a/tests/test_pd.py
+++ b/tests/test_pd.py
@@ -301,7 +301,7 @@ def test_pd_wavefront_exports_only_completed_chunks():
assert boundaries == [3, 6, 9, 12, 14]
assert request.loops_done == engine.model.config.total_ut_steps
- assert request.phase.value == "transferring"
+ assert request.status.value == "transferring"
engine.close()
@@ -477,9 +477,10 @@ def test_engine_yields_while_waiting_for_remote_kv(async_scheduling):
from types import SimpleNamespace
from unittest.mock import Mock
+ from vllm_rlt.core.execution import WorkKey
from vllm_rlt.core.scheduler import ScheduledRequest, Scheduler, SchedulerOutput
from vllm_rlt.engine.preemption import PreemptionManager
- from vllm_rlt.request import Request, RequestPhase
+ from vllm_rlt.request import Request, RequestStatus
from vllm_rlt.worker.output import ModelRunnerOutput
engine = object.__new__(LLMEngine)
@@ -491,7 +492,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(), eos_token_id=2)
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 = []
@@ -509,14 +510,12 @@ 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))
ticket = Mock()
- ticket.batch = SchedulerOutput(
- RequestPhase.CODA, [ScheduledRequest(active, phase=RequestPhase.CODA)]
- )
+ ticket.batch = SchedulerOutput([ScheduledRequest(active, WorkKey.DECODE_CODA)])
ticket.ready.return_value = True
ticket.collect.return_value = ModelRunnerOutput(
token_ids=[3], output_indices=(0,), depths=(4,)
@@ -527,7 +526,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 c31d4aeb..fb138422 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():
@@ -84,9 +85,7 @@ def test_incremental_growth_matches_reserved_generation():
assert outputs[0].exit_depths == outputs[1].exit_depths
-@pytest.mark.parametrize(
- "stage", [RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT]
-)
+@pytest.mark.parametrize("stage", [WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT])
def test_incremental_split_prefill_reserves_prompt_growth(stage):
prompt = list(range(1, 18))
e = LLMEngine(
@@ -105,13 +104,13 @@ 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:
for _ in range(8):
- if request.phase is RequestPhase.PREFILLING_RECURRENT:
+ if e.scheduler.pending_work(request) is WorkKey.PROMPT_RECURRENT:
break
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)
@@ -329,7 +328,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)
@@ -353,7 +351,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
@@ -366,7 +364,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]
@@ -393,10 +391,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 165159ea..c08e790b 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,19 +39,64 @@ def params(max_tokens=5):
return SamplingParams(max_tokens=max_tokens, ignore_eos=True)
-@pytest.mark.parametrize("wavefront", [False, True])
-def test_batch_selection_is_never_preempted_by_a_later_item(monkeypatch, wavefront):
+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_pending_work_index_tracks_enqueue_selection_and_removal():
+ e = engine()
+ e.add_request("r", [1, 2], params())
+ e.scheduler._admit()
+ request = e.scheduler.requests["r"]
+ assert e.scheduler.pending_work(request) is WorkKey.PROMPT_ATOMIC
+ assert e.scheduler.pending_works == {"r": WorkKey.PROMPT_ATOMIC}
+
+ batch = e.scheduler.take_works((WorkKey.PROMPT_ATOMIC,))
+ assert batch is not None
+ assert e.scheduler.pending_work(request) is None
+ assert not e.scheduler.pending_works
+
+ e.scheduler.enqueue_work(request, WorkKey.DECODE_CODA)
+ assert e.scheduler.remove_pending_work(request) is WorkKey.DECODE_CODA
+ assert not e.scheduler.pending_works
+ assert all("r" not in queue for queue in e.scheduler.work_queues.values())
+ 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(wavefront_chunked_prefill=wavefront)
+ e = engine()
for name in ("a", "b"):
e.add_request(name, [1, 2], params())
e.scheduler._admit()
monkeypatch.setattr(e.cache_manager, "ensure_capacity", lambda rid, frontier: rid != "b")
- phase = RequestPhase.PREFILLING_PRELUDE if wavefront else RequestPhase.PREFILL
- batch = e.scheduler._take(phase)
+ 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 == phase
+ assert e.scheduler.requests["a"].status == RequestStatus.ACTIVE
+
+
+def test_batch_selection_reports_a_duplicate_request_in_one_queue():
+ e = engine(max_num_seqs=3)
+ for request_id in ("a", "b"):
+ e.add_request(request_id, [1, 2], params())
+ e.scheduler._admit()
+ e.scheduler.work_queues[WorkKey.PROMPT_ATOMIC].append("a")
+
+ with pytest.raises(RuntimeError, match="queued more than once"):
+ e.scheduler.take_works((WorkKey.PROMPT_ATOMIC,))
+ e.close()
def test_wavefront_does_not_schedule_a_candidate_preempted_during_selection(monkeypatch):
@@ -65,13 +111,13 @@ def test_wavefront_does_not_schedule_a_candidate_preempted_during_selection(monk
lambda rid, frontier: rid == "a" and not e.cache_manager.has_allocation("b"),
)
- batch = e.scheduler._take(RequestPhase.PREFILLING_PRELUDE)
+ batch = e.scheduler.take_works((WorkKey.PROMPT_PRELUDE,))
assert {item.request.request_id for item in batch.requests} == {"a"}
assert e.preemption.preemptions == 1
- assert list(e.scheduler.queues[RequestPhase.WAITING]) == ["b"]
- assert not e.scheduler.queues[RequestPhase.PREFILLING_PRELUDE]
- assert e.scheduler.requests["b"].phase == RequestPhase.WAITING
+ assert list(e.scheduler.waiting) == ["b"]
+ assert not e.scheduler.work_queues[WorkKey.PROMPT_PRELUDE]
+ assert e.scheduler.requests["b"].status == RequestStatus.WAITING
e.close()
@@ -82,7 +128,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 = []
@@ -93,9 +139,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"}))]
@@ -106,12 +152,40 @@ 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
+def test_waiting_resident_victim_can_be_offloaded_without_changing_queue_position():
+ e = engine(max_num_seqs=1, policy="priority")
+ e.add_request("low", [1, 2], SamplingParams(max_tokens=3, priority=10))
+ e.scheduler._admit()
+ e.add_request("high", [3, 4], SamplingParams(max_tokens=1, priority=0))
+ low = e.scheduler.requests["low"]
+ high = e.scheduler.requests["high"]
+
+ assert e.scheduler.preempt(high, priority_only=True)
+ assert low.status == RequestStatus.WAITING
+ assert e.cache_manager.has_allocation("low")
+ waiting = list(e.scheduler.waiting)
+ assert e.preemption.snapshots["low"]["kv"] is None
+
+ assert e.scheduler.preempt(high, waiting_only=True)
+ assert list(e.scheduler.waiting) == waiting
+ assert not e.cache_manager.has_allocation("low")
+ assert e.preemption.snapshots["low"]["kv"] is not None
+ assert e.preemption.preemptions == 1
+
+ e.abort_request("high")
+ e.scheduler._admit()
+ assert low.status == RequestStatus.ACTIVE
+ assert e.scheduler.has_admission_slot(low)
+ assert e.preemption.resumptions == 1
+ e.close()
+
+
def test_reused_request_id_is_not_protected_by_the_previous_request():
"""Cancelling a request must not leave its ID protected for a new one."""
e = engine(max_num_seqs=1)
@@ -125,14 +199,14 @@ 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])
def test_refill_runs_boundary_restored_during_admission(boundary):
"""A sole resumed boundary must run in the admission step that restores it."""
e = engine(max_num_seqs=1)
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:
@@ -141,21 +215,21 @@ def test_refill_runs_boundary_restored_during_admission(boundary):
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
@@ -169,7 +243,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:
@@ -182,23 +256,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 b2593c83..cb98d1de 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, RequestStatus
from vllm_rlt.worker.sampler import Sampler
from vllm_rlt.worker.speculative import greedy_accept, rejection_sample
@@ -134,9 +135,7 @@ def record(hidden, batch, cache, **kwargs):
lambda module, inputs, output: head_outputs.append(output.detach().clone())
)
try:
- runner.submit(
- SchedulerOutput(RequestPhase.SPECULATIVE, [ScheduledRequest(r, 3, 4)])
- ).collect()
+ runner.submit(SchedulerOutput([ScheduledRequest(r, WorkKey.SPECULATIVE, 3, 4)])).collect()
finally:
hook.remove()
m.recurrent_prepared = core
@@ -352,7 +351,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())
@@ -366,10 +365,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
@@ -386,7 +385,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
@@ -532,7 +531,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
@@ -694,7 +693,7 @@ def test_abort_suspended_speculative_request_allows_id_reuse():
e.add_request("high", [7, 8, 9], high)
e.step()
assert "low" in e.preemption.snapshots
- assert suspended.phase == RequestPhase.WAITING
+ assert suspended.status == RequestStatus.WAITING
assert e.abort_request("low").finished
assert "low" not in e.preemption.snapshots
e.add_request("low", [5, 6], low)
@@ -709,11 +708,8 @@ def test_abort_suspended_speculative_request_allows_id_reuse():
@pytest.mark.parametrize("chunk_size", [1, 3, 8])
@pytest.mark.parametrize("budget", [2, 7])
-@pytest.mark.parametrize("wavefront", [False, True])
-def test_split_prefill_and_speculation_share_forward_and_match_atomic(
- chunk_size, budget, wavefront
-):
- from vllm_rlt.request import Operation
+def test_split_prefill_and_speculation_share_forward_and_match_atomic(chunk_size, budget):
+ from vllm_rlt.core.execution import Operation
def run(wavefront):
e = engine(
@@ -756,7 +752,7 @@ def test_common_submit_entry_preserves_prefill_and_speculative_results(wavefront
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)
@@ -768,13 +764,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()
@@ -794,7 +791,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)
@@ -818,7 +815,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
@@ -854,11 +851,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().speculative 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 3c36fb5d..ef67c346 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 (
MixedPolicy,
NoRefillPolicy,
@@ -25,7 +26,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
@@ -60,10 +61,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(
@@ -96,8 +96,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(
@@ -132,11 +140,19 @@ def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(wavefront
mid, final = (e.scheduler.requests[rid] for rid in ("mid", "final"))
offset = 4 if wavefront else 0
batch = SchedulerOutput(
- RequestPhase.PREFILLING_RECURRENT,
[
- ScheduledRequest(mid, 0, 2, depth=1, chunk_start=0, prefill_hidden_row_offset=offset),
+ ScheduledRequest(
+ mid,
+ WorkKey.PROMPT_RECURRENT,
+ 0,
+ 2,
+ depth=1,
+ chunk_start=0,
+ prefill_hidden_row_offset=offset,
+ ),
ScheduledRequest(
final,
+ WorkKey.PROMPT_RECURRENT,
0,
2,
depth=3,
@@ -181,7 +197,7 @@ def test_final_depth_prefill_scatter_retains_only_coda_hidden_per_item(wavefront
def test_wavefront_final_depth_does_not_store_any_chunk_hidden(monkeypatch, mixed):
from types import SimpleNamespace
- from vllm_rlt.request import ExecutionBatch, StateSpan
+ from vllm_rlt.core.execution import ExecutionBatch, StateSpan
from vllm_rlt.worker.model_runner import BatchInputs, PreparedExecution
e = make_prefill_engine(
@@ -271,10 +287,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(
@@ -319,10 +334,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
@@ -385,8 +407,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)
]
@@ -468,8 +489,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
@@ -514,8 +543,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
@@ -578,8 +610,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
@@ -630,7 +665,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
@@ -655,20 +690,20 @@ def test_split_prefill_steps_one_depth_per_chunk_and_keeps_the_prompt_ordered():
break
assert all(item.request is request for item in batch.requests)
for item in batch.requests:
- if batch.phase == RequestPhase.PREFILLING_PRELUDE:
+ if item.work == WorkKey.PROMPT_PRELUDE:
prelude_positions.extend(
range(item.token_start, item.token_start + item.token_count)
)
else:
- assert batch.phase == RequestPhase.PREFILLING_RECURRENT
+ assert item.work == WorkKey.PROMPT_RECURRENT
assert batch.loop_count == 1
for position in range(item.token_start, item.token_start + item.token_count):
key = (position, item.depth)
recurrent_visits[key] = recurrent_visits.get(key, 0) + 1
e.scheduler._apply_wavefront_prefill_batch_progress(
- batch.phase, batch.requests, loop_count=batch.loop_count
+ batch.requests[0].work, batch.requests, loop_count=batch.loop_count
)
- if request.phase == RequestPhase.CODA:
+ if e.scheduler.pending_work(request) == WorkKey.DECODE_CODA:
break
assert sorted(prelude_positions) == list(range(len(request.prompt_token_ids)))
@@ -701,14 +736,15 @@ def test_engine_split_prefill_transition_advances_the_same_frontier_contract():
break
assert len(batch.requests) == 1
scheduled_request = batch.requests[0]
- progress = e.scheduler._apply_prefill_progress(
+ progress = e.scheduler._apply_split_prefill_progress(
+ request,
+ scheduled_request.work,
scheduled_request,
- batch.phase,
loop_count=batch.loop_count,
)
trace.append(
(
- batch.phase,
+ scheduled_request.work,
scheduled_request.token_start,
scheduled_request.token_count,
scheduled_request.depth,
@@ -719,12 +755,12 @@ def test_engine_split_prefill_transition_advances_the_same_frontier_contract():
break
assert trace == [
- (RequestPhase.PREFILLING_PRELUDE, 0, 2, None, TOTAL_DEPTHS),
- (RequestPhase.PREFILLING_PRELUDE, 2, 2, None, TOTAL_DEPTHS),
- (RequestPhase.PREFILLING_PRELUDE, 4, 1, None, TOTAL_DEPTHS),
+ (WorkKey.PROMPT_PRELUDE, 0, 2, None, TOTAL_DEPTHS),
+ (WorkKey.PROMPT_PRELUDE, 2, 2, None, TOTAL_DEPTHS),
+ (WorkKey.PROMPT_PRELUDE, 4, 1, None, TOTAL_DEPTHS),
]
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(
@@ -765,7 +801,7 @@ def test_wavefront_prefill_batches_ready_depths_and_reuses_retired_ring_slots():
request = e.scheduler.requests["wave"]
prelude = e.scheduler.schedule()
- assert prelude.phase is RequestPhase.PREFILLING_PRELUDE
+ assert prelude.works == frozenset((WorkKey.PROMPT_PRELUDE,))
assert len(prelude.requests) == TOTAL_DEPTHS
assert [(item.chunk_start, item.prefill_hidden_row_offset) for item in prelude.requests] == [
(0, 0),
@@ -780,13 +816,13 @@ def test_wavefront_prefill_batches_ready_depths_and_reuses_retired_ring_slots():
for _ in range(20):
batch = e.scheduler.schedule()
assert batch is not None
- if batch.phase == RequestPhase.PREFILLING_PRELUDE:
+ if WorkKey.PROMPT_PRELUDE in batch.works:
reused = batch.requests[0]
assert (reused.chunk_start, reused.prefill_hidden_row_offset) == (8, 0)
else:
recurrent_trace.append([(item.token_start, item.depth) for item in batch.requests])
e._update(batch, None)
- if request.phase is RequestPhase.CODA:
+ if e.scheduler.pending_work(request) is WorkKey.DECODE_CODA:
break
assert recurrent_trace[:4] == [
@@ -798,7 +834,7 @@ def test_wavefront_prefill_batches_ready_depths_and_reuses_retired_ring_slots():
assert reused is not None
assert request.num_prefilled_tokens == len(request.prompt_token_ids)
assert request.prefill_reclaimable_tokens == len(request.prompt_token_ids)
- assert request.phase is RequestPhase.CODA
+ assert e.scheduler.pending_work(request) is WorkKey.DECODE_CODA
e.close()
@@ -849,7 +885,10 @@ def test_wavefront_prefill_preserves_hidden_and_all_depth_kv(budget):
break
engine.step()
for request in engine.scheduler.requests.values():
- if request.phase == RequestPhase.CODA and request.request_id not in completed:
+ if (
+ engine.scheduler.pending_work(request) is WorkKey.DECODE_CODA
+ and request.request_id not in completed
+ ):
tensors = [request.hidden_state.clone()]
for depth in range(TOTAL_DEPTHS):
for layer in range(engine.model.config.num_hidden_layers):
@@ -871,7 +910,7 @@ def test_wavefront_prefill_preserves_hidden_and_all_depth_kv(budget):
torch.testing.assert_close(actual, expected, atol=4e-5, rtol=4e-5)
-def test_wavefront_mixed_prelude_prioritizes_decode_before_prompt_window():
+def test_wavefront_mixed_co_batches_compatible_prelude_work():
import torch
scheduler = SchedulerConfig(
@@ -886,18 +925,16 @@ def test_wavefront_mixed_prelude_prioritizes_decode_before_prompt_window():
e.add_request("prompt", [2, 3, 4, 5, 6, 7], SamplingParams(max_tokens=2))
e.scheduler._admit()
decode = e.scheduler.requests["decode"]
- e.scheduler.queues[RequestPhase.PREFILLING_PRELUDE].remove("decode")
+ e.scheduler.remove_pending_work(decode)
decode.generated_token_ids.append(9)
decode.hidden_state = torch.zeros(e.model.config.hidden_size)
- e.scheduler.enqueue(decode, RequestPhase.PRELUDE)
+ e.scheduler.enqueue_work(decode, WorkKey.DECODE_PRELUDE)
batch = e.scheduler.schedule()
- assert batch.is_mixed
+ assert batch.works == frozenset((WorkKey.DECODE_PRELUDE, WorkKey.PROMPT_PRELUDE))
assert batch.num_tokens == 6
- assert batch.phase_for(batch.requests[0]) is RequestPhase.PRELUDE
- prompt_items = [
- item for item in batch.requests if item.phase is RequestPhase.PREFILLING_PRELUDE
- ]
+ assert batch.requests[0].work is WorkKey.DECODE_PRELUDE
+ prompt_items = [item for item in batch.requests if item.work is WorkKey.PROMPT_PRELUDE]
assert [item.chunk_start for item in prompt_items] == [0, 2, 4]
assert sum(item.token_count for item in prompt_items) == 5
e.close()
@@ -1029,7 +1066,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:
@@ -1039,22 +1076,22 @@ def hold_split_ticket(batch):
held_batches = []
def record_batch(batch):
- if batch.phase in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT):
+ if batch.has_split_prefill_work:
held_batches.append(batch)
return hold_split_ticket(batch)
e.executor.submit = record_batch
for _ in range(32):
e.step()
- if request.phase is RequestPhase.CODA:
+ 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 0 < len(held) < 15 # Wavefront batches some chunk-depth tasks together.
assert all(
- batch.loop_count == (0 if batch.phase == RequestPhase.PREFILLING_PRELUDE else 1)
+ batch.loop_count == (0 if WorkKey.PROMPT_PRELUDE in batch.works else 1)
for batch in held_batches
)
assert e.model_runner.split_prefill_submissions == len(held)
@@ -1133,18 +1170,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_reclaimable_tokens == request.num_prefilled_tokens
assert all(left == request.num_prefilled_tokens for left, _ in request.prefill_depth_ranges)
assert e.cache_manager.num_used_blocks == 0
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
@@ -1184,7 +1221,7 @@ def test_priority_admission_preempts_split_prefill_and_matches_reference(wavefro
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],
@@ -1236,7 +1273,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)
@@ -1364,10 +1401,11 @@ 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")
+ parked_work = e.scheduler.remove_pending_work(split_request)
+ assert parked_work is WorkKey.PROMPT_RECURRENT
e.add_request(
"decode",
@@ -1378,7 +1416,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
@@ -1387,7 +1425,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 None
assert split_request.num_prefilled_tokens == 0
# WCPB stages prompt chunks in a prelude wavefront: with
# total_ut_steps=4, prefill_chunk_size=2 and a 4-token budget, both
@@ -1402,7 +1440,8 @@ 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)
+ assert e.scheduler.pending_work(split_request) is WorkKey.PROMPT_RECURRENT
finished = {}
for _ in range(200):
for output in e.step():
@@ -1505,7 +1544,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"]
@@ -1521,7 +1560,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)
@@ -1577,14 +1616,13 @@ 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.clear()
- recurrent.append("a")
+ request_b = e.scheduler.requests["b"]
+ e.scheduler.remove_pending_work(request_b)
e.step() # Advance only request a to depth one.
- recurrent.append("b")
+ e.scheduler.enqueue_work(request_b, WorkKey.PROMPT_RECURRENT)
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()
@@ -1605,19 +1643,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.remove_pending_work(request)
assert e.scheduler._prelude_window_open(request) is False
@@ -1633,9 +1671,7 @@ 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
@@ -1649,7 +1685,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
@@ -1662,32 +1698,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()
@@ -1705,9 +1745,7 @@ 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
@@ -1758,14 +1796,14 @@ def test_staged_prefill_policy_is_selected_by_scheduler_requirements():
)
-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,
)
@@ -1876,8 +1914,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
@@ -1913,7 +1950,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 6290592a..209f34d1 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")
@@ -189,7 +190,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")
@@ -222,7 +223,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()
@@ -240,13 +241,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()
@@ -390,7 +391,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()
@@ -583,7 +584,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
@@ -613,7 +614,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)
@@ -652,7 +654,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:
@@ -691,7 +693,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..df0aa14b 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()
@@ -184,10 +186,11 @@ def test_commit_validation_and_late_write_rejection():
engine.close()
-def test_ack_alone_cannot_release_unsubmitted_or_locally_pending_export():
+def test_ack_alone_cannot_release_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()
@@ -214,6 +217,27 @@ def test_ack_alone_cannot_release_unsubmitted_or_locally_pending_export():
engine.close()
+def test_unsubmitted_export_does_not_consume_compute_slot_or_retire_on_ack():
+ engine = make_stage_engine()
+ connector = make_connector(engine, slots=1)
+ transfer = reserve(connector, "unsubmitted", receiving=False)
+ assert transfer.request.status == RequestStatus.WAITING
+ assert not engine.scheduler.has_admission_slot(transfer.request)
+ assert engine.scheduler.has_admission_capacity()
+
+ connector.acknowledge(transfer.tid)
+ connector.progress()
+ assert transfer.tid in connector.transfers
+ assert not transfer.request.compute_released
+ assert engine.cache_manager.has_transfer_lease(transfer.tid)
+
+ connector.cancel(transfer.tid)
+ connector.progress()
+ assert transfer.tid not in connector.transfers
+ assert engine.cache_manager.num_used_blocks == 0
+ engine.close()
+
+
def test_connector_controls_protect_pending_remote_access():
engine = make_stage_engine(role="decode")
connector = make_connector(engine)
@@ -256,15 +280,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 +300,20 @@ 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 +338,18 @@ 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..ab5fdd7f
--- /dev/null
+++ b/vllm_rlt/core/execution.py
@@ -0,0 +1,112 @@
+"""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
+ prefill_hidden_row_offset: 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 dfcdf0df..cecdb6c9 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 62e0cdee..4f80941e 100644
--- a/vllm_rlt/core/scheduler.py
+++ b/vllm_rlt/core/scheduler.py
@@ -5,25 +5,28 @@
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, requires_staged_prefill
-from vllm_rlt.request import FinishReason, Request, RequestPhase
+from vllm_rlt.request import FinishReason, Request, RequestStatus
from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput
-# Items of these phases cover a contiguous span of prompt tokens, so their KV
+# Items of these works 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
@@ -32,7 +35,6 @@ class ScheduledRequest:
depth: int | None = None
chunk_start: int | None = None
completes_prompt: bool = False
- phase: RequestPhase | None = None
# Row offset in the request's prefill hidden-state buffer. None lets the runner
# derive token_start - chunk_start.
prefill_hidden_row_offset: int | None = None
@@ -40,8 +42,6 @@ class ScheduledRequest:
@dataclass(frozen=True)
class SchedulerOutput:
- # None means a mixed batch; each scheduled request carries its own phase.
- phase: RequestPhase | None
requests: list[ScheduledRequest]
# Actual number of recurrent stages in this submission. None is reserved for
# direct callers that construct a batch without passing through Scheduler.
@@ -54,19 +54,26 @@ def __post_init__(self):
raise ValueError("loop_count must be a nonnegative integer")
@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
- 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 has_split_prefill_work(self) -> bool:
+ return any(
+ item.work in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)
+ for item in self.requests
+ )
+
+ @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:
@@ -91,9 +98,7 @@ class SchedulerUpdate:
outputs: list[Request] = field(default_factory=list)
exited: list[Request] = field(default_factory=list)
finished: list[Request] = field(default_factory=list)
- prefill: list[tuple[ScheduledRequest, RequestPhase, PrefillProgress]] = field(
- default_factory=list
- )
+ prefill: list[tuple[ScheduledRequest, WorkKey, PrefillProgress]] = field(default_factory=list)
next_prelude: list[int] = field(default_factory=list)
retain_signal: list[int] = field(default_factory=list)
committed_tokens: int = 0
@@ -155,11 +160,15 @@ def __init__(
self.eos_token_ids = (eos_token_id,)
self.exit_mode = exit_mode
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.pending_works: dict[str, 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
self.resume_callback = None
+ self.has_preemption_snapshot = None
# Optional connector readiness and export predicates. Never mutate queues or blocks.
self.execution_ready = None
self.exports_request = None
@@ -176,36 +185,105 @@ 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 enqueue(self, request: Request, phase: RequestPhase):
- request.phase = phase
- self.queues[phase].append(request.request_id)
+ 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 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 request.request_id in self.pending_works:
+ raise ValueError(f"request {request.request_id!r} already has pending work")
+ self.work_queues[work].append(request.request_id)
+ self.pending_works[request.request_id] = work
+
+ def pending_work(self, request: Request) -> WorkKey | None:
+ return self.pending_works.get(request.request_id)
+
+ def remove_pending_work(self, request: Request) -> WorkKey | None:
+ pending = self.pending_works.pop(request.request_id, None)
+ if pending is not None:
+ self.work_queues[pending].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]
@@ -223,7 +301,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()
@@ -239,10 +317,24 @@ def _is_preemption_candidate(self, request, requesting_request, *, priority_only
"""
if (
request is requesting_request
- or request.phase == RequestPhase.RECEIVING
+ or request.status == RequestStatus.RECEIVING
+ or (
+ request.status == RequestStatus.WAITING
+ and (
+ priority_only
+ or not self.cache_manager.has_allocation(request.request_id)
+ or self.has_preemption_snapshot is None
+ or not self.has_preemption_snapshot(request.request_id)
+ )
+ )
+ or request.status == RequestStatus.TRANSFERRING
+ or request.status == RequestStatus.FINISHED
or (
- request.phase == RequestPhase.WAITING
- and (priority_only or not self.cache_manager.has_allocation(request.request_id))
+ request.status == RequestStatus.ACTIVE
+ and (
+ request.request_id not in self.admission_slots
+ or request.request_id not in self.pending_works
+ )
)
or request.num_output_placeholders
or request.num_speculative_rounds
@@ -268,7 +360,7 @@ def _select_preemption_victim(
candidates = (
request
for request in self.requests.values()
- if (not waiting_only or request.phase == RequestPhase.WAITING)
+ if (not waiting_only or request.status == RequestStatus.WAITING)
and self._is_preemption_candidate(
request, requesting_request, priority_only=priority_only, excluded=excluded
)
@@ -301,39 +393,28 @@ def preempt(
)
if victim is None:
return False
+ was_waiting = victim.status == RequestStatus.WAITING
self.suspend_callback(victim, offload_kv=not priority_only)
- if victim.phase == RequestPhase.WAITING:
- return True
- for queue in self.queues.values():
- while victim.request_id in queue:
- queue.remove(victim.request_id)
- self.enqueue(victim, RequestPhase.WAITING)
+ if not was_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. The victim retains KV until an allocation actually needs it.
- 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
@@ -343,7 +424,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:
@@ -357,14 +439,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:
@@ -373,7 +454,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
@@ -402,7 +485,10 @@ 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.
@@ -427,7 +513,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):
@@ -466,10 +552,11 @@ def _commit_admission(self, request: Request, plan: AdmissionPlan) -> bool:
return True
def _enter_prefill(self, request):
+ self.activate_request(request)
if self.staged_prefill:
- self.enqueue(request, RequestPhase.PREFILLING_PRELUDE)
+ self.enqueue_work(request, WorkKey.PROMPT_PRELUDE)
else:
- self.enqueue(request, RequestPhase.PREFILL)
+ self.enqueue_work(request, WorkKey.PROMPT_ATOMIC)
def _admit(self):
"""Admit waiting requests: order -> slot -> restore/allocate -> defer/enter.
@@ -478,13 +565,9 @@ def _admit(self):
admission increments earlier blocked requests' bypass counts; reaching
the limit stops further bypasses so active requests can drain.
"""
- if not self.queues[RequestPhase.WAITING]:
+ if not self.waiting:
return
- active_count = sum(
- r.phase not in (RequestPhase.WAITING, RequestPhase.RECEIVING) and not r.compute_released
- for r in self.requests.values()
- )
- if active_count >= self.config.max_num_seqs and (
+ if not self.has_admission_capacity() and (
self.config.policy != "priority" or self.suspend_callback is None
):
return
@@ -496,18 +579,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):
@@ -516,18 +598,17 @@ 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)
admitted = self._commit_admission(request, plan)
while not admitted and self.preempt(request, waiting_only=True):
# Recompute prefix claims and growth reservations after
# releasing a resident allocation; the previous plan is stale.
- plan = self._plan_admission(request, active_count)
+ plan = self._plan_admission(request)
admitted = self._commit_admission(request, plan)
if not admitted:
deferred.append(request_id)
@@ -535,7 +616,6 @@ def _admit(self):
break
continue
- active_count += 1
for blocked_id in deferred:
self.requests[blocked_id].admission_bypasses += 1
if any(
@@ -546,35 +626,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 execution range without advancing request progress.
+ """Snapshot one work item's execution range without advancing progress.
Prompt length 10, prefilled 4, chunk 4, budget 6 -> range [4, 8).
- Wavefront expands a queued request into multiple ranges in take_phases().
- Decode phases contribute one position; its loop may differ by request.
+ 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_reclaimable_tokens
+ request, work, start, 0, depth=0, chunk_start=request.prefill_reclaimable_tokens
)
first_token_depth = request.prefill_depth(start)
# One run per depth frontier: every token of the run already completed
@@ -591,6 +669,7 @@ def _make_scheduled_request(
)
return ScheduledRequest(
request,
+ work,
start,
count,
depth=first_token_depth,
@@ -600,18 +679,18 @@ 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 _make_wavefront_task(
- self, request, phase, token_budget, *, start, depth=None
+ self, request, work, token_budget, *, start, depth=None
) -> ScheduledRequest | None:
"""Snapshot one chunk within the bounded per-request wavefront ring."""
chunk_size = self.config.prefill_chunk_size
@@ -625,7 +704,7 @@ def _make_wavefront_task(
slot = ((chunk_start - origin) // chunk_size) % total_depths
prefill_hidden_row_offset = slot * chunk_size + chunk_offset
- if phase == RequestPhase.PREFILLING_PRELUDE:
+ if work == WorkKey.PROMPT_PRELUDE:
window_end = min(
len(request.prompt_token_ids),
request.prefill_reclaimable_tokens + total_depths * chunk_size,
@@ -644,12 +723,13 @@ def _make_wavefront_task(
return None
return ScheduledRequest(
request,
+ work,
start,
count,
depth=depth,
chunk_start=chunk_start,
completes_prompt=(
- phase == RequestPhase.PREFILLING_RECURRENT
+ work == WorkKey.PROMPT_RECURRENT
and depth == total_depths - 1
and start + count == len(request.prompt_token_ids)
),
@@ -688,30 +768,30 @@ def _prelude_window_open(self, request: Request) -> bool:
return request.num_preluded_tokens < window_end
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]
- if self.wavefront_chunked_prefill and phase in (
- RequestPhase.PREFILLING_PRELUDE,
- RequestPhase.PREFILLING_RECURRENT,
+ for work in works:
+ queue = self.work_queues[work]
+ if self.wavefront_chunked_prefill and work in (
+ WorkKey.PROMPT_PRELUDE,
+ WorkKey.PROMPT_RECURRENT,
):
- token_budget = self._take_wavefront_phase(
- phase,
+ token_budget = self._take_wavefront_work(
+ work,
queue,
token_budget,
scheduled_requests,
@@ -723,111 +803,111 @@ def take_phases(self, phases, *, eligible=None):
while queue and token_budget and len(selected) < self.config.max_num_seqs and remaining:
remaining -= 1
request = self.requests[queue.popleft()]
+ if request.request_id in selected:
+ raise RuntimeError(
+ f"request {request.request_id!r} is queued more than once in a batch"
+ )
if eligible is not None and not eligible(request):
queue.append(request.request_id)
continue
- if phase == RequestPhase.PREFILLING_PRELUDE and not self._prelude_window_open(
- request
- ):
+ 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):
queue.append(request.request_id)
continue
# Selecting A must prevent B's later capacity check from evicting A.
+ self.pending_works.pop(request.request_id, None)
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}
- prompt_phases = selected_phases.intersection(self.policy.PREFILL_PHASES)
- recorded = (
- RequestPhase.RECURRENT
- if RequestPhase.RECURRENT in selected_phases
- else next(iter(prompt_phases))
- if prompt_phases
- else scheduled_requests[0].phase
- )
- self.policy.record_batch(recorded)
loop_count = self._planned_loop_count(scheduled_requests)
scheduled_requests = [
replace(
item,
- completes_prompt=(
- self._scheduled_prefill_completes(item, loop_count)
- if item.phase
- in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT)
- else item.completes_prompt
- ),
+ completes_prompt=self._scheduled_prefill_completes(item, loop_count)
+ if item.work in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)
+ else item.completes_prompt,
)
for item in scheduled_requests
]
- return SchedulerOutput(
- next(iter(selected_phases)) if len(selected_phases) == 1 else None,
- scheduled_requests,
- loop_count,
- )
+ self.policy.record_batch(item.work for item in scheduled_requests)
+ return SchedulerOutput(scheduled_requests, loop_count)
def _planned_loop_count(self, scheduled_requests):
- phases = {item.phase for item in scheduled_requests}
- if phases == {RequestPhase.PREFILL}:
+ works = {item.work for item in scheduled_requests}
+ if works == {WorkKey.PROMPT_ATOMIC}:
return self.total_ut_steps
- if RequestPhase.SPECULATIVE in phases or RequestPhase.CODA in phases:
+ if WorkKey.SPECULATIVE in works or WorkKey.DECODE_CODA in works:
return 0
- if RequestPhase.PRELUDE in phases:
+ if WorkKey.DECODE_PRELUDE in works and WorkKey.PROMPT_PRELUDE in works:
+ return 0
+ if WorkKey.DECODE_PRELUDE in works and not works.intersection(
+ (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)
+ ):
return 0
if self.wavefront_chunked_prefill:
- return 1 if RequestPhase.PREFILLING_RECURRENT in phases else 0
- if RequestPhase.PREFILLING_PRELUDE in phases:
+ return int(WorkKey.PROMPT_RECURRENT in works or WorkKey.DECODE_RECURRENT in works)
+ if WorkKey.PROMPT_PRELUDE in works:
# In mixed mode, stage just the embeddings when recurrent decode work
# is ready so the next depth can join that batch. Without a reason to
# yield, PRELUDE can execute the full prompt depth in this submission.
- if self.config.mode == "mixed" and self.has_ready(RequestPhase.RECURRENT):
+ if self.config.mode == "mixed" and self.has_ready_work(WorkKey.DECODE_RECURRENT):
return 0
return self.total_ut_steps
- if RequestPhase.PREFILLING_RECURRENT in phases:
+ if WorkKey.PROMPT_RECURRENT in works:
if self.config.mode == "mixed" and (
- RequestPhase.RECURRENT in phases or self.has_ready(RequestPhase.RECURRENT)
+ WorkKey.DECODE_RECURRENT in works or self.has_ready_work(WorkKey.DECODE_RECURRENT)
):
return 1
return min(self._max_contiguous_prefill_loops(item) for item in scheduled_requests)
- if RequestPhase.RECURRENT in phases:
+ if WorkKey.DECODE_RECURRENT in works:
return 1
- raise ValueError(f"cannot derive recurrent boundary for phases: {phases}")
+ return 0
def _batch_loop_count(self, batch):
if batch.loop_count is not None:
return batch.loop_count
- phases = {batch.phase_for(item) for item in batch.requests}
- if RequestPhase.PREFILL in phases:
+ works = batch.works
+ if WorkKey.PROMPT_ATOMIC in works:
return self.total_ut_steps
- if phases.intersection((RequestPhase.PRELUDE, RequestPhase.CODA, RequestPhase.SPECULATIVE)):
+ if works.intersection(
+ (
+ WorkKey.PROMPT_PRELUDE,
+ WorkKey.DECODE_PRELUDE,
+ WorkKey.DECODE_CODA,
+ WorkKey.SPECULATIVE,
+ )
+ ):
+ if WorkKey.PROMPT_RECURRENT in works or WorkKey.DECODE_RECURRENT in works:
+ return 1
return 0
- if phases.intersection((RequestPhase.PREFILLING_RECURRENT, RequestPhase.RECURRENT)):
+ if works.intersection((WorkKey.PROMPT_RECURRENT, WorkKey.DECODE_RECURRENT)):
return 1
return 0
@@ -844,8 +924,8 @@ def _max_contiguous_prefill_loops(self, scheduled_request):
return max(1, loops)
def _scheduled_prefill_completes(self, scheduled_request, loop_count):
- phase = scheduled_request.phase
- depth = scheduled_request.depth if phase == RequestPhase.PREFILLING_RECURRENT else 0
+ work = scheduled_request.work
+ depth = scheduled_request.depth if work == WorkKey.PROMPT_RECURRENT else 0
return (
loop_count > 0
and depth + loop_count == self.total_ut_steps
@@ -853,8 +933,8 @@ def _scheduled_prefill_completes(self, scheduled_request, loop_count):
== len(scheduled_request.request.prompt_token_ids)
)
- def _take_wavefront_phase(
- self, phase, queue, token_budget, scheduled_requests, selected, eligible
+ def _take_wavefront_work(
+ self, work, queue, token_budget, scheduled_requests, selected, eligible
):
"""Round-robin each request's chunks/depths into the shared token budget."""
remaining = len(queue)
@@ -863,17 +943,15 @@ def _take_wavefront_phase(
while queue and remaining and len(candidates) < owner_limit and token_budget:
remaining -= 1
request = self.requests[queue.popleft()]
+ queue.append(request.request_id)
if eligible is not None and not eligible(request):
- queue.append(request.request_id)
continue
- if phase == RequestPhase.PREFILLING_PRELUDE and not self._prelude_window_open(request):
- queue.append(request.request_id)
+ if work == WorkKey.PROMPT_PRELUDE and not self._prelude_window_open(request):
continue
- if phase == RequestPhase.PREFILLING_RECURRENT and not any(
+ if work == WorkKey.PROMPT_RECURRENT and not any(
start < min(end, request.num_preluded_tokens)
for start, end in request.prefill_depth_ranges
):
- queue.append(request.request_id)
continue
candidates.append(request)
@@ -882,20 +960,20 @@ def _take_wavefront_phase(
added = set()
blocked = set()
- if phase == RequestPhase.PREFILLING_PRELUDE:
+ if work == WorkKey.PROMPT_PRELUDE:
cursors = {request.request_id: request.num_preluded_tokens for request in candidates}
for _ in range(len(candidates[0].prefill_depth_ranges)):
for request in candidates:
request_id = request.request_id
if token_budget <= 0:
break
- if request_id in blocked or request.phase != phase:
+ if request_id in blocked or request.status != RequestStatus.ACTIVE:
continue
start = cursors[request_id]
if start >= len(request.prompt_token_ids):
continue
scheduled_request = self._make_wavefront_task(
- request, phase, token_budget, start=start
+ request, work, token_budget, start=start
)
if scheduled_request is None:
continue
@@ -903,7 +981,7 @@ def _take_wavefront_phase(
if not self._ensure_execution_capacity(request, frontier, selected=selected):
blocked.add(request_id)
continue
- scheduled_requests.append(replace(scheduled_request, phase=phase))
+ scheduled_requests.append(scheduled_request)
selected.add(request_id)
added.add(request_id)
cursors[request_id] += scheduled_request.token_count
@@ -915,14 +993,14 @@ def _take_wavefront_phase(
request_id = request.request_id
if token_budget <= 0:
break
- if request_id in blocked or request.phase != phase:
+ if request_id in blocked or request.status != RequestStatus.ACTIVE:
continue
start, end = request.prefill_depth_ranges[depth]
end = min(end, request.num_preluded_tokens)
if start >= end:
continue
scheduled_request = self._make_wavefront_task(
- request, phase, token_budget, depth=depth, start=start
+ request, work, token_budget, depth=depth, start=start
)
if scheduled_request is None:
continue
@@ -936,15 +1014,21 @@ def _take_wavefront_phase(
token_budget -= scheduled_request.token_count
tasks.sort(key=lambda item: (item[0], item[1].token_start))
- scheduled_requests.extend(
- replace(scheduled_request, phase=phase) for _, scheduled_request in tasks
- )
-
- for request in candidates:
- # Growing another candidate may have suspended this request and put
- # it in WAITING. It must not also return to its old prompt queue.
- if request.request_id not in added and request.phase == phase:
- queue.append(request.request_id)
+ scheduled_requests.extend(scheduled_request for _, scheduled_request in tasks)
+
+ # Remove selected and suspended candidates in one queue pass. Other
+ # candidates retain their pending work and queue position.
+ removed = {
+ request.request_id
+ for request in candidates
+ if request.request_id in added or request.status != RequestStatus.ACTIVE
+ }
+ if removed:
+ queue_copy = [request_id for request_id in queue if request_id not in removed]
+ queue.clear()
+ queue.extend(queue_copy)
+ for request_id in removed:
+ self.pending_works.pop(request_id, None)
return token_budget
def schedule(self, *, prefer_recurrent=False) -> SchedulerOutput | None:
@@ -959,14 +1043,17 @@ def schedule(self, *, prefer_recurrent=False) -> SchedulerOutput | None:
return None
return self.policy.schedule(self, prefer_recurrent=prefer_recurrent)
- def _prompt_destination(self, request: Request) -> RequestPhase:
+ def _prompt_destination(self, request: Request):
exporting = self.exports_request is not None and self.exports_request(request)
- return RequestPhase.TRANSFERRING if exporting else RequestPhase.CODA
+ if exporting:
+ self.suspend_for_transfer(request)
+ else:
+ self.enqueue_work(request, WorkKey.DECODE_CODA)
def _apply_split_prefill_progress(
self,
request: Request,
- phase: RequestPhase,
+ work: WorkKey,
scheduled_request: ScheduledRequest,
*,
loop_count: int | None = None,
@@ -981,26 +1068,24 @@ def _apply_split_prefill_progress(
start = scheduled_request.token_start
end = start + scheduled_request.token_count
if loop_count is None:
- loop_count = 0 if phase == RequestPhase.PREFILLING_PRELUDE else 1
-
- if phase == RequestPhase.PREFILLING_PRELUDE:
- request.prefill_reclaimable_tokens = start
+ loop_count = 0 if work == WorkKey.PROMPT_PRELUDE else 1
+ if work == WorkKey.PROMPT_PRELUDE:
+ request.prefill_reclaimable_tokens = min(request.prefill_reclaimable_tokens, start)
request.num_preluded_tokens = end
depth = 0
else:
depth = scheduled_request.depth
if loop_count <= 0:
- if phase != RequestPhase.PREFILLING_PRELUDE:
+ if work != WorkKey.PROMPT_PRELUDE:
raise ValueError("recurrent prefill progress requires at least one loop")
- self.enqueue(request, RequestPhase.PREFILLING_RECURRENT)
+ self.enqueue_work(request, WorkKey.PROMPT_RECURRENT)
return PrefillProgress(False, False)
-
boundaries = request.prefill_depth_ranges
if depth + loop_count > total_depths:
raise ValueError("prefill execution range exceeds model loop depth")
for offset in range(loop_count):
current_depth = depth + offset
- if phase == RequestPhase.PREFILLING_RECURRENT and (
+ if work == WorkKey.PROMPT_RECURRENT and (
boundaries[current_depth][0] != start or boundaries[current_depth][1] < end
):
raise RuntimeError("prefill execution range crosses an incomplete depth frontier")
@@ -1013,7 +1098,7 @@ def _apply_split_prefill_progress(
final_depth_reached = depth + loop_count == total_depths
chunk_complete = final_depth_reached and end == request.num_preluded_tokens
if not final_depth_reached or not chunk_complete:
- self.enqueue(request, RequestPhase.PREFILLING_RECURRENT)
+ self.enqueue_work(request, WorkKey.PROMPT_RECURRENT)
return PrefillProgress(False, False)
self.cache_manager.publish_prefix(
@@ -1026,23 +1111,17 @@ def _apply_split_prefill_progress(
prompt_complete = end == len(request.prompt_token_ids)
if prompt_complete:
request.loops_done = total_depths
- next_phase = self._prompt_destination(request)
+ self._prompt_destination(request)
else:
- next_phase = RequestPhase.PREFILLING_PRELUDE
- self.enqueue(request, next_phase)
+ self.enqueue_work(request, WorkKey.PROMPT_PRELUDE)
return PrefillProgress(True, prompt_complete)
def _apply_prefill_progress(
- self,
- scheduled_request: ScheduledRequest,
- phase: RequestPhase,
- *,
- loop_count: int | None = None,
- event=None,
+ self, scheduled_request: ScheduledRequest, work: WorkKey, *, loop_count=None, event=None
) -> PrefillProgress:
"""One prompt frontier update for local and remotely consumed outputs."""
request = scheduled_request.request
- 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(
@@ -1050,17 +1129,17 @@ def _apply_prefill_progress(
)
if final:
request.loops_done = self.total_ut_steps
- self.enqueue(
- request, self._prompt_destination(request) if final else RequestPhase.PREFILL
- )
+ self._prompt_destination(request)
+ else:
+ self.enqueue_work(request, WorkKey.PROMPT_ATOMIC)
return PrefillProgress(True, final)
return self._apply_split_prefill_progress(
- request, phase, scheduled_request, loop_count=loop_count, event=event
+ request, work, scheduled_request, loop_count=loop_count, event=event
)
def _apply_wavefront_prefill_batch_progress(
self,
- phase: RequestPhase,
+ work: WorkKey,
scheduled_requests: list[ScheduledRequest],
*,
loop_count: int = 1,
@@ -1073,14 +1152,14 @@ def _apply_wavefront_prefill_batch_progress(
request = scheduled_requests[0].request
if any(item.request is not request for item in scheduled_requests):
raise ValueError("wavefront progress spans must share one request")
- if phase not in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT):
- raise ValueError(f"invalid wavefront prefill phase: {phase}")
- expected_loops = 0 if phase == RequestPhase.PREFILLING_PRELUDE else 1
+ if work not in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT):
+ raise ValueError(f"invalid wavefront prefill work: {work}")
+ expected_loops = 0 if work == WorkKey.PROMPT_PRELUDE else 1
if loop_count != expected_loops:
raise ValueError("wavefront prefill requires one recurrent depth per scheduling step")
transfer_start = request.prefill_reclaimable_tokens
- if phase == RequestPhase.PREFILLING_PRELUDE:
+ if work == WorkKey.PROMPT_PRELUDE:
frontier = request.num_preluded_tokens
for item in scheduled_requests:
if item.token_start != frontier:
@@ -1089,10 +1168,10 @@ def _apply_wavefront_prefill_batch_progress(
request.num_preluded_tokens = frontier
progress = PrefillProgress(False, False, transfer_start)
representative = scheduled_requests[0]
- next_phase = (
- RequestPhase.PREFILLING_PRELUDE
+ next_work = (
+ WorkKey.PROMPT_PRELUDE
if self._prelude_window_open(request)
- else RequestPhase.PREFILLING_RECURRENT
+ else WorkKey.PROMPT_RECURRENT
)
else:
chunk_complete = False
@@ -1132,15 +1211,16 @@ def _apply_wavefront_prefill_batch_progress(
)
if prompt_complete:
request.loops_done = self.total_ut_steps
- next_phase = self._prompt_destination(request)
+ next_work = self._prompt_destination(request)
else:
- next_phase = (
- RequestPhase.PREFILLING_PRELUDE
+ next_work = (
+ WorkKey.PROMPT_PRELUDE
if self._prelude_window_open(request)
- else RequestPhase.PREFILLING_RECURRENT
+ else WorkKey.PROMPT_RECURRENT
)
- self.enqueue(request, next_phase)
+ if next_work is not None and request.status == RequestStatus.ACTIVE:
+ self.enqueue_work(request, next_work)
return representative, progress
def _wavefront_prefill_groups(self, batch: SchedulerOutput):
@@ -1148,14 +1228,14 @@ def _wavefront_prefill_groups(self, batch: SchedulerOutput):
if not self.wavefront_chunked_prefill:
return groups
for item in batch.requests:
- phase = batch.phase_for(item)
+ work = item.work
request = item.request
if (
- phase in (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT)
+ work in (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)
and self.requests.get(request.request_id) is request
- and request.phase != RequestPhase.FINISHED
+ and request.status != RequestStatus.FINISHED
):
- groups.setdefault((request.request_id, phase), []).append(item)
+ groups.setdefault((request.request_id, work), []).append(item)
return groups
def _should_exit(self, request: Request) -> bool:
@@ -1221,15 +1301,15 @@ def update_from_output(
# Compare object identity so an old result cannot update the new request.
if (
self.requests.get(request.request_id) is not request
- or request.phase == RequestPhase.FINISHED
+ or request.status == RequestStatus.FINISHED
):
continue
- phase = batch.phase_for(scheduled_request)
- if self.wavefront_chunked_prefill and phase in (
- RequestPhase.PREFILLING_PRELUDE,
- RequestPhase.PREFILLING_RECURRENT,
+ work = scheduled_request.work
+ if self.wavefront_chunked_prefill and work in (
+ WorkKey.PROMPT_PRELUDE,
+ WorkKey.PROMPT_RECURRENT,
):
- grouped = wavefront_groups.pop((request.request_id, phase), None)
+ grouped = wavefront_groups.pop((request.request_id, work), None)
if grouped is not None:
event = (
request_events.get(request.request_id)
@@ -1237,14 +1317,14 @@ def update_from_output(
else result.completion
)
representative, progress = self._apply_wavefront_prefill_batch_progress(
- phase, grouped, loop_count=self._batch_loop_count(batch), event=event
+ work, grouped, loop_count=self._batch_loop_count(batch), event=event
)
- update.prefill.append((representative, phase, progress))
+ update.prefill.append((representative, work, progress))
continue
- if phase in (
- RequestPhase.PREFILL,
- RequestPhase.PREFILLING_PRELUDE,
- RequestPhase.PREFILLING_RECURRENT,
+ if work in (
+ WorkKey.PROMPT_ATOMIC,
+ WorkKey.PROMPT_PRELUDE,
+ WorkKey.PROMPT_RECURRENT,
):
event = (
request_events.get(request.request_id)
@@ -1252,18 +1332,15 @@ def update_from_output(
else result.completion
)
progress = self._apply_prefill_progress(
- scheduled_request,
- phase,
- loop_count=self._batch_loop_count(batch),
- event=event,
+ scheduled_request, work, loop_count=self._batch_loop_count(batch), event=event
)
- update.prefill.append((scheduled_request, phase, progress))
- elif phase == RequestPhase.PRELUDE:
+ update.prefill.append((scheduled_request, work, progress))
+ elif work == WorkKey.DECODE_PRELUDE:
request.loops_done = 0
request.remaining_probability = 1.0
request.pending_exit_depth = None
- self.enqueue(request, RequestPhase.RECURRENT)
- elif phase == RequestPhase.RECURRENT:
+ self.enqueue_work(request, WorkKey.DECODE_RECURRENT)
+ elif work == WorkKey.DECODE_RECURRENT:
request.loops_done += 1
if self.exit_mode == "trace":
should_exit = self._trace_exit(request)
@@ -1284,8 +1361,10 @@ def update_from_output(
request.pending_exit_depth = request.loops_done + 1
if should_exit:
update.exited.append(request)
- self.enqueue(request, RequestPhase.CODA if should_exit else RequestPhase.RECURRENT)
- elif phase == RequestPhase.CODA:
+ self.enqueue_work(request, WorkKey.DECODE_CODA)
+ else:
+ self.enqueue_work(request, WorkKey.DECODE_RECURRENT)
+ elif work == WorkKey.DECODE_CODA:
if delivered:
if not result.output_indices or result.output_indices[row] != len(
request.generated_token_ids
@@ -1298,14 +1377,12 @@ def update_from_output(
if self._record_token(request, result.token_ids[row], depth):
update.finished.append(request)
elif not delivered:
- self.enqueue(
- request,
- RequestPhase.SPECULATIVE
- if self.speculative_config
- else RequestPhase.PRELUDE,
+ next_work = (
+ WorkKey.SPECULATIVE if self.speculative_config else WorkKey.DECODE_PRELUDE
)
+ self.enqueue_work(request, next_work)
update.outputs.append(request)
- elif phase == RequestPhase.SPECULATIVE:
+ elif work == WorkKey.SPECULATIVE:
spec = result.speculative[row]
if result.frontiers is not None:
request.num_output_placeholders -= (
@@ -1340,8 +1417,8 @@ def update_from_output(
else scheduled_request.token_start + emitted,
)
request.loops_done = 0
- if request.request_id not in self.queues[RequestPhase.SPECULATIVE]:
- self.enqueue(request, RequestPhase.SPECULATIVE)
+ if self.pending_work(request) is None:
+ self.enqueue_work(request, WorkKey.SPECULATIVE)
update.outputs.append(request)
return update
@@ -1360,15 +1437,15 @@ def update_after_submission(
request = scheduled_request.request
if (
self.requests.get(request.request_id) is not request
- or request.phase == RequestPhase.FINISHED
+ or request.status == RequestStatus.FINISHED
):
continue
- phase = batch.phase_for(scheduled_request)
- if self.wavefront_chunked_prefill and phase in (
- RequestPhase.PREFILLING_PRELUDE,
- RequestPhase.PREFILLING_RECURRENT,
+ work = scheduled_request.work
+ if self.wavefront_chunked_prefill and work in (
+ WorkKey.PROMPT_PRELUDE,
+ WorkKey.PROMPT_RECURRENT,
):
- grouped = wavefront_groups.pop((request.request_id, phase), None)
+ grouped = wavefront_groups.pop((request.request_id, work), None)
if grouped is not None:
event = (
request_events.get(request.request_id)
@@ -1376,14 +1453,14 @@ def update_after_submission(
else completion
)
representative, progress = self._apply_wavefront_prefill_batch_progress(
- phase, grouped, loop_count=self._batch_loop_count(batch), event=event
+ work, grouped, loop_count=self._batch_loop_count(batch), event=event
)
- update.prefill.append((representative, phase, progress))
+ update.prefill.append((representative, work, progress))
continue
- if phase in (
- RequestPhase.PREFILL,
- RequestPhase.PREFILLING_PRELUDE,
- RequestPhase.PREFILLING_RECURRENT,
+ if work in (
+ WorkKey.PROMPT_ATOMIC,
+ WorkKey.PROMPT_PRELUDE,
+ WorkKey.PROMPT_RECURRENT,
):
event = (
request_events.get(request.request_id)
@@ -1391,31 +1468,28 @@ def update_after_submission(
else completion
)
progress = self._apply_prefill_progress(
- scheduled_request,
- phase,
- loop_count=self._batch_loop_count(batch),
- event=event,
+ scheduled_request, work, loop_count=self._batch_loop_count(batch), event=event
)
- update.prefill.append((scheduled_request, phase, progress))
- elif phase == RequestPhase.PRELUDE:
+ update.prefill.append((scheduled_request, work, progress))
+ elif work == WorkKey.DECODE_PRELUDE:
request.loops_done = 0
request.remaining_probability = 1.0
request.pending_exit_depth = None
- self.enqueue(request, RequestPhase.RECURRENT)
- elif phase == RequestPhase.CODA:
+ self.enqueue_work(request, WorkKey.DECODE_RECURRENT)
+ elif work == WorkKey.DECODE_CODA:
request.num_output_placeholders += 1
if request.num_scheduled_outputs < request.sampling_params.max_tokens:
update.next_prelude.append(row)
- self.enqueue(request, RequestPhase.PRELUDE)
- elif phase == RequestPhase.SPECULATIVE:
+ self.enqueue_work(request, WorkKey.DECODE_PRELUDE)
+ elif work == WorkKey.SPECULATIVE:
request.num_speculative_rounds += 1
request.num_output_placeholders += (
1
if self.speculative_config.mode == "pipeline"
else scheduled_request.token_count
)
- self.enqueue(request, RequestPhase.SPECULATIVE)
- elif phase == RequestPhase.RECURRENT:
+ self.enqueue_work(request, WorkKey.SPECULATIVE)
+ elif work == WorkKey.DECODE_RECURRENT:
consume_signal = (
request.loops_done > 0
and request.sampling_params.exit_threshold < 1
@@ -1446,5 +1520,8 @@ def update_after_submission(
and request.sampling_params.exit_threshold < 1
):
update.retain_signal.append(row)
- self.enqueue(request, RequestPhase.CODA if should_exit else RequestPhase.RECURRENT)
+ self.enqueue_work(
+ request,
+ WorkKey.DECODE_CODA if should_exit else WorkKey.DECODE_RECURRENT,
+ )
return update
diff --git a/vllm_rlt/core/scheduling_policy.py b/vllm_rlt/core/scheduling_policy.py
index 39b405ec..4e32073f 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,14 @@ 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 +33,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 +58,16 @@ class PromptOnlyPolicy(SchedulingPolicy):
def __init__(self, config, *, staged_prefill):
super().__init__(config)
- self.PREFILL_PHASES = (
- (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT)
+ self.PREFILL_WORKS = (
+ (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)
if staged_prefill
- else (RequestPhase.PREFILL,)
+ else (WorkKey.PROMPT_ATOMIC,)
)
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 +77,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 +115,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,55 +146,52 @@ class SpeculativePolicy(SchedulingPolicy):
def __init__(self, config):
super().__init__(config)
- self.PREFILL_PHASES = (
- (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT)
+ self.PREFILL_WORKS = (
+ (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)
if requires_staged_prefill(config, produces_tokens=True)
- else (RequestPhase.PREFILL,)
+ else (WorkKey.PROMPT_ATOMIC,)
)
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 for staged prompt work and recurrent depth boundaries."""
- PREFILL_PHASES = (RequestPhase.PREFILLING_PRELUDE, RequestPhase.PREFILLING_RECURRENT)
+ PREFILL_WORKS = (WorkKey.PROMPT_PRELUDE, WorkKey.PROMPT_RECURRENT)
class MixedPolicy(SplitPrefillPolicy):
- """Choose compatible prompt/generation work; execution is policy-agnostic."""
+ """Combine compatible prompt and decode work while sharing refill ordering."""
+
+ 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 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)
- )
- if prefer_recurrent and scheduler.has_ready(RequestPhase.RECURRENT):
- return scheduler._take(RequestPhase.RECURRENT)
- if (
- scheduler.wavefront_chunked_prefill
- and scheduler.has_ready(RequestPhase.PRELUDE)
- and scheduler.has_ready(RequestPhase.PREFILLING_PRELUDE)
+ if self.has_ready_exact_work(scheduler, WorkKey.DECODE_RECURRENT) and (
+ self.has_ready_exact_work(scheduler, WorkKey.PROMPT_RECURRENT)
):
- return scheduler.take_phases((RequestPhase.PRELUDE, RequestPhase.PREFILLING_PRELUDE))
- batch = super().schedule(scheduler, prefer_recurrent=prefer_recurrent)
- return batch
+ return scheduler.take_works(RECURRENT_WORKS)
+ if prefer_recurrent and self.has_ready_exact_work(scheduler, WorkKey.DECODE_RECURRENT):
+ return scheduler.take_works((WorkKey.DECODE_RECURRENT,))
+ return super().schedule(scheduler, prefer_recurrent=prefer_recurrent)
def create_scheduling_policy(
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 f8be4a75..99c6664f 100644
--- a/vllm_rlt/engine/preemption.py
+++ b/vllm_rlt/engine/preemption.py
@@ -20,21 +20,27 @@ def discard_snapshot(self, request_id: str) -> None:
"""Discard saved state; normal request cleanup releases any resident KV."""
self.snapshots.pop(request_id, None)
+ def has_snapshot(self, request_id: str) -> bool:
+ return request_id in self.snapshots
+
def suspend(self, victim, *, offload_kv=True):
"""Release a runner slot, optionally reclaiming KV capacity as well.
An already suspended resident request may be offloaded under later KV
- pressure. Its saved phase and RNG must survive that conversion.
+ pressure. Its saved work and RNG must survive that conversion.
"""
e = self.engine
cache = e.cache_manager
previous = self.snapshots.get(victim.request_id)
if previous is None:
e.model_runner.synchronize()
- e._rewind_split_prefill(victim)
- phase, hidden, token = victim.phase, victim.hidden_state, victim.input_token_tensor
+ work = e.scheduler.pending_work(victim)
+ if work is None:
+ raise ValueError("a preemption victim must have pending work")
+ work = e._rewind_split_prefill(victim, work)
+ hidden, token = victim.hidden_state, victim.input_token_tensor
else:
- phase, hidden, token = previous["phase"], previous["hidden"], previous["token"]
+ work, hidden, token = previous["work"], previous["hidden"], previous["token"]
device = torch.device("cpu") if offload_kv else cache.device
runner = e.speculative_runner
@@ -46,8 +52,8 @@ def suspend(self, victim, *, offload_kv=True):
# Build a new record before releasing resources. In particular, a failed
# offload must leave the prior resident state available for resumption.
snapshot = dict(
+ work=work,
pipeline=pipeline,
- phase=phase,
kv=cache.snapshot(victim.request_id) if offload_kv else None,
hidden=None if hidden is None else hidden.to(device).clone(),
token=None if token is None else token.to(device).clone(),
@@ -68,7 +74,7 @@ def suspend(self, victim, *, offload_kv=True):
self.snapshots[victim.request_id] = snapshot
def resume(self, request):
- """Return the saved phase, None without saved state, 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
@@ -91,4 +97,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 a83d5101..a2958488 100644
--- a/vllm_rlt/engine/stage_engine.py
+++ b/vllm_rlt/engine/stage_engine.py
@@ -10,6 +10,7 @@
SchedulerConfig,
resolve_cuda_graph_rows,
)
+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 Scheduler
@@ -18,7 +19,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
@@ -262,6 +263,7 @@ def __init__(
if scheduler_config.enable_preemption:
self.scheduler.suspend_callback = self.preemption.suspend
self.scheduler.resume_callback = self.preemption.resume
+ self.scheduler.has_preemption_snapshot = self.preemption.has_snapshot
@property
def connector(self):
@@ -380,7 +382,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(
@@ -411,7 +413,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]
@@ -503,30 +505,27 @@ 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 _notify_connector_prefill(self, scheduled_request, phase, progress):
+ def _notify_connector_prefill(self, scheduled_request, work, progress):
request = scheduled_request.request
if self.connector is not None and self.connector.exports(request):
self.connector.prefill_submitted(
request,
- phase=phase,
+ work=work,
start=(
scheduled_request.token_start
- if phase == RequestPhase.PREFILL
+ if work == WorkKey.PROMPT_ATOMIC
else progress.transfer_start
if progress.transfer_start is not None
else scheduled_request.chunk_start
@@ -565,8 +564,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):
@@ -581,12 +579,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
+ self.scheduler.release_admission_slot(request)
request.compute_released = True
return True
@@ -596,7 +595,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:
@@ -617,8 +616,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)
@@ -631,16 +630,13 @@ 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 _rewind_split_prefill(self, request):
- """Rewind an in-flight chunk; the caller rebuilds its queue membership."""
- if 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_window_origin = frontier
@@ -650,11 +646,11 @@ def _rewind_split_prefill(self, request):
self.model.config.total_ut_steps,
start=frontier,
)
- request.phase = RequestPhase.PREFILLING_PRELUDE
+ return WorkKey.PROMPT_PRELUDE
def _apply_update(self, update) -> list[RequestOutput]:
- for scheduled_request, phase, progress in update.prefill:
- self._notify_connector_prefill(scheduled_request, phase, progress)
+ for scheduled_request, work, progress in update.prefill:
+ self._notify_connector_prefill(scheduled_request, work, progress)
if update.exited:
self.model_runner.finalize_many(update.exited)
if self.speculative_runner is not None:
@@ -680,11 +676,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
if self.speculative_runner is not None:
@@ -720,7 +712,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
@@ -741,7 +733,7 @@ def _step_async_speculative(self):
if (
self.pipeline_speculative
and self.scheduler.config.enable_preemption
- and self.scheduler.queues[RequestPhase.WAITING]
+ and self.scheduler.waiting
):
while self._pending_speculative:
outputs.extend(self._collect_speculative())
@@ -752,7 +744,7 @@ 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)
self._apply_update(self.scheduler.update_after_submission(batch))
@@ -793,9 +785,10 @@ def _step_async(self):
# If the preceding core is still running, boundary work can overlap
# the independent next core. Once it has completed, refill first:
# splitting on a briefly pending coda readback fragments the batch.
- prefer_recurrent=self._overlap_boundary
+ prefer_recurrent=self.scheduler.config.mode != "mixed"
+ and 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
@@ -807,7 +800,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.
@@ -826,7 +819,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 = (
@@ -835,11 +828,11 @@ def _step_async(self):
)
signals = {}
for row, scheduled_request in enumerate(batch.requests):
- phase = batch.phase_for(scheduled_request)
+ work = scheduled_request.work
request = scheduled_request.request
- if phase == RequestPhase.PRELUDE:
+ if work == WorkKey.DECODE_PRELUDE:
self._pending_exit_signals.pop(request.request_id, None)
- elif phase == RequestPhase.RECURRENT:
+ elif work == WorkKey.DECODE_RECURRENT:
# Submission is already enqueued. Only now may the preceding score
# block CPU delivery; max-depth/trace decisions require no readback.
previous = self._pending_exit_signals.pop(request.request_id, None)
diff --git a/vllm_rlt/pd/connector.py b/vllm_rlt/pd/connector.py
index 121a6538..29df1b23 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,12 +184,12 @@ 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 or (
- phase == RequestPhase.PREFILLING_RECURRENT
+ if work == WorkKey.PROMPT_PRELUDE or (
+ work == WorkKey.PROMPT_RECURRENT
and self.engine.scheduler.config.wavefront_chunked_prefill
):
transfer.transfer_overlaps += self.transport.busy(transfer.tid)
@@ -373,7 +374,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
)
@@ -411,7 +412,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 47f58ad1..47980356 100644
--- a/vllm_rlt/pd/engine.py
+++ b/vllm_rlt/pd/engine.py
@@ -17,7 +17,7 @@
from vllm_rlt.models import prepare_model
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
@@ -396,7 +396,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 ef784437..5325cfc9 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.
@@ -108,94 +96,7 @@ 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 lets the runner derive token_start - chunk_start. Otherwise this is a
- # row offset into the request's prefill hidden-state buffer, including its ring slot.
- chunk_start: int | None = None
- retain_chunk: bool = False
- retain_last: bool = True
- 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
- # Row offset in the request's prefill hidden-state buffer. None lets the runner
- # derive token_start - chunk_start.
- prefill_hidden_row_offset: int | None = None
-
-
-@dataclass(frozen=True)
-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)
+@dataclass
class RequestOutput:
request_id: str
prompt_token_ids: list[int]
@@ -212,7 +113,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 6b708df2..d4b72857 100644
--- a/vllm_rlt/worker/async_speculative.py
+++ b/vllm_rlt/worker/async_speculative.py
@@ -8,9 +8,10 @@
import numpy as np
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.cuda_graph import (
CodaGraphs,
DeviceRecurrentGraphs,
diff --git a/vllm_rlt/worker/executor.py b/vllm_rlt/worker/executor.py
index bad146d4..35d4d28d 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,43 +20,43 @@ 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
planned_loops = scheduled_batch.loop_count
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 = planned_loops
+ if loop_count is None:
+ 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
)
- execution_loops = (
- planned_loops
- if planned_loops is not None
- else total_depths
- if phase == RequestPhase.PREFILL
- else int(operation == Operation.RECURRENT)
- )
spans.append(
StateSpan(
request=request,
@@ -64,24 +64,28 @@ 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 depth + execution_loops < total_depths,
+ retain_chunk=chunk
+ and (work == WorkKey.PROMPT_PRELUDE or depth + loop_count < 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]
)
- if prompt_input and (not chunk or execution_loops > 0)
+ if prompt_input and (not chunk or loop_count > 0)
else None,
prefill_hidden_row_offset=(
scheduled_request.prefill_hidden_row_offset if chunk else None
),
)
)
- operations.add(operation)
- loop_counts.add(execution_loops)
- if len(operations) != 1 or len(loop_counts) != 1:
- raise ValueError("batch requests must share an operation and execution boundary")
- return ExecutionBatch(operations.pop(), tuple(spans), loop_counts.pop())
+ 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 ada23d20..5dc7bbb0 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,
@@ -1180,7 +1181,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
@@ -1188,7 +1189,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 = (
@@ -1254,7 +1255,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 85f0e235..ff4a89d9 100644
--- a/vllm_rlt/worker/speculative.py
+++ b/vllm_rlt/worker/speculative.py
@@ -9,8 +9,9 @@
import torch
+from vllm_rlt.core.execution import Operation
from vllm_rlt.core.scheduler import SchedulerOutput
-from vllm_rlt.request import Operation, Request
+from vllm_rlt.request import Request
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