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 5, 08:30 UTC
[LMCache] vLLM Block Allocation Event (#38856)
Signed-off-by: yuwei <yuwei@dev.local> Co-authored-by: yuwei <yuwei@dev.local>
Yuwei An committed
Apr 9, 2026 at 17:30 UTC
467d3247c3bd3b0ef7056b39bbdb692b4b2957fa
1 file changed
+72
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+72
@@ -29,7 +29,18 @@ try:
29
LMCacheMPWorkerAdapter,
30
LoadStoreOp,
31
)
32
+
33
+ try:
34
+ from lmcache.v1.multiprocess.custom_types import RequestAllocationRecord
35
+ except ImportError:
36
+ from lmcache.v1.multiprocess.custom_types import (
37
+ BlockAllocationRecord as RequestAllocationRecord,
38
+ )
39
except ImportError:
40
+ from lmcache.v1.multiprocess.custom_types import (
41
+ BlockAllocationRecord as RequestAllocationRecord,
42
+ )
43
+
44
from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_integration import (
45
LMCacheMPSchedulerAdapter,
46
LMCacheMPWorkerAdapter,
@@ -837,6 +848,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):
848
if len(metadata) > 0:
849
logger.debug("Final connector metadata: %s", metadata)
850
851
+ # Report block allocation deltas to LMCache for observability
852
+ self._report_block_allocation_deltas(scheduler_output)
853
+
854
return metadata
855
856
def update_connector_output(self, connector_output: KVConnectorOutput):
@@ -1007,6 +1021,64 @@ class LMCacheMPConnector(KVConnectorBase_V1):
1021
if r_meta is not None:
1022
metadata.add_request_metadata(r_meta)
1023
1024
+ def _report_block_allocation_deltas(
1025
+ self,
1026
+ scheduler_output: SchedulerOutput,
1027
+ ) -> None:
1028
+ """Gather per-request block allocation deltas and report to LMCache.
1029
+
1030
+ For new requests: all allocated_block_ids and token_ids are new.
1031
+ For cached requests: only newly appended block_ids and token_ids.
1032
+ """
1033
+ records: list[RequestAllocationRecord] = []
1034
+
1035
+ # New requests: send all tokens covering all allocated blocks so
1036
+ # the L0 metrics subscriber can correctly map each block to its
1037
+ # actual token content (not just the newly-scheduled slice).
1038
+ for new_request in scheduler_output.scheduled_new_reqs:
1039
+ tracker = self.request_trackers.get(new_request.req_id)
1040
+ if tracker is None:
1041
+ continue
1042
+ num_blocks = len(tracker.allocated_block_ids)
1043
+ total_tokens = num_blocks * self.vllm_block_size
1044
+ records.append(
1045
+ RequestAllocationRecord(
1046
+ req_id=new_request.req_id,
1047
+ new_block_ids=list(tracker.allocated_block_ids),
1048
+ new_token_ids=list(tracker.all_token_ids[:total_tokens]),
1049
+ )
1050
+ )
1051
+
1052
+ # Cached requests: only the newly added blocks and their full
1053
+ # token content. We send all tokens covered by the new blocks
1054
+ # (not just the tokens scheduled this step) so the L0 subscriber
1055
+ # can correctly identify block content.
1056
+ cached_reqs = scheduler_output.scheduled_cached_reqs
1057
+ for idx, request_id in enumerate(cached_reqs.req_ids):
1058
+ new_block_ids = reformat_block_ids(cached_reqs.new_block_ids[idx])
1059
+ if not new_block_ids:
1060
+ continue
1061
+ tracker = self.request_trackers.get(request_id)
1062
+ if tracker is None:
1063
+ continue
1064
+ # The new blocks sit at the end of the request's block list.
1065
+ # Compute the token range they cover.
1066
+ total_blocks = len(tracker.allocated_block_ids)
1067
+ num_new_blocks = len(new_block_ids)
1068
+ start_token = (total_blocks - num_new_blocks) * self.vllm_block_size
1069
+ end_token = total_blocks * self.vllm_block_size
1070
+ new_token_ids = list(tracker.all_token_ids[start_token:end_token])
1071
+ records.append(
1072
+ RequestAllocationRecord(
1073
+ req_id=request_id,
1074
+ new_block_ids=new_block_ids,
1075
+ new_token_ids=new_token_ids,
1076
+ )
1077
+ )
1078
+
1079
+ if records:
1080
+ self.scheduler_adapter.report_block_allocations(records)
1081
+
1082
def _get_request_tracker(self, request_id: str) -> LMCacheMPRequestTracker:
1083
assert request_id in self.request_trackers, (
1084
f"Request tracker for request_id {request_id} not found. "