diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 4f82e828d3..caff9c9a49 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -731,3 +731,63 @@ def test_async_loads_both_admitted_when_pool_fits(): for req in reqs: assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def test_async_load_reserves_blocks_for_promotion_margin(): + """An async load is not admitted unless the blocks its own promotion will + need are still free. + + A parked load is allocated without lookahead slots, but promotion pads it + to ``1 + num_spec_tokens`` and asks for the lookahead margin on top. Without + reserving that margin, two loads can be admitted that together consume the + whole pool; the head then fails ``allocate_slots`` forever while the load + behind it is never reached, and with nothing running no block is ever freed. + + req_a (4 blocks) and req_b (3 blocks) exactly fill the 7 usable blocks, so + req_b must be held back in WAITING holding no blocks, leaving req_a room to + be promoted and run. + """ + vllm_config = create_vllm_config(num_speculative_tokens=3) + BLOCK_SIZE = vllm_config.cache_config.block_size + scheduler = create_scheduler(vllm_config, num_blocks=8) # usable = 7 + + req_a = create_request( + request_id=1, + block_size=BLOCK_SIZE, + num_tokens=BLOCK_SIZE * 4, + do_remote_prefill=True, + max_tokens=1, + ) + req_b = create_request( + request_id=2, + block_size=BLOCK_SIZE, + num_tokens=BLOCK_SIZE * 3, + do_remote_prefill=True, + max_tokens=1, + ) + scheduler.add_request(req_a) + scheduler.add_request(req_b) + + # Both get a full external hit, so each would hold its whole prompt. + with patch.object( + scheduler.connector, + "get_num_new_matched_tokens", + side_effect=[(BLOCK_SIZE * 4, True), (BLOCK_SIZE * 3, True)], + ): + scheduler.schedule() + + assert req_a.status == RequestStatus.WAITING_FOR_REMOTE_KVS + assert req_b.status == RequestStatus.WAITING + req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[ + 0 + ].req_to_blocks + assert req_b.request_id not in req_to_blocks + + # req_a's load lands: it must be promotable and actually get scheduled. + scheduler.update_from_output( + scheduler.schedule(), + create_model_runner_output([], finished_recving={req_a.request_id}), + ) + scheduler_output = scheduler.schedule() + assert req_a.status == RequestStatus.RUNNING + assert scheduler_output.num_scheduled_tokens[req_a.request_id] > 0 diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 80b97c6a13..85ada84b0d 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -17,6 +17,7 @@ from vllm.config import ( KVTransferConfig, ModelConfig, SchedulerConfig, + SpeculativeConfig, VllmConfig, ) from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory @@ -106,6 +107,7 @@ def create_vllm_config( kv_connector_module_path: str | None = None, kv_role: str = "kv_consumer", disable_hybrid_kv_cache_manager: bool | None = None, + num_speculative_tokens: int | None = None, ) -> VllmConfig: """Initialize VllmConfig For Testing.""" model_config = ModelConfig( @@ -141,6 +143,11 @@ def create_vllm_config( kv_load_failure_policy=kv_load_failure_policy, ) attention_config = AttentionConfig(backend=attention_backend) + speculative_config = ( + SpeculativeConfig(model="ngram", num_speculative_tokens=num_speculative_tokens) + if num_speculative_tokens is not None + else None + ) return VllmConfig( scheduler_config=scheduler_config, model_config=model_config, @@ -148,6 +155,7 @@ def create_vllm_config( kv_transfer_config=kv_transfer_config, device_config=DeviceConfig("cpu"), attention_config=attention_config, + speculative_config=speculative_config, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 08f1191111..19ea411aa1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -31,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.multimodal.utils import get_mm_features_in_window +from vllm.utils.math_utils import cdiv from vllm.v1.core.encoder_cache_manager import ( EncoderCacheManager, EncoderDecoderCacheManager, @@ -1140,9 +1141,13 @@ class Scheduler(SchedulerInterface): if load_kv_async: # An async load holds its blocks for the whole transfer with # no forward progress and isn't preemptible here. Admit it - # only if it fits in (free - other in-flight reservations), to - # avoid deadlock and predictable preemptions. - reserved_blocks = self._inflight_prefill_reserved_blocks() + # only if it fits in (free - other in-flight reservations) + # plus its own spec decode step blocks, to avoid deadlock and + # predictable preemptions. + reserved_blocks = ( + self._inflight_prefill_reserved_blocks() + + self._spec_decode_step_blocks() + ) new_blocks = self.kv_cache_manager.allocate_slots( request, @@ -2892,9 +2897,9 @@ class Scheduler(SchedulerInterface): return delay_free or partial_tail_delay, kv_xfer_params def _request_remaining_blocks(self, request: Request) -> int: - """Blocks `request` still needs to allocate to hold its full sequence.""" + """Blocks `request` still needs to hold its full sequence and be promoted.""" full_num_tokens = min(request.num_tokens, self.max_model_len) - return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( + num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( request_id=request.request_id, num_tokens=full_num_tokens, new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, @@ -2904,6 +2909,20 @@ class Scheduler(SchedulerInterface): num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) + return num_blocks + self._spec_decode_step_blocks() + + def _spec_decode_step_blocks(self) -> int: + """Number of blocks for the extra KV slots a spec decode step needs. + + When using async kv load, scheduler must reserve enough blocks for + full sequence + the spec decode step, otherwise request cannot be + able to run after the async_load if we are out of kv blocks. + """ + if not self.num_spec_tokens: + return 0 + return cdiv( + 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size + ) def _inflight_prefill_reserved_blocks(self) -> int: """Num blocks in-flight prefills still need to finish (their reservation)."""