diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 27107aa6b0..168a0f8978 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -353,6 +353,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.finished_req_ids = set() scheduler.finished_req_ids_dict = None scheduler.grammar_compile_error_reqs = set() + scheduler.encoder_cache_mismatch_reqs = set() scheduler.vllm_config = Mock() scheduler.aux_output_connector = None scheduler.return_sampling_mask = False diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index 3948f7c6e3..3c61f3367f 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -114,6 +114,7 @@ def test_eviction_when_cache_is_full(): # 'x' should have been evicted. assert "x" not in manager.cached + assert "x" not in manager.num_embeds_by_hash assert "req1" not in manager.request_cached_ids assert "x" in manager.get_freed_mm_hashes() @@ -273,6 +274,31 @@ def test_encoder_cache_with_is_embed_mask(): assert savings_ratio == 12.5 +def test_encoder_cache_tracks_embedding_count_by_identifier(): + class MockRequestWithMask(MockRequest): + def get_num_encoder_embeds(self, input_id: int) -> int: + return self.mm_features[input_id].mm_position.get_num_embeds() + + is_embed = torch.tensor([False, True, True, False, True, False]) + request = MockRequestWithMask("owner", ["img1"], [len(is_embed)]) + request.mm_features[0] = MultiModalFeatureSpec( + data=None, + modality="image", + identifier="img1", + mm_position=PlaceholderRange(offset=0, length=6, is_embed=is_embed), + ) + + manager = EncoderCacheManager(cache_size=10) + manager.allocate(request, 0) + assert manager.get_cached_num_encoder_embeds(request, 0) == 3 + + manager.free(request) + assert manager.get_cached_num_encoder_embeds(request, 0) == 3 + + manager.reset() + assert manager.get_cached_num_encoder_embeds(request, 0) is None + + def test_encoder_cache_mask_based_retrieval(): class MockRequestWithMask(MockRequest): def get_num_encoder_embeds(self, input_id: int) -> int: @@ -336,6 +362,7 @@ def test_reset_clears_all_state(): manager.reset() assert len(manager.cached) == 0 + assert len(manager.num_embeds_by_hash) == 0 assert len(manager.request_cached_ids) == 0 assert len(manager.freeable) == 0 assert len(manager.freed) == 0 diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 062b38c2ef..505eef5919 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3896,6 +3896,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.finished_req_ids_dict = None scheduler.aux_output_connector = None scheduler.grammar_compile_error_reqs = set() + scheduler.encoder_cache_mismatch_reqs = set() scheduler.vllm_config = Mock() scheduler.return_sampling_mask = False scheduler.recompute_kv_load_failures = False @@ -6089,6 +6090,201 @@ def test_unavailable_encoder_input_fails_the_request_as_retryable(): assert [o.finish_reason for o in engine_outputs] == [FinishReason.ERROR] +def _create_mm_request_with_embedding_positions( + request_id: str, identifier: str, is_embed: list[bool] +) -> Request: + return create_requests( + num_requests=1, + num_tokens=len(is_embed), + mm_hashes_list=[[identifier]], + mm_positions=[ + [ + PlaceholderRange( + offset=0, + length=len(is_embed), + is_embed=torch.tensor(is_embed), + ) + ] + ], + req_ids=[request_id], + )[0] + + +def test_encoder_cache_accepts_matching_embed_count(): + """An actively referenced entry is reused when embedding counts match.""" + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") + identifier = "reused-client-uuid" + owner = _create_mm_request_with_embedding_positions( + "owner", identifier, [True, True, True, False, False, False] + ) + compatible = _create_mm_request_with_embedding_positions( + "compatible", identifier, [True, False, True, False, True, False, False] + ) + scheduler.encoder_cache_manager.allocate(owner, 0) + scheduler.add_request(compatible) + + scheduler_output = scheduler.schedule() + + assert compatible.request_id in scheduler_output.num_scheduled_tokens + assert compatible.request_id not in scheduler.encoder_cache_mismatch_reqs + assert compatible.request_id not in scheduler_output.scheduled_encoder_inputs + + +def test_encoder_cache_rejects_mismatched_embed_count(caplog_vllm): + """A completed request's freeable entry rejects an incompatible reuse.""" + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") + identifier = "reused-client-uuid" + owner = _create_mm_request_with_embedding_positions( + "owner", identifier, [True, True, True, False, False, False] + ) + mismatched = _create_mm_request_with_embedding_positions( + "mismatched", identifier, [True, True, True, True, False, False] + ) + scheduler.encoder_cache_manager.allocate(owner, 0) + scheduler.encoder_cache_manager.free(owner) + freeable_slots_before = scheduler.encoder_cache_manager.num_freeable_slots + scheduler.add_request(mismatched) + + scheduler_output = scheduler.schedule() + assert not scheduler_output.num_scheduled_tokens + assert mismatched.request_id in scheduler.encoder_cache_mismatch_reqs + assert ( + "multimodal identifier 'reused-client-uuid' holds 3 embeddings" + in caplog_vllm.text + and "expects 4" in caplog_vllm.text + ) + + outputs = scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput(req_ids=[], req_id_to_index={}), + ) + + assert mismatched.status == RequestStatus.FINISHED_ERROR + assert not scheduler.encoder_cache_mismatch_reqs + assert mismatched.request_id not in scheduler.requests + assert scheduler.encoder_cache_manager.cached[identifier] == set() + assert ( + mismatched.request_id not in scheduler.encoder_cache_manager.request_cached_ids + ) + assert scheduler.encoder_cache_manager.freeable[identifier] == 3 + assert scheduler.encoder_cache_manager.num_freeable_slots == freeable_slots_before + output = outputs[mismatched.client_index].outputs[0] + assert output.request_id == mismatched.request_id + assert output.finish_reason == FinishReason.ERROR + + healthy = create_requests(num_requests=1, req_ids=["healthy"], num_tokens=6)[0] + scheduler.add_request(healthy) + next_output = scheduler.schedule() + assert healthy.request_id in next_output.num_scheduled_tokens + + +def test_encoder_cache_embed_count_mismatch_restores_prior_hit_state(): + """A prior hit is released when a later mismatch fails the request.""" + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") + manager = scheduler.encoder_cache_manager + valid_identifier = "valid-reuse" + mismatch_identifier = "mismatched-reuse" + valid_owner = _create_mm_request_with_embedding_positions( + "valid-owner", valid_identifier, [True, True, False] + ) + mismatch_owner = _create_mm_request_with_embedding_positions( + "mismatch-owner", mismatch_identifier, [True, True, True, False] + ) + manager.allocate(valid_owner, 0) + manager.free(valid_owner) + manager.allocate(mismatch_owner, 0) + manager.free(mismatch_owner) + freeable_slots_before = manager.num_freeable_slots + + request = create_requests( + num_requests=1, + num_tokens=7, + mm_hashes_list=[[valid_identifier, mismatch_identifier]], + mm_positions=[ + [ + PlaceholderRange( + offset=0, + length=3, + is_embed=torch.tensor([True, True, False]), + ), + PlaceholderRange( + offset=3, + length=4, + is_embed=torch.tensor([True, True, True, True]), + ), + ] + ], + req_ids=["partially-mutated"], + )[0] + scheduler.add_request(request) + + scheduler_output = scheduler.schedule() + assert not scheduler_output.num_scheduled_tokens + assert request.request_id in scheduler.encoder_cache_mismatch_reqs + assert request.request_id in manager.cached[valid_identifier] + assert valid_identifier not in manager.freeable + assert manager.num_freeable_slots == freeable_slots_before - 2 + + outputs = scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput(req_ids=[], req_id_to_index={}), + ) + + assert request.status == RequestStatus.FINISHED_ERROR + assert request.request_id not in scheduler.requests + assert request.request_id not in manager.request_cached_ids + assert manager.cached[valid_identifier] == set() + assert manager.freeable[valid_identifier] == 2 + assert manager.freeable[mismatch_identifier] == 3 + assert manager.num_freeable_slots == freeable_slots_before + output = outputs[request.client_index].outputs[0] + assert output.finish_reason == FinishReason.ERROR + + +def test_encoder_cache_rejects_mismatched_embed_count_within_request(): + """One request cannot reuse an identifier with a different embed count.""" + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") + identifier = "reused-within-request" + request = create_requests( + num_requests=1, + num_tokens=8, + mm_hashes_list=[[identifier, identifier]], + mm_positions=[ + [ + PlaceholderRange( + offset=0, + length=4, + is_embed=torch.tensor([True, True, False, False]), + ), + PlaceholderRange( + offset=4, + length=4, + is_embed=torch.tensor([True, True, True, False]), + ), + ] + ], + req_ids=["mismatched-within-request"], + )[0] + scheduler.add_request(request) + + scheduler_output = scheduler.schedule() + assert not scheduler_output.num_scheduled_tokens + assert request.request_id in scheduler.encoder_cache_mismatch_reqs + assert identifier not in scheduler.encoder_cache_manager.cached + + outputs = scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput(req_ids=[], req_id_to_index={}), + ) + + assert request.status == RequestStatus.FINISHED_ERROR + assert not scheduler.encoder_cache_mismatch_reqs + assert request.request_id not in scheduler.requests + output = outputs[request.client_index].outputs[0] + assert output.request_id == request.request_id + assert output.finish_reason == FinishReason.ERROR + + def test_free_encoder_inputs_defers_for_eagle_lookahead(): """With EAGLE speculative decoding, the encoder input is retained one extra position so the drafter's +1 look-ahead mm-embedding gather (which reads one diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index f069ce7453..24cf609d22 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -59,6 +59,8 @@ class EncoderCacheManager: reference the cached entry. If the set is empty, the entry exists but is not referenced by any request and is eligible for reclamation. + num_embeds_by_hash: Mapping from mm_hash to the number of encoder + embeddings stored for that cache entry. freeable: List of tuples (mm_hash, num_encoder_embeds) representing entries whose no current running request is needed and that can be freed to make space when needed. @@ -80,6 +82,8 @@ class EncoderCacheManager: # mm_hash of mm_data => ids of requests that reference the mm_data self.cached: dict[str, set[str]] = {} + # mm_hash of mm_data => number of cached encoder embeddings + self.num_embeds_by_hash: dict[str, int] = {} # request_id => set of input_ids cached for that request self.request_cached_ids: dict[str, set[int]] = {} @@ -94,6 +98,7 @@ class EncoderCacheManager: Called when model weights are updated to invalidate stale embeddings. """ self.cached.clear() + self.num_embeds_by_hash.clear() self.request_cached_ids.clear() self.freeable.clear() self.freed.clear() @@ -130,6 +135,13 @@ class EncoderCacheManager: self.request_cached_ids.setdefault(request.request_id, set()).add(input_id) return True + def get_cached_num_encoder_embeds( + self, request: Request, input_id: int + ) -> int | None: + """Return the encoder-embedding count stored for a cached input.""" + mm_hash = request.mm_features[input_id].identifier + return self.num_embeds_by_hash.get(mm_hash) + def can_allocate( self, request: Request, @@ -188,6 +200,7 @@ class EncoderCacheManager: while num_embeds > self.num_free_slots: mm_hash, num_free_embeds = self.freeable.popitem(last=False) del self.cached[mm_hash] + del self.num_embeds_by_hash[mm_hash] self.freed.append(mm_hash) self.num_free_slots += num_free_embeds return True @@ -209,6 +222,12 @@ class EncoderCacheManager: self.cached[mm_hash] = set() num_encoder_embeds = request.get_num_encoder_embeds(input_id) + cached_num_encoder_embeds = self.num_embeds_by_hash.setdefault( + mm_hash, num_encoder_embeds + ) + # _try_schedule_encoder_inputs rejects local-hit and same-step + # identifier/count mismatches before allocation. + assert cached_num_encoder_embeds == num_encoder_embeds # NOTE: Encoder cache should always have enough space for encoder inputs # that are scheduled since eviction takes place at can_allocate(). @@ -258,7 +277,7 @@ class EncoderCacheManager: return self.cached[mm_hash].discard(req_id) if not self.cached[mm_hash]: - num_encoder_embeds = request.get_num_encoder_embeds(input_id) + num_encoder_embeds = self.num_embeds_by_hash[mm_hash] self.freeable[mm_hash] = num_encoder_embeds self.num_freeable_slots += num_encoder_embeds diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 59ba590d9d..ce832a18f6 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -234,6 +234,7 @@ class Scheduler(SchedulerInterface): # Grammar compilation failures to finish as per-request errors in # update_from_output. self.grammar_compile_error_reqs: set[str] = set() + self.encoder_cache_mismatch_reqs: set[str] = set() # Encoder-related. # Calculate encoder cache size if applicable @@ -1714,6 +1715,29 @@ class Scheduler(SchedulerInterface): num_output_tokens=num_output_tokens, ) + def _reject_on_encoder_cache_embed_mismatch( + self, + request: Request, + input_id: int, + identifier: str, + cached_num_encoder_embeds: int, + requested_num_encoder_embeds: int, + ) -> bool: + if cached_num_encoder_embeds == requested_num_encoder_embeds: + return False + logger.warning( + "Encoder cache entry for multimodal identifier %r holds %d " + "embeddings, but request %s input %d expects %d. A client-provided " + "multimodal UUID may have been reused for different content.", + identifier, + cached_num_encoder_embeds, + request.request_id, + input_id, + requested_num_encoder_embeds, + ) + self.encoder_cache_mismatch_reqs.add(request.request_id) + return True + def _try_schedule_encoder_inputs( self, request: Request, @@ -1752,7 +1776,7 @@ class Scheduler(SchedulerInterface): # NOTE: since scheduler operates on the request level (possibly with # multiple encoder inputs per request), we need to create temporary # trackers for accounting at the encoder input level. - mm_hashes_to_schedule = set() + mm_hashes_to_schedule: dict[str, int] = {} num_embeds_to_schedule = 0 encoder_window_end = ( @@ -1794,10 +1818,35 @@ class Scheduler(SchedulerInterface): # We are not using the encoder cache for encoder-decoder models, # yet. if item_identifier in mm_hashes_to_schedule: + scheduled_num_encoder_embeds = mm_hashes_to_schedule[ + item_identifier + ] + if self._reject_on_encoder_cache_embed_mismatch( + request, + i, + item_identifier, + scheduled_num_encoder_embeds, + num_encoder_embeds, + ): + return [], 0, encoder_compute_budget, [] # The same encoder input has already been scheduled in the # current step. continue + cached_num_encoder_embeds = ( + self.encoder_cache_manager.get_cached_num_encoder_embeds(request, i) + ) + if cached_num_encoder_embeds is not None and ( + self._reject_on_encoder_cache_embed_mismatch( + request, + i, + item_identifier, + cached_num_encoder_embeds, + num_encoder_embeds, + ) + ): + return [], 0, encoder_compute_budget, [] + if self.encoder_cache_manager.check_and_update_cache(request, i): # The encoder input is already computed and cached from a # previous step. @@ -1857,14 +1906,14 @@ class Scheduler(SchedulerInterface): if self.ec_connector is not None and self.ec_connector.has_cache_item( item_identifier ): - mm_hashes_to_schedule.add(item_identifier) + mm_hashes_to_schedule[item_identifier] = num_encoder_embeds external_load_encoder_input.append(i) num_embeds_to_schedule += num_encoder_embeds continue num_embeds_to_schedule += num_encoder_embeds encoder_compute_budget -= num_encoder_embeds - mm_hashes_to_schedule.add(item_identifier) + mm_hashes_to_schedule[item_identifier] = num_encoder_embeds encoder_inputs_to_schedule.append(i) return ( @@ -2185,6 +2234,8 @@ class Scheduler(SchedulerInterface): error_req_ids = set(self.grammar_compile_error_reqs) self.grammar_compile_error_reqs.clear() + error_req_ids.update(self.encoder_cache_mismatch_reqs) + self.encoder_cache_mismatch_reqs.clear() if failed_kv_load_req_ids and not self.recompute_kv_load_failures: error_req_ids.update(failed_kv_load_req_ids) if self.ec_connector is not None: