diff --git a/docs/scheduler.md b/docs/scheduler.md index 6f1c99e1..6ced3330 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -208,16 +208,17 @@ arcs on the right indicate separate paths, not joins. | Diagram step | What is shown | | --- | --- | | 1. Trigger | A: failed token-capacity growth. B: no active slot during admission. Available capacity/slots bypass preemption; disabled preemption or non-priority B stops the attempt. | -| 2. Protection | Exclude the requester, already-selected requests, pending outputs, WAITING/RECEIVING requests and transfer leases; B also excludes equal/higher-priority victims. An empty candidate set leads to the A/B defer/stop outcomes. | -| 3. Ranking | Lower service priority first, then larger estimated footprint; pick one victim. | +| 2. Protection | Exclude the requester, already-selected requests, pending work, RECEIVING requests and transfer leases; B also excludes WAITING and equal/higher-priority victims. Capacity reclamation can offload a WAITING resident. | +| 3. Ranking | Lower service priority first, then larger estimated footprint; priority admission also checks whether each candidate could free sufficient KV before selecting it. | | 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. | +| 5. Requester | A retries after each safe victim until it can select or must requeue. 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. | | 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 -behavior, ordering ties and current recovery limitations. The diagram covers -preemption policy; it does not expand the separate connector import protocol. +behavior, ordering ties and current recovery limitations. The diagram does not +show the priority-admission KV preflight, the failed-admission guard or the +separate connector import protocol. #### What happens to each kind of data? @@ -225,10 +226,10 @@ preemption policy; it does not expand the separate connector import protocol. | --- | --- | --- | --- | | 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` | | 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` | +| KV tensors | Priority slot preemption retains resident KV. Under capacity pressure, copy allocated pages to CPU and release the allocation; shared/cached pages need not all become free. | Reuse the resident allocation, or allocate fresh blocks and restore saved pages. | `KVCacheManager.snapshot()` / `restore()` | +| KV logical metadata: capacity, page count, written state per plane | Retained with resident KV, or deep-copied into `KVSnapshot` when offloaded. | Reuse resident metadata or restore it 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. An offloaded allocation is released. | A restored allocation establishes new mappings; subsequent preparation derives 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, clone to the cache device for a resident suspension, or to CPU for offload; clear their references on the Request. | Attach the saved tensors to the same Request on the cache device. | `PreemptionManager` | There is no dedicated queue containing snapshots. WAITING contains request IDs; `snapshots[request_id]` contains `phase`, `kv`, `hidden` and `token`. Persistent @@ -249,17 +250,19 @@ rejects preemption; it requires refill scheduling with preemption disabled. | Situation | Action in the current implementation | | --- | --- | -| A selected candidate needs a larger token frontier and `ensure_capacity()` fails | Attempt one safe victim, then retry allocation once for this candidate. | -| Active-request slots are full under priority admission | Attempt one strictly lower-priority victim to free a slot. This uses the same snapshot path even if the immediate shortage is a slot rather than KV. | +| A selected candidate needs a larger token frontier and `ensure_capacity()` fails | Reclaim safe victims until allocation succeeds or no candidate remains. | +| Active-request slots are full under priority admission | Consider strictly lower-priority victims in rank order; skip one if admission cannot fit even with its reclaimable KV and eligible waiting resident KV. A prior failed attempt can also suppress another preemption until the budget improves. | | Active slots are full under FCFS, or no victim can free a priority slot | Stop that admission scan. | -| A fresh request passes the slot check but fails its allocation budget | Defer it under admission scan/bypass rules. This branch does not independently loop over victims to make its allocation fit. | -| Restoring a suspended request cannot allocate KV | Keep its snapshot and defer it; restoration does not itself recursively preempt other requests. The preceding priority slot check can still preempt. | +| A fresh request passes the slot check but fails its allocation budget | Reclaim waiting resident allocations and retry; if still blocked, defer it under admission scan/bypass rules. | +| Restoring a suspended request cannot allocate KV | Reclaim waiting resident allocations and retry; if still blocked, retain the snapshot and defer it under the same bypass limit. | -Before ranking victims, exclude the requesting Request itself, WAITING and -RECEIVING requests, requests with pending output placeholders, requests already -selected for the current batch, and allocations with transfer leases. Priority -admission additionally excludes equal/higher-priority requests. The selected -batch exclusion set is empty during admission, which runs before batch selection. +Before ranking victims, exclude the requesting Request itself, RECEIVING +requests, requests with pending output placeholders or speculative rounds, +requests already selected for the current batch, and allocations with transfer +leases. A WAITING request is eligible only when it still has a resident +allocation and the caller is reclaiming capacity, not for priority slot +preemption. Priority admission additionally excludes equal/higher-priority +requests. Admission runs before batch selection, so its exclusion set is empty. #### Current policy: ranking and retry @@ -275,31 +278,53 @@ That footprint is not the exact number of reclaimable physical pages and does not account for transfer cost. Exact ties follow candidate iteration order. Memory-pressure selection uses this ranking even under FCFS; only priority admission requires the victim to be strictly lower priority than the requester. +Priority admission tries the next ranked victim when the previous one cannot +provide a feasible slot and KV budget. -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 -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 -own capacity checks. With no selected work, no batch is returned. +If no safe victim exists and allocation still fails, batch construction returns +the requesting ID to its phase queue and considers other candidates within the +scan/token/sequence limits. A later scheduling attempt can try again. With no +selected work, no batch is returned. #### Current policy: suspend, restore and cleanup | 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. | -| Relinquish resources | Suspend runner resources while retaining RNG; poll cache prefixes; free the victim's cache allocation; clear hidden/token references; store the snapshot. | +| Suspend device state | Synchronize the runner's configured streams; rewind unfinished split prefill; save the phase and optional hidden/token tensors. KV is copied to CPU when offloaded. | +| Relinquish resources | Suspend runner resources while retaining RNG; release KV only when offloading; clear hidden/token references; store the snapshot. | | Move the request | Remove the victim's old queue entries 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. | | 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 -has `max_admission_bypasses` accounting to prevent unlimited bypass of blocked -requests. Successful restoration takes an early continuation: it does not use -the common bypass-counter update. Blocked restoration also takes an early -continuation rather than checking the fresh-allocation bypass limit. These are -current branch behaviors, not a guarantee of starvation-free restoration. +Deferred IDs are put back at the front in their existing order. Both fresh +admissions and restorations count as bypasses of earlier blocked requests, and +both fresh and restored requests stop further bypasses at +`max_admission_bypasses`. + +Before a priority slot preemption, the scheduler computes an optimistic +admission budget for each eligible victim. It excludes that victim's future +prompt-growth reservation and counts blocks that would become available after +releasing the victim and eligible waiting resident allocations. Shared pages are +counted only once, and transfer-leased pages cannot be reclaimed. Fresh requests +include their prompt/capacity budget and watermark; offloaded snapshots need +their saved pages, while resident snapshots need no new KV. Prefix reuse counts +both available and pending prefix publications as possible hits and accounts for +published pages that are or would become cache-only after reclamation. This +optimistic check may still permit an attempt that later fails; actual admission +verifies and claims the prefix. If even this budget is insufficient, the victim +is not suspended or rewound. An unknown custom resume callback also skips the +preflight rejection. + +If a request first preempts a lower-priority active request for a slot, but its +own fresh allocation or snapshot restoration still fails, the scheduler records +the available budget: free blocks minus active prompt growth reservations. While +slots remain full, it will not repeat that request's priority preemption until +this budget improves. Successful admission removes the record; request completion +or cancellation clears the records so capacity and eligibility can be evaluated +again. This lets active work make progress instead of repeatedly rewinding it for +an admission that cannot yet succeed. It can temporarily delay strict-priority +service and is not a general starvation-freedom guarantee for all workloads. #### Implemented recovery versus possible extensions @@ -311,7 +336,7 @@ current branch behaviors, not a guarantee of starvation-free restoration. | Choose KV swap versus discard-and-recompute | Not implemented as a selectable recovery policy. Full-depth replay of generated tokens is not assumed equivalent for RLT. | | Other-GPU / external-storage snapshot backend | Not implemented by `PreemptionManager`. | | Asynchronous pooled snapshot transfer | Not implemented by this path; the normal input-transfer buffer pools do not make snapshots asynchronous. | -| Victim ranking by exact reclaimed bytes, swap cost or prior preemptions | Not implemented. There is no cost-based victim score or cooldown against repeated preemption. | +| Victim ranking by exact reclaimed bytes, swap cost or prior preemptions | Not implemented. The preflight can skip an infeasible victim but does not change the victim score. | Insufficient capacity and absence of a victim are explicit defer outcomes. Snapshot/copy/restore exceptions are not converted by `PreemptionManager` into a diff --git a/tests/test_kv_cache.py b/tests/test_kv_cache.py index d49c59f2..be467c08 100644 --- a/tests/test_kv_cache.py +++ b/tests/test_kv_cache.py @@ -12,6 +12,60 @@ def make_cache(**kwargs): return KVCacheManager(**parameters) +def test_reclaimable_blocks_account_for_shared_prefix_references(): + cache = make_cache(num_blocks=20, max_loops=2, enable_prefix_caching=True) + assert cache.allocate("first", 5) + cache.mark_imported_prefix("first", 5) + cache.publish_prefix("first", [1, 2, 3, 4, 5], 5) + prefix = cache.lookup_prefix([1, 2, 3, 4, 6]) + assert cache.allocate("second", 5, prefix=prefix) + free_before = cache.num_free_blocks + + assert cache.prefix_admission_bounds([1, 2, 3, 4, 6]) == (4, 0) + assert cache.prefix_admission_bounds([1, 2, 3, 4, 6], release_requests=["first", "second"]) == ( + 4, + 4, + ) + assert cache.reclaimable_blocks(["first"]) == 2 + assert cache.reclaimable_blocks(["first", "second"]) == 8 + assert cache.num_free_blocks == free_before + + cache.pin_transfer("first", "transfer") + assert cache.reclaimable_blocks(["first", "second"]) == 2 + cache.unpin_transfer("first", "transfer") + cache.free("first") + assert cache.num_free_blocks == free_before + 2 + assert cache.reclaimable_blocks(["second"]) == 6 + cache.free("second") + assert cache.num_free_blocks == free_before + 8 + assert cache.prefix_admission_bounds([1, 2, 3, 4, 6]) == (4, 4) + + +def test_prefix_admission_bounds_include_pending_publications_without_polling(): + class PendingEvent: + def __init__(self): + self.queries = 0 + + def query(self): + self.queries += 1 + return False + + cache = make_cache(num_blocks=20, max_loops=2, enable_prefix_caching=True) + assert cache.allocate("first", 5) + cache.mark_imported_prefix("first", 5) + event = PendingEvent() + cache.publish_prefix("first", [1, 2, 3, 4, 5], 5, event=event) + + assert cache.prefix_admission_bounds([1, 2, 3, 4, 6]) == (4, 0) + assert cache.prefix_admission_bounds([1, 9, 3, 4, 6]) == (0, 0) + assert event.queries == 0 + assert not cache.lookup_prefix([1, 2, 3, 4, 6]) + assert event.queries == 1 + assert cache.prefix_admission_bounds([1, 2, 3, 4, 6]) == (4, 0) + cache.free("first") + assert cache.prefix_admission_bounds([1, 2, 3, 4, 6]) == (0, 0) + + def test_admission_is_atomic_and_counts_all_depths(): cache = make_cache(num_blocks=12) assert cache.required_blocks(3) == 6 diff --git a/tests/test_prefill_livelock.py b/tests/test_prefill_livelock.py new file mode 100644 index 00000000..ce1a6a95 --- /dev/null +++ b/tests/test_prefill_livelock.py @@ -0,0 +1,698 @@ +"""Regression coverage for priority admission under staged prefill pressure.""" + +from collections import defaultdict + +import pytest + +from tests.helpers import tiny_ouro_model +from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.request import Request, RequestPhase + +REQUESTS = ( + (0, "r0", [24, 25, 9, 13], 1, 1), + (2, "r1", [6], 2, 2), + (2, "r2", [16, 52, 33], 1, 1), + (3, "r3", [26, 42, 2, 30], 5, 1), + (5, "r4", [30, 25, 32, 37, 13, 58, 54, 26], 3, 1), + (5, "r5", [32, 15], 5, 0), + (6, "r6", [45], 3, 2), + (8, "r7", [34, 27, 31, 58, 59], 2, 1), + (10, "r8", [47, 8, 43, 17, 7, 53, 5], 5, 2), + (11, "r9", [40, 54, 25, 7, 43, 4, 22], 1, 1), +) + +REVIEWER_REQUESTS = ( + (0, "r0", [6], 4, 1), + (4, "r1", [24, 54], 4, 0), + (5, "r2", [48, 52, 43], 6, 1), + (6, "r3", [17, 39, 14, 39, 3], 3, 1), + (8, "r4", [44, 11, 28, 41, 26, 52, 47, 56, 33, 24], 4, 1), + (8, "r5", [29, 33, 18, 58, 3, 56, 2, 24, 30], 6, 2), + (9, "r6", [59, 25, 28, 58, 57, 34], 2, 1), + (14, "r7", [36, 12, 16], 3, 2), + (18, "r8", [2, 12, 21, 12], 6, 0), + (18, "r9", [33, 33, 24], 3, 2), +) + + +def run_workload( + requests, *, mode="refill", preemption=True, wavefront=True, asynchronous=False, delayed=False +): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(40, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + mode=mode, + policy="priority", + enable_preemption=preemption, + max_num_seqs=4, + max_num_batched_tokens=6, + prefill_chunk_size=4, + wavefront_chunked_prefill=wavefront, + ), + exit_config=ExitConfig("ouro_delayed" if delayed else "ouro"), + execution_config=ExecutionConfig(async_scheduling=asynchronous), + ) + arrivals = defaultdict(list) + for step, request_id, prompt, max_tokens, priority in requests: + arrivals[step].append((request_id, prompt, max_tokens, priority)) + finished = {} + engine_steps = 0 + try: + for step in range(500): + for request_id, prompt, max_tokens, priority in arrivals[step]: + engine.add_request( + request_id, + prompt, + SamplingParams( + max_tokens=max_tokens, + ignore_eos=True, + priority=priority, + max_loops=4, + exit_threshold=0.6, + ), + ) + if not engine.has_unfinished_requests(): + if step > requests[-1][0]: + break + continue + outputs = engine.step() + engine_steps += 1 + for output in outputs: + if output.finished: + finished[output.request_id] = ( + output.token_ids, + output.exit_depths, + output.finish_reason, + ) + assert len(finished) == len(requests), ( + f"only {len(finished)}/{len(requests)} requests finished in 500 steps; " + f"preemptions={engine.preemption.preemptions if engine.preemption else 0}" + ) + assert set(finished) == {request_id for _, request_id, *_ in requests} + assert not engine.has_unfinished_requests() + assert engine.cache_manager.num_used_blocks == 0 + if engine.preemption is not None: + assert not engine.preemption.snapshots + return finished, engine.preemption.preemptions if engine.preemption else 0, engine_steps + finally: + engine.close() + + +@pytest.mark.parametrize("requests", [REQUESTS, REVIEWER_REQUESTS], ids=["seed-8", "reviewer"]) +@pytest.mark.parametrize("mode,wavefront", [("refill", True), ("mixed", True), ("mixed", False)]) +def test_priority_staged_prefill_drains_under_cache_pressure(requests, mode, wavefront): + reference, _, _ = run_workload(requests, preemption=False) + actual, preemptions, engine_steps = run_workload(requests, mode=mode, wavefront=wavefront) + assert preemptions > 0 + assert engine_steps < 500 + assert actual == reference + + +def test_priority_wavefront_prefill_drains_with_async_scheduling(): + reference, _, _ = run_workload(REQUESTS, preemption=False, delayed=True) + actual, preemptions, engine_steps = run_workload(REQUESTS, asynchronous=True, delayed=True) + assert preemptions > 0 + assert engine_steps < 500 + assert actual == reference + + +@pytest.mark.parametrize("prefix_caching", [False, True]) +def test_impossible_priority_preemption_preserves_prefill_progress(prefix_caching): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig( + 24, 2, incremental_allocation=True, enable_prefix_caching=prefix_caching + ), + scheduler_config=SchedulerConfig( + mode="refill", + policy="priority", + enable_preemption=True, + max_num_seqs=2, + max_num_batched_tokens=8, + prefill_chunk_size=4, + wavefront_chunked_prefill=True, + ), + ) + try: + for request_id in ("a", "b"): + engine.add_request(request_id, [1, 2, 3, 4], SamplingParams(max_tokens=1, priority=2)) + engine.step() + before = { + request_id: ( + engine.scheduler.requests[request_id].phase, + engine.scheduler.requests[request_id].num_preluded_tokens, + [row[:] for row in engine.scheduler.requests[request_id].prefill_depth_ranges], + ) + for request_id in ("a", "b") + } + engine.add_request("high", [1] * 10, SamplingParams(max_tokens=1, priority=0)) + + for _ in range(3): + engine.scheduler._admit() + + assert engine.preemption.preemptions == 0 + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + for request_id, (phase, preluded, ranges) in before.items(): + request = engine.scheduler.requests[request_id] + assert request.phase == phase + assert request.num_preluded_tokens == preluded + assert request.prefill_depth_ranges == ranges + + finished = set() + for _ in range(100): + finished.update(output.request_id for output in engine.step() if output.finished) + if not engine.has_unfinished_requests(): + break + assert finished == {"a", "b", "high"} + assert engine.cache_manager.num_used_blocks == 0 + finally: + engine.close() + + +def test_nonincremental_priority_preemption_reserves_output_capacity(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(32, 2, incremental_allocation=False), + scheduler_config=SchedulerConfig( + policy="priority", enable_preemption=True, max_num_seqs=2, prefill_chunk_size=4 + ), + ) + try: + for request_id in ("a", "b"): + engine.add_request(request_id, [1] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.scheduler._admit() + engine.add_request("high", [2] * 10, SamplingParams(max_tokens=5, priority=0)) + + engine.scheduler._admit() + + assert engine.preemption.preemptions == 0 + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + finished = set() + for _ in range(100): + finished.update(output.request_id for output in engine.step() if output.finished) + if not engine.has_unfinished_requests(): + break + assert finished == {"a", "b", "high"} + assert engine.cache_manager.num_used_blocks == 0 + finally: + engine.close() + + +def test_impossible_priority_restoration_does_not_preempt(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(32, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + mode="refill", + policy="priority", + enable_preemption=True, + max_num_seqs=3, + max_num_batched_tokens=10, + prefill_chunk_size=10, + wavefront_chunked_prefill=True, + ), + ) + try: + engine.add_request("high", [1] * 10, SamplingParams(max_tokens=1, priority=0)) + engine.step() + assert engine.scheduler.preempt(Request("pressure", [1], SamplingParams(max_tokens=1))) + for request_id, priority in (("a", -1), ("b", -1), ("c", 2)): + engine.add_request(request_id, [2] * 4, SamplingParams(max_tokens=1, priority=priority)) + engine.scheduler._admit() + phase = engine.scheduler.requests["c"].phase + + for _ in range(3): + engine.scheduler._admit() + + assert engine.preemption.preemptions == 1 + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["c"].phase == phase + finally: + engine.close() + + +def test_resident_priority_restoration_needs_no_new_kv(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(12, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + policy="priority", + enable_preemption=True, + max_num_seqs=2, + prefill_chunk_size=2, + wavefront_chunked_prefill=True, + ), + ) + try: + engine.add_request("high", [1, 2], SamplingParams(max_tokens=1, priority=0)) + engine.add_request("low", [3, 4], SamplingParams(max_tokens=1, priority=2)) + engine.scheduler._admit() + pressure = Request("pressure", [1], SamplingParams(max_tokens=1, priority=-1)) + assert engine.scheduler.preempt(pressure, priority_only=True, excluded=frozenset({"low"})) + assert engine.preemption.snapshots["high"]["kv"] is None + engine.add_request("other", [5, 6], SamplingParams(max_tokens=1, priority=-1)) + + engine.scheduler._admit() + + assert engine.cache_manager.num_free_blocks == 0 + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILLING_PRELUDE + assert engine.scheduler.requests["low"].phase == RequestPhase.WAITING + assert engine.preemption.resumptions == 1 + finally: + engine.close() + + +def test_offloaded_priority_restoration_can_reclaim_victim_kv(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(20, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + policy="priority", enable_preemption=True, max_num_seqs=2, prefill_chunk_size=4 + ), + ) + try: + engine.add_request("high", [1] * 4, SamplingParams(max_tokens=1, priority=0)) + engine.add_request("low", [2] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.scheduler._admit() + pressure = Request("pressure", [1], SamplingParams(max_tokens=1, priority=-2)) + assert engine.scheduler.preempt(pressure, excluded=frozenset({"low"})) + assert engine.preemption.snapshots["high"]["kv"] is not None + engine.add_request("other", [3] * 4, SamplingParams(max_tokens=1, priority=-1)) + + engine.scheduler._admit() + + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILL + assert engine.scheduler.requests["low"].phase == RequestPhase.WAITING + assert engine.preemption.snapshots["low"]["kv"] is not None + assert engine.preemption.resumptions == 1 + finally: + engine.close() + + +def test_priority_preemption_can_reclaim_waiting_resident_kv(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(24, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + policy="priority", enable_preemption=True, max_num_seqs=2, prefill_chunk_size=4 + ), + ) + try: + engine.add_request("active", [1] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.add_request("resident", [2] * 4, SamplingParams(max_tokens=1, priority=3)) + engine.scheduler._admit() + pressure = Request("pressure", [1], SamplingParams(max_tokens=1, priority=-2)) + assert engine.scheduler.preempt( + pressure, priority_only=True, excluded=frozenset({"active"}) + ) + engine.add_request("other", [3] * 2, SamplingParams(max_tokens=1, priority=-1)) + engine.scheduler._admit() + engine.add_request("high", [4] * 8, SamplingParams(max_tokens=1, priority=0)) + + engine.scheduler._admit() + + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILL + assert engine.preemption.snapshots["resident"]["kv"] is not None + assert engine.preemption.snapshots["active"]["kv"] is not None + assert engine.preemption.preemptions == 2 + finally: + engine.close() + + +def test_priority_preemption_tries_next_victim_if_first_cannot_make_room(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(20, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + policy="priority", enable_preemption=True, max_num_seqs=2, prefill_chunk_size=2 + ), + ) + try: + engine.add_request("long", [1] * 8, SamplingParams(max_tokens=1, priority=2)) + engine.add_request("short", [2] * 2, SamplingParams(max_tokens=1, priority=3)) + engine.scheduler._admit() + engine.add_request("high", [3] * 8, SamplingParams(max_tokens=1, priority=0)) + + engine.scheduler._admit() + + assert engine.scheduler.requests["long"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["short"].phase == RequestPhase.PREFILL + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILL + assert engine.preemption.preemptions == 1 + finally: + engine.close() + + +def test_priority_preemption_waives_watermark_when_victim_is_only_active_request(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(16, 2, incremental_allocation=True, watermark_ratio=0.25), + scheduler_config=SchedulerConfig( + policy="priority", enable_preemption=True, max_num_seqs=1, prefill_chunk_size=4 + ), + ) + try: + engine.add_request("low", [1] * 4, SamplingParams(max_tokens=1, priority=1)) + engine.scheduler._admit() + engine.add_request("high", [2] * 8, SamplingParams(max_tokens=1, priority=0)) + + engine.scheduler._admit() + + assert engine.scheduler.requests["low"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILL + assert engine.preemption.preemptions == 1 + finally: + engine.close() + + +def test_priority_preemption_can_reuse_shared_prefix(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(20, 2, incremental_allocation=True, enable_prefix_caching=True), + scheduler_config=SchedulerConfig( + policy="priority", enable_preemption=True, max_num_seqs=2, prefill_chunk_size=4 + ), + ) + try: + cache = engine.cache_manager + assert cache.allocate("seed", 5) + cache.mark_imported_prefix("seed", 5) + cache.publish_prefix("seed", [1, 2, 3, 4, 5], 5) + cache.free("seed") + engine.add_request("low", [1, 2, 3, 4, 7], SamplingParams(max_tokens=1, priority=2)) + engine.add_request("other", [8, 9], SamplingParams(max_tokens=1, priority=-1)) + engine.scheduler._admit() + engine.add_request("high", [1, 2, 3, 4, 6], SamplingParams(max_tokens=1, priority=0)) + + engine.scheduler._admit() + + assert engine.scheduler.requests["low"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILL + assert engine.scheduler.requests["high"].num_prefilled_tokens == 4 + assert engine.preemption.snapshots["low"]["kv"] is None + finally: + engine.close() + + +def test_cache_only_prefix_claim_prevents_futile_priority_preemption(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(24, 2, incremental_allocation=True, enable_prefix_caching=True), + scheduler_config=SchedulerConfig( + policy="priority", enable_preemption=True, max_num_seqs=2, prefill_chunk_size=4 + ), + ) + try: + cache = engine.cache_manager + assert cache.allocate("seed", 5) + cache.mark_imported_prefix("seed", 5) + cache.publish_prefix("seed", [1, 2, 3, 4, 5], 5) + cache.free("seed") + engine.add_request("a", [6] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.add_request("b", [7] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.scheduler._admit() + prompt = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + engine.add_request("high", prompt, SamplingParams(max_tokens=1, priority=0)) + + assert cache.prefix_admission_bounds(prompt) == (8, 8) + engine.scheduler._admit() + + assert engine.preemption.preemptions == 0 + assert engine.scheduler.requests["a"].phase == RequestPhase.PREFILL + assert engine.scheduler.requests["b"].phase == RequestPhase.PREFILL + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + finally: + engine.close() + + +def test_priority_preemption_avoids_victim_that_would_orphan_shared_prefix(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(24, 2, incremental_allocation=True, enable_prefix_caching=True), + scheduler_config=SchedulerConfig( + mode="refill", + policy="priority", + enable_preemption=True, + max_num_seqs=2, + max_num_batched_tokens=8, + prefill_chunk_size=4, + wavefront_chunked_prefill=True, + ), + ) + try: + cache = engine.cache_manager + assert cache.allocate("seed", 5) + cache.mark_imported_prefix("seed", 5) + cache.publish_prefix("seed", [1, 2, 3, 4, 5], 5) + cache.free("seed") + engine.add_request("a", [1, 2, 3, 4, 6], SamplingParams(max_tokens=1, priority=2)) + engine.add_request("b", [7] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.scheduler._admit() + prompt = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + engine.add_request("high", prompt, SamplingParams(max_tokens=1, priority=0)) + + assert cache.prefix_admission_bounds(prompt) == (8, 0) + assert cache.prefix_admission_bounds(prompt, release_requests=["a"]) == (8, 8) + engine.scheduler._admit() + assert engine.preemption.preemptions == 1 + assert engine.scheduler.requests["a"].phase == RequestPhase.PREFILLING_PRELUDE + assert engine.scheduler.requests["b"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILLING_PRELUDE + + finished = set() + for _ in range(100): + finished.update(output.request_id for output in engine.step() if output.finished) + if not engine.has_unfinished_requests(): + break + assert finished == {"a", "b", "high"} + assert cache.num_used_blocks == 0 + finally: + engine.close() + + +def test_failed_priority_admission_does_not_repeat_preemption(): + class PendingEvent: + def query(self): + return False + + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(24, 2, incremental_allocation=True, enable_prefix_caching=True), + scheduler_config=SchedulerConfig( + mode="refill", + policy="priority", + enable_preemption=True, + max_num_seqs=2, + max_num_batched_tokens=8, + prefill_chunk_size=4, + wavefront_chunked_prefill=True, + ), + ) + try: + engine.add_request("a", [1] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.add_request("b", [2] * 4, SamplingParams(max_tokens=1, priority=2)) + engine.step() + engine.cache_manager.publish_prefix("a", [1] * 4, 4, event=PendingEvent()) + engine.add_request("high", [1] * 10, SamplingParams(max_tokens=1, priority=0)) + + engine.scheduler._admit() + assert engine.preemption.preemptions == 1 + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["b"].phase == RequestPhase.PREFILLING_RECURRENT + for _ in range(5): + engine.scheduler._admit() + assert engine.preemption.preemptions == 1 + assert engine.scheduler.requests["b"].phase == RequestPhase.PREFILLING_RECURRENT + + finished = set() + for _ in range(100): + for output in engine.step(): + if output.finished: + finished.add(output.request_id) + if not engine.has_unfinished_requests(): + break + assert finished == {"a", "b", "high"} + assert not engine.scheduler._blocked_priority_preemptions + assert not engine.preemption.snapshots + assert engine.cache_manager.num_used_blocks == 0 + finally: + engine.close() + + +def test_failed_priority_restoration_does_not_repeat_preemption(monkeypatch): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(40, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + mode="refill", + policy="priority", + enable_preemption=True, + max_num_seqs=3, + max_num_batched_tokens=10, + prefill_chunk_size=10, + wavefront_chunked_prefill=True, + ), + ) + try: + engine.add_request("high", [1] * 10, SamplingParams(max_tokens=1, priority=0)) + engine.step() + assert engine.scheduler.preempt(Request("pressure", [1], SamplingParams(max_tokens=1))) + resume = engine.scheduler.resume_callback + monkeypatch.setattr( + engine.scheduler, + "resume_callback", + lambda request: False if request.request_id == "high" else resume(request), + ) + for request_id, priority in (("a", -1), ("b", -1), ("c", 2)): + engine.add_request(request_id, [2] * 4, SamplingParams(max_tokens=1, priority=priority)) + + engine.scheduler._admit() + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["high"].admission_bypasses == 1 + engine.scheduler._admit() + assert engine.preemption.preemptions == 2 + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + for _ in range(5): + engine.scheduler._admit() + assert engine.preemption.preemptions == 2 + + monkeypatch.setattr(engine.scheduler, "resume_callback", resume) + finished = set() + for _ in range(100): + for output in engine.step(): + if output.finished: + finished.add(output.request_id) + if not engine.has_unfinished_requests(): + break + assert finished == {"a", "b", "c", "high"} + assert not engine.scheduler._blocked_priority_preemptions + assert not engine.preemption.snapshots + assert engine.cache_manager.num_used_blocks == 0 + finally: + engine.close() + + +def test_priority_admission_retries_after_capacity_improves(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(40, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + mode="refill", + policy="priority", + enable_preemption=True, + max_num_seqs=2, + max_num_batched_tokens=8, + prefill_chunk_size=4, + wavefront_chunked_prefill=True, + ), + ) + try: + assert engine.cache_manager.allocate("external", 8) + for request_id in ("a", "b"): + engine.add_request(request_id, [1, 2, 3, 4], SamplingParams(max_tokens=1, priority=2)) + engine.step() + engine.add_request("high", [1] * 16, SamplingParams(max_tokens=1, priority=0)) + engine.scheduler._admit() + assert engine.preemption.preemptions == 0 + assert engine.scheduler.requests["high"].phase == RequestPhase.WAITING + engine.scheduler._admit() + assert all( + engine.scheduler.requests[request_id].phase != RequestPhase.WAITING + for request_id in ("a", "b") + ) + + engine.cache_manager.free("external") + engine.scheduler._admit() + assert engine.preemption.preemptions == 1 + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILLING_PRELUDE + finally: + engine.close() + + +def test_abort_clears_failed_priority_admission_for_reused_id(monkeypatch): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(24, 2, incremental_allocation=True, enable_prefix_caching=True), + scheduler_config=SchedulerConfig( + policy="priority", + enable_preemption=True, + max_num_seqs=2, + max_num_batched_tokens=8, + prefill_chunk_size=4, + wavefront_chunked_prefill=True, + ), + ) + try: + monkeypatch.setattr( + engine.cache_manager, + "prefix_admission_bounds", + lambda tokens, **kwargs: (((len(tokens) - 1) // 2) * 4, 0), + ) + for request_id in ("a", "b"): + engine.add_request(request_id, [1, 2, 3, 4], SamplingParams(max_tokens=1, priority=2)) + engine.step() + engine.add_request("high", [1] * 10, SamplingParams(max_tokens=1, priority=0)) + engine.scheduler._admit() + assert "high" in engine.scheduler._blocked_priority_preemptions + + engine.abort_request("high") + assert not engine.scheduler._blocked_priority_preemptions + engine.add_request("high", [3], SamplingParams(max_tokens=1, priority=0)) + engine.scheduler._admit() + assert engine.scheduler.requests["high"].phase == RequestPhase.PREFILLING_PRELUDE + finally: + engine.close() + + +def test_resumption_counts_as_an_admission_bypass(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(24, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + policy="fcfs", + enable_preemption=True, + max_num_seqs=3, + max_admission_bypasses=1, + prefill_chunk_size=4, + ), + ) + try: + engine.add_request("resumed", [2], SamplingParams(max_tokens=1, priority=10)) + engine.add_request("running-1", [1] * 4, SamplingParams(max_tokens=1)) + engine.add_request("running-2", [1] * 4, SamplingParams(max_tokens=1)) + engine.scheduler._admit() + engine.add_request("long", [1] * 9, SamplingParams(max_tokens=1)) + assert engine.scheduler.preempt(engine.scheduler.requests["long"]) + + engine.scheduler._admit() + assert engine.scheduler.requests["long"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["long"].admission_bypasses == 1 + assert engine.scheduler.requests["resumed"].phase == RequestPhase.PREFILL + + engine.abort_request("running-1") + engine.scheduler._admit() + assert engine.scheduler.requests["long"].phase == RequestPhase.WAITING + engine.abort_request("running-2") + engine.scheduler._admit() + assert engine.scheduler.requests["long"].phase == RequestPhase.PREFILL + finally: + engine.close() diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..5a20e731 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -8,7 +8,7 @@ import hashlib import struct -from collections import OrderedDict +from collections import Counter, OrderedDict from collections.abc import Callable, Sequence from copy import deepcopy from dataclasses import dataclass, field @@ -276,6 +276,37 @@ def lookup_prefix(self, tokens): found.append(blocks) return tuple(found) + def prefix_admission_bounds( + self, tokens: Sequence[int], *, release_requests: Sequence[str] = () + ) -> tuple[int, int]: + """Return possible reuse and a lower bound for exclusive prefix claims. + + Pending publications are counted even if their events or written ranges + are not ready. The claim bound only covers the contiguous prefix already + published, projected after releasing the specified allocations. + """ + if not self.enable_prefix_caching: + return 0, 0 + released = self._released_block_counts(release_requests) + pending_keys = set() + for request_id, allocation, pending_tokens, length, _ in self._pending_prefixes: + if self._allocations.get(request_id) is not allocation: + continue + publishable = min(length, len(pending_tokens) - 1) + pending_keys.update(self._prefix_keys(pending_tokens[:publishable])) + pages = claims = 0 + published_prefix = True + for key in self._prefix_keys(tokens[:-1]): + blocks = self._prefixes.get(key) + if blocks is None and key not in pending_keys: + break + if blocks is not None and published_prefix: + claims += sum(self._refs[block] - released[block] == 1 for block in blocks) + if blocks is None: + published_prefix = False + pages += 1 + return pages * self.storage_depths, claims + def allocate(self, request_id: str, max_tokens: int, *, initial_tokens=None, prefix=()) -> bool: """Reserve a logical capacity; acquire physical pages for the current frontier.""" if not isinstance(request_id, str) or not request_id: @@ -399,6 +430,32 @@ def allocated_blocks(self, request_id: str) -> int: """How many physical blocks this request holds, counting every plane.""" return len(self._get_allocation(request_id).block_tables[0]) * self.storage_depths + def _released_block_counts(self, request_ids: Sequence[str]) -> Counter: + allocations = (self._get_allocation(request_id) for request_id in set(request_ids)) + return Counter( + block + for allocation in allocations + if not allocation.transfer_leases + for table in allocation.block_tables + for block in table + ) + + def reclaimable_blocks(self, request_ids: Sequence[str]) -> int: + """Count additional available blocks if these allocations are released. + + Shared pages count only once, and cached pages become available once the + last non-cache reference is released. Transfer-leased allocations remain + pinned and contribute nothing until the lease is released. + """ + released = self._released_block_counts(request_ids) + cached = { + block for blocks in self._prefixes.values() for block in blocks if block in released + } + return sum( + self._refs[block] == count or (self._refs[block] == count + 1 and block in cached) + for block, count in released.items() + ) + def get_block_table(self, request_id: str, depth: int) -> tuple[int, ...]: self._validate_depth(depth) return self._get_allocation(request_id).block_tables[self._plane(depth)] diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 12aa040c..b98cf279 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -156,10 +156,12 @@ def __init__( self.exit_mode = exit_mode self.requests: dict[str, Request] = {} self.queues: dict[RequestPhase, deque[str]] = {s: deque() for s in RequestPhase} + self._blocked_priority_preemptions: dict[str, int] = {} # 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.resume_blocks_callback = None # Optional connector readiness and export predicates. Never mutate queues or blocks. self.execution_ready = None self.exports_request = None @@ -206,6 +208,7 @@ def finish(self, request: Request, reason: FinishReason): request.num_output_placeholders = 0 request.num_speculative_rounds = 0 self.requests.pop(request.request_id) + self._blocked_priority_preemptions.clear() def abort(self, request_id: str) -> Request: request = self.requests[request_id] @@ -317,18 +320,98 @@ def _ensure_active_slot(self, request: Request, active_count: int) -> int | None 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. + from preemption here. A previous unsuccessful priority preemption is + retried only when the available admission budget improves. """ if active_count < self.config.max_num_seqs: return active_count - if ( - self.config.policy == "priority" - and self.suspend_callback is not None - and self.preempt(request, priority_only=True, excluded=frozenset()) + blocked_budget = self._blocked_priority_preemptions.get(request.request_id) + if blocked_budget is not None and self._available_admission_budget() <= blocked_budget: + return None + if self.config.policy != "priority" or self.suspend_callback is None: + return None + excluded = set() + while victim := self._select_preemption_victim( + request, priority_only=True, excluded=frozenset(excluded) ): - return active_count - 1 + if self._priority_preemption_can_admit(request, victim, active_count): + if self.preempt(request, priority_only=True, excluded=frozenset(excluded)): + return active_count - 1 + return None + excluded.add(victim.request_id) return None + def _priority_preemption_can_admit( + self, request: Request, victim: Request, active_count: int + ) -> bool: + """Reject a slot preemption only if even optimistic KV capacity is too small. + + A victim keeps resident KV initially; admission may subsequently offload + it and any eligible waiting allocations. Prefix hits are optimistic, so + an uncertain preflight can proceed and the failed-preemption guard still + protects against a later allocation failure. + """ + cache = self.cache_manager + if request.phase == RequestPhase.RECEIVING or ( + cache.has_allocation(request.request_id) + and cache.has_transfer_lease(request.request_id) + ): + return True + if self.resume_callback is not None and self.resume_blocks_callback is None: + return True + restore_blocks = ( + self.resume_blocks_callback(request) + if self.resume_blocks_callback is not None + else None + ) + if restore_blocks == 0: + return True + reclaim_ids = [ + waiting.request_id + for waiting in self.requests.values() + if waiting.phase == RequestPhase.WAITING + and self._is_preemption_candidate( + waiting, request, priority_only=False, excluded=frozenset() + ) + ] + if cache.has_allocation(victim.request_id): + reclaim_ids.append(victim.request_id) + available = cache.num_free_blocks + cache.reclaimable_blocks(reclaim_ids) + if restore_blocks is not None: + return restore_blocks <= available + prompt_tokens = len(request.prompt_token_ids) + budget_tokens = ( + prompt_tokens + if cache.incremental_allocation + else prompt_tokens + request.sampling_params.max_tokens - 1 + ) + prefix_blocks, cached_claim = cache.prefix_admission_bounds( + request.prompt_token_ids, release_requests=reclaim_ids + ) + remaining_growth = self._reserved_growth_blocks( + reserve_outputs=False, excluded=frozenset({victim.request_id}) + ) + headroom = cache.watermark_blocks if active_count > 1 else 0 + return ( + cache.required_blocks(budget_tokens) + - prefix_blocks + + cached_claim + + remaining_growth + + headroom + <= available + ) + + def _available_admission_budget(self) -> int: + return self.cache_manager.num_free_blocks - self._reserved_growth_blocks( + reserve_outputs=False + ) + + def _record_failed_priority_preemption(self, request: Request) -> None: + """Do not repeatedly roll back active work without more KV capacity.""" + budget = self._available_admission_budget() + previous = self._blocked_priority_preemptions.get(request.request_id, budget) + self._blocked_priority_preemptions[request.request_id] = max(budget, previous) + def _try_resume(self, request: Request) -> ResumeResult: """Translate the existing restoration callback's three-way result. @@ -346,7 +429,7 @@ def _try_resume(self, request: Request) -> ResumeResult: self.enqueue(request, restored) return ResumeResult.RESTORED - def _reserved_growth_blocks(self, *, reserve_outputs: bool) -> int: + def _reserved_growth_blocks(self, *, reserve_outputs: bool, excluded=frozenset()) -> int: """Budget future growth of existing requests without allocating it. Example: a request needs 12 blocks eventually but owns 4. Without @@ -357,6 +440,8 @@ def _reserved_growth_blocks(self, *, reserve_outputs: bool) -> int: cache = self.cache_manager reserved = 0 for request in self.requests.values(): + if request.request_id in excluded: + continue # Imported/reserved requests already own pages before compute admission. # Without preemption their future decode growth must remain affordable too. has_allocation = cache.has_allocation(request.request_id) @@ -475,8 +560,8 @@ def _admit(self): """Admit waiting requests: order -> slot -> restore/allocate -> defer/enter. A blocked long request may let a short one pass. Each successful fresh - admission increments earlier blocked requests' bypass counts; reaching - the limit stops further bypasses so active requests can drain. + admission or resumption increments earlier blocked requests' bypass + counts; reaching the limit lets active requests drain. """ if not self.queues[RequestPhase.WAITING]: return @@ -499,6 +584,7 @@ def _admit(self): available_count = self._ensure_active_slot(request, active_count) if available_count is None: break + slot_preempted = available_count < active_count active_count = available_count request_id = waiting.popleft() request = self.requests[request_id] @@ -515,26 +601,30 @@ def _admit(self): self._enter_prefill(request) 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) - 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) - admitted = self._commit_admission(request, plan) - if not admitted: + if slot_preempted: + self._record_failed_priority_preemption(request) deferred.append(request_id) if request.admission_bypasses >= self.config.max_admission_bypasses: break continue + if resume_result == ResumeResult.NOT_PREEMPTED: + plan = self._plan_admission(request, active_count) + 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) + admitted = self._commit_admission(request, plan) + if not admitted: + if slot_preempted: + self._record_failed_priority_preemption(request) + deferred.append(request_id) + if request.admission_bypasses >= self.config.max_admission_bypasses: + break + continue + self._blocked_priority_preemptions.pop(request_id, None) active_count += 1 for blocked_id in deferred: self.requests[blocked_id].admission_bypasses += 1 diff --git a/vllm_rlt/engine/preemption.py b/vllm_rlt/engine/preemption.py index f8be4a75..1e0ed0ff 100644 --- a/vllm_rlt/engine/preemption.py +++ b/vllm_rlt/engine/preemption.py @@ -20,6 +20,14 @@ 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 resume_required_blocks(self, request) -> int | None: + """Return restore allocation size, or None when no snapshot exists.""" + state = self.snapshots.get(request.request_id) + if state is None: + return None + kv = state["kv"] + return 0 if kv is None else kv.pages * self.engine.cache_manager.storage_depths + def suspend(self, victim, *, offload_kv=True): """Release a runner slot, optionally reclaiming KV capacity as well. diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..8a1f9ebb 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -262,6 +262,7 @@ def __init__( if scheduler_config.enable_preemption: self.scheduler.suspend_callback = self.preemption.suspend self.scheduler.resume_callback = self.preemption.resume + self.scheduler.resume_blocks_callback = self.preemption.resume_required_blocks @property def connector(self):