A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 05:00 UTC
Support clear mm and encoder cache (#33452)
Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roger Wang <hey@rogerw.io>
jma99_2333 committed
Jan 31, 2026 at 07:22 UTC
22d9a056d5c4849af372f9506bb93cad7da755f6
15 files changed
+212
-1
docs/usage/security.md
+1
@@ -172,6 +172,7 @@ These endpoints are **only available when the environment variable `VLLM_SERVER_
172
- `/server_info` - Get detailed server configuration
173
- `/reset_prefix_cache` - Reset prefix cache (can disrupt service)
174
- `/reset_mm_cache` - Reset multimodal cache (can disrupt service)
175
+- `/reset_encoder_cache` - Reset encoder cache (can disrupt service)
176
- `/sleep` - Put engine to sleep (causes denial of service)
177
- `/wake_up` - Wake engine from sleep
178
- `/is_sleeping` - Check if engine is sleeping
tests/v1/core/test_encoder_cache_manager.py
+89
-1
@@ -4,7 +4,10 @@ import pytest
4
import torch
5
6
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
7
-from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
7
+from vllm.v1.core.encoder_cache_manager import (
8
+ EncoderCacheManager,
9
+ EncoderDecoderCacheManager,
10
+)
11
12
pytestmark = pytest.mark.cpu_test
13
@@ -247,3 +250,88 @@ def test_encoder_cache_mask_based_retrieval():
250
251
assert num_embeds_before == 0
252
assert num_embeds_in_range == 2
253
+
254
+
255
+def test_reset_clears_all_state():
256
+ """Test that reset() clears all cached entries and restores capacity."""
257
+ manager = EncoderCacheManager(cache_size=20)
258
+
259
+ req1 = MockRequest("req1", ["img1", "img2"], [5, 3])
260
+ req2 = MockRequest("req2", ["img3"], [4])
261
+
262
+ manager.allocate(req1, 0)
263
+ manager.allocate(req1, 1)
264
+ manager.allocate(req2, 0)
265
+ manager.free_encoder_input(req1, 0)
266
+
267
+ req3 = MockRequest("req3", ["img4"], [10])
268
+ manager.free_encoder_input(req1, 1)
269
+ manager.free_encoder_input(req2, 0)
270
+ manager.can_allocate(req3, 0, int(1e9), 0)
271
+ manager.allocate(req3, 0)
272
+
273
+ assert len(manager.cached) > 0
274
+ assert manager.num_free_slots < 20
275
+
276
+ manager.reset()
277
+
278
+ assert len(manager.cached) == 0
279
+ assert len(manager.freeable) == 0
280
+ assert len(manager.freed) == 0
281
+ assert manager.num_free_slots == 20
282
+ assert manager.num_freeable_slots == 20
283
+
284
+
285
+def test_reset_allows_fresh_allocations():
286
+ manager = EncoderCacheManager(cache_size=10)
287
+
288
+ req1 = MockRequest("req1", ["img1"], [10])
289
+ manager.allocate(req1, 0)
290
+ assert manager.num_free_slots == 0
291
+
292
+ manager.reset()
293
+
294
+ req2 = MockRequest("req2", ["img2"], [8])
295
+ assert manager.can_allocate(req2, 0, int(1e9), 0)
296
+ manager.allocate(req2, 0)
297
+
298
+ assert manager.num_free_slots == 2
299
+ assert "img2" in manager.cached
300
+ assert "img1" not in manager.cached
301
+
302
+
303
+def test_encoder_decoder_cache_manager_reset():
304
+ manager = EncoderDecoderCacheManager(cache_size=20)
305
+
306
+ req1 = MockRequest("req1", ["img1"], [5])
307
+ req2 = MockRequest("req2", ["img2"], [3])
308
+
309
+ manager.allocate(req1, 0)
310
+ manager.allocate(req2, 0)
311
+ manager.free(req1)
312
+ manager.get_freed_mm_hashes()
313
+
314
+ assert manager.num_free_slots < 20
315
+
316
+ manager.reset()
317
+
318
+ assert len(manager.allocated) == 0
319
+ assert len(manager.to_free) == 0
320
+ assert manager.num_free_slots == 20
321
+
322
+
323
+def test_encoder_decoder_cache_manager_reset_allows_fresh_allocations():
324
+ manager = EncoderDecoderCacheManager(cache_size=10)
325
+
326
+ req1 = MockRequest("req1", ["img1"], [10])
327
+ manager.allocate(req1, 0)
328
+ assert manager.num_free_slots == 0
329
+
330
+ manager.reset()
331
+
332
+ req2 = MockRequest("req2", ["img2"], [8])
333
+ assert manager.can_allocate(req2, 0, int(1e9), 0)
334
+ manager.allocate(req2, 0)
335
+
336
+ assert manager.num_free_slots == 2
337
+ assert "img2" in manager.allocated
vllm/engine/protocol.py
+5
@@ -113,6 +113,11 @@ class EngineClient(ABC):
113
"""Reset the multi-modal cache"""
114
...
115
116
+ @abstractmethod
117
+ async def reset_encoder_cache(self) -> None:
118
+ """Reset the encoder cache"""
119
+ ...
120
+
121
@abstractmethod
122
async def reset_prefix_cache(
123
self, reset_running_requests: bool = False, reset_connector: bool = False
vllm/entrypoints/serve/cache/api_router.py
+11
@@ -55,6 +55,17 @@ async def reset_mm_cache(raw_request: Request):
55
return Response(status_code=200)
56
57
58
+@router.post("/reset_encoder_cache")
59
+async def reset_encoder_cache(raw_request: Request):
60
+ """
61
+ Reset the encoder cache. Note that we currently do not check if the
62
+ encoder cache is successfully reset in the API server.
63
+ """
64
+ logger.info("Resetting encoder cache...")
65
+ await engine_client(raw_request).reset_encoder_cache()
66
+ return Response(status_code=200)
67
+
68
+
69
def attach_router(app: FastAPI):
70
if not envs.VLLM_SERVER_DEV_MODE:
71
return
vllm/v1/core/encoder_cache_manager.py
+18
@@ -77,6 +77,18 @@ class EncoderCacheManager:
77
self.freeable: OrderedDict[str, int] = OrderedDict()
78
self.freed: list[str] = []
79
80
+ def reset(self) -> None:
81
+ """Reset the encoder cache to its initial state.
82
+
83
+ This clears all cached encoder outputs and resets capacity tracking.
84
+ Called when model weights are updated to invalidate stale embeddings.
85
+ """
86
+ self.cached.clear()
87
+ self.freeable.clear()
88
+ self.freed.clear()
89
+ self.num_free_slots = self.cache_size
90
+ self.num_freeable_slots = self.cache_size
91
+
92
def check_and_update_cache(self, request: Request, input_id: int) -> bool:
93
"""Check if encoder output for a specific multimodal input is cached.
94
@@ -360,6 +372,12 @@ class EncoderDecoderCacheManager(EncoderCacheManager):
372
self.allocated: list[str] = []
373
self.to_free: list[str] = []
374
375
+ def reset(self) -> None:
376
+ """Reset the encoder cache to its initial state."""
377
+ self.num_free_slots = self.cache_size
378
+ self.allocated.clear()
379
+ self.to_free.clear()
380
+
381
def check_and_update_cache(self, request: Request, input_id: int) -> bool:
382
return False
383
vllm/v1/core/sched/interface.py
+9
@@ -183,6 +183,15 @@ class SchedulerInterface(ABC):
183
"""
184
raise NotImplementedError
185
186
+ @abstractmethod
187
+ def reset_encoder_cache(self) -> None:
188
+ """Reset the encoder cache to invalidate all cached encoder outputs.
189
+
190
+ This should be called when model weights are updated to ensure
191
+ stale vision embeddings are not reused.
192
+ """
193
+ raise NotImplementedError
194
+
195
@abstractmethod
196
def get_request_counts(self) -> tuple[int, int]:
197
"""Returns (num_running_reqs, num_waiting_reqs)."""
vllm/v1/core/sched/scheduler.py
+17
@@ -1763,6 +1763,14 @@ class Scheduler(SchedulerInterface):
1763
1764
return True
1765
1766
+ def reset_encoder_cache(self) -> None:
1767
+ """Reset the encoder cache to invalidate all cached encoder outputs.
1768
+
1769
+ This should be called when model weights are updated to ensure
1770
+ stale vision embeddings are not reused.
1771
+ """
1772
+ self.encoder_cache_manager.reset()
1773
+
1774
def make_stats(
1775
self,
1776
spec_decoding_stats: SpecDecodingStats | None = None,
@@ -1788,6 +1796,7 @@ class Scheduler(SchedulerInterface):
1796
num_running_reqs=len(self.running),
1797
num_waiting_reqs=len(self.waiting),
1798
kv_cache_usage=self.kv_cache_manager.usage,
1799
+ encoder_cache_usage=self._get_encoder_cache_usage(),
1800
prefix_cache_stats=prefix_cache_stats,
1801
connector_prefix_cache_stats=connector_prefix_cache_stats,
1802
kv_cache_eviction_events=eviction_events,
@@ -1797,6 +1806,14 @@ class Scheduler(SchedulerInterface):
1806
perf_stats=perf_stats,
1807
)
1808
1809
+ def _get_encoder_cache_usage(self) -> float:
1810
+ """Get encoder cache usage as a fraction (0.0 to 1.0)."""
1811
+ ecm = self.encoder_cache_manager
1812
+ if ecm.cache_size == 0:
1813
+ return 0.0
1814
+ used_slots = ecm.cache_size - ecm.num_free_slots
1815
+ return used_slots / ecm.cache_size
1816
+
1817
def make_spec_decoding_stats(
1818
self,
1819
spec_decoding_stats: SpecDecodingStats | None,
vllm/v1/engine/async_llm.py
+3
@@ -882,6 +882,9 @@ class AsyncLLM(EngineClient):
882
reset_running_requests, reset_connector
883
)
884
885
+ async def reset_encoder_cache(self) -> None:
886
+ await self.engine_core.reset_encoder_cache_async()
887
+
888
async def sleep(self, level: int = 1) -> None:
889
await self.reset_prefix_cache()
890
await self.engine_core.sleep_async(level)
vllm/v1/engine/core.py
+20
@@ -565,6 +565,26 @@ class EngineCore:
565
reset_running_requests, reset_connector
566
)
567
568
+ def reset_encoder_cache(self) -> None:
569
+ """Reset the encoder cache to invalidate all cached encoder outputs.
570
+
571
+ This should be called when model weights are updated to ensure
572
+ stale vision embeddings computed with old weights are not reused.
573
+ Clears both the scheduler's cache manager and the GPU model runner's cache.
574
+ """
575
+ # NOTE: Since this is mainly for debugging, we don't attempt to
576
+ # re-sync the internal caches (P0 sender, P1 receiver)
577
+ if self.scheduler.has_unfinished_requests():
578
+ logger.warning(
579
+ "Resetting the encoder cache when requests are "
580
+ "in progress may lead to desynced internal caches."
581
+ )
582
+
583
+ # Reset the scheduler's encoder cache manager (logical state)
584
+ self.scheduler.reset_encoder_cache()
585
+ # Reset the GPU model runner's encoder cache (physical storage)
586
+ self.model_executor.reset_encoder_cache()
587
+
588
def sleep(self, level: int = 1):
589
self.model_executor.sleep(level)
590
vllm/v1/engine/core_client.py
+15
@@ -144,6 +144,9 @@ class EngineCoreClient(ABC):
144
) -> bool:
145
raise NotImplementedError
146
147
+ def reset_encoder_cache(self) -> None:
148
+ raise NotImplementedError
149
+
150
def sleep(self, level: int = 1) -> None:
151
raise NotImplementedError
152
@@ -216,6 +219,9 @@ class EngineCoreClient(ABC):
219
) -> bool:
220
raise NotImplementedError
221
222
+ async def reset_encoder_cache_async(self) -> None:
223
+ raise NotImplementedError
224
+
225
async def sleep_async(self, level: int = 1) -> None:
226
raise NotImplementedError
227
@@ -300,6 +306,9 @@ class InprocClient(EngineCoreClient):
306
reset_running_requests, reset_connector
307
)
308
309
+ def reset_encoder_cache(self) -> None:
310
+ self.engine_core.reset_encoder_cache()
311
+
312
def sleep(self, level: int = 1) -> None:
313
self.engine_core.sleep(level)
314
@@ -765,6 +774,9 @@ class SyncMPClient(MPClient):
774
"reset_prefix_cache", reset_running_requests, reset_connector
775
)
776
777
+ def reset_encoder_cache(self) -> None:
778
+ self.call_utility("reset_encoder_cache")
779
+
780
def add_lora(self, lora_request: LoRARequest) -> bool:
781
return self.call_utility("add_lora", lora_request)
782
@@ -973,6 +985,9 @@ class AsyncMPClient(MPClient):
985
"reset_prefix_cache", reset_running_requests, reset_connector
986
)
987
988
+ async def reset_encoder_cache_async(self) -> None:
989
+ await self.call_utility_async("reset_encoder_cache")
990
+
991
async def sleep_async(self, level: int = 1) -> None:
992
await self.call_utility_async("sleep", level)
993
vllm/v1/engine/llm_engine.py
+8
@@ -332,6 +332,14 @@ class LLMEngine:
332
reset_running_requests, reset_connector
333
)
334
335
+ def reset_encoder_cache(self) -> None:
336
+ """Reset the encoder cache to invalidate all cached encoder outputs.
337
+
338
+ This should be called when model weights are updated to ensure
339
+ stale vision embeddings computed with old weights are not reused.
340
+ """
341
+ self.engine_core.reset_encoder_cache()
342
+
343
def sleep(self, level: int = 1):
344
self.engine_core.sleep(level)
345
vllm/v1/executor/abstract.py
+4
@@ -294,6 +294,10 @@ class Executor(ABC):
294
"""Reset the multi-modal cache in each worker."""
295
self.collective_rpc("reset_mm_cache")
296
297
+ def reset_encoder_cache(self) -> None:
298
+ """Reset the encoder cache in each worker to clear cached encoder outputs."""
299
+ self.collective_rpc("reset_encoder_cache")
300
+
301
def sleep(self, level: int = 1):
302
if self.is_sleeping:
303
logger.warning("Executor is already sleeping.")
vllm/v1/metrics/stats.py
+1
@@ -173,6 +173,7 @@ class SchedulerStats:
173
current_wave: int = 0
174
175
kv_cache_usage: float = 0.0
176
+ encoder_cache_usage: float = 0.0
177
178
prefix_cache_stats: PrefixCacheStats = field(default_factory=PrefixCacheStats)
179
connector_prefix_cache_stats: PrefixCacheStats | None = None
vllm/v1/worker/gpu_model_runner.py
+8
@@ -720,6 +720,14 @@ class GPUModelRunner(
720
if self.mm_budget:
721
self.mm_budget.reset_cache()
722
723
+ def reset_encoder_cache(self) -> None:
724
+ """Clear the GPU-side encoder cache storing vision embeddings.
725
+
726
+ This should be called when model weights are updated to ensure
727
+ stale embeddings computed with old weights are not reused.
728
+ """
729
+ self.encoder_cache.clear()
730
+
731
@torch.inference_mode()
732
def init_fp8_kv_scales(self) -> None:
733
"""
vllm/v1/worker/gpu_worker.py
+3
@@ -539,6 +539,9 @@ class Worker(WorkerBase):
539
def reset_mm_cache(self) -> None:
540
self.model_runner.reset_mm_cache()
541
542
+ def reset_encoder_cache(self) -> None:
543
+ self.model_runner.reset_encoder_cache()
544
+
545
def get_model(self) -> nn.Module:
546
return self.model_runner.get_model()
547