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

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

[Prefix Cache] Include lora_name in BlockStored event for deterministic KV-cache reconstruction (#27577)

Signed-off-by: Sage Ahrac <sagiahrak@gmail.com> Co-authored-by: Sage <80211083+sagiahrac@users.noreply.github.com>

Sage committed Dec 30, 2025 at 02:17 UTC 39512aba72537d364e0352141cf8c0523115ec09

8 files changed +49

examples/online_serving/kv_events_subscriber.py

@@ -28,8 +28,14 @@ class BlockStored(KVCacheEvent):

28 parent_block_hash: ExternalBlockHash | None

29 token_ids: list[int]

30 block_size: int

31 +

32 lora_id: int | None

33 + """Deprecated: use `lora_name` for KV block key hash.

34 + Retained for backward compatibility.

35 + """

36 +

37 medium: str | None

38 + lora_name: str | None

41 class BlockRemoved(KVCacheEvent):

tests/v1/engine/test_engine_core_client.py

@@ -733,6 +733,7 @@ def test_kv_cache_events(

733 )

734 assert event.parent_block_hash is None, "Parent block hash should be None"

735 assert event.lora_id is None, "Lora id should be None"

736 + assert event.lora_name is None, "Lora name should be None"

737 assert len(event.token_ids) == num_blocks * block_size, (

738 "Token ids should be the same as the custom tokens"

739 )

tests/v1/kv_connector/unit/test_lmcache_connector.py

+30

@@ -25,6 +25,7 @@ def mock_lmcache_engine_event():

25 lora_id,

26 block_size,

27 medium,

28 + lora_name,

29 ):

30 self.block_hashes = block_hashes

31 self.parent_block_hash = parent_block_hash

@@ -32,6 +33,7 @@ def mock_lmcache_engine_event():

33 self.lora_id = lora_id

34 self.block_size = block_size

35 self.medium = medium

36 + self.lora_name = lora_name

38 return MockEvent(

39 block_hashes=["hash1", "hash2"],

@@ -40,6 +42,7 @@ def mock_lmcache_engine_event():

42 lora_id=None,

43 block_size=16,

44 medium="GPU",

45 + lora_name=None,

46 )

@@ -109,6 +112,7 @@ class TestGetKVConnectorKVCacheEvents:

112 assert events[0].lora_id is None

113 assert events[0].block_size == 16

114 assert events[0].medium == "GPU"

115 + assert events[0].lora_name is None

116

117 def test_converts_multiple_events(self, mock_connector):

118 """Test conversion of multiple events from lmcache engine format."""

@@ -121,6 +125,7 @@ class TestGetKVConnectorKVCacheEvents:

125 self.lora_id = None

126 self.block_size = 16

127 self.medium = "GPU"

128 + self.lora_name = None

129

130 events = [MockEvent(i) for i in range(5)]

131 mock_connector._lmcache_engine.get_kv_events.return_value = events

@@ -150,6 +155,7 @@ class TestGetKVConnectorKVCacheEvents:

155 self.lora_id = 42

156 self.block_size = 32

157 self.medium = "DISK"

158 + self.lora_name = "lora_example"

159

160 mock_connector._lmcache_engine.get_kv_events.return_value = [

161 MockEventWithLora()

@@ -166,6 +172,7 @@ class TestGetKVConnectorKVCacheEvents:

172 assert event.lora_id == 42

173 assert event.block_size == 32

174 assert event.medium == "DISK"

175 + assert event.lora_name == "lora_example"

176

177 def test_handles_none_parent_block_hash(self, mock_connector):

178 """Test handling of events with None parent_block_hash."""

@@ -178,6 +185,7 @@ class TestGetKVConnectorKVCacheEvents:

185 self.lora_id = None

186 self.block_size = 16

187 self.medium = "GPU"

188 + self.lora_name = None

189

190 mock_connector._lmcache_engine.get_kv_events.return_value = [

191 MockEventNoParent()

@@ -223,6 +231,7 @@ class TestUpdateConnectorOutput:

231 block_size=16,

232 lora_id=None,

233 medium="GPU",

234 + lora_name=None,

235 )

236 kv_events.add_events([event])

237

@@ -243,6 +252,7 @@ class TestUpdateConnectorOutput:

252 block_size=16,

253 lora_id=None,

254 medium="GPU",

255 + lora_name=None,

256 )

257 existing_events.add_events([event1])

258 existing_events.add_events([event1]) # Simulate 2 workers reporting

@@ -258,6 +268,7 @@ class TestUpdateConnectorOutput:

268 block_size=16,

269 lora_id=None,

270 medium="GPU",

271 + lora_name=None,

272 )

273 new_events.add_events([event2])

274

@@ -288,6 +299,7 @@ class TestUpdateConnectorOutput:

299 block_size=16,

300 lora_id=None,

301 medium="GPU",

302 + lora_name=None,

303 )

304 new_events.add_events([event])

305

@@ -309,6 +321,7 @@ class TestUpdateConnectorOutput:

321 block_size=16,

322 lora_id=None,

323 medium="GPU",

324 + lora_name=None,

325 )

326 events1.add_events([event1])

327 output1 = KVConnectorOutput(kv_cache_events=events1)

@@ -323,6 +336,7 @@ class TestUpdateConnectorOutput:

336 block_size=16,

337 lora_id=None,

338 medium="GPU",

339 + lora_name=None,

340 )

341 events2.add_events([event2])

342 output2 = KVConnectorOutput(kv_cache_events=events2)

@@ -337,6 +351,7 @@ class TestUpdateConnectorOutput:

351 block_size=16,

352 lora_id=None,

353 medium="GPU",

354 + lora_name=None,

355 )

356 events3.add_events([event3])

357 output3 = KVConnectorOutput(kv_cache_events=events3)

@@ -358,6 +373,7 @@ class TestUpdateConnectorOutput:

373 block_size=16,

374 lora_id=None,

375 medium="GPU",

376 + lora_name=None,

377 )

378 events1.add_events([event1])

379 output1 = KVConnectorOutput(kv_cache_events=events1)

@@ -397,6 +413,7 @@ class TestTakeEvents:

413 block_size=16,

414 lora_id=None,

415 medium="GPU",

416 + lora_name=None,

417 )

418 event2 = BlockStored(

419 block_hashes=["hash2"],

@@ -405,6 +422,7 @@ class TestTakeEvents:

422 block_size=16,

423 lora_id=None,

424 medium="GPU",

425 + lora_name=None,

426 )

427 kv_events.add_events([event1, event2])

428 mock_connector._kv_cache_events = kv_events

@@ -431,6 +449,7 @@ class TestTakeEvents:

449 block_size=16,

450 lora_id=None,

451 medium="GPU",

452 + lora_name=None,

453 )

454 uncommon_event = BlockStored(

455 block_hashes=["hash_uncommon"],

@@ -439,6 +458,7 @@ class TestTakeEvents:

458 block_size=16,

459 lora_id=None,

460 medium="GPU",

461 + lora_name=None,

462 )

463

464 # All 3 workers report common_event

@@ -469,6 +489,7 @@ class TestTakeEvents:

489 block_size=16,

490 lora_id=None,

491 medium="GPU",

492 + lora_name=None,

493 )

494 kv_events1.add_events([event1])

495 mock_connector._kv_cache_events = kv_events1

@@ -491,6 +512,7 @@ class TestTakeEvents:

512 block_size=16,

513 lora_id=None,

514 medium="GPU",

515 + lora_name=None,

516 )

517 kv_events2.add_events([event2])

518 mock_connector._kv_cache_events = kv_events2

@@ -510,6 +532,7 @@ class TestTakeEvents:

532 block_size=16,

533 lora_id=None,

534 medium="GPU",

535 + lora_name=None,

536 )

537 event2 = BlockStored(

538 block_hashes=["hash2"],

@@ -518,6 +541,7 @@ class TestTakeEvents:

541 block_size=16,

542 lora_id=None,

543 medium="GPU",

544 + lora_name=None,

545 )

546

547 # Worker 1 reports event1

@@ -572,6 +596,7 @@ class TestIntegrationScenarios:

596 self.lora_id = None

597 self.block_size = 16

598 self.medium = "GPU"

599 + self.lora_name = None

600

601 # Worker 1

602 mock_connector._lmcache_engine.get_kv_events.return_value = [

@@ -628,6 +653,7 @@ class TestIntegrationScenarios:

653 self.lora_id = None

654 self.block_size = 16

655 self.medium = "GPU"

656 + self.lora_name = None

657

658 for cycle in range(3):

659 # Get events

@@ -667,6 +693,7 @@ class TestIntegrationScenarios:

693 block_size=16,

694 lora_id=None,

695 medium="GPU",

696 + lora_name=None,

697 )

698

699 worker1_unique_event = BlockStored(

@@ -676,6 +703,7 @@ class TestIntegrationScenarios:

703 block_size=16,

704 lora_id=None,

705 medium="GPU",

706 + lora_name=None,

707 )

708

709 worker2_unique_event = BlockStored(

@@ -685,6 +713,7 @@ class TestIntegrationScenarios:

713 block_size=16,

714 lora_id=None,

715 medium="GPU",

716 + lora_name=None,

717 )

718

719 worker3_unique_event = BlockStored(

@@ -694,6 +723,7 @@ class TestIntegrationScenarios:

723 block_size=16,

724 lora_id=None,

725 medium="GPU",

726 + lora_name=None,

727 )

728

729 # Create events for each worker

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -528,6 +528,7 @@ def test_offloading_connector(request_runner):

528 assert event.token_ids == []

529 assert event.parent_block_hash is None

530 assert event.lora_id is None

531 + assert event.lora_name is None

532 event = events[1]

533 assert isinstance(event, BlockRemoved)

534 assert event.block_hashes == to_hashes([4, 5, 6])

vllm/distributed/kv_events.py

@@ -51,8 +51,14 @@ class BlockStored(KVCacheEvent):

51 parent_block_hash: ExternalBlockHash | None

52 token_ids: list[int]

53 block_size: int

54 +

55 lora_id: int | None

56 + """Deprecated: use `lora_name` for KV block key hash.

57 + Retained for backward compatibility.

58 + """

59 +

60 medium: str | None

61 + lora_name: str | None

63 def __hash__(self) -> int:

64 return hash(

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

@@ -234,6 +234,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):

234 lora_id=e.lora_id,

235 block_size=e.block_size,

236 medium=e.medium,

237 + lora_name=e.lora_name,

238 )

239 for e in events

240 ]

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

@@ -406,6 +406,7 @@ class OffloadingConnectorScheduler:

406 lora_id=None,

407 block_size=event.block_size,

408 medium=event.medium,

409 + lora_name=None,

410 )

411

412

vllm/v1/core/block_pool.py

@@ -290,6 +290,9 @@ class BlockPool:

290 if request.lora_request

291 else None,

292 medium=MEDIUM_GPU,

293 + lora_name=request.lora_request.name

294 + if request.lora_request

295 + else None,

296 )

297 )

298