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, 17:45 UTC
[KV Connector] Add missing method overrides to MultiConnector (#33292)
Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Seiji Eicher committed
Feb 6, 2026 at 09:58 UTC
aca596741673df92922565f08189d01ea228f605
2 files changed
+177
-11
tests/v1/kv_connector/unit/test_multi_connector.py
+134
-6
@@ -190,27 +190,35 @@ def test_multi_example_connector_consistency():
190
)
191
192
events = get_connector_events()
193
- # get_num_new_matched_tokens and update_state_after_alloc will be called
194
- # on each connector in turn.
195
- assert events["storage1-SCHEDULER"][:3] == [
193
+ # First event is set_xfer_handshake_metadata from initialization, then
194
+ # get_num_new_matched_tokens and update_state_after_alloc from generate().
195
+ assert events["storage1-SCHEDULER"][:4] == [
196
+ "set_xfer_handshake_metadata",
197
"get_num_new_matched_tokens 0",
198
"update_state_after_alloc num_blocks=[0] 0",
199
"build_connector_meta",
200
]
200
- assert events["storage1-WORKER"][:5] == [
201
+ # First three events are from initialization (register_kv_caches,
202
+ # set_host_xfer_buffer_ops, get_handshake_metadata), then generate() events.
203
+ assert events["storage1-WORKER"][:7] == [
204
"register_kv_caches",
205
+ "set_host_xfer_buffer_ops",
206
+ "get_handshake_metadata",
207
"bind_connector_metadata",
208
"start_load_kv",
209
"wait_for_layer_load",
210
"save_kv_layer",
211
]
207
- assert events["storage2-SCHEDULER"][:3] == [
212
+ assert events["storage2-SCHEDULER"][:4] == [
213
+ "set_xfer_handshake_metadata",
214
"get_num_new_matched_tokens 0",
215
"update_state_after_alloc num_blocks=[0] 0",
216
"build_connector_meta",
217
]
212
- assert events["storage2-WORKER"][:5] == [
218
+ assert events["storage2-WORKER"][:7] == [
219
"register_kv_caches",
220
+ "set_host_xfer_buffer_ops",
221
+ "get_handshake_metadata",
222
"bind_connector_metadata",
223
"start_load_kv",
224
"wait_for_layer_load",
@@ -297,6 +305,90 @@ def test_engine_id_conflict():
305
)
306
307
308
+def test_multi_connector_handle_preemptions_integration():
309
+ """
310
+ Integration test: verify MultiConnector delegates handle_preemptions
311
+ to all sub-connectors.
312
+
313
+ Uses TestExampleConnector which logs all method calls to temp files.
314
+ This test directly calls handle_preemptions on a MultiConnector with
315
+ TestExampleConnector sub-connectors and verifies the calls are logged.
316
+ """
317
+ from tests.v1.kv_connector.unit.utils import (
318
+ create_scheduler,
319
+ create_vllm_config,
320
+ )
321
+
322
+ storage_path = Path(tempfile.mkdtemp())
323
+
324
+ try:
325
+ # Configure MultiConnector with two TestExampleConnectors
326
+ kv_transfer_config = KVTransferConfig(
327
+ kv_connector="MultiConnector",
328
+ kv_role="kv_both",
329
+ kv_connector_extra_config={
330
+ "connectors": [
331
+ {
332
+ "kv_connector": "TestExampleConnector",
333
+ "kv_role": "kv_both",
334
+ "kv_connector_extra_config": {
335
+ "shared_storage_path": str(storage_path / "s1"),
336
+ "name": "preempt1",
337
+ },
338
+ "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",
339
+ },
340
+ {
341
+ "kv_connector": "TestExampleConnector",
342
+ "kv_role": "kv_both",
343
+ "kv_connector_extra_config": {
344
+ "shared_storage_path": str(storage_path / "s2"),
345
+ "name": "preempt2",
346
+ },
347
+ "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",
348
+ },
349
+ ]
350
+ },
351
+ )
352
+
353
+ vllm_config = create_vllm_config(
354
+ block_size=16,
355
+ max_num_batched_tokens=100,
356
+ kv_connector_extra_config=kv_transfer_config.kv_connector_extra_config,
357
+ )
358
+ vllm_config.kv_transfer_config = kv_transfer_config
359
+
360
+ # Create scheduler - this initializes the MultiConnector with SCHEDULER role
361
+ scheduler = create_scheduler(vllm_config, num_blocks=10)
362
+
363
+ # Clear any events from initialization
364
+ get_connector_events()
365
+
366
+ # Directly call handle_preemptions on the scheduler's connector
367
+ # Note: handle_preemptions is normally a worker-side method, but we're
368
+ # testing the delegation behavior of MultiConnector here.
369
+ # The connector attribute contains the KV connector.
370
+ assert scheduler.connector is not None, "Scheduler should have a connector"
371
+ preempted_req_ids = {"req-1", "req-2", "req-3"}
372
+ scheduler.connector.handle_preemptions(preempted_req_ids)
373
+
374
+ # Verify both connectors received the handle_preemptions call
375
+ events = get_connector_events()
376
+
377
+ # Both SCHEDULER-role connectors should have logged handle_preemptions
378
+ assert "handle_preemptions" in events.get("preempt1-SCHEDULER", []), (
379
+ f"preempt1-SCHEDULER should have handle_preemptions call. "
380
+ f"Got events: {events}"
381
+ )
382
+ assert "handle_preemptions" in events.get("preempt2-SCHEDULER", []), (
383
+ f"preempt2-SCHEDULER should have handle_preemptions call. "
384
+ f"Got events: {events}"
385
+ )
386
+
387
+ finally:
388
+ # Cleanup
389
+ shutil.rmtree(storage_path, ignore_errors=True)
390
+
391
+
392
class TestMultiConnectorStats:
393
"""Tests for MultiConnector stats reconstruction and operations."""
394
@@ -647,3 +739,39 @@ class TestMultiConnectorPreferCrossLayerBlocks:
739
MockConnector.__new__(MockConnector), # default False
740
]
741
assert mc.prefer_cross_layer_blocks is False
742
+
743
+
744
+def test_multi_connector_overrides_all_base_methods():
745
+ """
746
+ Ensure MultiConnector overrides all public methods from KVConnectorBase_V1.
747
+ """
748
+ # These are fine to inherit from KVConnectorBase_V1
749
+ # TODO(https://github.com/vllm-project/vllm/pull/31811): Remove
750
+ # get_kv_connector_kv_cache_events from INHERITED_OK once implemented.
751
+ INHERITED_OK = {
752
+ "role",
753
+ "has_connector_metadata",
754
+ "get_kv_connector_kv_cache_events",
755
+ }
756
+
757
+ base_members = {
758
+ name for name in dir(KVConnectorBase_V1) if not name.startswith("_")
759
+ } - KVConnectorBase_V1.__abstractmethods__
760
+
761
+ missing = [
762
+ name
763
+ for name in sorted(base_members)
764
+ if name not in INHERITED_OK and name not in MultiConnector.__dict__
765
+ ]
766
+
767
+ if missing:
768
+ pytest.fail(f"""
769
+MultiConnector does not override these KVConnectorBase_V1 methods: {missing}
770
+
771
+MultiConnector wraps other connectors and must delegate all methods.
772
+Please add overrides that delegate to self._connectors.
773
+
774
+Options:
775
+ 1. Add delegation in MultiConnector (preferred)
776
+ 2. Add to INHERITED_OK if the base implementation works correctly
777
+""")
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+43
-5
@@ -12,7 +12,9 @@ from vllm.config.kv_transfer import KVTransferConfig
12
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
13
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
14
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15
+ CopyBlocksOp,
16
KVConnectorBase_V1,
17
+ KVConnectorHandshakeMetadata,
18
KVConnectorMetadata,
19
KVConnectorRole,
20
)
@@ -272,11 +274,26 @@ class MultiConnector(KVConnectorBase_V1):
274
agg_block_ids |= c.get_block_ids_with_load_errors()
275
return agg_block_ids
276
275
- # TODO: Add a generic implementation of 'get_kv_connector_kv_cache_events' method
276
- # for the MultiConnector. It should be able to get events from multiple
277
- # connectors, handling the case where only a subset of the requested connectors
278
- # implements the 'get_kv_connector_kv_cache_events'
279
- # Follow on PR from https://github.com/vllm-project/vllm/pull/28309#pullrequestreview-3566351082
277
+ def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
278
+ """Set xPU-specific copy ops for all sub-connectors."""
279
+ for c in self._connectors:
280
+ c.set_host_xfer_buffer_ops(copy_operation)
281
+
282
+ def handle_preemptions(self, preempted_req_ids: set[str]):
283
+ """Handle preempted requests for all sub-connectors."""
284
+ for c in self._connectors:
285
+ c.handle_preemptions(preempted_req_ids)
286
+
287
+ def get_finished_count(self) -> int | None:
288
+ # TODO(https://github.com/vllm-project/vllm/issues/33400)
289
+ # Currently no connectors return non-None
290
+ return None
291
+
292
+ # TODO: Add a generic implementation of 'get_kv_connector_kv_cache_events'
293
+ # method for the MultiConnector. It should be able to get events from
294
+ # multiple connectors, handling the case where only a subset of the
295
+ # requested connectors implements the 'get_kv_connector_kv_cache_events'
296
+ # WIP: https://github.com/vllm-project/vllm/pull/31811
297
298
# ==============================
299
# Scheduler-side methods
@@ -332,6 +349,27 @@ class MultiConnector(KVConnectorBase_V1):
349
for c in self._connectors:
350
c.update_connector_output(connector_output)
351
352
+ def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
353
+ """
354
+ Get the KVConnector handshake metadata from sub-connectors.
355
+ Returns the first non-None metadata from sub-connectors.
356
+ """
357
+ for c in self._connectors:
358
+ metadata = c.get_handshake_metadata()
359
+ if metadata is not None:
360
+ return metadata
361
+ return None
362
+
363
+ def set_xfer_handshake_metadata(
364
+ self, metadata: dict[int, KVConnectorHandshakeMetadata]
365
+ ) -> None:
366
+ """
367
+ Set the KV connector handshake metadata for all sub-connectors.
368
+ This is needed to start the NIXL listener thread for NixlConnector.
369
+ """
370
+ for c in self._connectors:
371
+ c.set_xfer_handshake_metadata(metadata)
372
+
373
def request_finished(
374
self,
375
request: "Request",