diff --git a/.lil/changes/vllm-restore-admission.json b/.lil/changes/vllm-restore-admission.json new file mode 100644 index 000000000000..1d7fdf96e909 --- /dev/null +++ b/.lil/changes/vllm-restore-admission.json @@ -0,0 +1,20 @@ +{ + "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." + ], + "requires": [], + "pull_requests": [ + 922 + ] +} 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