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 6, 20:30 UTC
[KVConnector] Support worker -> scheduler metadata (#31964)
Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Or Ozeri committed
Mar 11, 2026 at 19:36 UTC
a1a3523a5647a58e00096ca7430e9f1ad4a50a97
6 files changed
+283
-29
tests/v1/kv_connector/unit/test_multi_connector.py
+176
-25
@@ -5,21 +5,27 @@ import shutil
5
import tempfile
6
from pathlib import Path
7
from typing import Any
8
+from unittest.mock import MagicMock
9
10
import pytest
11
12
+from tests.v1.kv_connector.unit.utils import create_vllm_config
13
from vllm import LLM, SamplingParams
14
from vllm.config import KVTransferConfig
15
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
16
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
17
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
18
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
19
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
20
MultiConnector,
21
MultiKVConnectorStats,
22
+ MultiKVConnectorWorkerMetadata,
23
)
24
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
25
NixlKVConnectorStats,
26
)
27
+from vllm.v1.kv_cache_interface import KVCacheConfig
28
+from vllm.v1.outputs import KVConnectorOutput, KVConnectorWorkerMetadata
29
30
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
31
@@ -40,7 +46,14 @@ class MockConnectorStats(KVConnectorStats):
46
47
48
class MockConnector(KVConnectorBase_V1):
43
- """Mock connector that implements build_kv_connector_stats for testing."""
49
+ """Mock connector for testing."""
50
+
51
+ def __new__(cls, *args, **kwargs):
52
+ # mock all KVConnectorBase_V1 functions
53
+ mock = MagicMock(spec_set=KVConnectorBase_V1)
54
+ # Override just build_kv_connector_stats
55
+ mock.build_kv_connector_stats = cls.build_kv_connector_stats
56
+ return mock
57
58
@classmethod
59
def build_kv_connector_stats(
@@ -70,16 +83,42 @@ class MockConnector(KVConnectorBase_V1):
83
pass
84
85
73
-class MockCrossLayerConnector(MockConnector):
74
- @property
75
- def prefer_cross_layer_blocks(self) -> bool:
76
- return True
77
-
78
-
86
# Register the mock connector
87
KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
88
89
90
+@pytest.fixture
91
+def mc() -> MultiConnector:
92
+ """MultiConnector using two mocked connectors"""
93
+ vllm_config = create_vllm_config()
94
+
95
+ mock_connector_config = {
96
+ "kv_connector": "MockConnector",
97
+ "kv_role": "kv_both",
98
+ "kv_connector_module_path": "tests.v1.kv_connector.unit.test_multi_connector",
99
+ }
100
+
101
+ vllm_config.kv_transfer_config = KVTransferConfig(
102
+ kv_connector="MultiConnector",
103
+ kv_role="kv_both",
104
+ kv_connector_extra_config={
105
+ "connectors": [mock_connector_config, mock_connector_config],
106
+ },
107
+ )
108
+
109
+ kv_cache_config = KVCacheConfig(
110
+ num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
111
+ )
112
+
113
+ mc = MultiConnector(
114
+ vllm_config=vllm_config,
115
+ role=KVConnectorRole.WORKER,
116
+ kv_cache_config=kv_cache_config,
117
+ )
118
+
119
+ return mc
120
+
121
+
122
# Helper function to compare directories recursively
123
def _compare_directories(dir1: Path, dir2: Path) -> bool:
124
"""Compares two directories recursively for identical content."""
@@ -715,24 +754,6 @@ class TestMultiConnectorStats:
754
assert not stats.is_empty()
755
756
718
-class TestMultiConnectorPreferCrossLayerBlocks:
719
- def test_all_connectors_prefer_cross_layer_blocks(self):
720
- mc = MultiConnector.__new__(MultiConnector)
721
- mc._connectors = [
722
- MockCrossLayerConnector.__new__(MockCrossLayerConnector),
723
- MockCrossLayerConnector.__new__(MockCrossLayerConnector),
724
- ]
725
- assert mc.prefer_cross_layer_blocks is True
726
-
727
- def test_mixed_connectors_do_not_prefer_cross_layer_blocks(self):
728
- mc = MultiConnector.__new__(MultiConnector)
729
- mc._connectors = [
730
- MockCrossLayerConnector.__new__(MockCrossLayerConnector),
731
- MockConnector.__new__(MockConnector), # default False
732
- ]
733
- assert mc.prefer_cross_layer_blocks is False
734
-
735
-
757
def test_multi_connector_overrides_all_base_methods():
758
"""
759
Ensure MultiConnector overrides all public methods from KVConnectorBase_V1.
@@ -767,3 +788,133 @@ Options:
788
1. Add delegation in MultiConnector (preferred)
789
2. Add to INHERITED_OK if the base implementation works correctly
790
""")
791
+
792
+
793
+def test_multi_connector_prefer_cross_layer_blocks(mc):
794
+ mc._connectors[0].prefer_cross_layer_blocks = False
795
+ mc._connectors[1].prefer_cross_layer_blocks = True
796
+ assert mc.prefer_cross_layer_blocks is False
797
+
798
+ mc._connectors[0].prefer_cross_layer_blocks = True
799
+ mc._connectors[1].prefer_cross_layer_blocks = True
800
+ assert mc.prefer_cross_layer_blocks is True
801
+
802
+
803
+def test_multi_connector_worker_metadata(mc):
804
+ class MockConnectorWorkerMetadata(KVConnectorWorkerMetadata):
805
+ def __init__(self, data: set[str]):
806
+ self.data = data
807
+
808
+ class MockConnectorWorkerMetadata0(MockConnectorWorkerMetadata):
809
+ def aggregate(
810
+ self, other: KVConnectorWorkerMetadata
811
+ ) -> KVConnectorWorkerMetadata:
812
+ assert isinstance(other, MockConnectorWorkerMetadata)
813
+ return MockConnectorWorkerMetadata0(data=self.data | other.data)
814
+
815
+ class MockConnectorWorkerMetadata1(MockConnectorWorkerMetadata):
816
+ def aggregate(
817
+ self, other: KVConnectorWorkerMetadata
818
+ ) -> KVConnectorWorkerMetadata:
819
+ assert isinstance(other, MockConnectorWorkerMetadata)
820
+ return MockConnectorWorkerMetadata1(data=self.data | other.data)
821
+
822
+ # -------------------- test build_worker_connector_meta -------------------
823
+
824
+ # both connectors return None
825
+ mc._connectors[0].build_connector_worker_meta.return_value = None
826
+ mc._connectors[1].build_connector_worker_meta.return_value = None
827
+ assert mc.build_connector_worker_meta() is None
828
+
829
+ # only first connector returns None
830
+ worker_meta1a = MockConnectorWorkerMetadata1({"1a"})
831
+ mc._connectors[0].build_connector_worker_meta.return_value = None
832
+ mc._connectors[1].build_connector_worker_meta.return_value = worker_meta1a
833
+ mc_worker_meta_none_1a = mc.build_connector_worker_meta()
834
+ assert isinstance(mc_worker_meta_none_1a, MultiKVConnectorWorkerMetadata)
835
+ assert mc_worker_meta_none_1a.metadata == (None, worker_meta1a)
836
+
837
+ # only second connector returns None
838
+ worker_meta0a = MockConnectorWorkerMetadata0({"0a"})
839
+ mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0a
840
+ mc._connectors[1].build_connector_worker_meta.return_value = None
841
+ mc_worker_meta_0a_none = mc.build_connector_worker_meta()
842
+ assert isinstance(mc_worker_meta_0a_none, MultiKVConnectorWorkerMetadata)
843
+ assert mc_worker_meta_0a_none.metadata == (worker_meta0a, None)
844
+
845
+ # both connectors do not return None
846
+ worker_meta0b = MockConnectorWorkerMetadata0({"0b"})
847
+ worker_meta1b = MockConnectorWorkerMetadata1({"1b"})
848
+ mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0b
849
+ mc._connectors[1].build_connector_worker_meta.return_value = worker_meta1b
850
+ mc_worker_meta_0b_1b = mc.build_connector_worker_meta()
851
+ assert isinstance(mc_worker_meta_0b_1b, MultiKVConnectorWorkerMetadata)
852
+ assert mc_worker_meta_0b_1b.metadata == (worker_meta0b, worker_meta1b)
853
+
854
+ # ----------------------------- test aggregate ----------------------------
855
+
856
+ # aggregate ({"0a"}, None) and (None, {"1a"}) -> ({"0a"}, {"1a"})
857
+ mc_worker_meta_0a_1a = mc_worker_meta_0a_none.aggregate(mc_worker_meta_none_1a)
858
+ assert isinstance(mc_worker_meta_0a_1a, MultiKVConnectorWorkerMetadata)
859
+ assert mc_worker_meta_0a_1a.metadata == (worker_meta0a, worker_meta1a)
860
+
861
+ # aggregate ({"0a"}, None) and ({"0b"}, None) -> ({"0a", "0b"}, None)
862
+ mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0b
863
+ mc._connectors[1].build_connector_worker_meta.return_value = None
864
+ mc_worker_meta_0b_none = mc.build_connector_worker_meta()
865
+ mc_worker_meta_0a_0b = mc_worker_meta_0a_none.aggregate(mc_worker_meta_0b_none)
866
+ assert isinstance(mc_worker_meta_0a_0b, MultiKVConnectorWorkerMetadata)
867
+ assert mc_worker_meta_0a_0b.metadata[1] is None
868
+ connector0_md = mc_worker_meta_0a_0b.metadata[0]
869
+ assert isinstance(connector0_md, MockConnectorWorkerMetadata0)
870
+ assert connector0_md.data == {"0a", "0b"}
871
+
872
+ # aggregate ({"0a"}, {"1a"}) and ({"0b"}, {"1b"}) -> ({"0a", "0b"}, {"1a", "1b"})
873
+ mc_worker_meta_01a_01b = mc_worker_meta_0a_1a.aggregate(mc_worker_meta_0b_1b)
874
+ assert isinstance(mc_worker_meta_01a_01b, MultiKVConnectorWorkerMetadata)
875
+ metadata = mc_worker_meta_01a_01b.metadata
876
+ assert len(metadata) == 2
877
+ connector0_md, connector1_md = metadata
878
+ assert isinstance(connector0_md, MockConnectorWorkerMetadata0)
879
+ assert isinstance(connector1_md, MockConnectorWorkerMetadata1)
880
+ assert connector0_md.data == {"0a", "0b"}
881
+ assert connector1_md.data == {"1a", "1b"}
882
+
883
+ # ---------------------- test update_connector_output ---------------------
884
+
885
+ def verify_worker_metadata(expected_metadata: MockConnectorWorkerMetadata | None):
886
+ def _verify_worker_metadata(connector_output: KVConnectorOutput):
887
+ worker_meta = connector_output.kv_connector_worker_meta
888
+ if expected_metadata is None:
889
+ assert worker_meta is None
890
+ return
891
+
892
+ assert isinstance(worker_meta, MockConnectorWorkerMetadata)
893
+ assert type(worker_meta) is type(expected_metadata)
894
+ assert expected_metadata.data == worker_meta.data
895
+
896
+ return _verify_worker_metadata
897
+
898
+ def assert_update_connector_output_called(mc: MultiConnector):
899
+ for c in mc._connectors:
900
+ c.update_connector_output.assert_called_once()
901
+ c.update_connector_output.reset_mock()
902
+
903
+ # no worker meta
904
+ kv_connector_output = KVConnectorOutput()
905
+ mc._connectors[0].update_connector_output.side_effect = verify_worker_metadata(None)
906
+ mc._connectors[1].update_connector_output.side_effect = verify_worker_metadata(None)
907
+ mc.update_connector_output(kv_connector_output)
908
+ assert_update_connector_output_called(mc)
909
+
910
+ # multi worker meta
911
+ kv_connector_output.kv_connector_worker_meta = mc_worker_meta_01a_01b
912
+ mc._connectors[0].update_connector_output.side_effect = verify_worker_metadata(
913
+ connector0_md
914
+ )
915
+ mc._connectors[1].update_connector_output.side_effect = verify_worker_metadata(
916
+ connector1_md
917
+ )
918
+ mc.update_connector_output(kv_connector_output)
919
+ assert_update_connector_output_called(mc)
920
+ assert kv_connector_output.kv_connector_worker_meta == mc_worker_meta_01a_01b
vllm/distributed/kv_transfer/kv_connector/utils.py
+13
@@ -85,6 +85,7 @@ class KVOutputAggregator:
85
finished_sending = set[str]()
86
finished_recving = set[str]()
87
aggregated_kv_connector_stats = None
88
+ aggregated_kv_connector_worker_meta = None
89
combined_kv_cache_events = None
90
invalid_block_ids = set[int]()
91
for model_runner_output in outputs:
@@ -127,6 +128,17 @@ class KVOutputAggregator:
128
aggregated_kv_connector_stats.aggregate(kv_connector_stats)
129
)
130
131
+ # Aggregate kv_connector_worker_meta from all workers.
132
+ if aggregated_kv_connector_worker_meta is None:
133
+ # Use the first worker's kv_connector_worker_meta as accumulator.
134
+ aggregated_kv_connector_worker_meta = kv_output.kv_connector_worker_meta
135
+ elif kv_connector_worker_meta := kv_output.kv_connector_worker_meta:
136
+ aggregated_kv_connector_worker_meta = (
137
+ aggregated_kv_connector_worker_meta.aggregate(
138
+ kv_connector_worker_meta
139
+ )
140
+ )
141
+
142
# Combine kv_cache_events from all workers.
143
if combined_kv_cache_events is None:
144
# Use the first worker's kv_cache events as start event list.
@@ -151,6 +163,7 @@ class KVOutputAggregator:
163
finished_recving=finished_recving or None,
164
kv_connector_stats=aggregated_kv_connector_stats or None,
165
kv_cache_events=combined_kv_cache_events or None,
166
+ kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None,
167
invalid_block_ids=invalid_block_ids,
168
expected_finished_count=self._expected_finished_count,
169
)
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+35
-2
@@ -36,6 +36,8 @@ The class provides the following primitives:
36
37
get_finished() - called with ids of finished requests, returns
38
ids of requests that have completed async sending/recving.
39
+ build_connector_worker_meta() - builds metadata to be sent
40
+ back to the scheduler-side connector
41
"""
42
43
import enum
@@ -137,13 +139,34 @@ class KVConnectorHandshakeMetadata(ABC): # noqa: B024
139
140
class KVConnectorMetadata(ABC): # noqa: B024
141
"""
140
- Abstract Metadata used to communicate between the
141
- Scheduler KVConnector and Worker KVConnector.
142
+ Abstract Metadata used to communicate
143
+ Scheduler KVConnector -> Worker KVConnector.
144
"""
145
146
pass
147
148
149
+class KVConnectorWorkerMetadata(ABC):
150
+ """
151
+ Abstract Metadata used to communicate back
152
+ Worker KVConnector -> Scheduler KVConnector.
153
+
154
+ Each worker can output its own metadata.
155
+ For a single engine step, all metadata objects returned by workers
156
+ will be aggregated using the `aggregate` method below, before
157
+ being passed to the Scheduler KVConnector.
158
+ """
159
+
160
+ @abstractmethod
161
+ def aggregate(
162
+ self, other: "KVConnectorWorkerMetadata"
163
+ ) -> "KVConnectorWorkerMetadata":
164
+ """
165
+ Aggregate metadata with another `KVConnectorWorkerMetadata` object.
166
+ """
167
+ pass
168
+
169
+
170
class KVConnectorBase_V1(ABC):
171
"""
172
Base class for KV connectors.
@@ -409,6 +432,16 @@ class KVConnectorBase_V1(ABC):
432
"""
433
return None
434
435
+ def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
436
+ """
437
+ Build the KVConnector worker metadata for this engine step.
438
+
439
+ Returns:
440
+ KVConnectorWorkerMetadata: the worker metadata.
441
+ None if no worker metadata is available.
442
+ """
443
+ return None
444
+
445
# ==============================
446
# Scheduler-side methods
447
# ==============================
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+52
-2
@@ -17,6 +17,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17
KVConnectorHandshakeMetadata,
18
KVConnectorMetadata,
19
KVConnectorRole,
20
+ KVConnectorWorkerMetadata,
21
)
22
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
23
KVConnectorPromMetrics,
@@ -45,6 +46,26 @@ class MultiKVConnectorMetadata(KVConnectorMetadata):
46
extra_async_saves: dict[str, int] | None = None
47
48
49
+@dataclass
50
+class MultiKVConnectorWorkerMetadata(KVConnectorWorkerMetadata):
51
+ metadata: tuple[KVConnectorWorkerMetadata | None, ...]
52
+
53
+ def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetadata:
54
+ assert isinstance(other, MultiKVConnectorWorkerMetadata)
55
+
56
+ assert len(self.metadata) == len(other.metadata)
57
+ metadata_list = []
58
+ for metadata1, metadata2 in zip(self.metadata, other.metadata):
59
+ if metadata1 is None:
60
+ metadata_list.append(metadata2)
61
+ elif metadata2 is None:
62
+ metadata_list.append(metadata1)
63
+ else:
64
+ metadata_list.append(metadata1.aggregate(metadata2))
65
+
66
+ return MultiKVConnectorWorkerMetadata(metadata=tuple(metadata_list))
67
+
68
+
69
@dataclass
70
class MultiKVConnectorStats(KVConnectorStats):
71
"""
@@ -304,6 +325,18 @@ class MultiConnector(KVConnectorBase_V1):
325
# Currently no connectors return non-None
326
return None
327
328
+ def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
329
+ metadata_list: list[KVConnectorWorkerMetadata | None] | None = None
330
+ for i, c in enumerate(self._connectors):
331
+ kv_connector_worker_meta = c.build_connector_worker_meta()
332
+ if metadata_list is None and kv_connector_worker_meta is not None:
333
+ metadata_list = [None] * i
334
+ if metadata_list is not None:
335
+ metadata_list.append(kv_connector_worker_meta)
336
+ if metadata_list is None:
337
+ return None
338
+ return MultiKVConnectorWorkerMetadata(metadata=tuple(metadata_list))
339
+
340
# TODO: Add a generic implementation of 'get_kv_connector_kv_cache_events'
341
# method for the MultiConnector. It should be able to get events from
342
# multiple connectors, handling the case where only a subset of the
@@ -361,8 +394,25 @@ class MultiConnector(KVConnectorBase_V1):
394
return metadata
395
396
def update_connector_output(self, connector_output: KVConnectorOutput):
364
- for c in self._connectors:
365
- c.update_connector_output(connector_output)
397
+ multi_connector_worker_meta: MultiKVConnectorWorkerMetadata | None = None
398
+ if connector_output.kv_connector_worker_meta is not None:
399
+ assert isinstance(
400
+ connector_output.kv_connector_worker_meta,
401
+ MultiKVConnectorWorkerMetadata,
402
+ )
403
+ multi_connector_worker_meta = connector_output.kv_connector_worker_meta
404
+
405
+ try:
406
+ for i, c in enumerate(self._connectors):
407
+ if multi_connector_worker_meta is not None:
408
+ # set the connector-specific worker metadata
409
+ connector_output.kv_connector_worker_meta = (
410
+ multi_connector_worker_meta.metadata[i]
411
+ )
412
+ c.update_connector_output(connector_output)
413
+ finally:
414
+ # restore kv_connector_worker_meta
415
+ connector_output.kv_connector_worker_meta = multi_connector_worker_meta
416
417
def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
418
"""
vllm/v1/outputs.py
+6
@@ -14,9 +14,13 @@ from vllm.v1.core.sched.output import SchedulerOutput
14
15
if TYPE_CHECKING:
16
from vllm.distributed.kv_events import KVConnectorKVEvents
17
+ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
18
+ KVConnectorWorkerMetadata,
19
+ )
20
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
21
else:
22
KVConnectorStats = object
23
+ KVConnectorWorkerMetadata = object
24
KVConnectorKVEvents = object
25
26
@@ -142,6 +146,7 @@ class KVConnectorOutput:
146
finished_recving: set[str] | None = None
147
kv_connector_stats: KVConnectorStats | None = None
148
kv_cache_events: KVConnectorKVEvents | None = None
149
+ kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None
150
# IDs of externally computed KV blocks that failed to load.
151
# Requests referencing these blocks should be rescheduled to recompute them
152
invalid_block_ids: set[int] = field(default_factory=set)
@@ -159,6 +164,7 @@ class KVConnectorOutput:
164
and not self.kv_connector_stats
165
and not self.kv_cache_events
166
and not self.invalid_block_ids
167
+ and not self.kv_connector_worker_meta
168
)
169
170
@classmethod
vllm/v1/worker/kv_connector_model_runner_mixin.py
+1
@@ -123,6 +123,7 @@ class KVConnectorModelRunnerMixin:
123
124
output.kv_connector_stats = kv_connector.get_kv_connector_stats()
125
output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events()
126
+ output.kv_connector_worker_meta = kv_connector.build_connector_worker_meta()
127
128
if not defer_finalize:
129
kv_connector.clear_connector_metadata()