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 11, 23:45 UTC
[KVConnector][Feature] Support KV connector cache reset via /reset_prefix_cache (#27170)
Signed-off-by: tovam <tovam@pliops.com> Signed-off-by: Tova Movshovitz <tovam@pliops.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Tova Movshovitz committed
Dec 5, 2025 at 20:33 UTC
adb315060cd9943b38e68141cd9a54421144ae64
11 files changed
+105
-24
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+14
@@ -573,3 +573,17 @@ class KVConnectorBase_V1(ABC):
573
expose connector transfer stats via Prometheus.
574
"""
575
return None
576
+
577
+ def reset_cache(self) -> bool | None:
578
+ """
579
+ Reset the connector's internal cache.
580
+
581
+ Returns:
582
+ bool: True if the cache was successfully reset, False otherwise.
583
+ """
584
+ logger.debug(
585
+ "Connector cache reset requested, but %s does not implement reset_cache().",
586
+ type(self).__name__,
587
+ )
588
+
589
+ return None
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+4
@@ -452,3 +452,7 @@ class MultiConnector(KVConnectorBase_V1):
452
per_engine_labelvalues,
453
prom_metrics,
454
)
455
+
456
+ def reset_cache(self) -> bool:
457
+ results = [c.reset_cache() is not False for c in self._connectors]
458
+ return all(results)
vllm/engine/protocol.py
+4
-2
@@ -116,8 +116,10 @@ class EngineClient(ABC):
116
...
117
118
@abstractmethod
119
- async def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
120
- """Reset the prefix cache"""
119
+ async def reset_prefix_cache(
120
+ self, reset_running_requests: bool = False, reset_connector: bool = False
121
+ ) -> bool:
122
+ """Reset the prefix cache and optionally any configured connector cache"""
123
...
124
125
@abstractmethod
vllm/entrypoints/llm.py
+6
-2
@@ -1491,8 +1491,12 @@ class LLM:
1491
def stop_profile(self) -> None:
1492
self.llm_engine.stop_profile()
1493
1494
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
1495
- return self.llm_engine.reset_prefix_cache(reset_running_requests)
1494
+ def reset_prefix_cache(
1495
+ self, reset_running_requests: bool = False, reset_connector: bool = False
1496
+ ) -> bool:
1497
+ return self.llm_engine.reset_prefix_cache(
1498
+ reset_running_requests, reset_connector
1499
+ )
1500
1501
def sleep(self, level: int = 1):
1502
"""
vllm/entrypoints/openai/api_server.py
+17
-4
@@ -663,14 +663,27 @@ if envs.VLLM_SERVER_DEV_MODE:
663
664
@router.post("/reset_prefix_cache")
665
async def reset_prefix_cache(
666
- raw_request: Request, reset_running_requests: bool = Query(default=False)
666
+ raw_request: Request,
667
+ reset_running_requests: bool = Query(default=False),
668
+ reset_external: bool = Query(default=False),
669
):
670
"""
669
- Reset the prefix cache. Note that we currently do not check if the
670
- prefix cache is successfully reset in the API server.
671
+ Reset the local prefix cache.
672
+
673
+ Optionally, if the query parameter `reset_external=true`
674
+ also resets the external (connector-managed) prefix cache.
675
+
676
+ Note that we currently do not check if the prefix cache
677
+ is successfully reset in the API server.
678
+
679
+ Example:
680
+ POST /reset_prefix_cache?reset_external=true
681
"""
682
logger.info("Resetting prefix cache...")
673
- await engine_client(raw_request).reset_prefix_cache(reset_running_requests)
683
+
684
+ await engine_client(raw_request).reset_prefix_cache(
685
+ reset_running_requests, reset_external
686
+ )
687
return Response(status_code=200)
688
689
@router.post("/reset_mm_cache")
vllm/v1/core/sched/interface.py
+3
-1
@@ -152,7 +152,9 @@ class SchedulerInterface(ABC):
152
return self.has_unfinished_requests() or self.has_finished_requests()
153
154
@abstractmethod
155
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
155
+ def reset_prefix_cache(
156
+ self, reset_running_requests: bool = False, reset_connector: bool = False
157
+ ) -> bool:
158
"""Reset the prefix cache for KV cache.
159
160
This is particularly required when the model weights are live-updated.
vllm/v1/core/sched/scheduler.py
+21
-1
@@ -1380,7 +1380,9 @@ class Scheduler(SchedulerInterface):
1380
def has_finished_requests(self) -> bool:
1381
return len(self.finished_req_ids) > 0
1382
1383
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
1383
+ def reset_prefix_cache(
1384
+ self, reset_running_requests: bool = False, reset_connector: bool = False
1385
+ ) -> bool:
1386
"""Reset the KV prefix cache.
1387
1388
If reset_running_requests is True, all the running requests will be
@@ -1418,8 +1420,26 @@ class Scheduler(SchedulerInterface):
1420
"the presence of running requests waiting for remote KV transfer, "
1421
"which is not supported yet."
1422
)
1423
+
1424
+ if reset_connector:
1425
+ reset_successful = self.reset_connector_cache() and reset_successful
1426
+
1427
return reset_successful
1428
1429
+ def reset_connector_cache(self) -> bool:
1430
+ if self.connector is None:
1431
+ logger.warning("reset_connector called but no KV connector is configured.")
1432
+ return False
1433
+
1434
+ if self.connector.reset_cache() is False:
1435
+ return False
1436
+
1437
+ if self.log_stats:
1438
+ assert self.connector_prefix_cache_stats is not None
1439
+ self.connector_prefix_cache_stats.reset = True
1440
+
1441
+ return True
1442
+
1443
def make_stats(
1444
self,
1445
spec_decoding_stats: SpecDecodingStats | None = None,
vllm/v1/engine/async_llm.py
+6
-2
@@ -749,8 +749,12 @@ class AsyncLLM(EngineClient):
749
self.input_processor.clear_mm_cache()
750
await self.engine_core.reset_mm_cache_async()
751
752
- async def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
753
- return await self.engine_core.reset_prefix_cache_async(reset_running_requests)
752
+ async def reset_prefix_cache(
753
+ self, reset_running_requests: bool = False, reset_connector: bool = False
754
+ ) -> bool:
755
+ return await self.engine_core.reset_prefix_cache_async(
756
+ reset_running_requests, reset_connector
757
+ )
758
759
async def sleep(self, level: int = 1) -> None:
760
await self.reset_prefix_cache()
vllm/v1/engine/core.py
+6
-2
@@ -503,8 +503,12 @@ class EngineCore:
503
504
self.model_executor.reset_mm_cache()
505
506
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
507
- return self.scheduler.reset_prefix_cache(reset_running_requests)
506
+ def reset_prefix_cache(
507
+ self, reset_running_requests: bool = False, reset_connector: bool = False
508
+ ) -> bool:
509
+ return self.scheduler.reset_prefix_cache(
510
+ reset_running_requests, reset_connector
511
+ )
512
513
def sleep(self, level: int = 1):
514
self.model_executor.sleep(level)
vllm/v1/engine/core_client.py
+18
-8
@@ -138,7 +138,9 @@ class EngineCoreClient(ABC):
138
def reset_mm_cache(self) -> None:
139
raise NotImplementedError
140
141
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
141
+ def reset_prefix_cache(
142
+ self, reset_running_requests: bool = False, reset_connector: bool = False
143
+ ) -> bool:
144
raise NotImplementedError
145
146
def sleep(self, level: int = 1) -> None:
@@ -209,7 +211,7 @@ class EngineCoreClient(ABC):
211
raise NotImplementedError
212
213
async def reset_prefix_cache_async(
212
- self, reset_running_requests: bool = False
214
+ self, reset_running_requests: bool = False, reset_connector: bool = False
215
) -> bool:
216
raise NotImplementedError
217
@@ -289,8 +291,12 @@ class InprocClient(EngineCoreClient):
291
def reset_mm_cache(self) -> None:
292
self.engine_core.reset_mm_cache()
293
292
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
293
- return self.engine_core.reset_prefix_cache(reset_running_requests)
294
+ def reset_prefix_cache(
295
+ self, reset_running_requests: bool = False, reset_connector: bool = False
296
+ ) -> bool:
297
+ return self.engine_core.reset_prefix_cache(
298
+ reset_running_requests, reset_connector
299
+ )
300
301
def sleep(self, level: int = 1) -> None:
302
self.engine_core.sleep(level)
@@ -753,8 +759,12 @@ class SyncMPClient(MPClient):
759
def reset_mm_cache(self) -> None:
760
self.call_utility("reset_mm_cache")
761
756
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
757
- return self.call_utility("reset_prefix_cache", reset_running_requests)
762
+ def reset_prefix_cache(
763
+ self, reset_running_requests: bool = False, reset_connector: bool = False
764
+ ) -> bool:
765
+ return self.call_utility(
766
+ "reset_prefix_cache", reset_running_requests, reset_connector
767
+ )
768
769
def add_lora(self, lora_request: LoRARequest) -> bool:
770
return self.call_utility("add_lora", lora_request)
@@ -958,10 +968,10 @@ class AsyncMPClient(MPClient):
968
await self.call_utility_async("reset_mm_cache")
969
970
async def reset_prefix_cache_async(
961
- self, reset_running_requests: bool = False
971
+ self, reset_running_requests: bool = False, reset_connector: bool = False
972
) -> bool:
973
return await self.call_utility_async(
964
- "reset_prefix_cache", reset_running_requests
974
+ "reset_prefix_cache", reset_running_requests, reset_connector
975
)
976
977
async def sleep_async(self, level: int = 1) -> None:
vllm/v1/engine/llm_engine.py
+6
-2
@@ -328,8 +328,12 @@ class LLMEngine:
328
self.input_processor.clear_mm_cache()
329
self.engine_core.reset_mm_cache()
330
331
- def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
332
- return self.engine_core.reset_prefix_cache(reset_running_requests)
331
+ def reset_prefix_cache(
332
+ self, reset_running_requests: bool = False, reset_connector: bool = False
333
+ ) -> bool:
334
+ return self.engine_core.reset_prefix_cache(
335
+ reset_running_requests, reset_connector
336
+ )
337
338
def sleep(self, level: int = 1):
339
self.engine_core.sleep(level)