@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 94a9ebcf31
Fork 0

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 13, 08:15 UTC

[KV connector][WIP] KV cache proxy based on LMCache multi-process mode (#27902)

Signed-off-by: ApostaC <yihua98@uchicago.edu>

Yihua Cheng committed Nov 12, 2025 at 12:25 UTC 94a9ebcf317b804a6a70d74de9f91804d314ff6a
4 files changed +1265 -2
vllm/distributed/kv_transfer/kv_connector/factory.py
+6
@@ -161,6 +161,12 @@ KVConnectorFactory.register_connector(
161 "LMCacheConnectorV1",
162 )
163
164 +KVConnectorFactory.register_connector(
165 + "LMCacheMPConnector",
166 + "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_mp_connector",
167 + "LMCacheMPConnector",
168 +)
169 +
170 KVConnectorFactory.register_connector(
171 "NixlConnector",
172 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector",
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py
+13 -2
@@ -2,6 +2,17 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 -from . import vllm_v1_adapter
5 +from . import multi_process_adapter, vllm_v1_adapter
6 +from .multi_process_adapter import (
7 + LMCacheMPSchedulerAdapter,
8 + LMCacheMPWorkerAdapter,
9 + LoadStoreOp,
10 +)
11
7 -__all__ = ["vllm_v1_adapter"]
12 +__all__ = [
13 + "vllm_v1_adapter",
14 + "multi_process_adapter",
15 + "LMCacheMPSchedulerAdapter",
16 + "LMCacheMPWorkerAdapter",
17 + "LoadStoreOp",
18 +]
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py new
+379
@@ -0,0 +1,379 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import os
4 +from collections.abc import Iterable
5 +from dataclasses import dataclass
6 +from itertools import islice
7 +from typing import Any
8 +
9 +import torch
10 +import zmq
11 +from lmcache.utils import _lmcache_nvtx_annotate, init_logger
12 +from lmcache.v1.multiprocess.custom_types import (
13 + CudaIPCWrapper,
14 + IPCCacheEngineKey,
15 + KVCache,
16 +)
17 +from lmcache.v1.multiprocess.mq import MessageQueueClient, MessagingFuture
18 +from lmcache.v1.multiprocess.protocol import RequestType, get_response_class
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +def wrap_kv_caches(kv_caches: dict[str, KVCache]) -> KVCache:
24 + logger.info("KV caches keys are %s", list(kv_caches.keys()))
25 + return [CudaIPCWrapper(tensor) for tensor in kv_caches.values()]
26 +
27 +
28 +def send_lmcache_request(
29 + mq_client: MessageQueueClient,
30 + request_type: RequestType,
31 + payloads: list[Any],
32 +) -> MessagingFuture[Any]:
33 + future = mq_client.submit_request(
34 + request_type, payloads, get_response_class(request_type)
35 + )
36 + return future
37 +
38 +
39 +def get_lmcache_chunk_size(
40 + mq_client: MessageQueueClient,
41 +) -> int:
42 + future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, [])
43 + chunk_size = future.result()
44 + return chunk_size
45 +
46 +
47 +def striding_block_hashes(
48 + block_hashes: list[bytes],
49 + blocks_in_chunk,
50 +) -> Iterable[bytes]:
51 + """Striding the block hashes to get the block hashes for each chunk.
52 + For example, if blocks_in_chunk is 16, then we will get the block hashes
53 + for the 16th, 32nd, 48th, ... blocks.
54 + """
55 + return islice(block_hashes, blocks_in_chunk - 1, None, blocks_in_chunk)
56 +
57 +
58 +@dataclass
59 +class LoadStoreOp:
60 + block_hashes: list[bytes]
61 + block_ids: list[int]
62 +
63 + def __len__(self) -> int:
64 + return len(self.block_hashes)
65 +
66 + def __post_init__(self):
67 + assert len(self.block_hashes) == len(self.block_ids), (
68 + "The number of block hashes should be equal to the number of block ids "
69 + f"But got {len(self.block_hashes)} and {len(self.block_ids)}"
70 + )
71 +
72 +
73 +StoreResult = bool
74 +RetrieveResult = list[bool]
75 +LookupResult = list[bool]
76 +
77 +
78 +class LMCacheMPSchedulerAdapter:
79 + def __init__(
80 + self,
81 + server_url: str,
82 + context: zmq.Context,
83 + model_name: str,
84 + world_size: int,
85 + kv_rank: int,
86 + vllm_block_size: int,
87 + ):
88 + """
89 + Args:
90 + server_url: The server URL for the LMCache message queue
91 + context: The ZMQ context
92 +
93 + model_name: The model name used for LMCache keys
94 + world_size: The world size used for LMCache keys
95 + kv_rank: The kv rank used for LMCache keys
96 + vllm_block_size: The block size used in vLLM
97 + """
98 + self.mq_client = MessageQueueClient(server_url, context)
99 +
100 + # Request futures
101 + self.lookup_futures: dict[str, MessagingFuture[LookupResult]] = {}
102 +
103 + self.model_name = model_name
104 + self.world_size = world_size
105 + self.worker_id = kv_rank
106 +
107 + # Read chunk size from lmcache
108 + self.chunk_size = get_lmcache_chunk_size(self.mq_client)
109 + assert self.chunk_size % vllm_block_size == 0, (
110 + "LMCache chunk size should be a multiple of vLLM block size"
111 + )
112 + self.blocks_in_chunk = self.chunk_size // vllm_block_size
113 +
114 + @_lmcache_nvtx_annotate
115 + def maybe_submit_lookup_request(self, request_id: str, block_hashes: list[bytes]):
116 + if request_id in self.lookup_futures:
117 + # Skip if there is already a lookup request
118 + return
119 +
120 + s = striding_block_hashes(block_hashes, self.blocks_in_chunk)
121 + keys = [self._create_key(block_hash) for block_hash in s]
122 + future = send_lmcache_request(
123 + self.mq_client,
124 + RequestType.LOOKUP,
125 + [keys, True],
126 + )
127 + self.lookup_futures[request_id] = future
128 +
129 + @_lmcache_nvtx_annotate
130 + def check_lookup_result(self, request_id: str) -> int | None:
131 + assert request_id in self.lookup_futures, (
132 + f"Lookup request for request_id={request_id} has not been submitted"
133 + )
134 +
135 + future = self.lookup_futures[request_id]
136 + if not future.query():
137 + return None
138 +
139 + result = future.result()
140 + num_chunks = sum(result)
141 + return num_chunks * self.chunk_size
142 +
143 + def num_blocks_per_chunk(self) -> int:
144 + """
145 + Returns:
146 + The number of vllm blocks in a LMCache data chunk
147 + """
148 + return self.blocks_in_chunk
149 +
150 + # Helper functions
151 + def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:
152 + """Convert a block hash to an IPC cache engine key"""
153 + return IPCCacheEngineKey(
154 + model_name=self.model_name,
155 + world_size=self.world_size,
156 + worker_id=self.worker_id,
157 + chunk_hash=block_hash,
158 + )
159 +
160 +
161 +class LMCacheMPWorkerAdapter:
162 + def __init__(
163 + self,
164 + server_url: str,
165 + context: zmq.Context,
166 + model_name: str,
167 + world_size: int,
168 + kv_rank: int,
169 + vllm_block_size: int,
170 + ):
171 + self.mq_client = MessageQueueClient(server_url, context)
172 +
173 + # Instance id for GPU worker
174 + self.instance_id = os.getpid()
175 +
176 + # Registered kv caches from vLLM
177 + self.kv_caches: dict[str, torch.Tensor] = {}
178 +
179 + # Request futures
180 + # request_id -> (future, other merged requests)
181 + self.store_futures: dict[
182 + str, tuple[MessagingFuture[StoreResult], list[str]]
183 + ] = {}
184 + self.retrieve_futures: dict[
185 + str, tuple[MessagingFuture[RetrieveResult], list[str]]
186 + ] = {}
187 +
188 + self.finished_stores: set[str] = set()
189 + self.previously_finished: set[str] = set()
190 +
191 + self.model_name = model_name
192 + self.world_size = world_size
193 + self.worker_id = kv_rank
194 +
195 + # Read chunk size from lmcache
196 + chunk_size = get_lmcache_chunk_size(self.mq_client)
197 + assert chunk_size % vllm_block_size == 0, (
198 + "LMCache chunk size should be a multiple of vLLM block size"
199 + )
200 + self.blocks_in_chunk = chunk_size // vllm_block_size
201 +
202 + def register_kv_caches(self, kv_caches: dict[str, KVCache]):
203 + # Register kv cache and send the request
204 + self.kv_caches = kv_caches
205 + logger.info("Registering kv caches")
206 + future = send_lmcache_request(
207 + self.mq_client,
208 + RequestType.REGISTER_KV_CACHE,
209 + [self.instance_id, wrap_kv_caches(kv_caches)],
210 + )
211 + future.result()
212 +
213 + @_lmcache_nvtx_annotate
214 + def submit_store_request(
215 + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
216 + ):
217 + keys = self._block_hashes_to_keys(op.block_hashes)
218 + future = send_lmcache_request(
219 + self.mq_client,
220 + RequestType.STORE,
221 + [keys, self.instance_id, op.block_ids, event.ipc_handle()],
222 + ).to_cuda_future()
223 + self.store_futures[request_id] = (future, [])
224 +
225 + @_lmcache_nvtx_annotate
226 + def submit_retrieve_request(
227 + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
228 + ):
229 + keys = self._block_hashes_to_keys(op.block_hashes)
230 + future = send_lmcache_request(
231 + self.mq_client,
232 + RequestType.RETRIEVE,
233 + [keys, self.instance_id, op.block_ids, event.ipc_handle()],
234 + ).to_cuda_future()
235 + self.retrieve_futures[request_id] = (future, [])
236 +
237 + @_lmcache_nvtx_annotate
238 + def batched_submit_store_requests(
239 + self,
240 + request_ids: list[str],
241 + ops: list[LoadStoreOp],
242 + event: torch.cuda.Event,
243 + ):
244 + keys = []
245 + block_ids = []
246 + for op in ops:
247 + keys.extend(self._block_hashes_to_keys(op.block_hashes))
248 + block_ids.extend(op.block_ids)
249 + future = send_lmcache_request(
250 + self.mq_client,
251 + RequestType.STORE,
252 + [keys, self.instance_id, block_ids, event.ipc_handle()],
253 + ).to_cuda_future()
254 + self.store_futures[request_ids[0]] = (future, request_ids[1:])
255 +
256 + @_lmcache_nvtx_annotate
257 + def batched_submit_retrieve_requests(
258 + self,
259 + request_ids: list[str],
260 + ops: list[LoadStoreOp],
261 + event: torch.cuda.Event,
262 + ):
263 + keys = []
264 + block_ids = []
265 + for op in ops:
266 + keys.extend(self._block_hashes_to_keys(op.block_hashes))
267 + block_ids.extend(op.block_ids)
268 + future = send_lmcache_request(
269 + self.mq_client,
270 + RequestType.RETRIEVE,
271 + [keys, self.instance_id, block_ids, event.ipc_handle()],
272 + ).to_cuda_future()
273 + self.retrieve_futures[request_ids[0]] = (future, request_ids[1:])
274 +
275 + @_lmcache_nvtx_annotate
276 + def get_finished(
277 + self, finished_req_ids: set[str]
278 + ) -> tuple[set[str] | None, set[str] | None]:
279 + finished_stores = set()
280 + finished_retrieves = set()
281 + for request_id, (future, other_reqs) in self.store_futures.items():
282 + if not future.query():
283 + continue
284 +
285 + result = future.result()
286 + finished_stores.add(request_id)
287 + finished_stores.update(other_reqs)
288 +
289 + if not result:
290 + # TODO: add error handling here
291 + logger.error(
292 + "Something went wrong when processing the "
293 + "store request for request_id=%s",
294 + request_id,
295 + )
296 +
297 + for request_id, (future, other_reqs) in self.retrieve_futures.items():
298 + if not future.query():
299 + continue
300 +
301 + result = future.result()
302 + finished_retrieves.add(request_id)
303 + finished_retrieves.update(other_reqs)
304 +
305 + if not all(result):
306 + # TODO: add error handing here
307 + logger.error(
308 + "Something went wrong when processing the "
309 + "retrieve request for request_id=%s, result=%s",
310 + request_id,
311 + result,
312 + )
313 + logger.info("Retrieve request for request_id=%s finished", request_id)
314 +
315 + # Remove the finished requests from the tracking dicts
316 + for request_id in finished_stores:
317 + self.store_futures.pop(request_id, None)
318 + for request_id in finished_retrieves:
319 + self.retrieve_futures.pop(request_id, None)
320 +
321 + # Update the internal states
322 + self.finished_stores.update(finished_stores)
323 +
324 + ret_stores = set()
325 + for req_id in finished_req_ids:
326 + if req_id in self.finished_stores or req_id in self.store_futures:
327 + self.previously_finished.add(req_id)
328 + else:
329 + ret_stores.add(req_id)
330 +
331 + # Calculate the final finished stores
332 + ret_stores.update(self._update_and_get_finished_store())
333 +
334 + return ret_stores, finished_retrieves
335 +
336 + def num_blocks_per_chunk(self) -> int:
337 + """
338 + Returns:
339 + The number of vllm blocks in a LMCache data chunk
340 + """
341 + return self.blocks_in_chunk
342 +
343 + def shutdown(self):
344 + # Unregister kv cache
345 + logger.info("Unregistering kv caches")
346 + send_lmcache_request(
347 + self.mq_client, RequestType.UNREGISTER_KV_CACHE, [self.instance_id]
348 + ).result()
349 +
350 + self.mq_client.close()
351 +
352 + # Helper functions
353 + def _update_and_get_finished_store(
354 + self,
355 + ) -> set[str]:
356 + """Converge the internal states about finished stores
357 + and returns the 'safe finished store request ids' back
358 + """
359 + safe_finished_s = self.finished_stores.intersection(self.previously_finished)
360 + self.finished_stores.difference_update(self.previously_finished)
361 + self.previously_finished.difference_update(safe_finished_s)
362 +
363 + return safe_finished_s
364 +
365 + def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:
366 + """Convert a block hash to an IPC cache engine key"""
367 + return IPCCacheEngineKey(
368 + model_name=self.model_name,
369 + world_size=self.world_size,
370 + worker_id=self.worker_id,
371 + chunk_hash=block_hash,
372 + )
373 +
374 + def _block_hashes_to_keys(
375 + self, block_hashes: list[bytes]
376 + ) -> list[IPCCacheEngineKey]:
377 + """Convert block hashes to IPC cache engine keys"""
378 + s = striding_block_hashes(block_hashes, self.blocks_in_chunk)
379 + return [self._create_key(block_hash) for block_hash in s]
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py new
+867
@@ -0,0 +1,867 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import enum
4 +from collections.abc import Iterable
5 +from dataclasses import dataclass, field
6 +from typing import TYPE_CHECKING, Any, Literal, Optional, cast
7 +
8 +import torch
9 +import zmq
10 +from lmcache.utils import init_logger as lmcache_init_logger
11 +
12 +from vllm.config import VllmConfig
13 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
14 + KVConnectorBase_V1,
15 + KVConnectorMetadata,
16 + KVConnectorRole,
17 +)
18 +from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_integration import (
19 + LMCacheMPSchedulerAdapter,
20 + LMCacheMPWorkerAdapter,
21 + LoadStoreOp,
22 +)
23 +from vllm.v1.core.sched.output import SchedulerOutput
24 +from vllm.v1.outputs import KVConnectorOutput
25 +from vllm.v1.utils import ConstantList
26 +
27 +if TYPE_CHECKING:
28 + from vllm.attention.backends.abstract import AttentionMetadata
29 + from vllm.config import VllmConfig
30 + from vllm.distributed.kv_events import KVCacheEvent
31 + from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
32 + KVConnectorPromMetrics,
33 + KVConnectorStats,
34 + PromMetric,
35 + PromMetricT,
36 + )
37 + from vllm.forward_context import ForwardContext
38 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
39 + from vllm.v1.core.kv_cache_utils import BlockHash
40 + from vllm.v1.kv_cache_interface import KVCacheConfig
41 + from vllm.v1.request import Request
42 +
43 +logger = lmcache_init_logger(__name__)
44 +
45 +
46 +# Helper functions
47 +def reformat_block_ids(block_ids: tuple[list[int], ...] | None) -> list[int]:
48 + if block_ids is None:
49 + return []
50 + assert isinstance(block_ids, tuple), (
51 + f"Expected block_ids to be a tuple of lists, but got {type(block_ids)}"
52 + )
53 +
54 + if len(block_ids) > 1:
55 + raise RuntimeError(
56 + "LMCacheMPConnector only works without hybrid kv cache manager. "
57 + "Please pass --disable-hybrid-kv-cache-manager when starting vllm"
58 + )
59 +
60 + return block_ids[0]
61 +
62 +
63 +def create_scheduler_adapter(
64 + server_url: str, zmq_context: zmq.Context, vllm_config: VllmConfig
65 +) -> LMCacheMPSchedulerAdapter:
66 + # TODO: have a helper function to calculate the correct rank and
67 + # world size for the MLA and other models
68 + return LMCacheMPSchedulerAdapter(
69 + server_url,
70 + zmq_context,
71 + vllm_config.model_config.model,
72 + vllm_config.parallel_config.world_size,
73 + vllm_config.parallel_config.rank,
74 + vllm_config.cache_config.block_size,
75 + )
76 +
77 +
78 +def create_worker_adapter(
79 + server_url: str, zmq_context: zmq.Context, vllm_config: VllmConfig
80 +) -> LMCacheMPWorkerAdapter:
81 + # TODO: have a helper function to calculate the correct rank and
82 + # world size for the MLA and other models
83 + return LMCacheMPWorkerAdapter(
84 + server_url,
85 + zmq_context,
86 + vllm_config.model_config.model,
87 + vllm_config.parallel_config.world_size,
88 + vllm_config.parallel_config.rank,
89 + vllm_config.cache_config.block_size,
90 + )
91 +
92 +
93 +def convert_block_hashes_to_bytes(
94 + block_hashes: list["BlockHash"],
95 +) -> list[bytes]:
96 + return cast(list[bytes], block_hashes)
97 +
98 +
99 +class LMCacheMPRequestState(enum.Enum):
100 + """
101 + State machine:
102 + PREFETCHING -- update_state_after_alloc --> WAITING_FOR_LOAD
103 + WAITING_FOR_LOAD -- process_loading_requests --> READY
104 + """
105 +
106 + PREFETCHING = enum.auto()
107 + WAITING_FOR_LOAD = enum.auto()
108 + READY = enum.auto()
109 +
110 +
111 +@dataclass
112 +class LMCacheMPRequestTracker:
113 + # NOTE: this class used vLLM data structures, should be part of
114 + # vLLM integration code
115 +
116 + request_id: str
117 +
118 + # Read-only lists to track the token ids and block hashes
119 + all_token_ids: ConstantList[int]
120 + block_hashes: ConstantList["BlockHash"]
121 +
122 + # Block ids and hashes will be updated at update_states_after_alloc and
123 + # during the generation
124 + allocated_block_ids: list[int] = field(default_factory=list)
125 +
126 + # Number of scheduled tokens in this request. We keep tracking this to
127 + # avoid saving half-full blocks.
128 + num_scheduled_tokens: int = 0
129 +
130 + # Number of blocks stored will be initialized when lookup the external
131 + # hit tokens and will be updated when processing new requests and cached
132 + # requests.
133 + num_stored_blocks: int = 0
134 +
135 + # Staging load operation -- save vllm and lmcache hit tokens during lookup
136 + num_vllm_hit_blocks: int = 0
137 + num_lmcache_hit_blocks: int = 0
138 +
139 + # Main state
140 + state: LMCacheMPRequestState = LMCacheMPRequestState.PREFETCHING
141 +
142 + def __init__(self, request: "Request"):
143 + self.request_id = request.request_id
144 + self.all_token_ids = request.all_token_ids
145 + self.block_hashes = ConstantList(request.block_hashes)
146 + self.allocated_block_ids = []
147 + self.num_stored_blocks = 0
148 + self.num_vllm_hit_blocks = 0
149 + self.num_lmcache_hit_blocks = 0
150 + self.state = LMCacheMPRequestState.PREFETCHING
151 +
152 + ####
153 + # Check the state of the request
154 + ####
155 + def needs_retrieve(self) -> bool:
156 + """Check whether the current request needs retrieve, will be used
157 + update_stage_after_alloc"""
158 + return (
159 + self.num_lmcache_hit_blocks > self.num_vllm_hit_blocks
160 + and self.state != LMCacheMPRequestState.READY
161 + )
162 +
163 + def is_ready_for_retrieving(self) -> bool:
164 + """Check whether the current request is ready for retrieving,
165 + will be used in process_loading_requests"""
166 + return (
167 + self.state == LMCacheMPRequestState.WAITING_FOR_LOAD
168 + and self.needs_retrieve()
169 + )
170 +
171 + ####
172 + # Update internal states
173 + ####
174 + def increase_num_scheduled_tokens(self, num_new_tokens: int):
175 + self.num_scheduled_tokens += num_new_tokens
176 +
177 + def increase_num_stored_blocks(self, num_new_blocks: int):
178 + """Increase the number of stored blocks for the current request
179 + This function will be called when processing the cached requests.
180 + """
181 + self.num_stored_blocks += num_new_blocks
182 +
183 + def update_block_ids(
184 + self,
185 + new_block_ids: list[int],
186 + ):
187 + """Update the block ids for the current request
188 + This function will be called when processing the cached requests.
189 + """
190 + self.allocated_block_ids.extend(new_block_ids)
191 +
192 + ####
193 + # For debugging
194 + ####
195 + def __repr__(self) -> str:
196 + return (
197 + f"LMCacheMPRequestTracker(request_id={self.request_id}, "
198 + f"num_tokens={len(self.all_token_ids)}, "
199 + f"num_block_hashes={len(self.block_hashes)}, "
200 + f"num_allocated_blocks={len(self.allocated_block_ids)}, "
201 + f"num_stored_blocks={self.num_stored_blocks}, "
202 + f"vllm_hit_blocks={self.num_vllm_hit_blocks}, "
203 + f"lmcache_hit_blocks={self.num_lmcache_hit_blocks}, "
204 + f"state={self.state})"
205 + )
206 +
207 + def __str__(self) -> str:
208 + return self.__repr__()
209 +
210 +
211 +@dataclass
212 +class LMCacheMPRequestMetadata:
213 + request_id: str
214 + direction: Literal["STORE", "RETRIEVE"]
215 + op: LoadStoreOp
216 +
217 + @staticmethod
218 + def GetStoreMetadata(
219 + tracker: LMCacheMPRequestTracker,
220 + blocks_in_chunk: int,
221 + vllm_block_size: int,
222 + ) -> "LMCacheMPRequestMetadata | None":
223 + """
224 + Generate the store metadata for the current request tracker.
225 +
226 + Args:
227 + tracker: The request tracker to generate the metadata from.
228 + blocks_in_chunk: the number of blocks in a LMCache data chunk
229 + """
230 + # Store the blocks that has block hashes
231 + # NOTE: the invariant here is that `num_stored_blocks` should
232 + # always be a multiple of `blocks_in_chunk`
233 + # TODO: This should be checked everytime we update the num_stored_blocks
234 + min_available_blocks = min(
235 + len(tracker.block_hashes),
236 + len(tracker.allocated_block_ids),
237 + tracker.num_scheduled_tokens // vllm_block_size,
238 + )
239 + num_staging_blocks = min_available_blocks - tracker.num_stored_blocks
240 + num_chunks = num_staging_blocks // blocks_in_chunk
241 +
242 + if num_chunks >= 1:
243 + start = tracker.num_stored_blocks
244 + end = start + num_chunks * blocks_in_chunk
245 + block_hashes = convert_block_hashes_to_bytes(
246 + tracker.block_hashes[start:end]
247 + )
248 + block_ids = tracker.allocated_block_ids[start:end]
249 +
250 + ret = LMCacheMPRequestMetadata(
251 + request_id=tracker.request_id,
252 + direction="STORE",
253 + op=LoadStoreOp(block_hashes=block_hashes, block_ids=block_ids),
254 + )
255 +
256 + # Update the request tracker
257 + tracker.increase_num_stored_blocks(end - start)
258 + return ret
259 +
260 + return None
261 +
262 + @staticmethod
263 + def GetRetrieveMetadata(
264 + tracker: LMCacheMPRequestTracker,
265 + blocks_in_chunk: int,
266 + ) -> "LMCacheMPRequestMetadata | None":
267 + """
268 + Generate the retrieve metadata for the current request tracker.
269 +
270 + Args:
271 + tracker: The request tracker to generate the metadata from.
272 + blocks_in_chunk: the number of blocks in a LMCache data chunk
273 + """
274 + if not tracker.is_ready_for_retrieving():
275 + return None
276 +
277 + # |---------------------|-----------------|----------------|
278 + # | num_vllm_hit_blocks |
279 + # | lmcache chunk 1 | lmcache chunk 2 |
280 + # | need to retrieve |
281 +
282 + start = tracker.num_vllm_hit_blocks // blocks_in_chunk * blocks_in_chunk
283 + end = tracker.num_lmcache_hit_blocks
284 + assert end % blocks_in_chunk == 0, (
285 + "The number of LMCache hit blocks should be a multiple of the "
286 + "number of blocks in a lmcache chunk. "
287 + )
288 + assert len(tracker.block_hashes) >= end, (
289 + "The number of block hashes should be greater than or equal to the "
290 + "number of LMCache hit blocks. "
291 + )
292 + if end > start:
293 + block_hashes = convert_block_hashes_to_bytes(
294 + tracker.block_hashes[start:end]
295 + )
296 + block_ids = tracker.allocated_block_ids[start:end]
297 +
298 + ret = LMCacheMPRequestMetadata(
299 + request_id=tracker.request_id,
300 + direction="RETRIEVE",
301 + op=LoadStoreOp(block_hashes=block_hashes, block_ids=block_ids),
302 + )
303 + return ret
304 +
305 + return None
306 +
307 +
308 +class LMCacheMPConnectorMetadata(KVConnectorMetadata):
309 + def __init__(self):
310 + super().__init__()
311 + self.requests: list[LMCacheMPRequestMetadata] = []
312 +
313 + def add_request_metadata(self, request_metadata: LMCacheMPRequestMetadata):
314 + self.requests.append(request_metadata)
315 +
316 + def __len__(self):
317 + return len(self.requests)
318 +
319 + # For debugging
320 + def __str__(self):
321 + request_strs = []
322 + for req_meta in self.requests:
323 + request_strs.append(
324 + f"RequestMetadata(request_id={req_meta.request_id}, "
325 + f"direction={req_meta.direction}, "
326 + f"num_blocks={len(req_meta.op)}, "
327 + f"block_ids={req_meta.op.block_ids})"
328 + )
329 + return "[" + "\n".join(request_strs) + "]"
330 +
331 + def __repr__(self):
332 + return self.__str__()
333 +
334 +
335 +class LMCacheMPConnector(KVConnectorBase_V1):
336 + """
337 + The connector for LMCache multi-process mode.
338 +
339 + Extra configs (kv_transfer_config.extra_config):
340 + - lmcache.mp.host: the host of the LMCache server.
341 + - lmcache.mp.port: the port of the LMCache server.
342 + """
343 +
344 + def __init__(
345 + self,
346 + vllm_config: "VllmConfig",
347 + role: KVConnectorRole,
348 + kv_cache_config: Optional["KVCacheConfig"] = None,
349 + ):
350 + super().__init__(vllm_config, role, kv_cache_config)
351 +
352 + assert vllm_config.kv_transfer_config is not None
353 + server_host = vllm_config.kv_transfer_config.get_from_extra_config(
354 + "lmcache.mp.host", "tcp://localhost"
355 + )
356 + server_port = vllm_config.kv_transfer_config.get_from_extra_config(
357 + "lmcache.mp.port", 5555
358 + )
359 +
360 + server_url = f"{server_host}:{server_port}"
361 + zmq_context = zmq.Context.instance()
362 + if self.role == KVConnectorRole.SCHEDULER:
363 + self.scheduler_adapter = create_scheduler_adapter(
364 + server_url, zmq_context, vllm_config
365 + )
366 + self.request_trackers: dict[str, LMCacheMPRequestTracker] = {}
367 + elif self.role == KVConnectorRole.WORKER:
368 + self.worker_adapter = create_worker_adapter(
369 + server_url, zmq_context, vllm_config
370 + )
371 + else:
372 + raise ValueError(f"Unknown KVConnectorRole: {self.role}")
373 +
374 + self.vllm_block_size = vllm_config.cache_config.block_size
375 +
376 + @property
377 + def role(self) -> KVConnectorRole:
378 + return self._role
379 +
380 + # ==============================
381 + # Worker-side methods
382 + # ==============================
383 +
384 + def _get_connector_metadata(self) -> KVConnectorMetadata:
385 + """Get the connector metadata.
386 +
387 + This function should only be called inside the connector.
388 +
389 + Returns:
390 + ConnectorMetadata: the connector metadata.
391 + """
392 +
393 + # Should only be called while set to valid metadata.
394 + assert self._connector_metadata is not None
395 + return self._connector_metadata
396 +
397 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
398 + """
399 + Initialize with the KV caches. Useful for pre-registering the
400 + KV Caches in the KVConnector (e.g. for NIXL).
401 +
402 + Args:
403 + kv_caches: dictionary of layer names, kv cache
404 + """
405 + logger.info("Registering kv caches!")
406 + self.worker_adapter.register_kv_caches(kv_caches)
407 + return
408 +
409 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
410 + """
411 + Start loading the KV cache from the connector to vLLM's paged
412 + KV buffer. This is called from the forward context before the
413 + forward pass to enable async loading during model execution.
414 +
415 + Args:
416 + forward_context (ForwardContext): the forward context.
417 + **kwargs: additional arguments for the load operation
418 +
419 + Note:
420 + The number of elements in kv_caches and layer_names should be
421 + the same.
422 +
423 + """
424 + metadata = self._get_connector_metadata()
425 + assert isinstance(metadata, LMCacheMPConnectorMetadata)
426 +
427 + with torch.cuda.stream(torch.cuda.current_stream()):
428 + event = torch.cuda.Event(interprocess=True)
429 + event.record()
430 +
431 + request_ids = []
432 + ops = []
433 +
434 + for meta in metadata.requests:
435 + if meta.direction != "RETRIEVE":
436 + continue
437 + request_ids.append(meta.request_id)
438 + ops.append(meta.op)
439 +
440 + if len(request_ids) > 0:
441 + logger.info(
442 + "HERE! SUBMITTING THE BATCHED RETRIEVE REQUESTS %s", request_ids
443 + )
444 + self.worker_adapter.batched_submit_retrieve_requests(
445 + request_ids, ops, event
446 + )
447 +
448 + def wait_for_layer_load(self, layer_name: str) -> None:
449 + """
450 + Block until the KV for a specific layer is loaded into vLLM's
451 + paged buffer. This is called from within attention layer to ensure
452 + async copying from start_load_kv is complete.
453 +
454 + This interface will be useful for layer-by-layer pipelining.
455 +
456 + Args:
457 + layer_name: the name of that layer
458 + """
459 + return
460 +
461 + def save_kv_layer(
462 + self,
463 + layer_name: str,
464 + kv_layer: torch.Tensor,
465 + attn_metadata: "AttentionMetadata",
466 + **kwargs: Any,
467 + ) -> None:
468 + """
469 + Start saving a layer of KV cache from vLLM's paged buffer
470 + to the connector. This is called from within attention layer to
471 + enable async copying during execution.
472 +
473 + Args:
474 + layer_name (str): the name of the layer.
475 + kv_layer (torch.Tensor): the paged KV buffer of the current
476 + layer in vLLM.
477 + attn_metadata (AttentionMetadata): the attention metadata.
478 + **kwargs: additional arguments for the save operation.
479 + """
480 + return
481 +
482 + def wait_for_save(self):
483 + """
484 + Block until all the save operations is done. This is called
485 + as the forward context exits to ensure that the async saving
486 + from save_kv_layer is complete before finishing the forward.
487 +
488 + This prevents overwrites of paged KV buffer before saving done.
489 + """
490 + metadata = self._get_connector_metadata()
491 + assert isinstance(metadata, LMCacheMPConnectorMetadata)
492 +
493 + with torch.cuda.stream(torch.cuda.current_stream()):
494 + event = torch.cuda.Event(interprocess=True)
495 + event.record()
496 +
497 + request_ids = []
498 + ops = []
499 + for meta in metadata.requests:
500 + if meta.direction != "STORE":
501 + continue
502 + request_ids.append(meta.request_id)
503 + ops.append(meta.op)
504 +
505 + if len(request_ids) > 0:
506 + self.worker_adapter.batched_submit_store_requests(request_ids, ops, event)
507 +
508 + def get_finished(
509 + self, finished_req_ids: set[str]
510 + ) -> tuple[set[str] | None, set[str] | None]:
511 + """
512 + Notifies worker-side connector ids of requests that have
513 + finished generating tokens on the worker.
514 + The scheduler process (via the Executors) will use this output
515 + to track which workers are done.
516 +
517 + Returns:
518 + ids of requests that have finished asynchronous transfer
519 + (requests that previously returned True from request_finished()),
520 + tuple of (sending/saving ids, recving/loading ids).
521 + The finished saves/sends req ids must belong to a set provided in a
522 + call to this method (this call or a prior one).
523 + """
524 + val = self.worker_adapter.get_finished(finished_req_ids)
525 + # logger.error("Finished req ids: %s, %s", val[0], val[1])
526 + return val
527 +
528 + def get_block_ids_with_load_errors(self) -> set[int]:
529 + """
530 + Get the set of block IDs that failed to load.
531 +
532 + Returns:
533 + Set of block IDs that encountered load errors.
534 + Empty set if no load errors occurred.
535 +
536 + Notes:
537 + - Applies to both sync- and async-loading requests.
538 + - Async loading: failed blocks may be reported in any forward pass
539 + up to and including the pass where the request ID is returned by
540 + `get_finished()`. Even if failures occur, the request must still
541 + be reported via `get_finished()`, and the failed block IDs must
542 + appear here no later than that same pass.
543 + - Sync loading: failed blocks should be reported in the forward
544 + pass in which they are detected.
545 + """
546 + # TODO: add error tracking
547 + return set()
548 +
549 + def shutdown(self):
550 + """
551 + Shutdown the connector. This is called when the worker process
552 + is shutting down to ensure that all the async operations are
553 + completed and the connector is cleaned up properly.
554 + """
555 + if hasattr(self, "worker_adapter"):
556 + self.worker_adapter.shutdown()
557 + return None
558 +
559 + def get_kv_connector_stats(self) -> Optional["KVConnectorStats"]:
560 + """
561 + Get the KV connector stats collected during the last interval.
562 + """
563 + return None
564 +
565 + # ==============================
566 + # Scheduler-side methods
567 + # ==============================
568 +
569 + def get_num_new_matched_tokens(
570 + self,
571 + request: "Request",
572 + num_computed_tokens: int,
573 + ) -> tuple[int | None, bool]:
574 + """
575 + Get number of new tokens that can be loaded from the
576 + external KV cache beyond the num_computed_tokens.
577 +
578 + Args:
579 + request (Request): the request object.
580 + num_computed_tokens (int): the number of locally
581 + computed tokens for this request
582 +
583 + Returns:
584 + A tuple with the following elements:
585 + - An optional number of tokens that can be loaded from the
586 + external KV cache beyond what is already computed.
587 + If None, it means that the connector needs more time to
588 + determine the number of matched tokens, and the scheduler
589 + should query for this request again later.
590 + - `True` if external KV cache tokens will be loaded
591 + asynchronously (between scheduler steps). Must be
592 + 'False' if the first element is 0.
593 +
594 + Notes:
595 + The connector should only consider the largest prefix of prompt-
596 + tokens for which KV cache is actually available at the time of the
597 + call. If the cache cannot be loaded for some tokens (e.g., due to
598 + connectivity issues or eviction), those tokens must not be taken
599 + into account.
600 + """
601 + tracker = self._get_or_create_request_tracker(request)
602 +
603 + self.scheduler_adapter.maybe_submit_lookup_request(
604 + request.request_id, convert_block_hashes_to_bytes(request.block_hashes)
605 + )
606 +
607 + ret = self.scheduler_adapter.check_lookup_result(request.request_id)
608 + if ret is None:
609 + return None, True
610 +
611 + if ret == 0:
612 + return 0, False
613 +
614 + assert (
615 + ret % (self.scheduler_adapter.num_blocks_per_chunk() * self.vllm_block_size)
616 + == 0
617 + )
618 +
619 + # Update num stored blocks for the tracker
620 + num_vllm_blocks = num_computed_tokens // self.vllm_block_size
621 + num_lmcache_blocks = ret // self.vllm_block_size
622 + tracker.increase_num_stored_blocks(num_lmcache_blocks)
623 +
624 + # Save the vllm and lmcache hit tokens
625 + tracker.num_vllm_hit_blocks = num_vllm_blocks
626 + tracker.num_lmcache_hit_blocks = num_lmcache_blocks
627 +
628 + need_to_load = max(0, ret - num_computed_tokens)
629 + logger.debug(
630 + "vLLM hit is: %d, Need to load is %d", num_computed_tokens, need_to_load
631 + )
632 + return need_to_load, need_to_load > 0
633 +
634 + def update_state_after_alloc(
635 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
636 + ):
637 + """
638 + Update KVConnector state after block allocation.
639 +
640 + If get_num_new_matched_tokens previously returned True for a
641 + request, this function may be called twice for that same request -
642 + first when blocks are allocated for the connector tokens to be
643 + asynchronously loaded into, and second when any additional blocks
644 + are allocated, after the load/transfer is complete.
645 +
646 + Args:
647 + request (Request): the request object.
648 + blocks (KVCacheBlocks): the blocks allocated for the request.
649 + num_external_tokens (int): the number of tokens that will be
650 + loaded from the external KV cache.
651 + """
652 + # NOTE: the `blocks` are NEW BLOCKS allocated for this request.
653 + tracker = self._get_request_tracker(request.request_id)
654 + block_ids = reformat_block_ids(blocks.get_block_ids())
655 +
656 + # No matter we need to retrieve or not, we need to update
657 + # the block ids into the tracker
658 + tracker.update_block_ids(block_ids)
659 +
660 + # Update the state of the tracker
661 + condition = tracker.needs_retrieve()
662 + if tracker.state == LMCacheMPRequestState.PREFETCHING:
663 + # If need to retrieve, change to WAITING_FOR_LOAD
664 + # Otherwise, change to READY
665 + tracker.state = (
666 + LMCacheMPRequestState.WAITING_FOR_LOAD
667 + if condition
668 + else LMCacheMPRequestState.READY
669 + )
670 +
671 + def build_connector_meta(
672 + self, scheduler_output: SchedulerOutput
673 + ) -> KVConnectorMetadata:
674 + """
675 + Build the connector metadata for this step.
676 +
677 + This function should NOT modify fields in the scheduler_output.
678 + Also, calling this function will reset the state of the connector.
679 +
680 + Args:
681 + scheduler_output (SchedulerOutput): the scheduler output object.
682 + """
683 + metadata = LMCacheMPConnectorMetadata()
684 +
685 + self._process_retrieve_requests(metadata)
686 + self._process_new_requests(scheduler_output, metadata)
687 + self._process_cached_requests(scheduler_output, metadata)
688 +
689 + if len(metadata) > 0:
690 + logger.debug("Final connector metadata: %s", metadata)
691 +
692 + return metadata
693 +
694 + def update_connector_output(self, connector_output: KVConnectorOutput):
695 + """
696 + Update KVConnector state from worker-side connectors output.
697 +
698 + Args:
699 + connector_output (KVConnectorOutput): the worker-side
700 + connectors output.
701 + """
702 + return
703 +
704 + def request_finished(
705 + self,
706 + request: "Request",
707 + block_ids: list[int],
708 + ) -> tuple[bool, dict[str, Any] | None]:
709 + """
710 + Called exactly once when a request has finished, before its blocks are
711 + freed.
712 +
713 + The connector may assumes responsibility for freeing the blocks
714 + asynchronously by returning True.
715 +
716 + Returns:
717 + True if the request is being saved/sent asynchronously and blocks
718 + should not be freed until the request_id is returned from
719 + get_finished().
720 + Optional KVTransferParams to be included in the request outputs
721 + returned by the engine.
722 + """
723 + return True, None
724 +
725 + def take_events(self) -> Iterable["KVCacheEvent"]:
726 + """
727 + Take the KV cache events from the connector.
728 +
729 + Yields:
730 + New KV cache events since the last call.
731 + """
732 + return ()
733 +
734 + @classmethod
735 + def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
736 + """
737 + Get the required KV cache layout for this connector.
738 + Args:
739 + vllm_config (VllmConfig): the vllm config.
740 +
741 + Returns:
742 + str: the required KV cache layout. e.g. HND, or NHD.
743 + None if the connector does not require a specific layout.
744 + """
745 +
746 + if cls is KVConnectorBase_V1:
747 + raise TypeError(
748 + "get_required_kvcache_layout should not be called "
749 + "on the abstract base class"
750 + )
751 + return None
752 +
753 + def get_finished_count(self) -> int | None:
754 + """
755 + Get the count of requests expected to complete send/receive operations
756 + via this connector. This method is used to initialize the
757 + KVOutputAggregator, overwriting the default world_size.
758 +
759 + Returns:
760 + int: expected sending or receiving completion count.
761 + """
762 + return None
763 +
764 + @classmethod
765 + def build_kv_connector_stats(
766 + cls, data: dict[str, Any] | None = None
767 + ) -> Optional["KVConnectorStats"]:
768 + """
769 + KVConnectorStats resolution method. This method allows dynamically
770 + registered connectors to return their own KVConnectorStats object,
771 + which can implement custom aggregation logic on the data dict.
772 + """
773 + return None
774 +
775 + @classmethod
776 + def build_prom_metrics(
777 + cls,
778 + vllm_config: "VllmConfig",
779 + metric_types: dict[type["PromMetric"], type["PromMetricT"]],
780 + labelnames: list[str],
781 + per_engine_labelvalues: dict[int, list[str]],
782 + ) -> Optional["KVConnectorPromMetrics"]:
783 + """
784 + Create a KVConnectorPromMetrics subclass which should register
785 + per-connector Prometheus metrics and implement observe() to
786 + expose connector transfer stats via Prometheus.
787 + """
788 + return None
789 +
790 + ##############################
791 + # Helper functions
792 + ##############################
793 + def _process_retrieve_requests(
794 + self,
795 + metadata: LMCacheMPConnectorMetadata,
796 + ) -> None:
797 + blocks_per_chunk = self.scheduler_adapter.num_blocks_per_chunk()
798 +
799 + for request_tracker in self.request_trackers.values():
800 + if request_tracker.state != LMCacheMPRequestState.WAITING_FOR_LOAD:
801 + continue
802 + r_metadata = LMCacheMPRequestMetadata.GetRetrieveMetadata(
803 + request_tracker, blocks_per_chunk
804 + )
805 + if r_metadata is not None:
806 + metadata.add_request_metadata(r_metadata)
807 + request_tracker.state = LMCacheMPRequestState.READY
808 +
809 + def _process_new_requests(
810 + self,
811 + scheduler_output: SchedulerOutput,
812 + metadata: LMCacheMPConnectorMetadata,
813 + ) -> None:
814 + blocks_per_chunk = self.scheduler_adapter.num_blocks_per_chunk()
815 +
816 + for new_request in scheduler_output.scheduled_new_reqs:
817 + request_tracker = self._get_request_tracker(new_request.req_id)
818 +
819 + num_new_tokens = scheduler_output.num_scheduled_tokens[new_request.req_id]
820 + request_tracker.increase_num_scheduled_tokens(num_new_tokens)
821 +
822 + r_meta = LMCacheMPRequestMetadata.GetStoreMetadata(
823 + request_tracker, blocks_per_chunk, self.vllm_block_size
824 + )
825 + if r_meta is not None:
826 + metadata.add_request_metadata(r_meta)
827 +
828 + def _process_cached_requests(
829 + self,
830 + scheduler_output: SchedulerOutput,
831 + metadata: LMCacheMPConnectorMetadata,
832 + ) -> None:
833 + blocks_per_chunk = self.scheduler_adapter.num_blocks_per_chunk()
834 +
835 + cached_reqs = scheduler_output.scheduled_cached_reqs
836 + for idx, request_id in enumerate(cached_reqs.req_ids):
837 + request_tracker = self._get_request_tracker(request_id)
838 +
839 + # Update block ids
840 + new_block_ids = reformat_block_ids(cached_reqs.new_block_ids[idx])
841 + request_tracker.update_block_ids(new_block_ids)
842 +
843 + # Update new scheduled tokens
844 + num_new_tokens = cached_reqs.num_computed_tokens[idx]
845 + request_tracker.increase_num_scheduled_tokens(num_new_tokens)
846 +
847 + r_meta = LMCacheMPRequestMetadata.GetStoreMetadata(
848 + request_tracker, blocks_per_chunk, self.vllm_block_size
849 + )
850 +
851 + if r_meta is not None:
852 + metadata.add_request_metadata(r_meta)
853 +
854 + def _get_request_tracker(self, request_id: str) -> LMCacheMPRequestTracker:
855 + assert request_id in self.request_trackers, (
856 + f"Request tracker for request_id {request_id} not found. "
857 + )
858 + return self.request_trackers[request_id]
859 +
860 + def _get_or_create_request_tracker(
861 + self, request: "Request"
862 + ) -> LMCacheMPRequestTracker:
863 + request_id = request.request_id
864 + if request_id not in self.request_trackers:
865 + new_tracker = LMCacheMPRequestTracker(request)
866 + self.request_trackers[request_id] = new_tracker
867 + return self.request_trackers[request_id]