[EC Connector] Non blocking EC Connector lookup (#41627)

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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."
+    )

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):
         """

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