@tduncan / Nvidia-Optimization-Stack-VLLM / commits / bbc4ddeca9

[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>

jackLei committed Sep 24, 2026 at 01:40 UTC bbc4ddeca9ece9851aad26f03afc4bab37c6e8df

5 files changed +298 -4

tests/v1/core/test_async_scheduler.py

@@ -353,6 +353,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():

353 scheduler.finished_req_ids = set()

354 scheduler.finished_req_ids_dict = None

355 scheduler.grammar_compile_error_reqs = set()

356 + scheduler.encoder_cache_mismatch_reqs = set()

357 scheduler.vllm_config = Mock()

358 scheduler.aux_output_connector = None

359 scheduler.return_sampling_mask = False

tests/v1/core/test_encoder_cache_manager.py

+27

@@ -114,6 +114,7 @@ def test_eviction_when_cache_is_full():

114

115 # 'x' should have been evicted.

116 assert "x" not in manager.cached

117 + assert "x" not in manager.num_embeds_by_hash

118 assert "req1" not in manager.request_cached_ids

119 assert "x" in manager.get_freed_mm_hashes()

120

@@ -273,6 +274,31 @@ def test_encoder_cache_with_is_embed_mask():

274 assert savings_ratio == 12.5

275

276

277 +def test_encoder_cache_tracks_embedding_count_by_identifier():

278 + class MockRequestWithMask(MockRequest):

279 + def get_num_encoder_embeds(self, input_id: int) -> int:

280 + return self.mm_features[input_id].mm_position.get_num_embeds()

281 +

282 + is_embed = torch.tensor([False, True, True, False, True, False])

283 + request = MockRequestWithMask("owner", ["img1"], [len(is_embed)])

284 + request.mm_features[0] = MultiModalFeatureSpec(

285 + data=None,

286 + modality="image",

287 + identifier="img1",

288 + mm_position=PlaceholderRange(offset=0, length=6, is_embed=is_embed),

289 + )

290 +

291 + manager = EncoderCacheManager(cache_size=10)

292 + manager.allocate(request, 0)

293 + assert manager.get_cached_num_encoder_embeds(request, 0) == 3

294 +

295 + manager.free(request)

296 + assert manager.get_cached_num_encoder_embeds(request, 0) == 3

297 +

298 + manager.reset()

299 + assert manager.get_cached_num_encoder_embeds(request, 0) is None

300 +

301 +

302 def test_encoder_cache_mask_based_retrieval():

303 class MockRequestWithMask(MockRequest):

304 def get_num_encoder_embeds(self, input_id: int) -> int:

@@ -336,6 +362,7 @@ def test_reset_clears_all_state():

362 manager.reset()

363

364 assert len(manager.cached) == 0

365 + assert len(manager.num_embeds_by_hash) == 0

366 assert len(manager.request_cached_ids) == 0

367 assert len(manager.freeable) == 0

368 assert len(manager.freed) == 0

tests/v1/core/test_scheduler.py

+196

@@ -3896,6 +3896,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():

3896 scheduler.finished_req_ids_dict = None

3897 scheduler.aux_output_connector = None

3898 scheduler.grammar_compile_error_reqs = set()

3899 + scheduler.encoder_cache_mismatch_reqs = set()

3900 scheduler.vllm_config = Mock()

3901 scheduler.return_sampling_mask = False

3902 scheduler.recompute_kv_load_failures = False

@@ -6089,6 +6090,201 @@ def test_unavailable_encoder_input_fails_the_request_as_retryable():

6090 assert [o.finish_reason for o in engine_outputs] == [FinishReason.ERROR]

6091

6092

6093 +def _create_mm_request_with_embedding_positions(

6094 + request_id: str, identifier: str, is_embed: list[bool]

6095 +) -> Request:

6096 + return create_requests(

6097 + num_requests=1,

6098 + num_tokens=len(is_embed),

6099 + mm_hashes_list=[[identifier]],

6100 + mm_positions=[

6101 + [

6102 + PlaceholderRange(

6103 + offset=0,

6104 + length=len(is_embed),

6105 + is_embed=torch.tensor(is_embed),

6106 + )

6107 + ]

6108 + ],

6109 + req_ids=[request_id],

6110 + )[0]

6111 +

6112 +

6113 +def test_encoder_cache_accepts_matching_embed_count():

6114 + """An actively referenced entry is reused when embedding counts match."""

6115 + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")

6116 + identifier = "reused-client-uuid"

6117 + owner = _create_mm_request_with_embedding_positions(

6118 + "owner", identifier, [True, True, True, False, False, False]

6119 + )

6120 + compatible = _create_mm_request_with_embedding_positions(

6121 + "compatible", identifier, [True, False, True, False, True, False, False]

6122 + )

6123 + scheduler.encoder_cache_manager.allocate(owner, 0)

6124 + scheduler.add_request(compatible)

6125 +

6126 + scheduler_output = scheduler.schedule()

6127 +

6128 + assert compatible.request_id in scheduler_output.num_scheduled_tokens

6129 + assert compatible.request_id not in scheduler.encoder_cache_mismatch_reqs

6130 + assert compatible.request_id not in scheduler_output.scheduled_encoder_inputs

6131 +

6132 +

6133 +def test_encoder_cache_rejects_mismatched_embed_count(caplog_vllm):

6134 + """A completed request's freeable entry rejects an incompatible reuse."""

6135 + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")

6136 + identifier = "reused-client-uuid"

6137 + owner = _create_mm_request_with_embedding_positions(

6138 + "owner", identifier, [True, True, True, False, False, False]

6139 + )

6140 + mismatched = _create_mm_request_with_embedding_positions(

6141 + "mismatched", identifier, [True, True, True, True, False, False]

6142 + )

6143 + scheduler.encoder_cache_manager.allocate(owner, 0)

6144 + scheduler.encoder_cache_manager.free(owner)

6145 + freeable_slots_before = scheduler.encoder_cache_manager.num_freeable_slots

6146 + scheduler.add_request(mismatched)

6147 +

6148 + scheduler_output = scheduler.schedule()

6149 + assert not scheduler_output.num_scheduled_tokens

6150 + assert mismatched.request_id in scheduler.encoder_cache_mismatch_reqs

6151 + assert (

6152 + "multimodal identifier 'reused-client-uuid' holds 3 embeddings"

6153 + in caplog_vllm.text

6154 + and "expects 4" in caplog_vllm.text

6155 + )

6156 +

6157 + outputs = scheduler.update_from_output(

6158 + scheduler_output,

6159 + ModelRunnerOutput(req_ids=[], req_id_to_index={}),

6160 + )

6161 +

6162 + assert mismatched.status == RequestStatus.FINISHED_ERROR

6163 + assert not scheduler.encoder_cache_mismatch_reqs

6164 + assert mismatched.request_id not in scheduler.requests

6165 + assert scheduler.encoder_cache_manager.cached[identifier] == set()

6166 + assert (

6167 + mismatched.request_id not in scheduler.encoder_cache_manager.request_cached_ids

6168 + )

6169 + assert scheduler.encoder_cache_manager.freeable[identifier] == 3

6170 + assert scheduler.encoder_cache_manager.num_freeable_slots == freeable_slots_before

6171 + output = outputs[mismatched.client_index].outputs[0]

6172 + assert output.request_id == mismatched.request_id

6173 + assert output.finish_reason == FinishReason.ERROR

6174 +

6175 + healthy = create_requests(num_requests=1, req_ids=["healthy"], num_tokens=6)[0]

6176 + scheduler.add_request(healthy)

6177 + next_output = scheduler.schedule()

6178 + assert healthy.request_id in next_output.num_scheduled_tokens

6179 +

6180 +

6181 +def test_encoder_cache_embed_count_mismatch_restores_prior_hit_state():

6182 + """A prior hit is released when a later mismatch fails the request."""

6183 + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")

6184 + manager = scheduler.encoder_cache_manager

6185 + valid_identifier = "valid-reuse"

6186 + mismatch_identifier = "mismatched-reuse"

6187 + valid_owner = _create_mm_request_with_embedding_positions(

6188 + "valid-owner", valid_identifier, [True, True, False]

6189 + )

6190 + mismatch_owner = _create_mm_request_with_embedding_positions(

6191 + "mismatch-owner", mismatch_identifier, [True, True, True, False]

6192 + )

6193 + manager.allocate(valid_owner, 0)

6194 + manager.free(valid_owner)

6195 + manager.allocate(mismatch_owner, 0)

6196 + manager.free(mismatch_owner)

6197 + freeable_slots_before = manager.num_freeable_slots

6198 +

6199 + request = create_requests(

6200 + num_requests=1,

6201 + num_tokens=7,

6202 + mm_hashes_list=[[valid_identifier, mismatch_identifier]],

6203 + mm_positions=[

6204 + [

6205 + PlaceholderRange(

6206 + offset=0,

6207 + length=3,

6208 + is_embed=torch.tensor([True, True, False]),

6209 + ),

6210 + PlaceholderRange(

6211 + offset=3,

6212 + length=4,

6213 + is_embed=torch.tensor([True, True, True, True]),

6214 + ),

6215 + ]

6216 + ],

6217 + req_ids=["partially-mutated"],

6218 + )[0]

6219 + scheduler.add_request(request)

6220 +

6221 + scheduler_output = scheduler.schedule()

6222 + assert not scheduler_output.num_scheduled_tokens

6223 + assert request.request_id in scheduler.encoder_cache_mismatch_reqs

6224 + assert request.request_id in manager.cached[valid_identifier]

6225 + assert valid_identifier not in manager.freeable

6226 + assert manager.num_freeable_slots == freeable_slots_before - 2

6227 +

6228 + outputs = scheduler.update_from_output(

6229 + scheduler_output,

6230 + ModelRunnerOutput(req_ids=[], req_id_to_index={}),

6231 + )

6232 +

6233 + assert request.status == RequestStatus.FINISHED_ERROR

6234 + assert request.request_id not in scheduler.requests

6235 + assert request.request_id not in manager.request_cached_ids

6236 + assert manager.cached[valid_identifier] == set()

6237 + assert manager.freeable[valid_identifier] == 2

6238 + assert manager.freeable[mismatch_identifier] == 3

6239 + assert manager.num_freeable_slots == freeable_slots_before

6240 + output = outputs[request.client_index].outputs[0]

6241 + assert output.finish_reason == FinishReason.ERROR

6242 +

6243 +

6244 +def test_encoder_cache_rejects_mismatched_embed_count_within_request():

6245 + """One request cannot reuse an identifier with a different embed count."""

6246 + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")

6247 + identifier = "reused-within-request"

6248 + request = create_requests(

6249 + num_requests=1,

6250 + num_tokens=8,

6251 + mm_hashes_list=[[identifier, identifier]],

6252 + mm_positions=[

6253 + [

6254 + PlaceholderRange(

6255 + offset=0,

6256 + length=4,

6257 + is_embed=torch.tensor([True, True, False, False]),

6258 + ),

6259 + PlaceholderRange(

6260 + offset=4,

6261 + length=4,

6262 + is_embed=torch.tensor([True, True, True, False]),

6263 + ),

6264 + ]

6265 + ],

6266 + req_ids=["mismatched-within-request"],

6267 + )[0]

6268 + scheduler.add_request(request)

6269 +

6270 + scheduler_output = scheduler.schedule()

6271 + assert not scheduler_output.num_scheduled_tokens

6272 + assert request.request_id in scheduler.encoder_cache_mismatch_reqs

6273 + assert identifier not in scheduler.encoder_cache_manager.cached

6274 +

6275 + outputs = scheduler.update_from_output(

6276 + scheduler_output,

6277 + ModelRunnerOutput(req_ids=[], req_id_to_index={}),

6278 + )

6279 +

6280 + assert request.status == RequestStatus.FINISHED_ERROR

6281 + assert not scheduler.encoder_cache_mismatch_reqs

6282 + assert request.request_id not in scheduler.requests

6283 + output = outputs[request.client_index].outputs[0]

6284 + assert output.request_id == request.request_id

6285 + assert output.finish_reason == FinishReason.ERROR

6286 +

6287 +

6288 def test_free_encoder_inputs_defers_for_eagle_lookahead():

6289 """With EAGLE speculative decoding, the encoder input is retained one extra

6290 position so the drafter's +1 look-ahead mm-embedding gather (which reads one

vllm/v1/core/encoder_cache_manager.py

+20 -1

@@ -59,6 +59,8 @@ class EncoderCacheManager:

59 reference the cached entry. If the set is empty, the entry exists

60 but is not referenced by any request and is eligible for

61 reclamation.

62 + num_embeds_by_hash: Mapping from mm_hash to the number of encoder

63 + embeddings stored for that cache entry.

64 freeable: List of tuples (mm_hash, num_encoder_embeds) representing entries

65 whose no current running request is needed and that can be freed to

66 make space when needed.

@@ -80,6 +82,8 @@ class EncoderCacheManager:

83 # mm_hash of mm_data => ids of requests that reference the mm_data

84 self.cached: dict[str, set[str]] = {}

85 + # mm_hash of mm_data => number of cached encoder embeddings

86 + self.num_embeds_by_hash: dict[str, int] = {}

87 # request_id => set of input_ids cached for that request

88 self.request_cached_ids: dict[str, set[int]] = {}

@@ -94,6 +98,7 @@ class EncoderCacheManager:

98 Called when model weights are updated to invalidate stale embeddings.

99 """

100 self.cached.clear()

101 + self.num_embeds_by_hash.clear()

102 self.request_cached_ids.clear()

103 self.freeable.clear()

104 self.freed.clear()

@@ -130,6 +135,13 @@ class EncoderCacheManager:

135 self.request_cached_ids.setdefault(request.request_id, set()).add(input_id)

136 return True

137

138 + def get_cached_num_encoder_embeds(

139 + self, request: Request, input_id: int

140 + ) -> int | None:

141 + """Return the encoder-embedding count stored for a cached input."""

142 + mm_hash = request.mm_features[input_id].identifier

143 + return self.num_embeds_by_hash.get(mm_hash)

144 +

145 def can_allocate(

146 self,

147 request: Request,

@@ -188,6 +200,7 @@ class EncoderCacheManager:

200 while num_embeds > self.num_free_slots:

201 mm_hash, num_free_embeds = self.freeable.popitem(last=False)

202 del self.cached[mm_hash]

203 + del self.num_embeds_by_hash[mm_hash]

204 self.freed.append(mm_hash)

205 self.num_free_slots += num_free_embeds

206 return True

@@ -209,6 +222,12 @@ class EncoderCacheManager:

222 self.cached[mm_hash] = set()

223

224 num_encoder_embeds = request.get_num_encoder_embeds(input_id)

225 + cached_num_encoder_embeds = self.num_embeds_by_hash.setdefault(

226 + mm_hash, num_encoder_embeds

227 + )

228 + # _try_schedule_encoder_inputs rejects local-hit and same-step

229 + # identifier/count mismatches before allocation.

230 + assert cached_num_encoder_embeds == num_encoder_embeds

231

232 # NOTE: Encoder cache should always have enough space for encoder inputs

233 # that are scheduled since eviction takes place at can_allocate().

@@ -258,7 +277,7 @@ class EncoderCacheManager:

277 return

278 self.cached[mm_hash].discard(req_id)

279 if not self.cached[mm_hash]:

261 - num_encoder_embeds = request.get_num_encoder_embeds(input_id)

280 + num_encoder_embeds = self.num_embeds_by_hash[mm_hash]

281 self.freeable[mm_hash] = num_encoder_embeds

282 self.num_freeable_slots += num_encoder_embeds

283

vllm/v1/core/sched/scheduler.py

+54 -3

@@ -234,6 +234,7 @@ class Scheduler(SchedulerInterface):

234 # Grammar compilation failures to finish as per-request errors in

235 # update_from_output.

236 self.grammar_compile_error_reqs: set[str] = set()

237 + self.encoder_cache_mismatch_reqs: set[str] = set()

238

239 # Encoder-related.

240 # Calculate encoder cache size if applicable

@@ -1714,6 +1715,29 @@ class Scheduler(SchedulerInterface):

1715 num_output_tokens=num_output_tokens,

1716 )

1717

1718 + def _reject_on_encoder_cache_embed_mismatch(

1719 + self,

1720 + request: Request,

1721 + input_id: int,

1722 + identifier: str,

1723 + cached_num_encoder_embeds: int,

1724 + requested_num_encoder_embeds: int,

1725 + ) -> bool:

1726 + if cached_num_encoder_embeds == requested_num_encoder_embeds:

1727 + return False

1728 + logger.warning(

1729 + "Encoder cache entry for multimodal identifier %r holds %d "

1730 + "embeddings, but request %s input %d expects %d. A client-provided "

1731 + "multimodal UUID may have been reused for different content.",

1732 + identifier,

1733 + cached_num_encoder_embeds,

1734 + request.request_id,

1735 + input_id,

1736 + requested_num_encoder_embeds,

1737 + )

1738 + self.encoder_cache_mismatch_reqs.add(request.request_id)

1739 + return True

1740 +

1741 def _try_schedule_encoder_inputs(

1742 self,

1743 request: Request,

@@ -1752,7 +1776,7 @@ class Scheduler(SchedulerInterface):

1776 # NOTE: since scheduler operates on the request level (possibly with

1777 # multiple encoder inputs per request), we need to create temporary

1778 # trackers for accounting at the encoder input level.

1755 - mm_hashes_to_schedule = set()

1779 + mm_hashes_to_schedule: dict[str, int] = {}

1780 num_embeds_to_schedule = 0

1781

1782 encoder_window_end = (

@@ -1794,10 +1818,35 @@ class Scheduler(SchedulerInterface):

1818 # We are not using the encoder cache for encoder-decoder models,

1819 # yet.

1820 if item_identifier in mm_hashes_to_schedule:

1821 + scheduled_num_encoder_embeds = mm_hashes_to_schedule[

1822 + item_identifier

1823 + ]

1824 + if self._reject_on_encoder_cache_embed_mismatch(

1825 + request,

1826 + i,

1827 + item_identifier,

1828 + scheduled_num_encoder_embeds,

1829 + num_encoder_embeds,

1830 + ):

1831 + return [], 0, encoder_compute_budget, []

1832 # The same encoder input has already been scheduled in the

1833 # current step.

1834 continue

1835

1836 + cached_num_encoder_embeds = (

1837 + self.encoder_cache_manager.get_cached_num_encoder_embeds(request, i)

1838 + )

1839 + if cached_num_encoder_embeds is not None and (

1840 + self._reject_on_encoder_cache_embed_mismatch(

1841 + request,

1842 + i,

1843 + item_identifier,

1844 + cached_num_encoder_embeds,

1845 + num_encoder_embeds,

1846 + )

1847 + ):

1848 + return [], 0, encoder_compute_budget, []

1849 +

1850 if self.encoder_cache_manager.check_and_update_cache(request, i):

1851 # The encoder input is already computed and cached from a

1852 # previous step.

@@ -1857,14 +1906,14 @@ class Scheduler(SchedulerInterface):

1906 if self.ec_connector is not None and self.ec_connector.has_cache_item(

1907 item_identifier

1908 ):

1860 - mm_hashes_to_schedule.add(item_identifier)

1909 + mm_hashes_to_schedule[item_identifier] = num_encoder_embeds

1910 external_load_encoder_input.append(i)

1911 num_embeds_to_schedule += num_encoder_embeds

1912 continue

1913

1914 num_embeds_to_schedule += num_encoder_embeds

1915 encoder_compute_budget -= num_encoder_embeds

1867 - mm_hashes_to_schedule.add(item_identifier)

1916 + mm_hashes_to_schedule[item_identifier] = num_encoder_embeds

1917 encoder_inputs_to_schedule.append(i)

1918

1919 return (

@@ -2185,6 +2234,8 @@ class Scheduler(SchedulerInterface):

2234

2235 error_req_ids = set(self.grammar_compile_error_reqs)

2236 self.grammar_compile_error_reqs.clear()

2237 + error_req_ids.update(self.encoder_cache_mismatch_reqs)

2238 + self.encoder_cache_mismatch_reqs.clear()

2239 if failed_kv_load_req_ids and not self.recompute_kv_load_failures:

2240 error_req_ids.update(failed_kv_load_req_ids)

2241 if self.ec_connector is not None: