[Bugfix][V1] Reject encoder-cache hits with mismatched embedding counts (#57696)
Signed-off-by: jackLei0901 <42642542+jackLei0901@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
jackLei
co-authored by
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
b44895cf93
commit
bbc4ddeca9
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user