From 4c5e4272e58dfc599e7e8d5fb0cb36406ec7574d Mon Sep 17 00:00:00 2001 From: Costa Tsaousis Date: Sun, 27 Sep 2026 22:34:04 +0000 Subject: [PATCH 1/7] Preserve external checkpoint restores through capacity-constrained admission Signed-off-by: Costa Tsaousis --- .lil/changes/vllm-restore-admission.json | 16 + tests/v1/core/test_boundary_admission.py | 416 ++++++++++++++++++++++- vllm/v1/core/kv_cache_manager.py | 308 +++++++++++++++-- vllm/v1/core/sched/scheduler.py | 60 +++- 4 files changed, 765 insertions(+), 35 deletions(-) create mode 100644 .lil/changes/vllm-restore-admission.json diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json new file mode 100644 index 000000000000..c6f2c86894e2 --- /dev/null +++ b/.lil/changes/vllm-restore-admission.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-restore-admission", + "category": "fix", + "summary": "External recurrent checkpoint restores can reserve execution capacity and retain ownership through admission.", + "models": ["all"], + "compatibility": "Use with the paired LMCache restore-admission change to wait for checkpoint capacity instead of recomputing cached prefixes.", + "details": [ + "Restores reserve GPU blocks and running slots before copies start, and keep restored pages owned until execution takes over.", + "Sliding-window and chunked-local continuation reservations count physical pages rather than skipped positions in the block table.", + "Requests with reserved capacity can progress past a blocked queue head; cancellation releases reservations after submitted copies drain.", + "Existing external-checkpoint callers retain their previous behavior unless they opt into admission reservations." + ], + "authors": ["ktsaou"], + "requires": [] +} diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index 02f02b449331..6f7c0c4f958f 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -11,12 +11,17 @@ from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request from vllm.utils.hashing import sha256 -from vllm.v1.core.boundary_checkpoint import BoundaryCheckpointCache +from vllm.v1.core.boundary_checkpoint import ( + BoundaryCheckpointCache, + boundary_checkpoint_slots, +) from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, KVCacheConfig, KVCacheGroupSpec, MambaSpec, MLAAttentionSpec, + SlidingWindowSpec, ) from vllm.v1.request import RequestStatus @@ -38,7 +43,7 @@ def compute_share_fixture_options(share: float | None) -> dict[str, Any]: return options -def manager(): +def manager(dcp=1, num_blocks=128): attention = MLAAttentionSpec( block_size=64, num_kv_heads=1, @@ -56,13 +61,15 @@ def manager(): groups = [KVCacheGroupSpec(["attention"], attention)] groups += [KVCacheGroupSpec([f"recurrent-{i}"], recurrent) for i in range(3)] return make_kv_cache_manager( - KVCacheConfig(num_blocks=128, kv_cache_tensors=[], kv_cache_groups=groups), + KVCacheConfig( + num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=groups + ), max_model_len=512, max_in_flight_tokens=32, enable_caching=True, use_eagle=True, num_prefill_lookahead=1, - dcp_world_size=1, + dcp_world_size=dcp, scheduler_block_size=64, hash_block_size=16, enable_boundary_checkpoints=True, @@ -81,6 +88,292 @@ def drain(cache): cache.block_pool.free_blocks(copies) +def reserve_import(cache, consumer, prefix=100): + return cache.reserve_external_boundary_checkpoint( + consumer, + prefix, + cache.boundary_checkpoint_page_positions(prefix), + draft_prefix_len=prefix, + kind="prompt", + num_ranks=2, + reserve_admission=True, + ) + + +@pytest.mark.parametrize("dcp", [1, 2]) +@pytest.mark.parametrize("prefix", [16, 100, 128]) +@pytest.mark.parametrize("chunk", [1, 16, 33, 127]) +def test_external_import_owns_execution_capacity_through_chunked_prefill( + dcp, prefix, chunk +): + """Other admissions cannot consume a restored request's remaining suffix.""" + cache = manager(dcp) + consumer = request("import", length=400) + free = cache.block_pool.get_num_free_blocks() + checkpoint = reserve_import(cache, consumer, prefix=prefix) + assert checkpoint is not None + credits = cache.external_boundary_reserved_blocks() + assert credits > 0 + assert not cache.acknowledge_external_boundary_checkpoint( + checkpoint.checkpoint_id, 0 + ) + assert cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, 1) + assert not cache.reset_prefix_cache() + pressure = cache.block_pool.get_new_blocks( + cache.block_pool.get_num_free_blocks() - credits + ) + assert cache.allocate_slots(request("competitor", "other"), 1) is None + blocks, hits, _ = cache.get_computed_blocks(consumer) + assert hits == prefix + assert ( + cache.allocate_slots( + consumer, + chunk, + num_new_computed_tokens=prefix, + new_computed_blocks=blocks, + num_lookahead_tokens=3, + full_sequence_must_fit=True, + ) + is not None + ) + consumer.num_computed_tokens = prefix + chunk + drain(cache) + assert cache.has_external_boundary_admission(consumer.request_id) + assert cache.external_boundary_reserved_blocks() > 0 + while consumer.num_computed_tokens < consumer.num_tokens: + count = min(chunk, consumer.num_tokens - consumer.num_computed_tokens) + assert cache.allocate_slots(consumer, count, num_lookahead_tokens=3) is not None + consumer.num_computed_tokens += count + drain(cache) + assert not cache.has_external_boundary_admission(consumer.request_id) + cache.free(consumer) + cache.block_pool.free_blocks(pressure) + assert cache.block_pool.get_num_free_blocks() == free + + +@pytest.mark.parametrize( + "attention_kind,retained", [("sliding", 0), ("sliding", 48), ("chunked", 0)] +) +@pytest.mark.parametrize( + "prefix,length", [(128, 400), (128, 128), (128, 129), (100, 100), (100, 101)] +) +def test_scheduler_preserves_recycling_import_capacity( + attention_kind, retained, prefix, length +): + """Null positions must not release credits while the live window moves.""" + from vllm.v1.core.sched.scheduler import Scheduler + from vllm.v1.outputs import ModelRunnerOutput + + template = make_scheduler( + enable_prefix_caching=True, + use_v2_model_runner=True, + max_num_batched_tokens=32, + long_prefill_token_threshold=16, + block_size=64, + max_model_len=512, + ) + config = template.vllm_config + config.cache_config.mamba_cache_mode = "align" + config.cache_config.num_gpu_blocks = 64 + attention_args = dict(block_size=64, num_kv_heads=1, head_size=1, dtype=torch.uint8) + attention = ( + SlidingWindowSpec( + **attention_args, sliding_window=128, extra_retained_tokens=retained + ) + if attention_kind == "sliding" + else ChunkedLocalAttentionSpec(**attention_args, attention_chunk_size=128) + ) + recurrent = MambaSpec( + block_size=16, shapes=((1,),), dtypes=(torch.uint8,), mamba_cache_mode="align" + ) + scheduler = Scheduler( + vllm_config=config, + kv_cache_config=KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], attention), + KVCacheGroupSpec(["recurrent"], recurrent), + ], + ), + structured_output_manager=template.structured_output_manager, + block_size=64, + hash_block_size=16, + ) + scheduler.use_v2_model_runner = True + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + consumer = request("recycling-import", length=length) + scheduler.add_request(consumer) + checkpoint = reserve_import(cache, consumer, prefix=prefix) + assert checkpoint is not None + if length <= prefix + 1: + # One attention page and two private recurrent states suffice for + # this suffix; imported read-only history need not be replaced. + assert cache.external_boundary_reserved_blocks() <= ( + 3 + 3 * len(boundary_checkpoint_slots(consumer)) + ) + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, rank) + pressure = [] + while consumer.num_computed_tokens < consumer.num_tokens: + spare = ( + cache.block_pool.get_num_free_blocks() + - cache.external_boundary_reserved_blocks() + ) + pressure.extend(cache.block_pool.get_new_blocks(max(0, spare))) + expected = min(16, max(1, length - max(prefix, consumer.num_computed_tokens))) + output = scheduler.schedule() + assert output.num_scheduled_tokens == {consumer.request_id: expected} + assert consumer.num_preemptions == 0 + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=[consumer.request_id], + req_id_to_index={consumer.request_id: 0}, + sampled_token_ids=[[]], + ), + ) + cache.free(consumer) + cache.block_pool.free_blocks(pressure) + assert cache.external_boundary_reserved_blocks() == 0 + assert cache.block_pool.get_num_free_blocks() == 63 + + +@pytest.mark.parametrize("cancel_at", ["copying", "ready"]) +def test_external_import_cancellation_releases_only_after_copy_drain(cancel_at): + cache = manager() + consumer = request("cancel") + free = cache.block_pool.get_num_free_blocks() + checkpoint = reserve_import(cache, consumer) + assert checkpoint is not None + if cancel_at == "ready": + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint( + checkpoint.checkpoint_id, rank + ) + cache.free(consumer) + if cancel_at == "copying": + assert cache.block_pool.get_num_free_blocks() < free + assert cache.external_boundary_reserved_blocks() > 0 + cache.discard_external_boundary_checkpoint(checkpoint.checkpoint_id) + cache.release_external_boundary_admission(consumer.request_id) + assert cache.external_boundary_reserved_blocks() == 0 + assert cache.block_pool.get_num_free_blocks() == free + + +def test_external_imports_reserve_slots_without_serializing_available_imports(): + cache = manager() + cache.set_external_boundary_admission_context(2, 0) + first, second, third = [request(name, name) for name in ("one", "two", "three")] + assert reserve_import(cache, first) is not None + assert reserve_import(cache, second) is not None + free = cache.block_pool.get_num_free_blocks() + assert reserve_import(cache, third) is None + assert not cache.can_admit_external_boundary_request(third.request_id) + assert cache.block_pool.get_num_free_blocks() == free + assert cache.can_admit_external_boundary_request(first.request_id) + + +def test_external_import_preflight_does_not_evict_cache_on_execution_pressure(): + cache = manager() + consumer = request("consumer") + cache.set_external_boundary_admission_context( + 1, cache.block_pool.get_num_free_blocks() + ) + free = cache.block_pool.get_num_free_blocks() + assert reserve_import(cache, consumer) is None + assert cache.block_pool.get_num_free_blocks() == free + assert cache.external_boundary_reserved_blocks() == 0 + + +def test_capacity_waiter_prevents_newer_admissions_from_stealing_released_space(): + cache = manager() + older, newer = request("older"), request("newer", "newer") + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert reserve_import(cache, older) is None + cache.block_pool.free_blocks(pressure) + assert not cache.can_admit_external_boundary_request(newer.request_id) + assert reserve_import(cache, newer) is None + assert reserve_import(cache, older) is not None + assert cache.can_admit_external_boundary_request(newer.request_id) + + +def test_cancelled_capacity_waiter_releases_admission_barrier(): + cache = manager() + consumer = request("cancel-waiter") + cache.set_external_boundary_admission_context( + 1, cache.block_pool.get_num_free_blocks() + ) + assert reserve_import(cache, consumer) is None + assert not cache.can_admit_external_boundary_request("later") + cache.release_external_boundary_admission(consumer.request_id) + assert cache.can_admit_external_boundary_request("later") + + +def test_import_preflight_preserves_watermark_and_classifies_impossible_restore(): + cache = manager() + consumer = request("watermark") + free = cache.block_pool.get_num_free_blocks() + cache.watermark_blocks = free + cache.set_external_boundary_admission_context(1, 0, has_scheduled_reqs=True) + assert reserve_import(cache, consumer) is None + assert cache.block_pool.get_num_free_blocks() == free + # The last runnable request finishing removes the normal admission watermark. + cache.set_external_boundary_admission_context(1, 0, has_scheduled_reqs=False) + assert reserve_import(cache, consumer) is not None + cache = manager(num_blocks=16) + free = cache.block_pool.get_num_free_blocks() + with pytest.raises(ValueError, match="exceed the GPU pool"): + reserve_import(cache, consumer) + assert cache.block_pool.get_num_free_blocks() == free + assert not cache.has_pending_external_boundary_admissions() + + +def test_eight_external_sessions_rotate_without_cold_prefill_or_leaked_blocks(): + """Aggregate checkpoints exceed the pool; each admitted turn still restores.""" + cache = manager(dcp=2, num_blocks=64) + free = cache.block_pool.get_num_free_blocks() + consumers = [request(f"session-{i}", str(i), length=400) for i in range(8)] + pages = sum(map(len, cache.boundary_checkpoint_page_positions(400))) + 1 + assert pages * len(consumers) > free + remaining = list(consumers) + completed = [] + while remaining: + admitted = [] + cache.set_external_boundary_admission_context(8, 0) + for consumer in remaining: + checkpoint = reserve_import(cache, consumer, prefix=400) + if checkpoint is not None: + admitted.append((consumer, checkpoint)) + assert admitted, "imports must make progress when previous owners drain" + for consumer, checkpoint in admitted: + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint( + checkpoint.checkpoint_id, rank + ) + blocks, hits, _ = cache.get_computed_blocks(consumer) + assert hits == consumer.num_tokens + assert ( + cache.allocate_slots( + consumer, + 1, + num_new_computed_tokens=hits, + new_computed_blocks=blocks, + num_lookahead_tokens=3, + full_sequence_must_fit=True, + ) + is not None + ) + drain(cache) + cache.free(consumer) + remaining.remove(consumer) + completed.append(consumer.request_id) + assert completed == [consumer.request_id for consumer in consumers] + assert cache.block_pool.get_num_free_blocks() == free + assert cache.external_boundary_reserved_blocks() == 0 + + def seed(cache, salt): producer = request("producer-" + salt, salt) cache.get_computed_blocks(producer) @@ -270,6 +563,121 @@ def make_scheduler(**kwargs): return create_scheduler(**kwargs) +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +@pytest.mark.parametrize("lanes", [1, 2]) +@pytest.mark.parametrize("limit", ["slots", "blocks"]) +def test_ready_import_behind_blocked_head_gets_saved_logits_step(policy, lanes, limit): + """Reserved imports remain reachable while unrelated requests wait.""" + from unittest.mock import Mock + + from tests.v1.core.utils import create_requests + + scheduler = make_scheduler( + enable_prefix_caching=True, + use_v2_model_runner=True, + async_scheduling=True, + max_num_seqs=2 if limit == "slots" else 3, + max_parallel_prefills=lanes, + scheduling_policy=policy, + ) + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + running, consumer = create_requests( + num_requests=2, + num_tokens=32, + req_ids=["running", "consumer"], + ) + (cold,) = create_requests(num_requests=1, num_tokens=160, req_ids=["cold"]) + scheduler.add_request(running) + assert scheduler.schedule().num_scheduled_tokens == {"running": 32} + checkpoint = cache.reserve_external_boundary_checkpoint( + consumer, + 32, + cache.boundary_checkpoint_page_positions(32), + draft_prefix_len=32, + kind="prompt", + num_ranks=2, + reserve_admission=True, + ) + assert checkpoint is not None + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, rank) + pressure = [] + if limit == "blocks": + pressure = cache.block_pool.get_new_blocks( + cache.block_pool.get_num_free_blocks() + - cache.external_boundary_reserved_blocks() + ) + scheduler.connector = Mock() + scheduler.connector.poll_boundary_checkpoint.return_value = True + scheduler.connector.boundary_checkpoint_external_tokens.side_effect = lambda req: ( + 32 if req is consumer else 0 + ) + scheduler.connector.get_num_new_matched_tokens.return_value = (0, False) + scheduler.add_request(cold) + scheduler.add_request(consumer) + output = scheduler.schedule() + assert output.boundary_logits_only + assert output.num_scheduled_tokens == {"consumer": 1} + assert consumer.prefill_stats.num_computed_tokens == 0 + assert consumer.prefill_stats.num_external_cached_tokens == 32 + assert cold.status == RequestStatus.WAITING + cache.block_pool.free_blocks(pressure) + + +def test_streaming_owner_can_resume_ahead_of_capacity_waiter(): + from unittest.mock import Mock + + from tests.v1.core.utils import create_requests + + scheduler = make_scheduler( + enable_prefix_caching=True, + use_v2_model_runner=True, + max_num_seqs=8, + num_blocks=12, + ) + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + (session,) = create_requests(num_requests=1, num_tokens=100, req_ids=["stream"]) + session.resumable = True + scheduler.add_request(session) + assert scheduler.schedule().num_scheduled_tokens == {"stream": 100} + session.num_in_flight_tokens = 0 + scheduler.running.remove(session) + assert not scheduler._handle_stopped_request(session) + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + (consumer,) = create_requests(num_requests=1, num_tokens=64, req_ids=["restore"]) + connector = Mock() + + def poll(req): + if req is not consumer: + return True + restored = cache.reserve_external_boundary_checkpoint( + req, + 64, + cache.boundary_checkpoint_page_positions(64), + draft_prefix_len=64, + kind="prompt", + num_ranks=2, + reserve_admission=True, + ) + assert restored is None + return False + + connector.poll_boundary_checkpoint.side_effect = poll + connector.boundary_checkpoint_external_tokens.return_value = 0 + connector.get_num_new_matched_tokens.return_value = (0, False) + scheduler.connector = connector + scheduler.add_request(consumer) + assert scheduler.schedule().num_scheduled_tokens == {} + assert not cache.can_admit_external_boundary_request(session.request_id) + (update,) = create_requests(num_requests=1, num_tokens=1, req_ids=["stream"]) + update.resumable = True + scheduler.add_request(update) + assert session.status == RequestStatus.WAITING + assert scheduler.schedule().num_scheduled_tokens == {"stream": 1} + + @pytest.mark.parametrize("release", ["move_victim", "finish_reader"]) @pytest.mark.parametrize("fairness", [None, 0.4]) def test_actual_scheduler_runs_decode_after_guard_defers_restore(fairness, release): diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index de3b56ff7183..a737fa9e6aa1 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -26,13 +26,18 @@ ) from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy -from vllm.v1.core.single_type_kv_cache_manager import MambaManager +from vllm.v1.core.single_type_kv_cache_manager import ( + MambaManager, + SingleTypeKVCacheManager, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, + ChunkedLocalAttentionSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, KVCacheConfig, MambaSpec, + SlidingWindowSpec, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, iter_layer_specs, @@ -43,6 +48,15 @@ logger = init_logger(__name__) +@dataclass +class _BoundaryImportAdmission: + request: Request + checkpoint: BoundaryCheckpoint + remaining_blocks: int + ready: bool = False + admitted: bool = False + + @dataclass class KVCacheBlocks: """ @@ -171,6 +185,7 @@ def __init__( # always supplies the real value at runtime. if max_in_flight_tokens is None: max_in_flight_tokens = max_model_len + self._boundary_max_in_flight_tokens = max_in_flight_tokens self.enable_caching = enable_caching self.enable_kv_cache_events = enable_kv_cache_events @@ -247,6 +262,11 @@ def __init__( else None ) self._boundary_imports: set[int] = set() + self._boundary_import_admissions: dict[str, _BoundaryImportAdmission] = {} + self._boundary_import_waiters: dict[str, None] = {} + self._external_admission_slots: int | None = None + self._external_admission_reserved_blocks = 0 + self._external_admission_has_scheduled_reqs = False self._boundary_allocations: dict[str, list[KVCacheBlock]] = {} self._boundary_readers: dict[str, BoundaryCheckpoint] = {} self._boundary_reader_horizons: dict[str, int] = {} @@ -339,7 +359,12 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int return self.empty_kv_cache_blocks, 0, 0 if boundary_cache is not None and boundary_cache.supports_request(request): - checkpoint = boundary_cache.find(request, request.num_tokens) + admission = self._boundary_import_admissions.get(request.request_id) + checkpoint = ( + admission.checkpoint + if admission is not None and admission.ready and not admission.admitted + else boundary_cache.find(request, request.num_tokens) + ) request.boundary_checkpoint = checkpoint if checkpoint is not None: checkpoint_blocks = tuple( @@ -597,9 +622,7 @@ def allocate_slots( boundary_blocks = 0 boundary_replay_managers = [] if request.use_boundary_checkpoints: - if request.request_id not in self._boundary_allocations: - checkpoint_slots = len(boundary_checkpoint_slots(request)) - boundary_blocks = checkpoint_slots * (self.num_kv_cache_groups + 1) + boundary_blocks = self._boundary_capture_blocks(request) if request.boundary_checkpoint is not None: boundary_blocks += sum( self.block_pool.blocks[i].ref_cnt == 0 @@ -610,15 +633,9 @@ def allocate_slots( and new_computed_blocks is not None and num_local_computed_tokens > 0 ): - boundary_replay_managers = [ - manager - for manager in self.coordinator.single_type_managers - if num_local_computed_tokens % manager.block_size == 0 - and not any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(manager.kv_cache_spec) - ) - ] + boundary_replay_managers = self._boundary_replay_managers( + num_local_computed_tokens + ) boundary_blocks += len(boundary_replay_managers) # The watermark is applied to waiting/preempted requests only, and only # when there's at least one request already scheduled. @@ -645,7 +662,10 @@ def allocate_slots( required_blocks = ( num_blocks_to_allocate + watermark_blocks + boundary_blocks ) - if required_blocks > self.block_pool.get_num_free_blocks(): + if required_blocks > ( + self.block_pool.get_num_free_blocks() + - self.external_boundary_reserved_blocks(request.request_id) + ): return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -681,7 +701,11 @@ def allocate_slots( # Keep `reserved_blocks` free for other in-flight sequences, and an # additional watermark of headroom for waiting/preempted admissions. - available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks + available_blocks = ( + self.block_pool.get_num_free_blocks() + - reserved_blocks + - self.external_boundary_reserved_blocks(request.request_id) + ) required_blocks = num_blocks_to_allocate + watermark_blocks + boundary_blocks if required_blocks > available_blocks: # Cannot allocate new blocks @@ -763,9 +787,14 @@ def allocate_slots( and request.request_id not in self._boundary_readers ): assert self.boundary_checkpoints is not None - checkpoint = self.boundary_checkpoints.acquire( - request.boundary_checkpoint.checkpoint_id - ) + admission = self._boundary_import_admissions.get(request.request_id) + if admission is not None and admission.ready and not admission.admitted: + checkpoint = admission.checkpoint + admission.admitted = True + else: + checkpoint = self.boundary_checkpoints.acquire( + request.boundary_checkpoint.checkpoint_id + ) if checkpoint is None: request.boundary_checkpoint = None return None @@ -813,6 +842,19 @@ def allocate_slots( by_slot.get(slot, (0,) * width) for slot in range(max(slots) + 1) ) + admission = self._boundary_import_admissions.get(request.request_id) + if admission is not None and admission.admitted: + if num_tokens_main_model >= request.num_tokens: + del self._boundary_import_admissions[request.request_id] + else: + admission.remaining_blocks = self._external_boundary_execution_blocks( + request, + num_tokens_main_model, + self.empty_kv_cache_blocks.blocks, + restoring=False, + ) + self._boundary_import_waiters.pop(request.request_id, None) + # P/D: delay caching blocks if we have to recv from # remote. Update state for locally cached blocks. if delay_cache_blocks: @@ -844,6 +886,7 @@ def free(self, request: Request) -> None: self.coordinator.free(request.request_id) def _pop_boundary_blocks(self, request: Request) -> list[KVCacheBlock]: + self.release_external_boundary_admission(request.request_id) blocks = self._boundary_allocations.pop(request.request_id, []) reader = self._boundary_readers.pop(request.request_id, None) self._boundary_reader_horizons.pop(request.request_id, None) @@ -967,6 +1010,7 @@ def reserve_external_boundary_checkpoint( draft_prefix_len: int, kind: BoundaryCheckpointKind, num_ranks: int, + reserve_admission: bool = False, ) -> BoundaryCheckpoint | None: """Reserve private destinations for an externally stored checkpoint. @@ -974,9 +1018,15 @@ def reserve_external_boundary_checkpoint( acknowledges its completed H2D copies. Destinations are pinned in the ordinary KV pool, including one page for target/draft auxiliary state. - Returns None if request-boundary caching is unavailable or the pool - cannot admit every live page. Raises ValueError for incompatible cache - geometry; callers must treat that as an external cache miss. + With reserve_admission, also reserve a running slot and continuation + credits. Successful publication retains the consumer pin until ordinary + allocation takes ownership. The scheduler must refresh admission context + before each candidate; cancellation must release the admission after all + copies drain. + + Returns None for temporary resource pressure or unavailable caching. + Raises ValueError for incompatible geometry or an import whose checkpoint + and execution state cannot fit even an otherwise empty pool. """ cache = self.boundary_checkpoints if ( @@ -988,6 +1038,10 @@ def reserve_external_boundary_checkpoint( if not 0 < num_tokens <= request.num_tokens: raise ValueError("Imported boundary must cover an existing request prefix") if not checkpoint_end_allowed(request, num_tokens): + if reserve_admission: + raise ValueError( + "Imported checkpoint ends inside an unauthenticated span" + ) return None if not 0 <= draft_prefix_len <= num_tokens: raise ValueError("Imported draft prefix exceeds the target prefix") @@ -999,7 +1053,37 @@ def reserve_external_boundary_checkpoint( if page_positions != self.boundary_checkpoint_page_positions(num_tokens): raise ValueError("External checkpoint does not cover every live cache page") count = sum(len(group) for group in page_positions) + 1 - if count > self.block_pool.get_num_free_blocks(): + execution_blocks = 0 + if reserve_admission: + if request.request_id in self._boundary_import_admissions: + return None + if not self.can_admit_external_boundary_request(request.request_id): + return None + # Estimation must not evict cached pages or mutate real block metadata. + pinned = KVCacheBlock(block_id=-1, ref_cnt=1) + prospective = tuple( + tuple( + pinned if i in positions else self.block_pool.null_block + for i in range(positions[-1] + 1) + ) + for positions in page_positions + ) + execution_blocks = self._external_boundary_execution_blocks( + request, + num_tokens, + prospective, + restoring=True, + ) + if count + execution_blocks > self.block_pool.num_gpu_blocks - 1: + raise ValueError("Checkpoint and execution state exceed the GPU pool") + reserved = self.external_boundary_reserved_blocks() + if reserve_admission: + reserved += self._external_admission_reserved_blocks + if self._external_admission_has_scheduled_reqs: + reserved += self.watermark_blocks + if count + execution_blocks + reserved > self.block_pool.get_num_free_blocks(): + if reserve_admission: + self._boundary_import_waiters.setdefault(request.request_id, None) return None allocation = self.block_pool.get_new_blocks(count) try: @@ -1020,6 +1104,11 @@ def reserve_external_boundary_checkpoint( ) cache.stage(request, checkpoint, num_ranks=num_ranks) self._boundary_imports.add(checkpoint.checkpoint_id) + if reserve_admission: + self._boundary_import_waiters.pop(request.request_id, None) + self._boundary_import_admissions[request.request_id] = ( + _BoundaryImportAdmission(request, checkpoint, execution_blocks) + ) return checkpoint finally: # stage() owns a separate pin on every dependency until publication @@ -1036,6 +1125,20 @@ def acknowledge_external_boundary_checkpoint( published = cache.acknowledge(checkpoint_id, rank) if not cache.is_pending(checkpoint_id): self._boundary_imports.remove(checkpoint_id) + for request_id, admission in tuple( + self._boundary_import_admissions.items() + ): + if admission.checkpoint.checkpoint_id != checkpoint_id: + continue + if published: + # No scheduler allocation may intervene between publication + # and acquiring the consumer's ownership. + acquired = cache.acquire(checkpoint_id) + assert acquired is not None + admission.ready = True + admission.request.boundary_checkpoint = admission.checkpoint + else: + del self._boundary_import_admissions[request_id] return published def discard_external_boundary_checkpoint(self, checkpoint_id: int) -> None: @@ -1048,6 +1151,165 @@ def discard_external_boundary_checkpoint(self, checkpoint_id: int) -> None: assert self.boundary_checkpoints is not None self.boundary_checkpoints.discard(checkpoint_id) self._boundary_imports.remove(checkpoint_id) + for request_id, admission in tuple( + self._boundary_import_admissions.items() + ): + if admission.checkpoint.checkpoint_id == checkpoint_id: + del self._boundary_import_admissions[request_id] + + def set_external_boundary_admission_context( + self, + available_slots: int, + reserved_blocks: int, + *, + has_scheduled_reqs: bool = False, + ) -> None: + """Refresh scheduler capacity immediately before polling a candidate. + + Args: + available_slots: Slots remaining after running and paused sessions. + reserved_blocks: Remaining needs of prefills without import tickets. + has_scheduled_reqs: Apply the normal admission watermark when true. + """ + self._external_admission_slots = available_slots + self._external_admission_reserved_blocks = reserved_blocks + self._external_admission_has_scheduled_reqs = has_scheduled_reqs + + def has_external_boundary_admission(self, request_id: str) -> bool: + """Whether this request already owns import continuation credits.""" + return request_id in self._boundary_import_admissions + + def external_boundary_reserved_blocks(self, owner: str | None = None) -> int: + """Return outstanding import credits, excluding the allocating owner.""" + return sum( + admission.remaining_blocks + for request_id, admission in self._boundary_import_admissions.items() + if request_id != owner + ) + + def can_admit_external_boundary_request(self, request_id: str) -> bool: + """Whether a candidate can use a running slot without stealing a ticket.""" + if request_id in self._boundary_import_admissions: + return True + if self._boundary_import_waiters and request_id != next( + iter(self._boundary_import_waiters) + ): + return False + return ( + self._external_admission_slots is None + or sum( + not admission.admitted + for admission in self._boundary_import_admissions.values() + ) + < self._external_admission_slots + ) + + def has_pending_external_boundary_admissions(self) -> bool: + """Whether an import still needs ordinary scheduler admission.""" + return any(not a.admitted for a in self._boundary_import_admissions.values()) + + def external_boundary_admission_ready(self, request_id: str) -> bool: + """Whether a waiting import has finished copying and owns its checkpoint.""" + admission = self._boundary_import_admissions.get(request_id) + return admission is not None and admission.ready and not admission.admitted + + def release_external_boundary_admission(self, request_id: str) -> None: + """Release a finished consumer's credits; in-flight copies must drain first.""" + self._boundary_import_waiters.pop(request_id, None) + admission = self._boundary_import_admissions.get(request_id) + if ( + admission is None + or admission.checkpoint.checkpoint_id in self._boundary_imports + ): + return + del self._boundary_import_admissions[request_id] + if admission.ready and not admission.admitted: + assert self.boundary_checkpoints is not None + self.boundary_checkpoints.release(admission.checkpoint) + + def _external_boundary_execution_blocks( + self, + request: Request, + prefix: int, + blocks: tuple[Sequence[KVCacheBlock], ...], + *, + restoring: bool, + ) -> int: + full = min(request.num_tokens + 1, self.max_model_len) + needed = self.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=min(full + self._boundary_restore_lookahead, self.max_model_len), + new_computed_blocks=blocks, + num_encoder_tokens=0, + total_computed_tokens=prefix, + num_local_computed_tokens=prefix, + num_tokens_main_model=full, + apply_admission_cap=True, + ) + reader = self._boundary_readers.get(request.request_id) + dependencies = set(reader.dependencies) if reader is not None else set() + for group, manager in zip(blocks, self.coordinator.single_type_managers): + spec = manager.kv_cache_spec + if isinstance(spec, MambaSpec) and spec.mamba_cache_mode == "align": + # Continuation overlaps running/previous states while the + # imported source stays pinned. + peak = ( + 2 + spec.num_speculative_blocks + spec.num_prefill_checkpoint_blocks + ) + elif isinstance(spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)): + # Logical tables retain null positions after a window moves; + # only private physical pages can satisfy the continuation peak. + peak = spec.max_admission_blocks_per_request( + self._boundary_max_in_flight_tokens, self.max_model_len + ) + else: + continue + owned = sum( + not block.is_null and block.block_id not in dependencies + for block in manager.req_to_blocks.get(request.request_id, ()) + ) + endpoint_need = manager.get_num_blocks_to_allocate( + request.request_id, + min(full + self._boundary_restore_lookahead, self.max_model_len), + group, + prefix, + prefix, + full, + apply_admission_cap=True, + ) + remaining = max(0, peak - owned) + if isinstance(spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)): + # A short suffix may never fill a complete private window. + suffix_need = manager.get_num_blocks_to_allocate( + request.request_id, + min(full + self._boundary_restore_lookahead, self.max_model_len), + group, + prefix, + prefix, + full, + ) + remaining = min(remaining, suffix_need) + needed += max(0, remaining - endpoint_need) + needed += self._boundary_capture_blocks(request) + if restoring and self.use_eagle: + needed += len(self._boundary_replay_managers(prefix)) + return needed + + def _boundary_capture_blocks(self, request: Request) -> int: + if request.request_id in self._boundary_allocations: + return 0 + return len(boundary_checkpoint_slots(request)) * (self.num_kv_cache_groups + 1) + + def _boundary_replay_managers(self, prefix: int) -> list[SingleTypeKVCacheManager]: + return [ + manager + for manager in self.coordinator.single_type_managers + if prefix % manager.block_size == 0 + and not any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(manager.kv_cache_spec) + ) + ] def remove_skipped_blocks( self, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f2150fd9b51d..a61b5fac1c9d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -782,14 +782,30 @@ def _has_waiting_boundary_logits(self) -> bool: queue = self._select_waiting_queue_for_scheduling() if queue is None: return False - request = queue.peek_request() - checkpoint = request.boundary_checkpoint - return ( + pending_imports = ( + self.kv_cache_manager.has_pending_external_boundary_admissions() + ) + candidates = itertools.chain( + (queue.peek_request(),), + ( + request + for request in ( + itertools.chain(self.skipped_waiting, self.waiting) + if pending_imports + else () + ) + if self.kv_cache_manager.external_boundary_admission_ready( + request.request_id + ) + ), + ) + return any( request.status in (RequestStatus.WAITING, RequestStatus.PREEMPTED) and request.num_computed_tokens == 0 - and checkpoint is not None - and checkpoint.num_tokens == request.num_tokens + and request.boundary_checkpoint is not None + and request.boundary_checkpoint.num_tokens == request.num_tokens and (request.num_stale_output_tokens == 0 or request.drop_stale_output) + for request in candidates ) def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: @@ -1373,6 +1389,32 @@ def schedule_running_requests( external_boundary_hit_tokens = 0 # Get already-cached tokens. + if ( + request.num_computed_tokens == 0 + and self.connector is not None + and self.kv_cache_manager.boundary_checkpoints is not None + ): + self.kv_cache_manager.set_external_boundary_admission_context( + self.max_num_running_reqs - num_running, + sum( + self._request_remaining_blocks(req) + for req in self._inflight_prefills + if not ( + self.kv_cache_manager.has_external_boundary_admission( + req.request_id + ) + ) + ), + has_scheduled_reqs=bool(self.running), + ) + if not self.kv_cache_manager.can_admit_external_boundary_request( + request.request_id + ): + request_queue.remove_request(request) + if prefill_interleave_step is not None: + prefill_interleave_step.mark_unavailable(request_id) + step_skipped_waiting.prepend_request(request) + continue if request.num_computed_tokens == 0: if ( self.connector is not None @@ -1753,11 +1795,13 @@ def schedule_running_requests( # manager if request.has_encoder_inputs: self.encoder_cache_manager.free(request) - if not interleaved_local_prefill: + if not interleaved_local_prefill and not ( + self.kv_cache_manager.has_pending_external_boundary_admissions() + ): break request_queue.remove_request(request) - assert prefill_interleave_step is not None - prefill_interleave_step.mark_unavailable(request_id) + if prefill_interleave_step is not None: + prefill_interleave_step.mark_unavailable(request_id) step_skipped_waiting.prepend_request(request) continue From ae02fd76f68a4e1127210133316dd1abb4d767b0 Mon Sep 17 00:00:00 2001 From: Costa Tsaousis Date: Sun, 27 Sep 2026 22:35:20 +0000 Subject: [PATCH 2/7] Link checkpoint admission release note to PR 922 Signed-off-by: Costa Tsaousis --- .lil/changes/vllm-restore-admission.json | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json index c6f2c86894e2..1d7fdf96e909 100644 --- a/.lil/changes/vllm-restore-admission.json +++ b/.lil/changes/vllm-restore-admission.json @@ -3,7 +3,9 @@ "id": "vllm-restore-admission", "category": "fix", "summary": "External recurrent checkpoint restores can reserve execution capacity and retain ownership through admission.", - "models": ["all"], + "models": [ + "all" + ], "compatibility": "Use with the paired LMCache restore-admission change to wait for checkpoint capacity instead of recomputing cached prefixes.", "details": [ "Restores reserve GPU blocks and running slots before copies start, and keep restored pages owned until execution takes over.", @@ -11,6 +13,8 @@ "Requests with reserved capacity can progress past a blocked queue head; cancellation releases reservations after submitted copies drain.", "Existing external-checkpoint callers retain their previous behavior unless they opt into admission reservations." ], - "authors": ["ktsaou"], - "requires": [] + "requires": [], + "pull_requests": [ + 922 + ] } From 3f94f17e19fcdf5998d2b4afa5bed4db7d3603ae Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:13:56 +0000 Subject: [PATCH 3/7] fix(kv-cache): hold import reservations only against new admissions Follow-up to #922 review findings on the admission reservations: - Import credits gated every allocation, so a running decode at a block boundary could be preempted (often the admitted consumer itself) to protect credits of a restore that had not started. Credits now apply to waiting and preempted admissions only; running requests keep growing and resolve pressure through ordinary preemption. - A capacity waiter blocked every other new admission, including requests the connector does not handle and higher-priority ones. A refused restore now only holds back admissions considered after it in the same scheduler step, and those still run when they fit beside the blocks it waits for. Requests ordered ahead of it are unaffected. - A capacity wait is bounded by VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S (default 60, 0 recomputes immediately). Once it expires, one WARNING records the wait, the reservation keeps returning None, and the scheduler admits the request to recompute its prompt even though the connector still defers it for capacity. - Cancelling an import mid-copy kept its credits and running slot until the copy drained. Only the staged pages must wait for the copy now; credits and the slot are released at once. Co-Authored-By: Claude Opus 5.5 --- tests/v1/core/test_boundary_admission.py | 219 ++++++++++++++++++++--- vllm/envs.py | 7 + vllm/v1/core/kv_cache_manager.py | 145 +++++++++++---- vllm/v1/core/sched/scheduler.py | 5 + 4 files changed, 320 insertions(+), 56 deletions(-) diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index 6f7c0c4f958f..d8f0df707d69 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -240,23 +240,30 @@ def test_scheduler_preserves_recycling_import_capacity( assert cache.block_pool.get_num_free_blocks() == 63 +def publish(cache, checkpoint): + for rank in range(2): + cache.acknowledge_external_boundary_checkpoint(checkpoint.checkpoint_id, rank) + + @pytest.mark.parametrize("cancel_at", ["copying", "ready"]) -def test_external_import_cancellation_releases_only_after_copy_drain(cancel_at): +def test_external_import_cancellation_releases_only_pages_after_copy_drain(cancel_at): + """Cancellation frees the slot and credits at once; pages wait for copies.""" cache = manager() + cache.set_external_boundary_admission_context(1, 0) consumer = request("cancel") free = cache.block_pool.get_num_free_blocks() checkpoint = reserve_import(cache, consumer) assert checkpoint is not None + assert not cache.can_admit_external_boundary_request("other") if cancel_at == "ready": - for rank in range(2): - cache.acknowledge_external_boundary_checkpoint( - checkpoint.checkpoint_id, rank - ) + publish(cache, checkpoint) cache.free(consumer) + assert cache.external_boundary_reserved_blocks() == 0 + assert cache.can_admit_external_boundary_request("other") if cancel_at == "copying": assert cache.block_pool.get_num_free_blocks() < free - assert cache.external_boundary_reserved_blocks() > 0 cache.discard_external_boundary_checkpoint(checkpoint.checkpoint_id) + # The connector releases again once the cancelled copy has drained. cache.release_external_boundary_admission(consumer.request_id) assert cache.external_boundary_reserved_blocks() == 0 assert cache.block_pool.get_num_free_blocks() == free @@ -287,28 +294,86 @@ def test_external_import_preflight_does_not_evict_cache_on_execution_pressure(): assert cache.external_boundary_reserved_blocks() == 0 -def test_capacity_waiter_prevents_newer_admissions_from_stealing_released_space(): +def refused_waiter_need(cache, waiter): + """Refuse a restore on a full pool and return the space it waits for.""" + probe = manager() + free = probe.block_pool.get_num_free_blocks() + assert reserve_import(probe, waiter) is not None + need = free - probe.block_pool.get_num_free_blocks() + need += probe.external_boundary_reserved_blocks() + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert reserve_import(cache, waiter) is None + return pressure, need + + +@pytest.mark.parametrize("released", ["waiter_only", "both"]) +def test_capacity_waiter_holds_back_only_admissions_that_would_take_its_space( + released, +): cache = manager() older, newer = request("older"), request("newer", "newer") - pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) - assert reserve_import(cache, older) is None - cache.block_pool.free_blocks(pressure) - assert not cache.can_admit_external_boundary_request(newer.request_id) - assert reserve_import(cache, newer) is None - assert reserve_import(cache, older) is not None + pressure, need = refused_waiter_need(cache, older) + cache.block_pool.free_blocks(pressure[: need * (2 if released == "both" else 1)]) assert cache.can_admit_external_boundary_request(newer.request_id) + # A later restore may use only the space beyond the waiter's need. + assert (reserve_import(cache, newer) is not None) == (released == "both") + assert reserve_import(cache, older) is not None + + +def test_capacity_waiter_does_not_hold_back_requests_considered_before_it(): + cache = manager() + older, newer = request("older"), request("newer", "newer") + pressure, need = refused_waiter_need(cache, older) + cache.new_step_starts() + cache.block_pool.free_blocks(pressure[:need]) + assert reserve_import(cache, newer) is not None + assert reserve_import(cache, older) is None def test_cancelled_capacity_waiter_releases_admission_barrier(): cache = manager() - consumer = request("cancel-waiter") - cache.set_external_boundary_admission_context( - 1, cache.block_pool.get_num_free_blocks() - ) - assert reserve_import(cache, consumer) is None - assert not cache.can_admit_external_boundary_request("later") + consumer, later = request("cancel-waiter"), request("later", "later") + pressure, need = refused_waiter_need(cache, consumer) + cache.block_pool.free_blocks(pressure[:need]) cache.release_external_boundary_admission(consumer.request_id) - assert cache.can_admit_external_boundary_request("later") + assert reserve_import(cache, later) is not None + + +def bounded_wait_clock(monkeypatch, max_wait): + """Drive restore waits from a fake clock and capture the manager's warnings.""" + import time + + from vllm.v1.core import kv_cache_manager + + clock = [100.0] + warnings = [] + monkeypatch.setattr(time, "monotonic", lambda: clock[0]) + monkeypatch.setattr( + kv_cache_manager.logger, "warning", lambda *args: warnings.append(args) + ) + monkeypatch.setenv("VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S", str(max_wait)) + return clock, warnings + + +def test_capacity_waiter_stops_waiting_after_bounded_wait(monkeypatch): + clock, warnings = bounded_wait_clock(monkeypatch, 30) + cache = manager() + waiter, later = request("waiter"), request("later", "later") + pressure, need = refused_waiter_need(cache, waiter) + clock[0] += 29.0 + assert reserve_import(cache, waiter) is None + assert not warnings and not cache.external_boundary_wait_expired("waiter") + clock[0] += 1.5 + assert reserve_import(cache, waiter) is None + assert cache.external_boundary_wait_expired(waiter.request_id) + assert len(warnings) == 1 and warnings[0][2] == pytest.approx(30.5) + # The request recomputes: it neither restores nor holds later requests back. + cache.block_pool.free_blocks(pressure[:need]) + assert reserve_import(cache, waiter) is None + assert reserve_import(cache, later) is not None + assert len(warnings) == 1 + cache.release_external_boundary_admission(waiter.request_id) + assert not cache.external_boundary_wait_expired(waiter.request_id) def test_import_preflight_preserves_watermark_and_classifies_impossible_restore(): @@ -670,7 +735,7 @@ def poll(req): scheduler.connector = connector scheduler.add_request(consumer) assert scheduler.schedule().num_scheduled_tokens == {} - assert not cache.can_admit_external_boundary_request(session.request_id) + assert consumer.request_id in cache._boundary_import_waiters (update,) = create_requests(num_requests=1, num_tokens=1, req_ids=["stream"]) update.resumable = True scheduler.add_request(update) @@ -678,6 +743,118 @@ def poll(req): assert scheduler.schedule().num_scheduled_tokens == {"stream": 1} +def import_scheduler(**options): + """A real scheduler with checkpoint caching and a restore connector mock. + + Tests install the connector after any setup step that processes output. + """ + from unittest.mock import Mock + + scheduler = make_scheduler( + **{ + "enable_prefix_caching": True, + "use_v2_model_runner": True, + "async_scheduling": True, + **options, + } + ) + cache = scheduler.kv_cache_manager + cache.boundary_checkpoints = BoundaryCheckpointCache(cache.block_pool) + connector = Mock() + connector.poll_boundary_checkpoint.return_value = True + connector.boundary_checkpoint_external_tokens.return_value = 0 + connector.get_num_new_matched_tokens.return_value = (0, False) + connector.has_pending_block_frees.return_value = False + return scheduler, cache, connector + + +def start_decode(scheduler, name="running"): + from tests.v1.core.utils import create_requests + + (running,) = create_requests(num_requests=1, num_tokens=16, req_ids=[name]) + scheduler.add_request(running) + assert scheduler.schedule().num_scheduled_tokens == {name: 16} + return running + + +def test_unadmitted_import_credits_do_not_preempt_running_decode(): + """Reservations gate new admissions, never the growth of running requests.""" + scheduler, cache, connector = import_scheduler() + running = start_decode(scheduler) + importer = request("import", "import", length=64) + checkpoint = reserve_import(cache, importer, prefix=16) + assert checkpoint is not None + publish(cache, checkpoint) + credits = cache.external_boundary_reserved_blocks() + assert credits > 0 + # The decode crosses a block boundary with only the import's credits free. + cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks() - credits) + scheduler.connector = connector + scheduler.add_request(importer) + output = scheduler.schedule() + assert output.num_scheduled_tokens["running"] == 1 + assert running.status == RequestStatus.RUNNING + assert running.num_preemptions == 0 + + +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +def test_capacity_waiter_does_not_block_requests_ahead_of_it(policy): + """A restore waiting for capacity cannot overtake requests ordered before it.""" + scheduler, cache, connector = import_scheduler( + num_blocks=64, scheduling_policy=policy + ) + earlier = request("earlier", "earlier", length=16) + waiter = request("waiter", "waiter", length=400) + lookup_done = [False] + + def poll(req): + if req is earlier: + return lookup_done[0] + assert reserve_import(cache, req, prefix=384) is None + return False + + connector.poll_boundary_checkpoint.side_effect = poll + scheduler.connector = connector + scheduler.add_request(earlier) + scheduler.add_request(waiter) + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert scheduler.schedule().num_scheduled_tokens == {} + assert waiter.request_id in cache._boundary_import_waiters + # Space returns for the earlier request, but not for the restore. + cache.block_pool.free_blocks(pressure[:16]) + lookup_done[0] = True + assert scheduler.schedule().num_scheduled_tokens == {"earlier": 16} + assert waiter.status == RequestStatus.WAITING + + +def test_restore_past_its_bounded_wait_is_admitted_to_recompute(monkeypatch): + """The connector keeps deferring the refused restore; the scheduler admits it.""" + clock, warnings = bounded_wait_clock(monkeypatch, 30) + scheduler, cache, connector = import_scheduler( + num_blocks=64, long_prefill_token_threshold=64 + ) + # A first chunk fits where the whole restore does not. + scheduler.scheduler_reserve_full_isl = False + waiter = request("waiter", "waiter", length=400) + + def poll(req): + assert reserve_import(cache, req, prefix=384) is None + return False + + connector.poll_boundary_checkpoint.side_effect = poll + scheduler.connector = connector + scheduler.add_request(waiter) + pressure = cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + assert scheduler.schedule().num_scheduled_tokens == {} + cache.block_pool.free_blocks(pressure[:20]) + assert scheduler.schedule().num_scheduled_tokens == {} + clock[0] += 31.0 + output = scheduler.schedule() + assert output.num_scheduled_tokens == {"waiter": 64} + assert len(warnings) == 1 and warnings[0][2] == pytest.approx(31.0) + assert not cache.external_boundary_wait_expired(waiter.request_id) + + @pytest.mark.parametrize("release", ["move_victim", "finish_reader"]) @pytest.mark.parametrize("fairness", [None, 0.4]) def test_actual_scheduler_runs_decode_after_guard_defers_restore(fairness, release): diff --git a/vllm/envs.py b/vllm/envs.py index 8052fb1d502b..88f38fc34505 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -174,6 +174,7 @@ VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS: bool = False + VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S: float = 60.0 VLLM_MLA_DISABLE: bool = False VLLM_RAY_PER_WORKER_GPUS: float = 1.0 VLLM_RAY_BUNDLE_INDICES: str = "" @@ -1508,6 +1509,11 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS": lambda: ( os.getenv("VLLM_SCHEDULER_UNCAP_PREFILL_ONLY_STEPS", "0") == "1" ), + # Seconds an external checkpoint restore may wait for GPU capacity before + # its request recomputes the prompt instead. 0 recomputes immediately. + "VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S": lambda: max( + 0.0, float(os.getenv("VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S", "60")) + ), # If set, vLLM will disable the MLA attention optimizations. "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), # If set, vLLM will pick up the provided Flash Attention MLA @@ -2531,6 +2537,7 @@ def compile_factors() -> dict[str, object]: "VLLM_LOGGING_COLOR", "VLLM_LOG_STATS_INTERVAL", "VLLM_REQUEST_STALL_WARNING_S", + "VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S", "VLLM_LOG_REQUEST_TIMELINE", "VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "VLLM_TUNED_CONFIG_FOLDER", diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index a737fa9e6aa1..8ae6a51d5916 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -2,10 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools +import time from collections.abc import Iterable, Sequence from dataclasses import dataclass from typing import Literal, overload +import vllm.envs as envs from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger from vllm.utils.math_utils import cdiv @@ -57,6 +59,12 @@ class _BoundaryImportAdmission: admitted: bool = False +@dataclass +class _BoundaryImportWaiter: + need: int + since: float + + @dataclass class KVCacheBlocks: """ @@ -263,7 +271,14 @@ def __init__( ) self._boundary_imports: set[int] = set() self._boundary_import_admissions: dict[str, _BoundaryImportAdmission] = {} - self._boundary_import_waiters: dict[str, None] = {} + self._boundary_import_waiters: dict[str, _BoundaryImportWaiter] = {} + # Requests whose restore outwaited its capacity bound. They recompute + # the prompt instead of restoring until they are admitted. + self._boundary_import_expired: set[str] = set() + # The first restore refused for capacity in this scheduler step and + # the blocks it needs. Admissions considered after it must leave them. + self._boundary_import_barrier: tuple[str, int] | None = None + self._boundary_import_max_wait = envs.VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S self._external_admission_slots: int | None = None self._external_admission_reserved_blocks = 0 self._external_admission_has_scheduled_reqs = False @@ -639,11 +654,17 @@ def allocate_slots( boundary_blocks += len(boundary_replay_managers) # The watermark is applied to waiting/preempted requests only, and only # when there's at least one request already scheduled. - if has_scheduled_reqs and request.status in ( + new_admission = request.status in ( RequestStatus.WAITING, RequestStatus.PREEMPTED, - ): + ) + if has_scheduled_reqs and new_admission: watermark_blocks = self.watermark_blocks + # Import reservations hold back new admissions only. Running requests + # keep growing and resolve pressure through ordinary preemption. + import_reserve = ( + self._blocks_reserved_for_imports(request) if new_admission else 0 + ) if full_sequence_must_fit: # First check and fail if the full request sequence won't fit. @@ -662,10 +683,7 @@ def allocate_slots( required_blocks = ( num_blocks_to_allocate + watermark_blocks + boundary_blocks ) - if required_blocks > ( - self.block_pool.get_num_free_blocks() - - self.external_boundary_reserved_blocks(request.request_id) - ): + if required_blocks > self.block_pool.get_num_free_blocks() - import_reserve: return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -702,9 +720,7 @@ def allocate_slots( # Keep `reserved_blocks` free for other in-flight sequences, and an # additional watermark of headroom for waiting/preempted admissions. available_blocks = ( - self.block_pool.get_num_free_blocks() - - reserved_blocks - - self.external_boundary_reserved_blocks(request.request_id) + self.block_pool.get_num_free_blocks() - reserved_blocks - import_reserve ) required_blocks = num_blocks_to_allocate + watermark_blocks + boundary_blocks if required_blocks > available_blocks: @@ -854,6 +870,7 @@ def allocate_slots( restoring=False, ) self._boundary_import_waiters.pop(request.request_id, None) + self._boundary_import_expired.discard(request.request_id) # P/D: delay caching blocks if we have to recv from # remote. Update state for locally cached blocks. @@ -1019,12 +1036,15 @@ def reserve_external_boundary_checkpoint( ordinary KV pool, including one page for target/draft auxiliary state. With reserve_admission, also reserve a running slot and continuation - credits. Successful publication retains the consumer pin until ordinary - allocation takes ownership. The scheduler must refresh admission context - before each candidate; cancellation must release the admission after all - copies drain. + credits that later new admissions must leave free. Successful + publication retains the consumer pin until ordinary allocation takes + ownership. The scheduler must refresh admission context whenever it + admits a request; cancellation releases credits and the slot at once, + while staged pages stay pinned until their copies drain. Returns None for temporary resource pressure or unavailable caching. + A restore refused for VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S keeps getting + None, and the scheduler admits its request to recompute the prompt. Raises ValueError for incompatible geometry or an import whose checkpoint and execution state cannot fit even an otherwise empty pool. """ @@ -1055,7 +1075,10 @@ def reserve_external_boundary_checkpoint( count = sum(len(group) for group in page_positions) + 1 execution_blocks = 0 if reserve_admission: - if request.request_id in self._boundary_import_admissions: + if ( + request.request_id in self._boundary_import_admissions + or request.request_id in self._boundary_import_expired + ): return None if not self.can_admit_external_boundary_request(request.request_id): return None @@ -1076,14 +1099,17 @@ def reserve_external_boundary_checkpoint( ) if count + execution_blocks > self.block_pool.num_gpu_blocks - 1: raise ValueError("Checkpoint and execution state exceed the GPU pool") - reserved = self.external_boundary_reserved_blocks() - if reserve_admission: - reserved += self._external_admission_reserved_blocks + reserved = ( + self._blocks_reserved_for_imports(request) + + self._external_admission_reserved_blocks + ) if self._external_admission_has_scheduled_reqs: reserved += self.watermark_blocks + else: + reserved = self.external_boundary_reserved_blocks() if count + execution_blocks + reserved > self.block_pool.get_num_free_blocks(): if reserve_admission: - self._boundary_import_waiters.setdefault(request.request_id, None) + self._wait_for_import_capacity(request, count + execution_blocks) return None allocation = self.block_pool.get_new_blocks(count) try: @@ -1164,7 +1190,9 @@ def set_external_boundary_admission_context( *, has_scheduled_reqs: bool = False, ) -> None: - """Refresh scheduler capacity immediately before polling a candidate. + """Refresh scheduler capacity before polling candidates. + + The scheduler refreshes it again after every admission in its pass. Args: available_slots: Slots remaining after running and paused sessions. @@ -1187,14 +1215,54 @@ def external_boundary_reserved_blocks(self, owner: str | None = None) -> int: if request_id != owner ) + def _blocks_reserved_for_imports(self, request: Request) -> int: + """Blocks a new admission must leave free for imports ahead of it.""" + reserved = self.external_boundary_reserved_blocks(request.request_id) + barrier = self._boundary_import_barrier + if ( + barrier is not None + and barrier[0] != request.request_id + and request.num_computed_tokens == 0 + and request.request_id not in self._boundary_import_admissions + ): + reserved += barrier[1] + return reserved + + def _wait_for_import_capacity(self, request: Request, need: int) -> None: + """Record a capacity refusal and bound how long the restore may wait.""" + now = time.monotonic() + waiter = self._boundary_import_waiters.get(request.request_id) + if waiter is None: + waiter = _BoundaryImportWaiter(need, now) + self._boundary_import_waiters[request.request_id] = waiter + waited = now - waiter.since + if waited >= self._boundary_import_max_wait: + self._forget_import_waiter(request.request_id) + self._boundary_import_expired.add(request.request_id) + logger.warning( + "Checkpoint restore for request %s waited %.1f s for %d GPU " + "blocks; recomputing its prompt instead", + request.request_id, + waited, + need, + ) + return + waiter.need = need + # Later candidates in this step are behind the waiter; earlier ones, + # and requests considered first in the next step, are not. + if self._boundary_import_barrier is None: + self._boundary_import_barrier = (request.request_id, need) + + def _forget_import_waiter(self, request_id: str) -> None: + self._boundary_import_waiters.pop(request_id, None) + barrier = self._boundary_import_barrier + if barrier is not None and barrier[0] == request_id: + self._boundary_import_barrier = None + def can_admit_external_boundary_request(self, request_id: str) -> bool: """Whether a candidate can use a running slot without stealing a ticket.""" if request_id in self._boundary_import_admissions: return True - if self._boundary_import_waiters and request_id != next( - iter(self._boundary_import_waiters) - ): - return False return ( self._external_admission_slots is None or sum( @@ -1213,19 +1281,25 @@ def external_boundary_admission_ready(self, request_id: str) -> bool: admission = self._boundary_import_admissions.get(request_id) return admission is not None and admission.ready and not admission.admitted + def external_boundary_wait_expired(self, request_id: str) -> bool: + """Whether a restore outwaited its capacity bound, so the prompt recomputes.""" + return request_id in self._boundary_import_expired + def release_external_boundary_admission(self, request_id: str) -> None: - """Release a finished consumer's credits; in-flight copies must drain first.""" - self._boundary_import_waiters.pop(request_id, None) - admission = self._boundary_import_admissions.get(request_id) - if ( - admission is None - or admission.checkpoint.checkpoint_id in self._boundary_imports - ): + """Release an import's credits and slot, or end its capacity wait. + + A copy that is still in flight keeps only its staged pages pinned, until + the connector acknowledges or discards it. + """ + self._forget_import_waiter(request_id) + self._boundary_import_expired.discard(request_id) + admission = self._boundary_import_admissions.pop(request_id, None) + if admission is None or not admission.ready or admission.admitted: return - del self._boundary_import_admissions[request_id] - if admission.ready and not admission.admitted: - assert self.boundary_checkpoints is not None - self.boundary_checkpoints.release(admission.checkpoint) + assert self.boundary_checkpoints is not None + self.boundary_checkpoints.release(admission.checkpoint) + if admission.request.boundary_checkpoint is admission.checkpoint: + admission.request.boundary_checkpoint = None def _external_boundary_execution_blocks( self, @@ -1655,4 +1729,5 @@ def finalize_partial_tail_offloads( def new_step_starts(self) -> None: """Notify the coordinator that a new step is starting.""" + self._boundary_import_barrier = None self.coordinator.new_step_starts() diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a61b5fac1c9d..35100581b80d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1420,6 +1420,11 @@ def schedule_running_requests( self.connector is not None and self.kv_cache_manager.boundary_checkpoints is not None and not self.connector.poll_boundary_checkpoint(request) + # The connector still defers a restore that outwaited + # its capacity bound; the request recomputes instead. + and not self.kv_cache_manager.external_boundary_wait_expired( + request_id + ) ): request_queue.remove_request(request) if prefill_interleave_step is not None: From 0dd394dfb15dfd0f2b4ce78d9558dc13acf287a7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:14:05 +0000 Subject: [PATCH 4/7] fix(scheduler): keep queue order and step cost while imports are pending Follow-up to #922 review findings on the waiting pass: - A failed allocation no longer stopped the waiting pass while any import was pending, so every later request could overtake a blocked queue head (FCFS and priority alike). The first failure stops ordinary admission again; only restored imports, which already own their capacity and slot, are still admitted past it. - With an import pending, every step walked the whole waiting queue (connector poll, checkpoint lookup and allocation per request), the saved-logits check scanned both queues (copying the priority heap), and the admission context summed in-flight prefills per candidate. Ready imports are now found through the allocator's admissions, the context is refreshed only after admissions, and priority queues test membership without the ordered heap copy. With 500 blocked 2K-token requests a step costs the same as without imports (0.06 ms instead of 15.7 ms FCFS / 28.3 ms priority). - A ready full-prefix import anywhere in the queue suppressed the running pass; if an older request was admitted instead, decodes lost the step. The import now gets its isolated saved-logits step first, and the running pass runs, without preempting new admissions, whenever that step is not taken. Co-Authored-By: Claude Opus 5.5 --- tests/v1/core/test_boundary_admission.py | 97 +++++++++ .../test_boundary_admission_continuation.py | 2 +- vllm/v1/core/kv_cache_manager.py | 8 + vllm/v1/core/sched/request_queue.py | 4 + vllm/v1/core/sched/scheduler.py | 206 ++++++++++++------ 5 files changed, 247 insertions(+), 70 deletions(-) diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index d8f0df707d69..cbf1adf902ab 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -777,6 +777,103 @@ def start_decode(scheduler, name="running"): return running +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +@pytest.mark.parametrize("state", ["copying", "ready"]) +def test_blocked_queue_head_admits_only_restored_imports_past_it(policy, state): + """Ordinary requests keep their queue order while imports are pending.""" + from tests.v1.core.utils import create_requests + + scheduler, cache, connector = import_scheduler( + num_blocks=64, scheduling_policy=policy + ) + importer = request("import", "import", length=48) + checkpoint = reserve_import(cache, importer, prefix=32) + assert checkpoint is not None + if state == "ready": + publish(cache, checkpoint) + (big,) = create_requests(num_requests=1, num_tokens=2000, req_ids=["big"]) + small = request("small", "small", length=16) + connector.poll_boundary_checkpoint.side_effect = lambda req: ( + req is not importer or state == "ready" + ) + scheduler.connector = connector + for req in (big, small, importer): + scheduler.add_request(req) + expected = {"import": 16} if state == "ready" else {} + assert scheduler.schedule().num_scheduled_tokens == expected + assert small.status == big.status == RequestStatus.WAITING + + +@pytest.mark.parametrize("policy", ["fcfs", "priority"]) +def test_pending_import_step_cost_does_not_scale_with_blocked_queue(policy): + """A copying import leaves a blocked 500-request queue at one lookup a step.""" + import time + + from tests.v1.core.utils import create_requests + + def blocked(with_import): + scheduler, cache, connector = import_scheduler( + num_blocks=256, scheduling_policy=policy + ) + queued = create_requests(num_requests=500, num_tokens=2048) + importer = request("import", "import", length=64) + if with_import: + assert reserve_import(cache, importer, prefix=32) is not None + cache.block_pool.get_new_blocks(cache.block_pool.get_num_free_blocks()) + connector.poll_boundary_checkpoint.side_effect = lambda req: ( + req is not importer + ) + scheduler.connector = connector + for req in (*queued, importer): + scheduler.add_request(req) + assert scheduler.schedule().num_scheduled_tokens == {} + connector.poll_boundary_checkpoint.reset_mock() + steps = 20 + start = time.perf_counter() + for _ in range(steps): + assert scheduler.schedule().num_scheduled_tokens == {} + elapsed = (time.perf_counter() - start) / steps + return elapsed, connector.poll_boundary_checkpoint.call_count / steps + + baseline, baseline_polls = blocked(with_import=False) + elapsed, polls = blocked(with_import=True) + assert polls == baseline_polls == 1 + assert elapsed < 10 * baseline + 1e-3 + + +@pytest.mark.parametrize("state", ["copying", "ready", "ready_but_skipped"]) +def test_ready_import_behind_older_request_keeps_running_decode(state): + """Running decodes skip a step only for a saved-logits step actually taken.""" + scheduler, cache, connector = import_scheduler() + running = start_decode(scheduler) + importer = request("import", "import", length=32) + checkpoint = reserve_import(cache, importer, prefix=32) + assert checkpoint is not None + if state != "copying": + publish(cache, checkpoint) + older = request("older", "older", length=16) + connector.poll_boundary_checkpoint.side_effect = lambda req: ( + req is not importer or state == "ready" + ) + connector.boundary_checkpoint_external_tokens.side_effect = lambda req: ( + 32 if req is importer else 0 + ) + scheduler.connector = connector + scheduler.add_request(older) + scheduler.add_request(importer) + output = scheduler.schedule() + if state == "ready": + assert output.boundary_logits_only + assert output.num_scheduled_tokens == {"import": 1} + output = scheduler.schedule() + assert output.num_scheduled_tokens["running"] == 1 + assert output.num_scheduled_tokens["older"] == 16 + else: + assert not output.boundary_logits_only + assert output.num_scheduled_tokens == {"running": 1, "older": 16} + assert running.num_preemptions == 0 + + def test_unadmitted_import_credits_do_not_preempt_running_decode(): """Reservations gate new admissions, never the growth of running requests.""" scheduler, cache, connector = import_scheduler() diff --git a/tests/v1/core/test_boundary_admission_continuation.py b/tests/v1/core/test_boundary_admission_continuation.py index 6e929a73f571..46b9d4c82b60 100644 --- a/tests/v1/core/test_boundary_admission_continuation.py +++ b/tests/v1/core/test_boundary_admission_continuation.py @@ -111,7 +111,7 @@ def test_nonisolated_continuation_deferral_serves_existing_decoder( ) base.queue_matching(scheduler, unrelated) base.victim_at_head(cache, checkpoint) - assert not scheduler._has_waiting_boundary_logits() + assert scheduler._waiting_boundary_logits_request() is None output = scheduler.schedule() assert not output.boundary_logits_only assert output.num_scheduled_tokens == {"first": 4} diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 8ae6a51d5916..4e99e1b4a82f 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -1285,6 +1285,14 @@ def external_boundary_wait_expired(self, request_id: str) -> bool: """Whether a restore outwaited its capacity bound, so the prompt recomputes.""" return request_id in self._boundary_import_expired + def ready_external_boundary_requests(self) -> list[Request]: + """Waiting consumers whose imports finished copying, in reservation order.""" + return [ + admission.request + for admission in self._boundary_import_admissions.values() + if admission.ready and not admission.admitted + ] + def release_external_boundary_admission(self, request_id: str) -> None: """Release an import's credits and slot, or end its capacity wait. diff --git a/vllm/v1/core/sched/request_queue.py b/vllm/v1/core/sched/request_queue.py index 38c7db94a45e..6bf56500d694 100644 --- a/vllm/v1/core/sched/request_queue.py +++ b/vllm/v1/core/sched/request_queue.py @@ -191,6 +191,10 @@ def __len__(self) -> int: """Get number of requests in queue.""" return len(self._heap) + def __contains__(self, request: object) -> bool: + """Check membership without the ordered iteration's heap copy.""" + return request in self._heap + def __iter__(self) -> Iterator[Request]: """Iterate over the queue according to priority policy.""" heap_copy = self._heap[:] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 35100581b80d..216e968629a3 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -56,6 +56,7 @@ ) from vllm.v1.core.sched.prefill_interleave import ( PrefillInterleaveController, + PrefillInterleaveStep, resolve_decode_refill_target, resolve_max_parallel_prefills, ) @@ -770,44 +771,68 @@ def _local_prefill_pressure(self, now: float) -> tuple[float, int]: pressure = max(pressure, 1.0 + age_seconds / expected_seconds) return pressure, backlog_tokens - def _has_waiting_boundary_logits(self) -> bool: - """Whether the queue head needs an isolated saved-logits step.""" + def _waiting_boundary_logits_request(self) -> Request | None: + """Return the waiting request that needs an isolated saved-logits step. + + This is the queue head or, failing that, a restored import that already + owns its reserved capacity. + """ if ( not (self.waiting or self.skipped_waiting) or self._pause_state != PauseState.UNPAUSED or len(self.running) + self.num_waiting_for_streaming_input >= self.max_num_running_reqs ): - return False + return None queue = self._select_waiting_queue_for_scheduling() if queue is None: - return False - pending_imports = ( - self.kv_cache_manager.has_pending_external_boundary_admissions() - ) - candidates = itertools.chain( - (queue.peek_request(),), + return None + return next( ( request for request in ( - itertools.chain(self.skipped_waiting, self.waiting) - if pending_imports - else () - ) - if self.kv_cache_manager.external_boundary_admission_ready( - request.request_id + queue.peek_request(), + *self._ready_import_requests(), ) + if request.status in (RequestStatus.WAITING, RequestStatus.PREEMPTED) + and request.num_computed_tokens == 0 + and request.boundary_checkpoint is not None + and request.boundary_checkpoint.num_tokens == request.num_tokens + and (request.num_stale_output_tokens == 0 or request.drop_stale_output) ), - ) - return any( - request.status in (RequestStatus.WAITING, RequestStatus.PREEMPTED) - and request.num_computed_tokens == 0 - and request.boundary_checkpoint is not None - and request.boundary_checkpoint.num_tokens == request.num_tokens - and (request.num_stale_output_tokens == 0 or request.drop_stale_output) - for request in candidates + None, ) + def _ready_import_requests(self) -> list[Request]: + """Restored imports that wait for admission, in scheduling order.""" + requests = self.kv_cache_manager.ready_external_boundary_requests() + if self.policy == SchedulingPolicy.PRIORITY: + requests.sort() + return requests + + def _ready_imports_to_consider( + self, + considered_ids: set[str], + prefill_interleave_step: PrefillInterleaveStep | None, + ) -> list[Request]: + """Restored imports not yet considered in this step's waiting pass.""" + return [ + request + for request in self._ready_import_requests() + if request.request_id not in considered_ids + and ( + prefill_interleave_step is None + or prefill_interleave_step.is_selected(request.request_id) + ) + ] + + def _waiting_queue_of(self, request: Request) -> RequestQueue | None: + """Return the waiting queue that holds the request, if any.""" + for queue in (self.skipped_waiting, self.waiting): + if request in queue: + return queue + return None + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.current_step += 1 # NOTE(woosuk) on the scheduling algorithm: @@ -1293,7 +1318,9 @@ def schedule_running_requests( # A full recurrent-cache hit cannot share a model step with forwards. # Its cache lookup in the waiting pass records that requirement. Give # it an isolated admission step instead of waiting for decodes to end. - reserve_boundary_logits_step = self._has_waiting_boundary_logits() + boundary_logits_request = self._waiting_boundary_logits_request() + reserve_boundary_logits_step = boundary_logits_request is not None + boundary_logits_taken = False if not reserve_boundary_logits_step: schedule_running_requests(initial_service_class) @@ -1310,6 +1337,15 @@ def schedule_running_requests( # Next, schedule the WAITING requests. if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: step_skipped_waiting = create_request_queue(self.policy) + external_admission = ( + self.connector is not None + and self.kv_cache_manager.boundary_checkpoints is not None + ) + admission_context_stale = True + # Set once ordinary admission stops: restored imports already own + # their capacity and slot, so only they may still be admitted. + ready_imports: list[Request] | None = None + considered_ids: set[str] = set() while (self.waiting or self.skipped_waiting) and token_budget > 0: if ( @@ -1323,20 +1359,36 @@ def schedule_running_requests( if num_running >= self.max_num_running_reqs: break - request_queue = self._select_waiting_queue_for_scheduling() - assert request_queue is not None - request = request_queue.peek_request() - if ( - prefill_interleave_step is not None - and self._request_has_local_prefill(request) - ): - waiting_selection = prefill_interleave_step.select_waiting_request( - (self.skipped_waiting, self.waiting) - ) - if waiting_selection is None: + selected: Request | None = None + if boundary_logits_request is not None: + selected, boundary_logits_request = boundary_logits_request, None + elif ready_imports is not None: + if not ready_imports: break - request_queue, request = waiting_selection + selected = ready_imports.pop(0) + if selected is None: + request_queue = self._select_waiting_queue_for_scheduling() + assert request_queue is not None + request = request_queue.peek_request() + if ( + prefill_interleave_step is not None + and self._request_has_local_prefill(request) + ): + waiting_selection = ( + prefill_interleave_step.select_waiting_request( + (self.skipped_waiting, self.waiting) + ) + ) + if waiting_selection is None: + break + request_queue, request = waiting_selection + else: + request_queue = self._waiting_queue_of(selected) + if request_queue is None: + continue + request = selected request_id = request.request_id + considered_ids.add(request_id) # try to promote blocked statuses while traversing skipped queue. if self._is_blocked_waiting_status( @@ -1389,31 +1441,23 @@ def schedule_running_requests( external_boundary_hit_tokens = 0 # Get already-cached tokens. - if ( - request.num_computed_tokens == 0 - and self.connector is not None - and self.kv_cache_manager.boundary_checkpoints is not None - ): - self.kv_cache_manager.set_external_boundary_admission_context( - self.max_num_running_reqs - num_running, - sum( - self._request_remaining_blocks(req) - for req in self._inflight_prefills - if not ( - self.kv_cache_manager.has_external_boundary_admission( - req.request_id - ) - ) - ), - has_scheduled_reqs=bool(self.running), - ) + if request.num_computed_tokens == 0 and external_admission: + if admission_context_stale: + # Only admissions change slots and in-flight needs. + self.kv_cache_manager.set_external_boundary_admission_context( + self.max_num_running_reqs - num_running, + self._unticketed_inflight_prefill_blocks(), + has_scheduled_reqs=bool(self.running), + ) + admission_context_stale = False if not self.kv_cache_manager.can_admit_external_boundary_request( - request.request_id + request_id ): - request_queue.remove_request(request) - if prefill_interleave_step is not None: - prefill_interleave_step.mark_unavailable(request_id) - step_skipped_waiting.prepend_request(request) + # Every remaining slot is reserved for a restored import. + if ready_imports is None: + ready_imports = self._ready_imports_to_consider( + considered_ids, prefill_interleave_step + ) continue if request.num_computed_tokens == 0: if ( @@ -1455,7 +1499,9 @@ def schedule_running_requests( and num_new_local_computed_tokens == request.num_tokens ) if boundary_logits_only and num_scheduled_tokens: - break + if ready_imports is None: + break + continue # Get externally-cached tokens if using a KVConnector. if self.connector is not None: @@ -1800,13 +1846,19 @@ def schedule_running_requests( # manager if request.has_encoder_inputs: self.encoder_cache_manager.free(request) - if not interleaved_local_prefill and not ( - self.kv_cache_manager.has_pending_external_boundary_admissions() - ): - break + if not interleaved_local_prefill: + # FCFS stops here; imports that already own their + # capacity cannot take it from this request. + if ready_imports is None: + ready_imports = self._ready_imports_to_consider( + considered_ids, prefill_interleave_step + ) + if not ready_imports: + break + continue request_queue.remove_request(request) - if prefill_interleave_step is not None: - prefill_interleave_step.mark_unavailable(request_id) + assert prefill_interleave_step is not None + prefill_interleave_step.mark_unavailable(request_id) step_skipped_waiting.prepend_request(request) continue @@ -1838,6 +1890,7 @@ def schedule_running_requests( ) request_queue.remove_request(request) + admission_context_stale = True if load_kv_async: # If loading async, allocate memory and put request # into the WAITING_FOR_REMOTE_KV state. @@ -1928,6 +1981,7 @@ def schedule_running_requests( # interleaving policy, including policies that do not # maintain a separate micro-prefill budget. token_budget = 0 + boundary_logits_taken = True break # re-queue requests skipped in this pass ahead of older skipped items. @@ -1939,10 +1993,15 @@ def schedule_running_requests( if not defer_prefills: self.prefill_capacity_bound = bool(self.waiting) - if reserve_boundary_logits_step and not num_scheduled_tokens: + if reserve_boundary_logits_step and not boundary_logits_taken: # Cache eviction, admission limits, or deferred input can prevent - # the isolated step. Running requests must still make progress. - schedule_running_requests(initial_service_class) + # the isolated step. Running requests must still make progress, + # without evicting requests admitted above. + schedule_running_requests( + initial_service_class, + allow_preemption=not num_scheduled_tokens, + enforce_lora_limit=bool(num_scheduled_tokens), + ) if ( prefill_interleave_step is not None @@ -3922,6 +3981,15 @@ def _inflight_prefill_reserved_blocks(self) -> int: self._request_remaining_blocks(req) for req in self._inflight_prefills ) + def _unticketed_inflight_prefill_blocks(self) -> int: + """Blocks in-flight prefills without import credits still need.""" + manager = self.kv_cache_manager + return sum( + self._request_remaining_blocks(req) + for req in self._inflight_prefills + if not manager.has_external_boundary_admission(req.request_id) + ) + def _update_waiting_for_remote_kv(self, request: Request) -> None: """ KV Connector: update request state after async recv is finished. From 7023056e0d7ce9aa978794e7e9affe28ab7f6ddd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:14:11 +0000 Subject: [PATCH 5/7] fix(scheduler): release restored imports on a forced prefix-cache reset reset_prefix_cache(reset_running_requests=True) preempts every running request so the pool can be cleared, but a restore that finished copying and was not yet admitted still pinned its pages, and the reset raised RuntimeError. pause_scheduler(mode="keep", clear_cache=True) and sleep(level>=1, mode="keep") reach this path. Such restores are now released with the running requests; their requests fall back to an ordinary lookup. A reset without preemption still reports failure while a restored import holds its pages. Co-Authored-By: Claude Opus 5.5 --- tests/v1/core/test_boundary_admission.py | 19 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 6 ++++++ 2 files changed, 25 insertions(+) diff --git a/tests/v1/core/test_boundary_admission.py b/tests/v1/core/test_boundary_admission.py index cbf1adf902ab..7d73321e1bad 100644 --- a/tests/v1/core/test_boundary_admission.py +++ b/tests/v1/core/test_boundary_admission.py @@ -952,6 +952,25 @@ def poll(req): assert not cache.external_boundary_wait_expired(waiter.request_id) +def test_reset_with_running_requests_releases_ready_unadmitted_import(): + """A cache reset drops a finished restore; its request then recomputes.""" + scheduler, cache, connector = import_scheduler() + running = start_decode(scheduler) + importer = request("import", "import", length=32) + checkpoint = reserve_import(cache, importer, prefix=16) + assert checkpoint is not None + publish(cache, checkpoint) + scheduler.connector = connector + scheduler.add_request(importer) + assert scheduler.reset_prefix_cache(reset_running_requests=True) + assert running.status == RequestStatus.PREEMPTED + assert importer.boundary_checkpoint is None + assert not cache.has_external_boundary_admission(importer.request_id) + assert cache.block_pool.get_num_free_blocks() == cache.block_pool.num_gpu_blocks - 1 + output = scheduler.schedule() + assert output.num_scheduled_tokens["import"] == 32 + + @pytest.mark.parametrize("release", ["move_victim", "finish_reader"]) @pytest.mark.parametrize("fairness", [None, 0.4]) def test_actual_scheduler_runs_decode_after_guard_defers_restore(fairness, release): diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 216e968629a3..9247348da340 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -3750,6 +3750,12 @@ def reset_prefix_cache( while self.running: request = self.running.pop() self._preempt_request(request, timestamp, drop_stale_output=True) + # Restores that finished copying but were never admitted hold their + # pages the same way; their requests fall back to ordinary lookup. + for request in self.kv_cache_manager.ready_external_boundary_requests(): + self.kv_cache_manager.release_external_boundary_admission( + request.request_id + ) # Clear scheduled request ids cache. Since we are forcing preemption # + resumption in the same step, we must act as if these requests were From f9082f2223d4a92ecb99f5df8a201c4ca70352cd Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:14:16 +0000 Subject: [PATCH 6/7] docs(release): describe fair, bounded checkpoint-restore admission Update the #922 fragment to the reviewed behavior: reservations hold back only new admissions, queue order and step cost are kept while imports are pending, saved-logits steps no longer cost decodes, cancellation frees credits at once, forced resets release restored imports, and VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S bounds capacity waits. Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-restore-admission.json | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json index 1d7fdf96e909..2d9b75986d23 100644 --- a/.lil/changes/vllm-restore-admission.json +++ b/.lil/changes/vllm-restore-admission.json @@ -2,15 +2,18 @@ "schema": "local-inference-release-change/v1", "id": "vllm-restore-admission", "category": "fix", - "summary": "External recurrent checkpoint restores can reserve execution capacity and retain ownership through admission.", + "summary": "External recurrent checkpoint restores wait for GPU capacity instead of recomputing cached prefixes, with fair scheduling and a bounded wait.", "models": [ "all" ], - "compatibility": "Use with the paired LMCache restore-admission change to wait for checkpoint capacity instead of recomputing cached prefixes.", + "compatibility": "Use with the paired LMCache restore-admission change. VLLM_CHECKPOINT_RESTORE_MAX_WAIT_S (default 60) bounds how long a restore may wait for capacity before its request recomputes the prompt; 0 recomputes immediately.", "details": [ - "Restores reserve GPU blocks and running slots before copies start, and keep restored pages owned until execution takes over.", + "Restores reserve GPU blocks and a running slot before copies start, and keep restored pages owned until execution takes over.", + "Reservations hold back only new admissions. Running requests keep growing, requests ordered ahead of a waiting restore are never blocked by it, and later requests are admitted when they fit beside it.", + "Queue order is kept: a request that does not fit still stops ordinary admission, and only restores that already own their capacity are admitted past it, without rescanning the waiting queue each step.", + "A finished full-prefix restore gets its isolated saved-logits step directly, and running decodes skip a step only when that step is taken.", + "Cancelling a restore frees its slot and reserved capacity at once; only the copied pages wait for in-flight transfers. A forced prefix-cache reset releases finished restores that were not yet admitted.", "Sliding-window and chunked-local continuation reservations count physical pages rather than skipped positions in the block table.", - "Requests with reserved capacity can progress past a blocked queue head; cancellation releases reservations after submitted copies drain.", "Existing external-checkpoint callers retain their previous behavior unless they opt into admission reservations." ], "requires": [], From 1ac6c2deb5b5cb1543b75a5e93176ffb7d048d00 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 28 Sep 2026 15:47:26 +0000 Subject: [PATCH 7/7] docs(release): link the restore-admission fragment to PR 929 Co-Authored-By: Claude Opus 5.5 --- .lil/changes/vllm-restore-admission.json | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json index 2d9b75986d23..ce684ef3819b 100644 --- a/.lil/changes/vllm-restore-admission.json +++ b/.lil/changes/vllm-restore-admission.json @@ -18,6 +18,10 @@ ], "requires": [], "pull_requests": [ - 922 + 929 + ], + "authors": [ + "ktsaou", + "Local Inference Lab" ] }