diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index db33b4c6df..7fa331747c 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -4349,3 +4349,128 @@ def test_eagle3_mm_encoder_cache_with_shift(): f"shifted_end={scheduled_end_with_shift}) overlapping MM at " f"{start_pos}. The fix must schedule encoder inputs." ) + + +@pytest.mark.parametrize("use_kv_connector", [False, True]) +def test_ec_connector_ensure_cache_available_defers_request(use_kv_connector): + """Test that ensure_cache_available() returning False defers the request. + + When the EC connector signals a prefetch is in progress (returns False), + the scheduler should: + 1. Not schedule the request (no KV cache or encoder cache allocated) + 2. Still schedule other requests behind the deferred one + 3. Schedule the deferred request on the next step when ensure_cache_available + returns True and has_cache_item returns True + """ + scheduler = create_scheduler( + model="llava-hf/llava-1.5-7b-hf", + enable_prefix_caching=True, + use_kv_connector=use_kv_connector, + use_ec_connector=True, + ec_role="ec_consumer", + ) + + NUM_TOKENS = 200 + NUM_ENCODER_TOKENS = 100 + + request_deferred = create_requests( + num_requests=1, + num_tokens=NUM_TOKENS, + mm_positions=[[PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS)]], + req_ids=["deferred"], + )[0] + + request_behind = create_requests( + num_requests=1, + num_tokens=20, + req_ids=["behind"], + )[0] + + # --- Step 1: ensure_cache_available returns False → request deferred --- + scheduler.ec_connector.ensure_cache_available = Mock(return_value=False) + + scheduler.add_request(request_deferred) + scheduler.add_request(request_behind) + output = scheduler.schedule() + + # ensure_cache_available must have been called with (request, num_computed_tokens=0) + # for a brand-new request that has no cached tokens yet. + scheduler.ec_connector.ensure_cache_available.assert_called_once_with( + request_deferred, 0 + ) + # Deferred request must NOT be scheduled + assert request_deferred.request_id not in output.num_scheduled_tokens + _assert_right_encoder_cache_allocated(scheduler, expected_total_allocated=0) + # No KV blocks allocated for the deferred request + for mgr in scheduler.kv_cache_manager.coordinator.single_type_managers: + assert request_deferred.request_id not in mgr.req_to_blocks + + # The text-only request behind the deferred one MUST still be scheduled + assert request_behind.request_id in output.num_scheduled_tokens + assert output.num_scheduled_tokens[request_behind.request_id] == 20 + + # --- Step 2: prefetch done, cache exists → request scheduled --- + # has_cache_item is called inside _try_schedule_encoder_inputs (not during + # deferral), so it is only relevant here in step 2. + scheduler.ec_connector.ensure_cache_available = Mock(return_value=True) + scheduler.ec_connector.has_cache_item = Mock(return_value=True) + + output = scheduler.schedule() + + # Now the deferred request should be scheduled + assert request_deferred.request_id in output.num_scheduled_tokens + assert output.num_scheduled_tokens[request_deferred.request_id] == NUM_TOKENS + _assert_right_encoder_cache_allocated(scheduler, requests=[request_deferred]) + # EC connector metadata should carry the deferred request's MM data + _assert_right_ec_connector_metadata( + output, mm_features_list=request_deferred.mm_features + ) + # No local encoder compute — all loaded externally + _assert_right_encoder_inputs(output, expected_total_reqs=0) + + +def test_ec_connector_pending_prefetch_only_checks_future_mm_features(): + """Test that future mm feature filtering only yields features beyond + the computed token frontier. + + Features already within num_computed_tokens (past/boundary) must be + filtered out; only features that extend beyond the frontier (future) should + be yielded so that connector implementations know which items to prefetch. + + Filter cases: + "past": end = 0 + 16 = 16 < 32 → filtered OUT + "boundary": end = 16 + 16 = 32 == 32 → filtered OUT (condition is >, not >=) + "future": end = 48 + 32 = 80 > 32 → yielded + """ + BLOCK_SIZE = 16 + NUM_COMPUTED_TOKENS = BLOCK_SIZE * 2 # 32 + NUM_TOKENS = BLOCK_SIZE * 8 # 128 + + HASH_PAST = "hash_past" + HASH_BOUNDARY = "hash_boundary" + HASH_FUTURE = "hash_future" + + request = create_requests( + num_requests=1, + num_tokens=NUM_TOKENS, + mm_hashes_list=[[HASH_PAST, HASH_BOUNDARY, HASH_FUTURE]], + mm_positions=[ + [ + PlaceholderRange(offset=0, length=BLOCK_SIZE), # end=16 (past) + PlaceholderRange(offset=16, length=BLOCK_SIZE), # end=32 (boundary) + PlaceholderRange(offset=48, length=BLOCK_SIZE * 2), # end=80 (future) + ] + ], + block_size=BLOCK_SIZE, + )[0] + + future_hashes = [ + f.identifier + for f in request.mm_features + if f.mm_position.offset + f.mm_position.length > NUM_COMPUTED_TOKENS + ] + + assert future_hashes == [HASH_FUTURE], ( + f"Expected only {HASH_FUTURE!r} from future mm feature filtering, " + f"got {future_hashes!r}. Past/boundary features must be filtered out." + ) diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 28370c8e1f..1d5f467027 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -211,6 +211,23 @@ class ECConnectorBase(ABC): """ pass + def ensure_cache_available( + self, request: "Request", num_computed_tokens: int + ) -> bool: + """ + Ensure encoder cache items are available for the given request. + May initiate asynchronous transfers for items not yet local. + + Args: + request: the request whose multimodal features to check. + num_computed_tokens: tokens already covered by cached KV blocks. + + Returns: + True if all items are ready or no transfer is needed. + False if any items are still in transit (request should be deferred). + """ + return True + @abstractmethod def update_state_after_alloc(self, request: "Request", index: int): """ diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b33135ff1e..ab932c8a6d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -629,6 +629,18 @@ class Scheduler(SchedulerInterface): ) assert num_computed_tokens <= request.num_tokens + # Skip request with pending mm encoding prefetches + if ( + self.ec_connector is not None + and request.mm_features + and not self.ec_connector.ensure_cache_available( + request, num_computed_tokens + ) + ): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + # Track first scheduled prefill, not post-preemption repeat prefills if request.prefill_stats is not None: assert num_computed_tokens <= request.num_prompt_tokens