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, 06:15 UTC
[KV connector][LMCache] Only record the cuda event when there are request to store/load (#30814)
Signed-off-by: ApostaC <yihua98@uchicago.edu>
Yihua Cheng committed
Dec 17, 2025 at 21:31 UTC
ec965569d94c09eb1c85d235319b24d1b795d048
2 files changed
+40
-17
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
+1
@@ -262,6 +262,7 @@ class LMCacheMPWorkerAdapter:
262
):
263
keys = []
264
block_ids = []
265
+
266
for op in ops:
267
keys.extend(self._block_hashes_to_keys(op.block_hashes))
268
block_ids.extend(op.block_ids)
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+39
-17
@@ -24,6 +24,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_integration import (
24
)
25
from vllm.v1.core.sched.output import SchedulerOutput
26
from vllm.v1.outputs import KVConnectorOutput
27
+from vllm.v1.request import RequestStatus
28
from vllm.v1.utils import ConstantList
29
30
if TYPE_CHECKING:
@@ -211,7 +212,7 @@ class LMCacheMPRequestTracker:
212
"""
213
self.num_stored_blocks += num_new_blocks
214
214
- def update_block_ids(
215
+ def append_block_ids(
216
self,
217
new_block_ids: list[int],
218
):
@@ -455,10 +456,6 @@ class LMCacheMPConnector(KVConnectorBase_V1):
456
metadata = self._get_connector_metadata()
457
assert isinstance(metadata, LMCacheMPConnectorMetadata)
458
458
- with torch.cuda.stream(torch.cuda.current_stream()):
459
- event = torch.cuda.Event(interprocess=True)
460
- event.record()
461
-
459
request_ids = []
460
ops = []
461
@@ -468,10 +465,14 @@ class LMCacheMPConnector(KVConnectorBase_V1):
465
request_ids.append(meta.request_id)
466
ops.append(meta.op)
467
471
- if len(request_ids) > 0:
472
- self.worker_adapter.batched_submit_retrieve_requests(
473
- request_ids, ops, event
474
- )
468
+ if len(request_ids) == 0:
469
+ return
470
+
471
+ with torch.cuda.stream(torch.cuda.current_stream()):
472
+ event = torch.cuda.Event(interprocess=True)
473
+ event.record()
474
+
475
+ self.worker_adapter.batched_submit_retrieve_requests(request_ids, ops, event)
476
477
def wait_for_layer_load(self, layer_name: str) -> None:
478
"""
@@ -518,10 +519,6 @@ class LMCacheMPConnector(KVConnectorBase_V1):
519
metadata = self._get_connector_metadata()
520
assert isinstance(metadata, LMCacheMPConnectorMetadata)
521
521
- with torch.cuda.stream(torch.cuda.current_stream()):
522
- event = torch.cuda.Event(interprocess=True)
523
- event.record()
524
-
522
request_ids = []
523
ops = []
524
for meta in metadata.requests:
@@ -530,8 +527,14 @@ class LMCacheMPConnector(KVConnectorBase_V1):
527
request_ids.append(meta.request_id)
528
ops.append(meta.op)
529
533
- if len(request_ids) > 0:
534
- self.worker_adapter.batched_submit_store_requests(request_ids, ops, event)
530
+ if len(request_ids) == 0:
531
+ return
532
+
533
+ with torch.cuda.stream(torch.cuda.current_stream()):
534
+ event = torch.cuda.Event(interprocess=True)
535
+ event.record()
536
+
537
+ self.worker_adapter.batched_submit_store_requests(request_ids, ops, event)
538
539
def get_finished(
540
self, finished_req_ids: set[str]
@@ -627,6 +630,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):
630
into account.
631
"""
632
tracker = self._get_or_create_request_tracker(request)
633
+ # TODO: support loading KV for preempted requests in the future
634
+ if request.status == RequestStatus.PREEMPTED:
635
+ return 0, False
636
637
self.scheduler_adapter.maybe_submit_lookup_request(
638
request.request_id, convert_block_hashes_to_bytes(request.block_hashes)
@@ -683,7 +689,7 @@ class LMCacheMPConnector(KVConnectorBase_V1):
689
690
# No matter we need to retrieve or not, we need to update
691
# the block ids into the tracker
686
- tracker.update_block_ids(block_ids)
692
+ tracker.append_block_ids(block_ids)
693
694
# Update the state of the tracker
695
condition = tracker.needs_retrieve()
@@ -866,7 +872,8 @@ class LMCacheMPConnector(KVConnectorBase_V1):
872
873
# Update block ids
874
new_block_ids = reformat_block_ids(cached_reqs.new_block_ids[idx])
869
- request_tracker.update_block_ids(new_block_ids)
875
+ if request_id not in cached_reqs.resumed_req_ids:
876
+ request_tracker.append_block_ids(new_block_ids)
877
878
# Update new scheduled tokens
879
num_new_tokens = cached_reqs.num_computed_tokens[idx]
@@ -889,6 +896,21 @@ class LMCacheMPConnector(KVConnectorBase_V1):
896
self, request: "Request"
897
) -> LMCacheMPRequestTracker:
898
request_id = request.request_id
899
+ # Remove the old trackers that is created before the preemption
900
+ if (
901
+ request.status == RequestStatus.PREEMPTED
902
+ and request_id in self.request_trackers
903
+ ):
904
+ tracker = self.request_trackers[request_id]
905
+
906
+ # NOTE: since this function may be called multiple times
907
+ # for a single request (because get_num_new_matched_tokens
908
+ # may be called multiple times) for the same request, we
909
+ # will only do the remove if the tracker is not in the "fresh"
910
+ # state, i.e., PREFETCHING
911
+ if tracker.state != LMCacheMPRequestState.PREFETCHING:
912
+ self.request_trackers.pop(request_id)
913
+
914
if request_id not in self.request_trackers:
915
new_tracker = LMCacheMPRequestTracker(request)
916
self.request_trackers[request_id] = new_tracker