@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f4417f8449
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 11, 15:00 UTC

[KVConnector] Add KV events to KV Connectors (#28309)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>

Martin Hickey committed Dec 11, 2025 at 14:30 UTC f4417f8449dc7a2cb890dbef659c0d1ce93432da
8 files changed +1036 -15
tests/v1/kv_connector/unit/test_lmcache_connector.py new
+756
@@ -0,0 +1,756 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from unittest.mock import MagicMock
4 +
5 +import pytest
6 +
7 +from vllm.distributed.kv_events import BlockStored
8 +from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import (
9 + LMCacheConnectorV1,
10 + LMCacheKVEvents,
11 +)
12 +from vllm.v1.outputs import KVConnectorOutput
13 +
14 +
15 +@pytest.fixture
16 +def mock_lmcache_engine_event():
17 + """Create a mock event object that mimics what the lmcache engine returns."""
18 +
19 + class MockEvent:
20 + def __init__(
21 + self,
22 + block_hashes,
23 + parent_block_hash,
24 + token_ids,
25 + lora_id,
26 + block_size,
27 + medium,
28 + ):
29 + self.block_hashes = block_hashes
30 + self.parent_block_hash = parent_block_hash
31 + self.token_ids = token_ids
32 + self.lora_id = lora_id
33 + self.block_size = block_size
34 + self.medium = medium
35 +
36 + return MockEvent(
37 + block_hashes=["hash1", "hash2"],
38 + parent_block_hash="parent_hash",
39 + token_ids=[1, 2, 3, 4],
40 + lora_id=None,
41 + block_size=16,
42 + medium="GPU",
43 + )
44 +
45 +
46 +@pytest.fixture
47 +def mock_connector():
48 + """Create a mock LMCacheConnectorV1 instance with mocked dependencies."""
49 + connector = MagicMock(spec=LMCacheConnectorV1)
50 + connector._kv_cache_events = None
51 + connector._lmcache_engine = MagicMock()
52 +
53 + # Make the methods use the real implementation
54 + connector.get_kv_connector_kv_cache_events = (
55 + LMCacheConnectorV1.get_kv_connector_kv_cache_events.__get__(
56 + connector, LMCacheConnectorV1
57 + )
58 + )
59 + connector.update_connector_output = (
60 + LMCacheConnectorV1.update_connector_output.__get__(
61 + connector, LMCacheConnectorV1
62 + )
63 + )
64 + connector.take_events = LMCacheConnectorV1.take_events.__get__(
65 + connector, LMCacheConnectorV1
66 + )
67 +
68 + return connector
69 +
70 +
71 +class TestGetKVConnectorKVCacheEvents:
72 + """Test get_kv_connector_kv_cache_events method."""
73 +
74 + def test_returns_none_when_no_events(self, mock_connector):
75 + """Test that None is returned when lmcache engine has no events."""
76 + mock_connector._lmcache_engine.get_kv_events.return_value = None
77 +
78 + result = mock_connector.get_kv_connector_kv_cache_events()
79 +
80 + assert result is None
81 + mock_connector._lmcache_engine.get_kv_events.assert_called_once()
82 +
83 + def test_returns_none_when_empty_list(self, mock_connector):
84 + """Test that None is returned when lmcache engine returns empty list."""
85 + mock_connector._lmcache_engine.get_kv_events.return_value = []
86 +
87 + result = mock_connector.get_kv_connector_kv_cache_events()
88 +
89 + assert result is None
90 +
91 + def test_converts_single_event(self, mock_connector, mock_lmcache_engine_event):
92 + """Test conversion of a single event from lmcache engine format."""
93 + mock_connector._lmcache_engine.get_kv_events.return_value = [
94 + mock_lmcache_engine_event
95 + ]
96 +
97 + result = mock_connector.get_kv_connector_kv_cache_events()
98 +
99 + assert result is not None
100 + assert isinstance(result, LMCacheKVEvents)
101 + assert result.get_number_of_workers() == 1
102 +
103 + events = result.get_all_events()
104 + assert len(events) == 1
105 + assert isinstance(events[0], BlockStored)
106 + assert events[0].block_hashes == ["hash1", "hash2"]
107 + assert events[0].parent_block_hash == "parent_hash"
108 + assert events[0].token_ids == [1, 2, 3, 4]
109 + assert events[0].lora_id is None
110 + assert events[0].block_size == 16
111 + assert events[0].medium == "GPU"
112 +
113 + def test_converts_multiple_events(self, mock_connector):
114 + """Test conversion of multiple events from lmcache engine format."""
115 +
116 + class MockEvent:
117 + def __init__(self, i):
118 + self.block_hashes = [f"hash{i}"]
119 + self.parent_block_hash = f"parent{i}"
120 + self.token_ids = [i]
121 + self.lora_id = None
122 + self.block_size = 16
123 + self.medium = "GPU"
124 +
125 + events = [MockEvent(i) for i in range(5)]
126 + mock_connector._lmcache_engine.get_kv_events.return_value = events
127 +
128 + result = mock_connector.get_kv_connector_kv_cache_events()
129 +
130 + assert result is not None
131 + assert isinstance(result, LMCacheKVEvents)
132 +
133 + converted_events = result.get_all_events()
134 + assert len(converted_events) == 5
135 +
136 + for i, event in enumerate(converted_events):
137 + assert isinstance(event, BlockStored)
138 + assert event.block_hashes == [f"hash{i}"]
139 + assert event.parent_block_hash == f"parent{i}"
140 + assert event.token_ids == [i]
141 +
142 + def test_preserves_event_attributes(self, mock_connector):
143 + """Test that all event attributes are correctly preserved."""
144 +
145 + class MockEventWithLora:
146 + def __init__(self):
147 + self.block_hashes = ["hash_a", "hash_b", "hash_c"]
148 + self.parent_block_hash = "parent_xyz"
149 + self.token_ids = [100, 200, 300]
150 + self.lora_id = 42
151 + self.block_size = 32
152 + self.medium = "DISK"
153 +
154 + mock_connector._lmcache_engine.get_kv_events.return_value = [
155 + MockEventWithLora()
156 + ]
157 +
158 + result = mock_connector.get_kv_connector_kv_cache_events()
159 +
160 + events = result.get_all_events()
161 + event = events[0]
162 +
163 + assert event.block_hashes == ["hash_a", "hash_b", "hash_c"]
164 + assert event.parent_block_hash == "parent_xyz"
165 + assert event.token_ids == [100, 200, 300]
166 + assert event.lora_id == 42
167 + assert event.block_size == 32
168 + assert event.medium == "DISK"
169 +
170 + def test_handles_none_parent_block_hash(self, mock_connector):
171 + """Test handling of events with None parent_block_hash."""
172 +
173 + class MockEventNoParent:
174 + def __init__(self):
175 + self.block_hashes = ["hash1"]
176 + self.parent_block_hash = None
177 + self.token_ids = [1, 2]
178 + self.lora_id = None
179 + self.block_size = 16
180 + self.medium = "GPU"
181 +
182 + mock_connector._lmcache_engine.get_kv_events.return_value = [
183 + MockEventNoParent()
184 + ]
185 +
186 + result = mock_connector.get_kv_connector_kv_cache_events()
187 +
188 + events = result.get_all_events()
189 + assert events[0].parent_block_hash is None
190 +
191 +
192 +class TestUpdateConnectorOutput:
193 + """Test update_connector_output method."""
194 +
195 + def test_does_nothing_when_kv_cache_events_is_none(self, mock_connector):
196 + """Test that method returns early when kv_cache_events is None."""
197 + connector_output = KVConnectorOutput(kv_cache_events=None)
198 +
199 + mock_connector.update_connector_output(connector_output)
200 +
201 + assert mock_connector._kv_cache_events is None
202 +
203 + def test_does_nothing_when_kv_cache_events_is_not_lmcache_kv_events(
204 + self, mock_connector
205 + ):
206 + """Test that method returns early when kv_cache_events is not
207 + LMCacheKVEvents."""
208 + # Create a mock object that is not LMCacheKVEvents
209 + fake_events = MagicMock()
210 + connector_output = KVConnectorOutput(kv_cache_events=fake_events)
211 +
212 + mock_connector.update_connector_output(connector_output)
213 +
214 + assert mock_connector._kv_cache_events is None
215 +
216 + def test_sets_kv_cache_events_when_none(self, mock_connector):
217 + """Test that _kv_cache_events is set when it was None."""
218 + kv_events = LMCacheKVEvents(num_workers=1)
219 + event = BlockStored(
220 + block_hashes=["hash1"],
221 + parent_block_hash=None,
222 + token_ids=[1, 2],
223 + block_size=16,
224 + lora_id=None,
225 + medium="GPU",
226 + )
227 + kv_events.add_events([event])
228 +
229 + connector_output = KVConnectorOutput(kv_cache_events=kv_events)
230 +
231 + mock_connector.update_connector_output(connector_output)
232 +
233 + assert mock_connector._kv_cache_events is kv_events
234 +
235 + def test_adds_events_when_kv_cache_events_already_exists(self, mock_connector):
236 + """Test that events are added when _kv_cache_events already exists."""
237 + # Set up existing events
238 + existing_events = LMCacheKVEvents(num_workers=2)
239 + event1 = BlockStored(
240 + block_hashes=["hash1"],
241 + parent_block_hash=None,
242 + token_ids=[1],
243 + block_size=16,
244 + lora_id=None,
245 + medium="GPU",
246 + )
247 + existing_events.add_events([event1])
248 + existing_events.add_events([event1]) # Simulate 2 workers reporting
249 +
250 + mock_connector._kv_cache_events = existing_events
251 +
252 + # Create new events to add
253 + new_events = LMCacheKVEvents(num_workers=1)
254 + event2 = BlockStored(
255 + block_hashes=["hash2"],
256 + parent_block_hash=None,
257 + token_ids=[2],
258 + block_size=16,
259 + lora_id=None,
260 + medium="GPU",
261 + )
262 + new_events.add_events([event2])
263 +
264 + connector_output = KVConnectorOutput(kv_cache_events=new_events)
265 +
266 + mock_connector.update_connector_output(connector_output)
267 +
268 + # Check that events were added
269 + all_events = mock_connector._kv_cache_events.get_all_events()
270 + assert len(all_events) == 3 # 2 from existing + 1 from new
271 + assert event1 in all_events
272 + assert event2 in all_events
273 +
274 + def test_increments_workers_when_kv_cache_events_already_exists(
275 + self, mock_connector
276 + ):
277 + """Test that worker count is incremented correctly."""
278 + # Set up existing events with 2 workers
279 + existing_events = LMCacheKVEvents(num_workers=2)
280 + mock_connector._kv_cache_events = existing_events
281 +
282 + # Create new events from 3 workers
283 + new_events = LMCacheKVEvents(num_workers=3)
284 + event = BlockStored(
285 + block_hashes=["hash1"],
286 + parent_block_hash=None,
287 + token_ids=[1],
288 + block_size=16,
289 + lora_id=None,
290 + medium="GPU",
291 + )
292 + new_events.add_events([event])
293 +
294 + connector_output = KVConnectorOutput(kv_cache_events=new_events)
295 +
296 + mock_connector.update_connector_output(connector_output)
297 +
298 + # Worker count should be 2 + 3 = 5
299 + assert mock_connector._kv_cache_events.get_number_of_workers() == 5
300 +
301 + def test_multiple_updates(self, mock_connector):
302 + """Test multiple consecutive updates."""
303 + # First update
304 + events1 = LMCacheKVEvents(num_workers=1)
305 + event1 = BlockStored(
306 + block_hashes=["hash1"],
307 + parent_block_hash=None,
308 + token_ids=[1],
309 + block_size=16,
310 + lora_id=None,
311 + medium="GPU",
312 + )
313 + events1.add_events([event1])
314 + output1 = KVConnectorOutput(kv_cache_events=events1)
315 + mock_connector.update_connector_output(output1)
316 +
317 + # Second update
318 + events2 = LMCacheKVEvents(num_workers=2)
319 + event2 = BlockStored(
320 + block_hashes=["hash2"],
321 + parent_block_hash=None,
322 + token_ids=[2],
323 + block_size=16,
324 + lora_id=None,
325 + medium="GPU",
326 + )
327 + events2.add_events([event2])
328 + output2 = KVConnectorOutput(kv_cache_events=events2)
329 + mock_connector.update_connector_output(output2)
330 +
331 + # Third update
332 + events3 = LMCacheKVEvents(num_workers=1)
333 + event3 = BlockStored(
334 + block_hashes=["hash3"],
335 + parent_block_hash=None,
336 + token_ids=[3],
337 + block_size=16,
338 + lora_id=None,
339 + medium="GPU",
340 + )
341 + events3.add_events([event3])
342 + output3 = KVConnectorOutput(kv_cache_events=events3)
343 + mock_connector.update_connector_output(output3)
344 +
345 + # Check final state
346 + all_events = mock_connector._kv_cache_events.get_all_events()
347 + assert len(all_events) == 3
348 + assert mock_connector._kv_cache_events.get_number_of_workers() == 4 # 1+2+1
349 +
350 + def test_updates_with_empty_events(self, mock_connector):
351 + """Test updating with empty event lists."""
352 + # First update with actual events
353 + events1 = LMCacheKVEvents(num_workers=1)
354 + event1 = BlockStored(
355 + block_hashes=["hash1"],
356 + parent_block_hash=None,
357 + token_ids=[1],
358 + block_size=16,
359 + lora_id=None,
360 + medium="GPU",
361 + )
362 + events1.add_events([event1])
363 + output1 = KVConnectorOutput(kv_cache_events=events1)
364 + mock_connector.update_connector_output(output1)
365 +
366 + # Second update with empty events
367 + events2 = LMCacheKVEvents(num_workers=2)
368 + # No events added
369 + output2 = KVConnectorOutput(kv_cache_events=events2)
370 + mock_connector.update_connector_output(output2)
371 +
372 + # Should still have the original event
373 + all_events = mock_connector._kv_cache_events.get_all_events()
374 + assert len(all_events) == 1
375 + assert mock_connector._kv_cache_events.get_number_of_workers() == 3
376 +
377 +
378 +class TestTakeEvents:
379 + """Test take_events method."""
380 +
381 + def test_yields_nothing_when_kv_cache_events_is_none(self, mock_connector):
382 + """Test that nothing is yielded when _kv_cache_events is None."""
383 + mock_connector._kv_cache_events = None
384 +
385 + events = list(mock_connector.take_events())
386 +
387 + assert events == []
388 +
389 + def test_yields_events_and_clears(self, mock_connector):
390 + """Test that events are yielded and then cleared."""
391 + # Set up events
392 + kv_events = LMCacheKVEvents(num_workers=1)
393 + event1 = BlockStored(
394 + block_hashes=["hash1"],
395 + parent_block_hash=None,
396 + token_ids=[1],
397 + block_size=16,
398 + lora_id=None,
399 + medium="GPU",
400 + )
401 + event2 = BlockStored(
402 + block_hashes=["hash2"],
403 + parent_block_hash=None,
404 + token_ids=[2],
405 + block_size=16,
406 + lora_id=None,
407 + medium="GPU",
408 + )
409 + kv_events.add_events([event1, event2])
410 + mock_connector._kv_cache_events = kv_events
411 +
412 + # Take events
413 + events = list(mock_connector.take_events())
414 +
415 + # Check that events were yielded
416 + assert len(events) == 2
417 + assert event1 in events
418 + assert event2 in events
419 +
420 + # Check that _kv_cache_events was cleared
421 + assert mock_connector._kv_cache_events is None
422 +
423 + def test_aggregates_before_yielding(self, mock_connector):
424 + """Test that events are aggregated before yielding."""
425 + # Set up events from multiple workers
426 + kv_events = LMCacheKVEvents(num_workers=3)
427 + common_event = BlockStored(
428 + block_hashes=["hash_common"],
429 + parent_block_hash=None,
430 + token_ids=[1],
431 + block_size=16,
432 + lora_id=None,
433 + medium="GPU",
434 + )
435 + uncommon_event = BlockStored(
436 + block_hashes=["hash_uncommon"],
437 + parent_block_hash=None,
438 + token_ids=[2],
439 + block_size=16,
440 + lora_id=None,
441 + medium="GPU",
442 + )
443 +
444 + # All 3 workers report common_event
445 + kv_events.add_events([common_event])
446 + kv_events.add_events([common_event])
447 + kv_events.add_events([common_event])
448 +
449 + # Only 1 worker reports uncommon_event
450 + kv_events.add_events([uncommon_event])
451 +
452 + mock_connector._kv_cache_events = kv_events
453 +
454 + # Take events
455 + events = list(mock_connector.take_events())
456 +
457 + # Only the common event should be yielded
458 + assert len(events) == 1
459 + assert events[0] == common_event
460 +
461 + def test_multiple_take_events_calls(self, mock_connector):
462 + """Test calling take_events multiple times."""
463 + # First call with events
464 + kv_events1 = LMCacheKVEvents(num_workers=1)
465 + event1 = BlockStored(
466 + block_hashes=["hash1"],
467 + parent_block_hash=None,
468 + token_ids=[1],
469 + block_size=16,
470 + lora_id=None,
471 + medium="GPU",
472 + )
473 + kv_events1.add_events([event1])
474 + mock_connector._kv_cache_events = kv_events1
475 +
476 + events1 = list(mock_connector.take_events())
477 + assert len(events1) == 1
478 + assert events1[0] == event1
479 + assert mock_connector._kv_cache_events is None
480 +
481 + # Second call with no events
482 + events2 = list(mock_connector.take_events())
483 + assert events2 == []
484 +
485 + # Third call after adding new events
486 + kv_events2 = LMCacheKVEvents(num_workers=1)
487 + event2 = BlockStored(
488 + block_hashes=["hash2"],
489 + parent_block_hash=None,
490 + token_ids=[2],
491 + block_size=16,
492 + lora_id=None,
493 + medium="GPU",
494 + )
495 + kv_events2.add_events([event2])
496 + mock_connector._kv_cache_events = kv_events2
497 +
498 + events3 = list(mock_connector.take_events())
499 + assert len(events3) == 1
500 + assert events3[0] == event2
501 +
502 + def test_yields_empty_after_aggregation_removes_all(self, mock_connector):
503 + """Test that nothing is yielded if aggregation removes all events."""
504 + # Set up events from 2 workers with no common events
505 + kv_events = LMCacheKVEvents(num_workers=2)
506 + event1 = BlockStored(
507 + block_hashes=["hash1"],
508 + parent_block_hash=None,
509 + token_ids=[1],
510 + block_size=16,
511 + lora_id=None,
512 + medium="GPU",
513 + )
514 + event2 = BlockStored(
515 + block_hashes=["hash2"],
516 + parent_block_hash=None,
517 + token_ids=[2],
518 + block_size=16,
519 + lora_id=None,
520 + medium="GPU",
521 + )
522 +
523 + # Worker 1 reports event1
524 + kv_events.add_events([event1])
525 + # Worker 2 reports event2
526 + kv_events.add_events([event2])
527 +
528 + mock_connector._kv_cache_events = kv_events
529 +
530 + # Take events
531 + events = list(mock_connector.take_events())
532 +
533 + # No common events, so nothing should be yielded
534 + assert events == []
535 + assert mock_connector._kv_cache_events is None
536 +
537 +
538 +class TestIntegrationScenarios:
539 + """Test integration scenarios."""
540 +
541 + def test_full_workflow(self, mock_connector, mock_lmcache_engine_event):
542 + """Test a complete workflow from getting events to taking them."""
543 + # Step 1: Get events from lmcache engine
544 + mock_connector._lmcache_engine.get_kv_events.return_value = [
545 + mock_lmcache_engine_event
546 + ]
547 + kv_events = mock_connector.get_kv_connector_kv_cache_events()
548 +
549 + assert kv_events is not None
550 + assert len(kv_events.get_all_events()) == 1
551 +
552 + # Step 2: Update connector output (simulate receiving from worker)
553 + output1 = KVConnectorOutput(kv_cache_events=kv_events)
554 + mock_connector.update_connector_output(output1)
555 +
556 + assert mock_connector._kv_cache_events is not None
557 +
558 + # Step 3: Take events
559 + taken_events = list(mock_connector.take_events())
560 +
561 + assert len(taken_events) == 1
562 + assert mock_connector._kv_cache_events is None
563 +
564 + def test_multiple_workers_workflow(self, mock_connector):
565 + """Test workflow with multiple workers."""
566 +
567 + class MockEvent:
568 + def __init__(self, hash_val):
569 + self.block_hashes = [hash_val]
570 + self.parent_block_hash = None
571 + self.token_ids = [1]
572 + self.lora_id = None
573 + self.block_size = 16
574 + self.medium = "GPU"
575 +
576 + # Worker 1
577 + mock_connector._lmcache_engine.get_kv_events.return_value = [
578 + MockEvent("hash_common"),
579 + MockEvent("hash_worker1"),
580 + ]
581 + kv_events1 = mock_connector.get_kv_connector_kv_cache_events()
582 + output1 = KVConnectorOutput(kv_cache_events=kv_events1)
583 + mock_connector.update_connector_output(output1)
584 +
585 + # Worker 2
586 + mock_connector._lmcache_engine.get_kv_events.return_value = [
587 + MockEvent("hash_common"),
588 + MockEvent("hash_worker2"),
589 + ]
590 + kv_events2 = mock_connector.get_kv_connector_kv_cache_events()
591 + output2 = KVConnectorOutput(kv_cache_events=kv_events2)
592 + mock_connector.update_connector_output(output2)
593 +
594 + # Take events (should only get common events)
595 + taken_events = list(mock_connector.take_events())
596 +
597 + # With aggregation, only events reported by both workers should be present
598 + # In this case, hash_common was reported by both
599 + event_hashes = [e.block_hashes[0] for e in taken_events]
600 + assert "hash_common" in event_hashes
601 +
602 + def test_empty_workflow(self, mock_connector):
603 + """Test workflow when there are no events at any stage."""
604 + # Get events returns None
605 + mock_connector._lmcache_engine.get_kv_events.return_value = None
606 + kv_events = mock_connector.get_kv_connector_kv_cache_events()
607 +
608 + assert kv_events is None
609 +
610 + # Update with None
611 + output = KVConnectorOutput(kv_cache_events=None)
612 + mock_connector.update_connector_output(output)
613 +
614 + # Take events
615 + taken_events = list(mock_connector.take_events())
616 +
617 + assert taken_events == []
618 + assert mock_connector._kv_cache_events is None
619 +
620 + def test_repeated_cycles(self, mock_connector):
621 + """Test multiple cycles of the complete workflow."""
622 +
623 + class MockEvent:
624 + def __init__(self, cycle_num):
625 + self.block_hashes = [f"hash_cycle_{cycle_num}"]
626 + self.parent_block_hash = None
627 + self.token_ids = [cycle_num]
628 + self.lora_id = None
629 + self.block_size = 16
630 + self.medium = "GPU"
631 +
632 + for cycle in range(3):
633 + # Get events
634 + mock_connector._lmcache_engine.get_kv_events.return_value = [
635 + MockEvent(cycle)
636 + ]
637 + kv_events = mock_connector.get_kv_connector_kv_cache_events()
638 +
639 + # Update
640 + output = KVConnectorOutput(kv_cache_events=kv_events)
641 + mock_connector.update_connector_output(output)
642 +
643 + # Take
644 + taken_events = list(mock_connector.take_events())
645 +
646 + # Verify
647 + assert len(taken_events) == 1
648 + assert taken_events[0].block_hashes[0] == f"hash_cycle_{cycle}"
649 + assert mock_connector._kv_cache_events is None
650 +
651 + def test_lmcache_kv_events_aggregation(self):
652 + """
653 + Test LMCacheKVEvents aggregation across TP ranks using
654 + KVOutputAggregator (used by MultiprocExecutor).
655 + """
656 + from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
657 + from vllm.v1.outputs import ModelRunnerOutput
658 +
659 + # Create KVOutputAggregator for 3 workers (simulating TP=3)
660 + aggregator = KVOutputAggregator(expected_finished_count=3)
661 +
662 + # Define common and unique events
663 + common_event = BlockStored(
664 + block_hashes=["hash_common"],
665 + parent_block_hash="parent_common",
666 + token_ids=[1, 2, 3],
667 + block_size=16,
668 + lora_id=None,
669 + medium="GPU",
670 + )
671 +
672 + worker1_unique_event = BlockStored(
673 + block_hashes=["hash_worker1"],
674 + parent_block_hash="parent_w1",
675 + token_ids=[4, 5],
676 + block_size=16,
677 + lora_id=None,
678 + medium="GPU",
679 + )
680 +
681 + worker2_unique_event = BlockStored(
682 + block_hashes=["hash_worker2"],
683 + parent_block_hash="parent_w2",
684 + token_ids=[6, 7],
685 + block_size=16,
686 + lora_id=None,
687 + medium="GPU",
688 + )
689 +
690 + worker3_unique_event = BlockStored(
691 + block_hashes=["hash_worker3"],
692 + parent_block_hash="parent_w3",
693 + token_ids=[8, 9],
694 + block_size=16,
695 + lora_id=None,
696 + medium="GPU",
697 + )
698 +
699 + # Create events for each worker
700 + # Worker 0: reports common event and its unique event
701 + worker0_events = LMCacheKVEvents(num_workers=1)
702 + worker0_events.add_events([common_event, worker1_unique_event])
703 +
704 + # Worker 1: reports common event and its unique event
705 + worker1_events = LMCacheKVEvents(num_workers=1)
706 + worker1_events.add_events([common_event, worker2_unique_event])
707 +
708 + # Worker 2: reports common event and its unique event
709 + worker2_events = LMCacheKVEvents(num_workers=1)
710 + worker2_events.add_events([common_event, worker3_unique_event])
711 +
712 + # Create ModelRunnerOutput instances for each worker
713 + worker_outputs = []
714 + for i, worker_events in enumerate(
715 + [worker0_events, worker1_events, worker2_events]
716 + ):
717 + output = ModelRunnerOutput(
718 + req_ids=[f"req_{i}"],
719 + req_id_to_index={f"req_{i}": 0},
720 + sampled_token_ids=[[123]], # dummy token
721 + logprobs=None,
722 + prompt_logprobs_dict={},
723 + pooler_output=[None],
724 + kv_connector_output=KVConnectorOutput(
725 + finished_sending=set([f"req_{i}_send"])
726 + if i < 2
727 + else None, # Workers 0,1 finished sending
728 + finished_recving=set([f"req_{i}_recv"])
729 + if i > 0
730 + else None, # Workers 1,2 finished receiving
731 + kv_cache_events=worker_events,
732 + ),
733 + )
734 + worker_outputs.append(output)
735 +
736 + # Use the real aggregation mechanism (like MultiprocExecutor.execute_model)
737 + aggregated_output = aggregator.aggregate(worker_outputs, output_rank=0)
738 + kv_cache_events = aggregated_output.kv_connector_output.kv_cache_events
739 +
740 + assert isinstance(kv_cache_events, LMCacheKVEvents)
741 +
742 + # After aggregation, events should be combined from all workers
743 + # The aggregator doesn't automatically aggregate events, so we need to call
744 + # aggregate() to get only common events
745 + kv_cache_events.aggregate()
746 + aggregated_events = kv_cache_events.get_all_events()
747 +
748 + # Only the common event should remain after aggregation
749 + # because it's the only event reported by all 3 workers
750 + assert len(aggregated_events) == 1
751 + assert aggregated_events[0] == common_event
752 +
753 + # Verify the common event properties
754 + assert aggregated_events[0].block_hashes == ["hash_common"]
755 + assert aggregated_events[0].parent_block_hash == "parent_common"
756 + assert aggregated_events[0].token_ids == [1, 2, 3]
vllm/distributed/kv_events.py
+129 -1
@@ -5,7 +5,7 @@ import queue
5 import threading
6 import time
7 from abc import ABC, abstractmethod
8 -from collections import deque
8 +from collections import Counter, deque
9 from collections.abc import Callable
10 from dataclasses import asdict
11 from itertools import count
@@ -54,11 +54,26 @@ class BlockStored(KVCacheEvent):
54 lora_id: int | None
55 medium: str | None
56
57 + def __hash__(self) -> int:
58 + return hash(
59 + (
60 + tuple(self.block_hashes),
61 + self.parent_block_hash,
62 + tuple(self.token_ids),
63 + self.block_size,
64 + self.lora_id,
65 + self.medium,
66 + )
67 + )
68 +
69
70 class BlockRemoved(KVCacheEvent):
71 block_hashes: list[ExternalBlockHash]
72 medium: str | None
73
74 + def __hash__(self) -> int:
75 + return hash((tuple(self.block_hashes), self.medium))
76 +
77
78 class AllBlocksCleared(KVCacheEvent):
79 pass
@@ -68,6 +83,119 @@ class KVEventBatch(EventBatch):
83 events: list[BlockStored | BlockRemoved | AllBlocksCleared]
84
85
86 +class KVEventAggregator:
87 + """
88 + Aggregates KV events across multiple workers.
89 + Tracks how many times each event appears and returns only those
90 + that were emitted by all workers.
91 + """
92 +
93 + __slots__ = ("_event_counter", "_num_workers")
94 +
95 + def __init__(self, num_workers: int) -> None:
96 + if num_workers <= 0:
97 + raise ValueError("num_workers must be greater than zero.")
98 + self._event_counter: Counter[KVCacheEvent] = Counter()
99 + self._num_workers: int = num_workers
100 +
101 + def add_events(self, events: list[KVCacheEvent]) -> None:
102 + """
103 + Add events from a worker batch.
104 +
105 + :param events: List of KVCacheEvent objects.
106 + """
107 + if not isinstance(events, list):
108 + raise TypeError("events must be a list of KVCacheEvent.")
109 + self._event_counter.update(events)
110 +
111 + def get_common_events(self) -> list[KVCacheEvent]:
112 + """
113 + Return events that appeared in all workers.
114 +
115 + :return: List of events present in all workers.
116 + """
117 + return [
118 + event
119 + for event, count in self._event_counter.items()
120 + if count == self._num_workers
121 + ]
122 +
123 + def get_all_events(self) -> list[KVCacheEvent]:
124 + """
125 + Return all events for all workers.
126 +
127 + :return: List of events for all workers.
128 + """
129 + return list(self._event_counter.elements())
130 +
131 + def clear_events(self) -> None:
132 + """
133 + Clear all tracked events.
134 + """
135 + self._event_counter.clear()
136 +
137 + def increment_workers(self, count: int = 1) -> None:
138 + """
139 + Increment the number of workers contributing events.
140 +
141 + :param count: Number to increment the workers by.
142 + """
143 + if count <= 0:
144 + raise ValueError("count must be positive.")
145 + self._num_workers += count
146 +
147 + def reset_workers(self) -> None:
148 + """
149 + Reset the number of workers to 1.
150 + """
151 + self._num_workers = 1
152 +
153 + def get_number_of_workers(self) -> int:
154 + """
155 + Return the number of workers.
156 +
157 + :return: int number of workers.
158 + """
159 + return self._num_workers
160 +
161 + def __repr__(self) -> str:
162 + return (
163 + f"<KVEventAggregator workers={self._num_workers}, "
164 + f"events={len(self._event_counter)}>"
165 + )
166 +
167 +
168 +class KVConnectorKVEvents(ABC):
169 + """
170 + Abstract base class for KV events.
171 + Acts as a container for KV events from the connector.
172 + """
173 +
174 + @abstractmethod
175 + def add_events(self, events: list[KVCacheEvent]) -> None:
176 + raise NotImplementedError
177 +
178 + @abstractmethod
179 + def aggregate(self) -> "KVConnectorKVEvents":
180 + raise NotImplementedError
181 +
182 + @abstractmethod
183 + def increment_workers(self, count: int = 1) -> None:
184 + raise NotImplementedError
185 +
186 + @abstractmethod
187 + def get_all_events(self) -> list[KVCacheEvent]:
188 + raise NotImplementedError
189 +
190 + @abstractmethod
191 + def get_number_of_workers(self) -> int:
192 + raise NotImplementedError
193 +
194 + @abstractmethod
195 + def clear_events(self) -> None:
196 + raise NotImplementedError
197 +
198 +
199 class EventPublisher(ABC):
200 """Lightweight publisher for EventBatch batches with data parallelism
201 support.
vllm/distributed/kv_transfer/kv_connector/utils.py
+15
@@ -78,6 +78,7 @@ class KVOutputAggregator:
78 finished_sending = set[str]()
79 finished_recving = set[str]()
80 aggregated_kv_connector_stats = None
81 + combined_kv_cache_events = None
82 invalid_block_ids = set[int]()
83 for model_runner_output in outputs:
84 assert model_runner_output is not None
@@ -119,6 +120,19 @@ class KVOutputAggregator:
120 aggregated_kv_connector_stats.aggregate(kv_connector_stats)
121 )
122
123 + # Combine kv_cache_events from all workers.
124 + if combined_kv_cache_events is None:
125 + # Use the first worker's kv_cache events as start event list.
126 + combined_kv_cache_events = kv_output.kv_cache_events
127 + elif kv_cache_events := kv_output.kv_cache_events:
128 + assert isinstance(
129 + combined_kv_cache_events,
130 + type(kv_cache_events),
131 + )
132 + worker_kv_cache_events = kv_cache_events.get_all_events()
133 + combined_kv_cache_events.add_events(worker_kv_cache_events)
134 + combined_kv_cache_events.increment_workers(1)
135 +
136 invalid_block_ids |= kv_output.invalid_block_ids
137
138 # select output of the worker specified by output_rank
@@ -129,6 +143,7 @@ class KVOutputAggregator:
143 finished_sending=finished_sending or None,
144 finished_recving=finished_recving or None,
145 kv_connector_stats=aggregated_kv_connector_stats or None,
146 + kv_cache_events=combined_kv_cache_events or None,
147 invalid_block_ids=invalid_block_ids,
148 expected_finished_count=self._expected_finished_count,
149 )
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+9 -1
@@ -49,7 +49,7 @@ from vllm.v1.outputs import KVConnectorOutput
49
50 if TYPE_CHECKING:
51 from vllm.config import VllmConfig
52 - from vllm.distributed.kv_events import KVCacheEvent
52 + from vllm.distributed.kv_events import KVCacheEvent, KVConnectorKVEvents
53 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
54 KVConnectorPromMetrics,
55 KVConnectorStats,
@@ -379,6 +379,14 @@ class KVConnectorBase_V1(ABC):
379 """
380 return None
381
382 + def get_kv_connector_kv_cache_events(self) -> Optional["KVConnectorKVEvents"]:
383 + """
384 + Get the KV connector kv cache events collected during the last interval.
385 + This function should be called by the model runner every time after the
386 + model execution and before cleanup.
387 + """
388 + return None
389 +
390 def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
391 """
392 Get the KVConnector handshake metadata for this connector.
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+114 -3
@@ -1,14 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable
4 from typing import TYPE_CHECKING, Any
5
6 import torch
6 -from lmcache.integration.vllm.vllm_v1_adapter import (
7 - LMCacheConnectorV1Impl as LMCacheConnectorLatestImpl,
8 -)
7
8 from vllm.attention.backends.abstract import AttentionMetadata
9 from vllm.config import VllmConfig
10 +from vllm.distributed.kv_events import (
11 + BlockStored,
12 + KVCacheEvent,
13 + KVConnectorKVEvents,
14 + KVEventAggregator,
15 +)
16 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17 KVConnectorBase_V1,
18 KVConnectorMetadata,
@@ -16,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
20 )
21 from vllm.logger import init_logger
22 from vllm.v1.core.sched.output import SchedulerOutput
23 +from vllm.v1.outputs import KVConnectorOutput
24
25 if TYPE_CHECKING:
26 from vllm.forward_context import ForwardContext
@@ -26,6 +31,44 @@ if TYPE_CHECKING:
31 logger = init_logger(__name__)
32
33
34 +class LMCacheKVEvents(KVConnectorKVEvents):
35 + """
36 + Concrete implementation of KVConnectorKVEvents using KVEventAggregator.
37 + """
38 +
39 + def __init__(self, num_workers: int) -> None:
40 + self._aggregator = KVEventAggregator(num_workers)
41 +
42 + def add_events(self, events: list[KVCacheEvent]) -> None:
43 + self._aggregator.add_events(events)
44 +
45 + def aggregate(self) -> "LMCacheKVEvents":
46 + """
47 + Aggregate KV events and retain only common events.
48 + """
49 + common_events = self._aggregator.get_common_events()
50 + self._aggregator.clear_events()
51 + self._aggregator.add_events(common_events)
52 + self._aggregator.reset_workers()
53 + return self
54 +
55 + def increment_workers(self, count: int = 1) -> None:
56 + self._aggregator.increment_workers(count)
57 +
58 + def get_all_events(self) -> list[KVCacheEvent]:
59 + return self._aggregator.get_all_events()
60 +
61 + def get_number_of_workers(self) -> int:
62 + return self._aggregator.get_number_of_workers()
63 +
64 + def clear_events(self) -> None:
65 + self._aggregator.clear_events()
66 + self._aggregator.reset_workers()
67 +
68 + def __repr__(self) -> str:
69 + return f"<LMCacheKVEvents events={self.get_all_events()}>"
70 +
71 +
72 class LMCacheConnectorV1(KVConnectorBase_V1):
73 def __init__(
74 self,
@@ -50,10 +93,17 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
93 cls = _adapter.LMCacheConnectorV1Impl
94 else:
95 logger.info("Initializing latest dev LMCache connector")
96 + # lazy import
97 + from lmcache.integration.vllm.vllm_v1_adapter import (
98 + LMCacheConnectorV1Impl as LMCacheConnectorLatestImpl,
99 + )
100 +
101 cls = LMCacheConnectorLatestImpl
102
103 self._lmcache_engine = cls(vllm_config, role, self)
104
105 + self._kv_cache_events: LMCacheKVEvents | None = None
106 +
107 # ==============================
108 # Worker-side methods
109 # ==============================
@@ -151,6 +201,31 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
201 # Fallback for older versions that don't support this method
202 return set()
203
204 + def get_kv_connector_kv_cache_events(self) -> LMCacheKVEvents | None:
205 + """
206 + Get the KV connector kv cache events collected during the last interval.
207 + """
208 +
209 + events = self._lmcache_engine.get_kv_events() # type: ignore [attr-defined]
210 + if not events:
211 + return None
212 +
213 + blocks: list[BlockStored] = [
214 + BlockStored(
215 + block_hashes=e.block_hashes,
216 + parent_block_hash=e.parent_block_hash,
217 + token_ids=e.token_ids,
218 + lora_id=e.lora_id,
219 + block_size=e.block_size,
220 + medium=e.medium,
221 + )
222 + for e in events
223 + ]
224 +
225 + lmcache_kv_events = LMCacheKVEvents(num_workers=1)
226 + lmcache_kv_events.add_events(blocks)
227 + return lmcache_kv_events
228 +
229 # ==============================
230 # Scheduler-side methods
231 # ==============================
@@ -198,6 +273,28 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
273 """
274 return self._lmcache_engine.build_connector_meta(scheduler_output)
275
276 + def update_connector_output(self, connector_output: KVConnectorOutput):
277 + """
278 + Update KVConnector state from worker-side connectors output.
279 +
280 + Args:
281 + connector_output (KVConnectorOutput): the worker-side
282 + connectors output.
283 + """
284 + # Get the KV events
285 + kv_cache_events = connector_output.kv_cache_events
286 + if not kv_cache_events or not isinstance(kv_cache_events, LMCacheKVEvents):
287 + return
288 +
289 + if self._kv_cache_events is None:
290 + self._kv_cache_events = kv_cache_events
291 + else:
292 + self._kv_cache_events.add_events(kv_cache_events.get_all_events())
293 + self._kv_cache_events.increment_workers(
294 + kv_cache_events.get_number_of_workers()
295 + )
296 + return
297 +
298 def request_finished(
299 self,
300 request: "Request",
@@ -214,3 +311,17 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
311 returned by the engine.
312 """
313 return self._lmcache_engine.request_finished(request, block_ids)
314 +
315 + def take_events(self) -> Iterable["KVCacheEvent"]:
316 + """
317 + Take the KV cache events from the connector.
318 +
319 + Yields:
320 + New KV cache events since the last call.
321 + """
322 + if self._kv_cache_events is not None:
323 + self._kv_cache_events.aggregate()
324 + kv_cache_events = self._kv_cache_events.get_all_events()
325 + yield from kv_cache_events
326 + self._kv_cache_events.clear_events()
327 + self._kv_cache_events = None
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+6
@@ -259,6 +259,12 @@ class MultiConnector(KVConnectorBase_V1):
259 agg_block_ids |= c.get_block_ids_with_load_errors()
260 return agg_block_ids
261
262 + # TODO: Add a generic implementation of 'get_kv_connector_kv_cache_events' method
263 + # for the MultiConnector. It should be able to get events from multiple
264 + # connectors, handling the case where only a subset of the requested connectors
265 + # implements the 'get_kv_connector_kv_cache_events'
266 + # Follow on PR from https://github.com/vllm-project/vllm/pull/28309#pullrequestreview-3566351082
267 +
268 # ==============================
269 # Scheduler-side methods
270 # ==============================
vllm/v1/outputs.py
+4
@@ -12,9 +12,11 @@ from vllm.compilation.cuda_graph import CUDAGraphStat
12 from vllm.v1.core.sched.output import SchedulerOutput
13
14 if TYPE_CHECKING:
15 + from vllm.distributed.kv_events import KVConnectorKVEvents
16 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
17 else:
18 KVConnectorStats = object
19 + KVConnectorKVEvents = object
20
21
22 class LogprobsLists(NamedTuple):
@@ -108,6 +110,7 @@ class KVConnectorOutput:
110 finished_sending: set[str] | None = None
111 finished_recving: set[str] | None = None
112 kv_connector_stats: KVConnectorStats | None = None
113 + kv_cache_events: KVConnectorKVEvents | None = None
114 # IDs of externally computed KV blocks that failed to load.
115 # Requests referencing these blocks should be rescheduled to recompute them
116 invalid_block_ids: set[int] = field(default_factory=set)
@@ -123,6 +126,7 @@ class KVConnectorOutput:
126 not self.finished_sending
127 and not self.finished_recving
128 and not self.kv_connector_stats
129 + and not self.kv_cache_events
130 and not self.invalid_block_ids
131 )
132
vllm/v1/worker/kv_connector_model_runner_mixin.py
+3 -10
@@ -22,7 +22,6 @@ from vllm.distributed.kv_transfer import (
22 has_kv_transfer_group,
23 )
24 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
25 -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
25 from vllm.forward_context import get_forward_context, set_forward_context
26 from vllm.logger import init_logger
27 from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
@@ -138,16 +137,10 @@ class KVConnectorModelRunnerMixin:
137 )
138 output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors()
139
141 - output.kv_connector_stats = (
142 - KVConnectorModelRunnerMixin.get_kv_connector_stats()
143 - )
144 - kv_connector.clear_connector_metadata()
140 + output.kv_connector_stats = kv_connector.get_kv_connector_stats()
141 + output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events()
142
146 - @staticmethod
147 - def get_kv_connector_stats() -> KVConnectorStats | None:
148 - if has_kv_transfer_group():
149 - return get_kv_transfer_group().get_kv_connector_stats()
150 - return None
143 + kv_connector.clear_connector_metadata()
144
145 @staticmethod
146 def use_uniform_kv_cache(