@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e94ec59733

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

[LMCache] Token Base IPC API (#34175)

Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>

Yuwei An committed Feb 9, 2026 at 17:18 UTC e94ec597334d9a3e9b0d04bc17152e2747c83d51

2 files changed +376 -90

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py

+344 -73

@@ -20,16 +20,42 @@ from lmcache.v1.multiprocess.protocol import RequestType, get_response_class

20 logger = init_logger(__name__)

23 -def wrap_kv_caches(kv_caches: dict[str, KVCache]) -> KVCache:

23 +def wrap_kv_caches(kv_caches: dict[str, torch.Tensor]) -> KVCache:

24 logger.info("KV caches keys are %s", list(kv_caches.keys()))

25 return [CudaIPCWrapper(tensor) for tensor in kv_caches.values()]

28 +def striding_block_hashes(

29 + block_hashes: list[bytes], blocks_in_chunk: int

30 +) -> Iterable[bytes]:

31 + """Extract chunk-level hashes from block hashes by striding.

32 +

33 + In hash-based vLLM, each vLLM block has its own hash. LMCache chunks

34 + span ``blocks_in_chunk`` consecutive blocks. The representative hash

35 + for a chunk is the hash of the **last** block in that chunk (because

36 + each block hash already encodes its prefix). So we start at index

37 + ``blocks_in_chunk - 1`` and stride by ``blocks_in_chunk``.

38 + """

39 + return islice(block_hashes, blocks_in_chunk - 1, None, blocks_in_chunk)

40 +

41 +

42 def send_lmcache_request(

43 mq_client: MessageQueueClient,

44 request_type: RequestType,

45 payloads: list[Any],

46 ) -> MessagingFuture[Any]:

47 + """

48 + Helper function to send the request to the LMCache multiprocess server

49 +

50 + Args:

51 + mq_client: The LMCache multiprocess mode message queue client

52 + request_type: The request type

53 + payloads: The request payloads

54 +

55 + Returns:

56 + A messaging future for the request

57 + """

58 +

59 future = mq_client.submit_request(

60 request_type, payloads, get_response_class(request_type)

61 )

@@ -39,40 +65,44 @@ def send_lmcache_request(

65 def get_lmcache_chunk_size(

66 mq_client: MessageQueueClient,

67 ) -> int:

42 - future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, [])

43 - chunk_size = future.result()

44 - return chunk_size

68 + """

69 + Helper function to get the LMCache chunk size from the server

71 + Args:

72 + mq_client: The LMCache multiprocess mode message queue client

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.

74 + Returns:

75 + An integer representing the LMCache chunk size

76 """

55 - return islice(block_hashes, blocks_in_chunk - 1, None, blocks_in_chunk)

77 + future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, [])

78 + chunk_size = future.result()

79 + return chunk_size

82 @dataclass

83 class LoadStoreOp:

60 - block_hashes: list[bytes]

84 block_ids: list[int]

85 + """Block ids for the load/store operation"""

63 - def __len__(self) -> int:

64 - return len(self.block_hashes)

87 + token_ids: list[int] | None = None

88 + """Token IDs for the load/store operation (token mode)"""

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 - )

90 + block_hashes: list[bytes] | None = None

91 + """Block hashes for the load/store operation (hash mode)"""

92 +

93 + start: int = 0

94 + """Start token index (token mode only)"""

95 +

96 + end: int = 0

97 + """End token index (token mode only)"""

98 +

99 + def __len__(self) -> int:

100 + return len(self.block_ids)

101

102

103 StoreResult = bool

104 RetrieveResult = list[bool]

75 -LookupResult = list[bool]

105 +LookupResult = int

106

107

108 class LMCacheMPSchedulerAdapter:

@@ -95,10 +125,6 @@ class LMCacheMPSchedulerAdapter:

125 kv_rank: The kv rank used for LMCache keys

126 vllm_block_size: The block size used in vLLM

127 """

98 - logger.warning(

99 - "Importing LMCacheMPSchedulerAdapter is deprecated. "

100 - "Please update your LMCache to the latest version."

101 - )

128 self.mq_client = MessageQueueClient(server_url, context)

129

130 # Request futures

@@ -116,22 +142,89 @@ class LMCacheMPSchedulerAdapter:

142 self.blocks_in_chunk = self.chunk_size // vllm_block_size

143

144 @_lmcache_nvtx_annotate

119 - def maybe_submit_lookup_request(self, request_id: str, block_hashes: list[bytes]):

145 + def maybe_submit_lookup_request(

146 + self,

147 + request_id: str,

148 + block_hashes: list[bytes] | None = None,

149 + token_ids: list[int] | None = None,

150 + ) -> None:

151 + """

152 + Submit a new lookup request to LMCache if there is no ongoing request.

153 +

154 + Supports both token-based and hash-based vLLM:

155 + - token_ids: token IDs (token-based vLLM) -> single token-mode key

156 + - block_hashes: block hashes (hash-based vLLM) -> strided hash-mode keys

157 +

158 + Exactly one of block_hashes or token_ids must be provided.

159 +

160 + Args:

161 + request_id: The ID of the lookup request. The same ID indicates it's

162 + from the same request

163 + block_hashes: Block hashes to lookup from LMCache (hash mode)

164 + token_ids: Token IDs to lookup from LMCache (token mode)

165 +

166 + Returns:

167 + None

168 +

169 + Notes:

170 + This function will have a side-effect: submitting a look up request to

171 + LMCache, which will essentially 'lock' the KV cache chunks in the LMCache

172 + for later retrieve operations.

173 + In the meantime, this function will record the lookup request, and the

174 + status of the look up request can be checked by `check_lookup_result`.

175 + """

176 if request_id in self.lookup_futures:

177 # Skip if there is already a lookup request

178 return

179

124 - s = striding_block_hashes(block_hashes, self.blocks_in_chunk)

125 - keys = [self._create_key(block_hash) for block_hash in s]

180 + assert (block_hashes is None) != (token_ids is None), (

181 + "Exactly one of block_hashes or token_ids must be provided"

182 + )

183 +

184 + if block_hashes is not None:

185 + # Hash mode: stride block hashes -> N hash-mode keys

186 + chunk_hashes = list(

187 + striding_block_hashes(block_hashes, self.blocks_in_chunk)

188 + )

189 + keys = [

190 + self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes

191 + ]

192 + else:

193 + # Token mode: truncate to chunk-aligned length

194 + assert token_ids is not None

195 + aligned_end = (len(token_ids) // self.chunk_size) * self.chunk_size

196 + if aligned_end == 0:

197 + return

198 + keys = [

199 + self._create_key(

200 + token_ids,

201 + start=0,

202 + end=aligned_end,

203 + request_id=request_id,

204 + ).no_worker_id_version()

205 + ]

206 +

207 future = send_lmcache_request(

208 self.mq_client,

209 RequestType.LOOKUP,

129 - [keys, True],

210 + [keys],

211 )

212 self.lookup_futures[request_id] = future

213

214 @_lmcache_nvtx_annotate

215 def check_lookup_result(self, request_id: str) -> int | None:

216 + """

217 + Check the result of a previously submitted lookup request.

218 +

219 + Args:

220 + request_id: The ID of the lookup request submitted in

221 + `maybe_submit_lookup_request`

222 +

223 + Returns:

224 + An integer representing the total number of tokens matched

225 + in LMCache (prefix matching), or

226 + None if the lookup request is not finished yet.

227 + """

228 assert request_id in self.lookup_futures, (

229 f"Lookup request for request_id={request_id} has not been submitted"

230 )

@@ -141,7 +234,7 @@ class LMCacheMPSchedulerAdapter:

234 return None

235

236 result = future.result()

144 - num_chunks = sum(result)

237 + num_chunks = result

238 return num_chunks * self.chunk_size

239

240 def num_blocks_per_chunk(self) -> int:

@@ -159,14 +252,47 @@ class LMCacheMPSchedulerAdapter:

252 """

253 self.lookup_futures.pop(request_id, None)

254

255 + def end_session(self, request_id: str) -> None:

256 + """

257 + Notify LMCache server to remove the session for a finished request.

258 + Args:

259 + request_id: The ID of the finished request.

260 + """

261 + send_lmcache_request(

262 + self.mq_client,

263 + RequestType.END_SESSION,

264 + [request_id],

265 + )

266 +

267 # Helper functions

163 - def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:

164 - """Convert a block hash to an IPC cache engine key"""

268 + def _create_key(

269 + self,

270 + token_ids: list[int],

271 + start: int = 0,

272 + end: int = 0,

273 + request_id: str | None = None,

274 + ) -> IPCCacheEngineKey:

275 + """Convert token IDs to an IPC cache engine key"""

276 return IPCCacheEngineKey(

277 model_name=self.model_name,

278 world_size=self.world_size,

279 worker_id=self.worker_id,

169 - chunk_hash=block_hash,

280 + token_ids=tuple(token_ids),

281 + start=start,

282 + end=end,

283 + request_id=request_id,

284 + )

285 +

286 + def _create_hash_key(

287 + self, chunk_hash: bytes, request_id: str | None = None

288 + ) -> IPCCacheEngineKey:

289 + """Create a hash-mode IPC cache engine key"""

290 + return IPCCacheEngineKey(

291 + model_name=self.model_name,

292 + world_size=self.world_size,

293 + worker_id=None,

294 + chunk_hash=chunk_hash,

295 + request_id=request_id,

296 )

297

298

@@ -180,10 +306,6 @@ class LMCacheMPWorkerAdapter:

306 kv_rank: int,

307 vllm_block_size: int,

308 ):

183 - logger.warning(

184 - "Importing LMCacheMPWorkerAdapter is deprecated. "

185 - "Please update your LMCache to the latest version."

186 - )

309 self.mq_client = MessageQueueClient(server_url, context)

310

311 # Instance id for GPU worker

@@ -201,7 +323,10 @@ class LMCacheMPWorkerAdapter:

323 str, tuple[MessagingFuture[RetrieveResult], list[str]]

324 ] = {}

325

326 + # The store requests that have finished execution in LMCache

327 self.finished_stores: set[str] = set()

328 + # The finished request ids that are passed via vLLM and also

329 + # have corresponding store requests submitted to LMCache before

330 self.previously_finished: set[str] = set()

331

332 self.model_name = model_name

@@ -215,7 +340,14 @@ class LMCacheMPWorkerAdapter:

340 )

341 self.blocks_in_chunk = chunk_size // vllm_block_size

342

218 - def register_kv_caches(self, kv_caches: dict[str, KVCache]):

343 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

344 + """

345 + Register the kv caches with LMCache server

346 +

347 + Args:

348 + kv_caches: A dict of kv caches to register. The keys are the

349 + layer names and the values are the corresponding tensors.

350 + """

351 # Register kv cache and send the request

352 self.kv_caches = kv_caches

353 logger.info("Registering kv caches")

@@ -230,7 +362,29 @@ class LMCacheMPWorkerAdapter:

362 def submit_store_request(

363 self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event

364 ):

233 - keys = self._block_hashes_to_keys(op.block_hashes)

365 + """

366 + Submit a KV cache store request to LMCache

367 +

368 + Args:

369 + request_id: The ID of the request

370 + op: The LoadStoreOp describing the store operation.

371 + event: The CUDA event that is recorded after the current

372 + model inference step

373 + """

374 + if op.block_hashes is not None:

375 + # Hash mode

376 + chunk_hashes = list(

377 + striding_block_hashes(op.block_hashes, self.blocks_in_chunk)

378 + )

379 + keys = [

380 + self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes

381 + ]

382 + else:

383 + # Token mode

384 + assert op.token_ids is not None

385 + keys = [

386 + self._create_key(op.token_ids, op.start, op.end, request_id=request_id)

387 + ]

388 future = send_lmcache_request(

389 self.mq_client,

390 RequestType.STORE,

@@ -242,7 +396,29 @@ class LMCacheMPWorkerAdapter:

396 def submit_retrieve_request(

397 self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event

398 ):

245 - keys = self._block_hashes_to_keys(op.block_hashes)

399 + """

400 + Submit a KV cache retrieve request to LMCache

401 +

402 + Args:

403 + request_id: The ID of the request

404 + op: The LoadStoreOp describing the retrieve operation.

405 + event: The CUDA event that is recorded after the current

406 + model inference step

407 + """

408 + if op.block_hashes is not None:

409 + # Hash mode

410 + chunk_hashes = list(

411 + striding_block_hashes(op.block_hashes, self.blocks_in_chunk)

412 + )

413 + keys = [

414 + self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes

415 + ]

416 + else:

417 + # Token mode

418 + assert op.token_ids is not None

419 + keys = [

420 + self._create_key(op.token_ids, op.start, op.end, request_id=request_id)

421 + ]

422 future = send_lmcache_request(

423 self.mq_client,

424 RequestType.RETRIEVE,

@@ -257,17 +433,47 @@ class LMCacheMPWorkerAdapter:

433 ops: list[LoadStoreOp],

434 event: torch.cuda.Event,

435 ):

260 - keys = []

261 - block_ids = []

262 - for op in ops:

263 - keys.extend(self._block_hashes_to_keys(op.block_hashes))

436 + """

437 + Submit a batched store request to LMCache

438 +

439 + Args:

440 + request_ids: The IDs of the requests

441 + ops: The LoadStoreOps describing the store operations. Should have

442 + the same length as request_ids

443 + event: The CUDA event that is recorded after the current

444 + model inference step

445 + """

446 + all_keys: list[IPCCacheEngineKey] = []

447 + block_ids: list[int] = []

448 + for request_id, op in zip(request_ids, ops, strict=False):

449 + if op.block_hashes is not None:

450 + chunk_hashes = list(

451 + striding_block_hashes(op.block_hashes, self.blocks_in_chunk)

452 + )

453 + keys = [

454 + self._create_hash_key(ch, request_id=request_id)

455 + for ch in chunk_hashes

456 + ]

457 + all_keys.extend(keys)

458 + else:

459 + assert op.token_ids is not None

460 + all_keys.append(

461 + self._create_key(

462 + op.token_ids, op.start, op.end, request_id=request_id

463 + )

464 + )

465 block_ids.extend(op.block_ids)

466 future = send_lmcache_request(

467 self.mq_client,

468 RequestType.STORE,

268 - [keys, self.instance_id, block_ids, event.ipc_handle()],

469 + [

470 + all_keys,

471 + self.instance_id,

472 + block_ids,

473 + event.ipc_handle(),

474 + ],

475 ).to_cuda_future()

270 - self.store_futures[request_ids[0]] = (future, request_ids[1:])

476 + self.store_futures[request_ids[0]] = (future, list(request_ids[1:]))

477

478 @_lmcache_nvtx_annotate

479 def batched_submit_retrieve_requests(

@@ -276,34 +482,83 @@ class LMCacheMPWorkerAdapter:

482 ops: list[LoadStoreOp],

483 event: torch.cuda.Event,

484 ):

279 - keys = []

280 - block_ids = []

485 + """

486 + Submit a batched retrieve request to LMCache

487

282 - for op in ops:

283 - keys.extend(self._block_hashes_to_keys(op.block_hashes))

488 + Args:

489 + request_ids: The IDs of the requests

490 + ops: The LoadStoreOps describing the retrieve operations. Should have

491 + the same length as request_ids

492 + event: The CUDA event that is recorded after the current

493 + model inference step

494 + """

495 + all_keys: list[IPCCacheEngineKey] = []

496 + block_ids: list[int] = []

497 + for request_id, op in zip(request_ids, ops, strict=False):

498 + if op.block_hashes is not None:

499 + chunk_hashes = list(

500 + striding_block_hashes(op.block_hashes, self.blocks_in_chunk)

501 + )

502 + keys = [

503 + self._create_hash_key(ch, request_id=request_id)

504 + for ch in chunk_hashes

505 + ]

506 + all_keys.extend(keys)

507 + else:

508 + assert op.token_ids is not None

509 + all_keys.append(

510 + self._create_key(

511 + op.token_ids, op.start, op.end, request_id=request_id

512 + )

513 + )

514 block_ids.extend(op.block_ids)

515 future = send_lmcache_request(

516 self.mq_client,

517 RequestType.RETRIEVE,

288 - [keys, self.instance_id, block_ids, event.ipc_handle()],

518 + [

519 + all_keys,

520 + self.instance_id,

521 + block_ids,

522 + event.ipc_handle(),

523 + ],

524 ).to_cuda_future()

290 - self.retrieve_futures[request_ids[0]] = (future, request_ids[1:])

525 + self.retrieve_futures[request_ids[0]] = (future, list(request_ids[1:]))

526

527 @_lmcache_nvtx_annotate

528 def get_finished(

294 - self, finished_req_ids: set[str]

529 + self, finished_req_ids_from_engine: set[str]

530 ) -> tuple[set[str] | None, set[str] | None]:

531 + """

532 + Check and get the finished store and retrieve requests.

533 +

534 + Args:

535 + finished_req_ids_from_engine: the set of request ids that are

536 + reported as finished from the vLLM engine side.

537 +

538 + Returns:

539 + A tuple of two sets:

540 + - The first set contains the finished store request ids. The returned

541 + store request ids MUST be seen before in the

542 + `finished_req_ids_from_engine`.

543 + - The second set contains the finished retrieve request ids.

544 +

545 + Notes:

546 + When enabling async scheduling in vLLM, the same request ID may appear

547 + multiple times in `finished_req_ids_from_engine`. The adapter should

548 + take care of deduplicating the request IDs and only return the request

549 + IDs that have not been returned before.

550 + """

551 finished_stores = set()

552 finished_retrieves = set()

298 - for request_id, (future, other_reqs) in self.store_futures.items():

299 - if not future.query():

553 + for request_id, (s_future, other_reqs) in self.store_futures.items():

554 + if not s_future.query():

555 continue

556

302 - result = future.result()

557 + s_result = s_future.result()

558 finished_stores.add(request_id)

559 finished_stores.update(other_reqs)

560

306 - if not result:

561 + if not s_result:

562 # TODO: add error handling here

563 logger.error(

564 "Something went wrong when processing the "

@@ -311,21 +566,21 @@ class LMCacheMPWorkerAdapter:

566 request_id,

567 )

568

314 - for request_id, (future, other_reqs) in self.retrieve_futures.items():

315 - if not future.query():

569 + for request_id, (r_future, other_reqs) in self.retrieve_futures.items():

570 + if not r_future.query():

571 continue

572

318 - result = future.result()

573 + r_result = r_future.result()

574 finished_retrieves.add(request_id)

575 finished_retrieves.update(other_reqs)

576

322 - if not all(result):

577 + if not all(r_result):

578 # TODO: add error handing here

579 logger.error(

580 "Something went wrong when processing the "

581 "retrieve request for request_id=%s, result=%s",

582 request_id,

328 - result,

583 + r_result,

584 )

585

586 # Remove the finished requests from the tracking dicts

@@ -338,7 +593,7 @@ class LMCacheMPWorkerAdapter:

593 self.finished_stores.update(finished_stores)

594

595 ret_stores = set()

341 - for req_id in finished_req_ids:

596 + for req_id in finished_req_ids_from_engine:

597 if req_id in self.finished_stores or req_id in self.store_futures:

598 self.previously_finished.add(req_id)

599 else:

@@ -357,7 +612,9 @@ class LMCacheMPWorkerAdapter:

612 return self.blocks_in_chunk

613

614 def shutdown(self):

360 - # Unregister kv cache

615 + """

616 + Shutdown the LMCache MP worker adapter

617 + """

618 logger.info("Unregistering kv caches")

619 send_lmcache_request(

620 self.mq_client, RequestType.UNREGISTER_KV_CACHE, [self.instance_id]

@@ -378,18 +635,32 @@ class LMCacheMPWorkerAdapter:

635

636 return safe_finished_s

637

381 - def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:

382 - """Convert a block hash to an IPC cache engine key"""

638 + def _create_key(

639 + self,

640 + token_ids: list[int],

641 + start: int = 0,

642 + end: int = 0,

643 + request_id: str | None = None,

644 + ) -> IPCCacheEngineKey:

645 + """Convert token IDs to an IPC cache engine key"""

646 return IPCCacheEngineKey(

647 model_name=self.model_name,

648 world_size=self.world_size,

649 worker_id=self.worker_id,

387 - chunk_hash=block_hash,

650 + token_ids=tuple(token_ids),

651 + start=start,

652 + end=end,

653 + request_id=request_id,

654 )

655

390 - def _block_hashes_to_keys(

391 - self, block_hashes: list[bytes]

392 - ) -> list[IPCCacheEngineKey]:

393 - """Convert block hashes to IPC cache engine keys"""

394 - s = striding_block_hashes(block_hashes, self.blocks_in_chunk)

395 - return [self._create_key(block_hash) for block_hash in s]

656 + def _create_hash_key(

657 + self, chunk_hash: bytes, request_id: str | None = None

658 + ) -> IPCCacheEngineKey:

659 + """Create a hash-mode IPC cache engine key"""

660 + return IPCCacheEngineKey(

661 + model_name=self.model_name,

662 + world_size=self.world_size,

663 + worker_id=self.worker_id,

664 + chunk_hash=chunk_hash,

665 + request_id=request_id,

666 + )

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py

+32 -17

@@ -3,7 +3,7 @@

3 import enum

4 from collections.abc import Iterable

5 from dataclasses import dataclass, field

6 -from typing import TYPE_CHECKING, Any, Literal, cast

6 +from typing import TYPE_CHECKING, Any, Literal

8 import torch

9 import zmq

@@ -130,12 +130,6 @@ def create_worker_adapter(

130 )

131

132

133 -def convert_block_hashes_to_bytes(

134 - block_hashes: list["BlockHash"],

135 -) -> list[bytes]:

136 - return cast(list[bytes], block_hashes)

137 -

138 -

133 class LMCacheMPRequestState(enum.Enum):

134 """

135 State machine:

@@ -266,6 +260,7 @@ class LMCacheMPRequestMetadata:

260 Args:

261 tracker: The request tracker to generate the metadata from.

262 blocks_in_chunk: the number of blocks in a LMCache data chunk

263 + vllm_block_size: the block size used in vLLM

264 """

265 # Store the blocks that has block hashes

266 # NOTE: the invariant here is that `num_stored_blocks` should

@@ -282,15 +277,21 @@ class LMCacheMPRequestMetadata:

277 if num_chunks >= 1:

278 start = tracker.num_stored_blocks

279 end = start + num_chunks * blocks_in_chunk

285 - block_hashes = convert_block_hashes_to_bytes(

286 - tracker.block_hashes[start:end]

287 - )

280 block_ids = tracker.allocated_block_ids[start:end]

281 + start_token_idx = start * vllm_block_size

282 + end_token_idx = end * vllm_block_size

283 + token_ids = list(tracker.all_token_ids)

284 + op = LoadStoreOp(

285 + token_ids=token_ids,

286 + block_ids=block_ids,

287 + start=start_token_idx,

288 + end=end_token_idx,

289 + )

290

291 ret = LMCacheMPRequestMetadata(

292 request_id=tracker.request_id,

293 direction="STORE",

293 - op=LoadStoreOp(block_hashes=block_hashes, block_ids=block_ids),

294 + op=op,

295 )

296

297 # Update the request tracker

@@ -303,6 +304,7 @@ class LMCacheMPRequestMetadata:

304 def GetRetrieveMetadata(

305 tracker: LMCacheMPRequestTracker,

306 blocks_in_chunk: int,

307 + vllm_block_size: int,

308 ) -> "LMCacheMPRequestMetadata | None":

309 """

310 Generate the retrieve metadata for the current request tracker.

@@ -310,6 +312,7 @@ class LMCacheMPRequestMetadata:

312 Args:

313 tracker: The request tracker to generate the metadata from.

314 blocks_in_chunk: the number of blocks in a LMCache data chunk

315 + vllm_block_size: the block size used in vLLM

316 """

317 if not tracker.is_ready_for_retrieving():

318 return None

@@ -330,15 +333,21 @@ class LMCacheMPRequestMetadata:

333 "number of LMCache hit blocks. "

334 )

335 if end > start:

333 - block_hashes = convert_block_hashes_to_bytes(

334 - tracker.block_hashes[start:end]

335 - )

336 block_ids = tracker.allocated_block_ids[start:end]

337 + start_token_idx = start * vllm_block_size

338 + end_token_idx = end * vllm_block_size

339 + token_ids = list(tracker.all_token_ids)

340 + op = LoadStoreOp(

341 + token_ids=token_ids,

342 + block_ids=block_ids,

343 + start=start_token_idx,

344 + end=end_token_idx,

345 + )

346

347 ret = LMCacheMPRequestMetadata(

348 request_id=tracker.request_id,

349 direction="RETRIEVE",

341 - op=LoadStoreOp(block_hashes=block_hashes, block_ids=block_ids),

350 + op=op,

351 )

352 return ret

353

@@ -643,7 +652,8 @@ class LMCacheMPConnector(KVConnectorBase_V1):

652 return 0, False

653

654 self.scheduler_adapter.maybe_submit_lookup_request(

646 - request.request_id, convert_block_hashes_to_bytes(request.block_hashes)

655 + request.request_id,

656 + token_ids=list(request.all_token_ids),

657 )

658

659 ret = self.scheduler_adapter.check_lookup_result(request.request_id)

@@ -766,6 +776,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):

776 """

777 # Clean up request tracker to prevent memory leak

778 self._cleanup_request_tracker(request.request_id)

779 + # Notify LMCache to end the session for this request

780 + self.scheduler_adapter.end_session(request.request_id)

781 +

782 return True, None

783

784 def take_events(self) -> Iterable["KVCacheEvent"]:

@@ -846,7 +859,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):

859 if request_tracker.state != LMCacheMPRequestState.WAITING_FOR_LOAD:

860 continue

861 r_metadata = LMCacheMPRequestMetadata.GetRetrieveMetadata(

849 - request_tracker, blocks_per_chunk

862 + request_tracker,

863 + blocks_per_chunk,

864 + vllm_block_size=self.vllm_block_size,

865 )

866 if r_metadata is not None:

867 metadata.add_request_metadata(r_metadata)