Repository navigation
Conversation
introduce RequestStatus for request lifecycle and WorkKey for pending work while keeping Operation as the model execution boundary. Move scheduling to WorkKey queues with explicit admission slots, and carry work identity through execution, preemption, async speculaive and PD paths. Update regression tests, scheduler documentation, and benchmark work traces. Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com> # Conflicts: # docs/scheduler.md # tests/test_pd.py # tests/test_split_prefill.py # vllm_rlt/core/scheduler.py # vllm_rlt/engine/stage_engine.py
…sfer_connector.py, vllm_rlt/core/scheduling_policy.py, vllm_rlt/worker/executor.py Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
Thanks — separating lifecycle (RequestStatus), pending work (WorkKey) and execution ownership (admission_slots) makes the scheduler much easier to follow. I checked the refactor against its base 4a35034 with CPU probes and on an RTX 5090.
Equivalence inside the diff looks good. For refill-atomic, no-refill, split-refill, priority + preemption (atomic and split) and synchronous speculation (atomic and split), the per-step batch composition is identical to the base, and so are the outputs. On every step of sync, CPU-async, mixed, preemption and random-abort runs, these invariants held:
len(admission_slots) <= max_num_seqs- every ACTIVE request holds a slot
- every slot belongs to a live ACTIVE, TRANSFERRING or FINISHED request
- at most one pending
WorkKeyper request - WAITING requests are only in
waiting - no slots or blocks remain at drain.
On the RTX 5090 (Triton), the GPU-marked suite at this head gives 245 passed, 53 failed, 24 skipped. The failures are 46 FlashAttention cases (both flash_attn and flash_attn_4 report that paged FA4 is not supported on SM120), 6 PD cases that need a second GPU, and test_cuda_memory_plan_reclaims_allocator_residue_on_recreation. That is the same environmental set I see on main-based code without this PR. CPU-async in the probes does not create AsyncState and does not use UVA or graphs, so the GPU async paths are only covered by the GPU suite.
Blocker: rebase onto main (#52, #128, #141)
vllm_rlt/core/scheduler.py and vllm_rlt/engine/preemption.py conflict with #52 ("yield execution slots without offloading resident KV"). The two designs fit together: #52's "yield the execution slot, keep KV" is exactly what admission_slots + park_waiting can express. However, three assumptions in this PR contradict main and will break after a naive resolution:
-
engine/preemption.pysuspend()raises whenpending_work(victim) is None. On main, capacity preemption may pick a victim that is already WAITING with resident KV (a priority victim that kept its KV) and only offloads it.PreemptionManager.suspendthen reuses the existing snapshot without syncing, rewinding, suspending the runner or counting a preemption. So:- when a snapshot exists, take work, hidden and token from it;
- call
pending_work/_rewind_split_prefill/model_runner.suspend/preemptions += 1only on the first suspension; - port main's
kv is Noneresume path.
-
_is_preemption_candidaterequires ACTIVE + slot + pending work. That rejects main's WAITING-but-resident victims. Here is the failure:- A high-priority arrival yields a low-priority victim, which keeps its KV.
- The arrival's allocation then fails.
preempt(..., waiting_only=True)finds no candidate, so the arrival is deferred.- The resident victim resumes into the freed slot: priority inversion, and possibly ping-pong.
The predicate needs a second case — WAITING with an allocation, a snapshot and
not priority_only— plus main'swaiting_onlyfilter. -
preempt()always callspark_waiting. For an already-WAITING victim this moves it to the tail ofwaiting; main returns early to preserve its queue position.
Main also changed the following, which need porting:
- the
_admitearly returns, expressed withhas_admission_capacity() - the retry loops in
_try_resume, in the admission commit and in_ensure_execution_capacity - the relaxed speculative + preemption check in
stage_engine.
Two places auto-merge without a conflict marker but break:
tests/test_speculative.py(main's new test) assertssuspended.phase == RequestPhase.WAITING.benchmarks/preemption.pycountsrequest.phase.value == "waiting".
Main's new speculative-preemption tests need to pass after the rebase:
test_priority_preempts_only_between_speculative_rounds_and_resumestest_speculative_preemption_preserves_request_rngtest_abort_suspended_speculative_request_allows_id_reuse
The A100 results in the description predate #52/#128/#141, so they do not cover the merged state.
Should fix
- MixedPolicy now co-batches PRELUDE work.
PRELUDE_WORKSgroupsDECODE_PRELUDEwithPROMPT_PRELUDE, so prompt chunks join top-priority decode-prelude batches. That happens before admission and before thedecode_duecheck, andrecord_batchcounts these batches as prefill.- On CPU the outputs are correct.
- No test asserts that such a batch forms: the mixed test only requires
len(works) == 2somewhere. - Under async scheduling, mixed prelude batches skip the
AsyncState.preparerouting inmodel_runner.py, which before this PR only ran for RECURRENT. - Please add a test that asserts a DECODE_PRELUDE + PROMPT_PRELUDE batch, covering CUDA, async and graphs, plus a boundary-KV model if mixed mode is allowed there.
- MixedPolicy ordering changed. The old Mixed policy took the mixed recurrent batch first whenever both recurrent queues were ready. The new one inherits refill order: DECODE_PRELUDE, then CODA, then
decode_due, and only then mixed recurrent.- The batch composition differs from the base. In one of four seeds, mixed recurrent batches went from 39 to 35, CODA batches from 26 to 36, and steps from 119 to 131; the other seeds improved slightly.
tests/test_mixed_execution.pynow clears the other queues to get the old selection, which suggests the test was adapted to the new order rather than kept.prefer_recurrentinstage_engine.pyalso now fires for mixed in-flight batches.- Please either restore the old priority, or document it and add a throughput/GSM8K comparison for mixed mode.
- Docs.
docs/scheduler.mdstill says Mixed combines only DECODE_RECURRENT + PROMPT_RECURRENT and that "PRELUDE and CODA are selected separately". The MixedPolicy row still describes the old rule. The SVGs (loop-scheduler.svg,stage-components.svgand others) anddocs/unified_execution.mdstill say "phase queues", "take_phases" or "saved phase". - A behavior change is folded into a test migration.
tests/test_transfer_connector.pynow calls_admit()before the export is suspended or acknowledged.- Previously, a never-admitted export cancelled into TRANSFERRING counted against
max_num_seqs. Now it does not. - That is probably a fix, but the original "unsubmitted export" scenario is no longer tested and the description doesn't mention it.
- Related:
pd/connector.pyusesnot has_admission_slotfor both "released after the final staging read" and "never admitted". This is not reachable today, but it is fragile.
- Previously, a never-admitted export cancelled into TRANSFERRING counted against
Nits
- Quadratic per-step cost.
pending_work/remove_pending_workscan all seven deques within, andenqueue_workcalls them on every enqueue. Re-enqueueing a full batch costs about 0.22 ms at 256 requests and 2.7 ms at 1024. Adict[request_id, WorkKey]index makes it O(1). - Selection. In the selection loop,
if request.request_id in selected: continuedrops the queue entry without re-appending it. It is dead under the uniqueness invariant, but it would strand a request silently if the invariant ever broke. Assert or re-append instead. - Weakened assertion. In
tests/test_scheduler_selection.py,batch.requests[0].work == WorkKey.PROMPT_ATOMICreplaced the meaningful "a was not moved to WAITING" check. Assertstatus == ACTIVE. - Public API.
vllm_rlt.requestno longer exportsRequestPhase,Operation,ExecutionBatchorStateSpan, with no aliases. Also,benchmarks/cdb_runtime.pytrace keys changed fromstagetoworks. Worth a note for downstream users. tests/test_engine.py.drain()usesnext(iter(batch.works)), which is arbitrary for mixed batches.
Pre-existing, not from this PR
Split prefill + policy="priority" + enable_preemption can livelock. I saw about 1958 consecutive PROMPT_PRELUDE batches with four requests stuck in prompt recurrent. The behaviour is identical on the base, this PR and main.
Repro: CacheConfig(40, 2, incremental_allocation=True), max_num_seqs=4, max_num_batched_tokens=6, prefill_chunk_size=4, atomic_prefill=False. I can open a separate issue if that helps.
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
|
@0z5a
Validation on the rebased revision:
The priority + split-prefill livelock remains outside this change. |
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
|
Following up with the complete reproduction configuration, and an update from rerunning the original workload today. The historical livelock is reproducible on the original base and reviewed head, but this same workload drains on today's main and rebased PR head. It should not be treated as an outstanding regression of the current #137 revision on the strength of my earlier comment. Version boundary and fresh resultsAll runs below are CPU, FP32,
The current API has removed Historical failing batch-label MD5: Exact trigger configuration
What actually repeats on the historical revisionThe last non-prelude batch is A representative trace is: At the cap, The observed cycle is repeated preemption/resumption and replay of the same split-prefill interval, with prelude work continually selected ahead of the ready prompt recurrent work. These observations explain the failure mode; they do not establish a minimal general-purpose scheduler fix. Controls on the old #137 head, changing one option at a time while preserving the same seed-generated arrivals:
Self-contained bounded reproducerSave the following as CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 PYTHONPATH="$PWD" timeout 120 python /absolute/path/livelock_exact.py > livelock-result.jsonIt intentionally stops after 2001 engine steps, emits arrivals, batch rows, progress checkpoints, preemption counters and final state as JSON, and closes the engine in Complete script"""Original seed-2 workload, with current API compatibility and bounded diagnostics."""
import argparse
import collections
import hashlib
import json
import random
import sys
import torch
from vllm_rlt import SamplingParams, SchedulerConfig, CacheConfig, ExecutionConfig
from vllm_rlt.engine.stage_engine import StageEngine
from vllm_rlt.models import OuroForCausalLM, OuroConfig
p = argparse.ArgumentParser()
p.add_argument('--blocks', type=int, default=40)
p.add_argument('--atomic', action='store_true')
p.add_argument('--no-preemption', action='store_true')
p.add_argument('--policy', default='priority')
p.add_argument('--steps', type=int, default=2001)
a = p.parse_args()
torch.set_num_threads(1)
torch.manual_seed(123)
random.seed(2)
cfg = OuroConfig(vocab_size=64,hidden_size=32,intermediate_size=64,num_hidden_layers=2,num_attention_heads=4,num_key_value_heads=2,head_dim=8,max_position_embeddings=128,total_ut_steps=4,bos_token_id=0,eos_token_id=0)
sc = dict(mode='refill', policy=a.policy, enable_preemption=not a.no_preemption, max_num_seqs=4,max_num_batched_tokens=6,prefill_chunk_size=4)
ec = {}
if 'atomic_prefill' in ExecutionConfig.__dataclass_fields__:
ec['atomic_prefill'] = a.atomic
else:
sc['wavefront_chunked_prefill'] = not a.atomic
def label(b):
if hasattr(b,'works'): return '+'.join(sorted(w.name for w in b.works))
m = {'prefilling_prelude':'PROMPT_PRELUDE','prefilling_recurrent':'PROMPT_RECURRENT','prelude':'DECODE_PRELUDE','recurrent':'DECODE_RECURRENT','coda':'DECODE_CODA','prefill':'PROMPT_ATOMIC'}
return '+'.join(sorted({m[(i.phase or b.phase).value] for i in b.requests}))
def state(e):
ans=[]
for rid,q in e.scheduler.requests.items():
work=e.scheduler.pending_work(q).name if hasattr(e.scheduler,'pending_work') and e.scheduler.pending_work(q) is not None else str(getattr(q,'phase',None))
ans.append(dict(id=rid,status=str(getattr(q,'status',getattr(q,'phase',None))),work=work,prefilled=q.num_prefilled_tokens,preluded=q.num_preluded_tokens,generated=len(q.generated_token_ids),ranges=[x[:] for x in q.prefill_depth_ranges]))
return ans
e=StageEngine(OuroForCausalLM(cfg), attention_backend='torch',cache_config=CacheConfig(a.blocks,2,incremental_allocation=True),scheduler_config=SchedulerConfig(**sc),execution_config=ExecutionConfig(**ec))
seq=[];trace=[];arrivals=[];finished={};step=0;checks={}
try:
prompts=[[random.randrange(1,60) for _ in range(random.randrange(1,12))] for _ in range(10)]
pending=list(enumerate(prompts))
while (pending or e.has_unfinished_requests()) and step<a.steps:
while pending and random.random()<0.5:
i,tokens=pending.pop(0)
params=SamplingParams(max_tokens=random.randrange(1,7),ignore_eos=True,priority=random.randrange(0,3),max_loops=4,exit_threshold=0.6)
arrivals.append(dict(step=step,id=f'r{i}',prompt=tokens,max_tokens=params.max_tokens,priority=params.priority))
e.add_request(f'r{i}',tokens,params)
if not e.has_unfinished_requests():continue
for out in e.step():
if out.finished:finished[out.request_id]=out.token_ids
b=e.last_schedule
if b is not None:
seq.append(label(b))
trace.append(dict(step=step,work=seq[-1],rows=[dict(id=i.request.request_id,start=i.token_start,count=i.token_count,depth=i.depth) for i in b.requests],preemptions=e.preemption.preemptions if e.preemption else 0,resumptions=e.preemption.resumptions if e.preemption else 0))
step+=1
if step in [50,100,200,1000,2000,2001]:checks[step]=state(e)
suffix=0
for item in reversed(seq):
if item!='PROMPT_PRELUDE':break
suffix+=1
print(json.dumps(dict(python=sys.version,torch=torch.__version__,args=vars(a),scheduler=sc,execution=ec,arrivals=arrivals,steps=step,finished=finished,state=state(e),checkpoints=checks,counts=dict(collections.Counter(seq)),sequence_md5=hashlib.md5(' '.join(seq).encode()).hexdigest(),prelude_suffix=suffix,preemptions=e.preemption.preemptions if e.preemption else 0,resumptions=e.preemption.resumptions if e.preemption else 0,blocks=e.cache_manager.num_used_blocks,trace=trace)))
finally:e.close()This follow-up is a CPU scheduler reproduction and current-head non-reproduction of the same workload. It is not GPU serving, throughput, async, graph or general preemption qualification. |
|
@0z5a Thanks for the complete reproduction details. I reran the original workload and also inspected the behavior change introduced by 1. The original workload no longer reproducesThe historical livelock is reproducible on:
Both revisions reach 2001 engine steps with only 2/10 requests finished and a long The same workload drains on:
Both current revisions finish 10/10 requests in 178 steps, with 16 preemptions, 16 resumptions, and zero occupied KV blocks at drain. Therefore, the original report is no longer reproducible as a regression of the current #137 head. 2. What changed in
|
|
@0z5a I wonder if it is convenient for you to open an issue to discuss it in detail or I'd be glad to : ). |
|
@LyxWxj Could you share the script for the remaining 10-request case? I'd like to reproduce it before changing the WCPB scheduling path. |
Required human contributor answers
1. What outcome are you trying to achieve?
I want to separate request lifecycle, scheduled work, and model execution.
StagePhasecurrently carries several of these meanings, which makes scheduling and execution ownership harder to follow and extend.2. What is the highest-leverage variable for that outcome?
The highest-leverage variable is how the scheduler represents and selects pending work. Explicit
WorkKeyvalues let the scheduler identify work independently from request lifecycle (RequestStatus) and model operations (Operation).3. How will this work move that variable rather than create busywork?
This change introduces lifecycle-only
RequestStatusand seven validWorkKeyvalues while keepingOperationas the model-call boundary. Scheduler queues and immutable scheduled work useWorkKey; compatible prompt and decode work can be selected into the same single-operation batch. Admission slots separately track execution-resource ownership, including asynchronous work that has logically finished but has not retired.4. Who is the customer or user, and what do they actually need?
The primary users are contributors maintaining or extending the vllm-rlt scheduler and engine. They need clear ownership of request lifecycle, pending work, and occupied execution resources, plus a reliable way to batch compatible prompt and decode work.
PR necessity and bugfix evidence
A separate PR is needed because the splite-prefill and prefill-decode mixed batching are already in
main, while theStage/WorkKeyredesign is not. This PR isolates that scheduler and execution-model change for review as a follow up for #118 .This is primarily an architecture and batching capability change, not a report of one production incident. The tests cover scheduler selection, mixed execution, preemption, asynchronous speculative work, and PD paths; no single production failure is claimed as the PR's justification.
Purpose
This PR separates request lifecycle from scheduled work and model execution:
RequestStatusrepresentsWAITING,RECEIVING,TRANSFERRING,ACTIVE, andFINISHED.WorkKeyrepresents the seven supported pending work types: atomic prompt, split prompt prelude/recurrent, decode prelude/recurrent/coda, and speculative work.Operationremains the model execution boundary:PRELUDE,RECURRENT, orCODA.WorkKey. A request can have at most one pending work item, and scheduled batches and asynchronous tickets preserve the selected work identity.admission_slotstracks execution-resource ownership independently of lifecycle status.activate_request()centralizes activation and slot reservation for admission, PD activation, and preemption resume.Operation. Decode readiness and fairness decisions use exact decode work queues.Atomic prefill remains the default. Mixed scheduling requires split prefill and does not support speculative work. Each request still has at most one queued
WorkKey; this change provides the scheduler basis for WCPB but does not add multiple concurrent prompt tasks per request.Test Plan
vllm-rlt Version / Commit: base
main4a350344027960a48384626aa646f1765fda6d4f;canhazgpu. FlashAttention-specific tests and the listed FA4/UVA patterns are excluded from this run; FA4 is unsupported on the available A100 SM80 device.git diff --checkon the changed files.Test Result
270 passed, 1 skipped, 1 failed, 843 deselectedin 241.91 seconds.tests/test_cdb_runtime.py::test_cuda_memory_plan_reclaims_allocator_residue_on_recreation: within the suite, the recreated KV budget differed by 190,840,832 bytes, exceeding the test's 16 MiB tolerance. The same test passed in isolation (1 passedin 1.70 seconds); the suite-level variance remains unresolved.6 passedin 40.26 seconds.tests/test_split_prefill.py::test_cuda_split_prefill_survives_decode_graph_capture_and_replay:1 passedin 1.82 seconds.python -m ruff check tests/test_split_prefill.pyandgit diff --checkpassed.pre-commit run --all-fileswere not run on this head. The excluded FlashAttention/FA4 paths were not validated on this A100.