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 4, 11:15 UTC
[KVConnector] MultiConnector SupportsHMA (#39571)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Apr 30, 2026 at 11:10 UTC
efdc95674db5c7b441d52ae02fa57e57c6bb3855
2 files changed
+217
-9
tests/v1/kv_connector/unit/test_multi_connector.py
+172
-2
@@ -8,13 +8,18 @@ from typing import Any
8
from unittest.mock import MagicMock
9
10
import pytest
11
+import torch
12
13
from tests.v1.kv_connector.unit.utils import create_vllm_config
14
from vllm import LLM, SamplingParams
15
from vllm.config import KVTransferConfig
16
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
17
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.base import (
19
+ KVConnectorBase_V1,
20
+ SupportsHMA,
21
+ supports_hma,
22
+)
23
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
24
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
25
MultiConnector,
@@ -83,8 +88,43 @@ class MockConnector(KVConnectorBase_V1):
88
pass
89
90
86
-# Register the mock connector
91
+class MockHMAConnector(KVConnectorBase_V1, SupportsHMA):
92
+ """Mock connector that supports HMA for testing."""
93
+
94
+ def __new__(cls, *args, **kwargs):
95
+ mock = MagicMock(spec_set=cls)
96
+ return mock
97
+
98
+ def start_load_kv(self, forward_context, **kwargs):
99
+ pass
100
+
101
+ def wait_for_layer_load(self, layer_name):
102
+ pass
103
+
104
+ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
105
+ pass
106
+
107
+ def wait_for_save(self):
108
+ pass
109
+
110
+ def build_connector_meta(self, scheduler_output):
111
+ return None
112
+
113
+ def get_num_new_matched_tokens(self, request, num_computed_tokens):
114
+ return (0, False)
115
+
116
+ def update_state_after_alloc(self, request, blocks, num_tokens) -> None:
117
+ pass
118
+
119
+ def request_finished_all_groups(self, request, block_ids):
120
+ return (False, None)
121
+
122
+
123
+# Register mock connectors
124
KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
125
+KVConnectorFactory.register_connector(
126
+ "MockHMAConnector", __name__, MockHMAConnector.__name__
127
+)
128
129
130
@pytest.fixture
@@ -920,3 +960,133 @@ def test_multi_connector_worker_metadata(mc):
960
mc.update_connector_output(kv_connector_output)
961
assert_update_connector_output_called(mc)
962
assert kv_connector_output.kv_connector_worker_meta == mc_worker_meta_01a_01b
963
+
964
+
965
+def _make_multi_connector(connector_names: list[str]) -> MultiConnector:
966
+ """Build a MultiConnector wrapping the given registered connectors."""
967
+ vllm_config = create_vllm_config()
968
+ connectors = [
969
+ {
970
+ "kv_connector": name,
971
+ "kv_role": "kv_both",
972
+ "kv_connector_module_path": "tests.v1.kv_connector.unit.test_multi_connector", # noqa: E501
973
+ }
974
+ for name in connector_names
975
+ ]
976
+ vllm_config.kv_transfer_config = KVTransferConfig(
977
+ kv_connector="MultiConnector",
978
+ kv_role="kv_both",
979
+ kv_connector_extra_config={"connectors": connectors},
980
+ )
981
+ kv_cache_config = KVCacheConfig(
982
+ num_blocks=0,
983
+ kv_cache_tensors=[],
984
+ kv_cache_groups=[],
985
+ )
986
+ return MultiConnector(
987
+ vllm_config=vllm_config,
988
+ role=KVConnectorRole.WORKER,
989
+ kv_cache_config=kv_cache_config,
990
+ )
991
+
992
+
993
+def test_multi_connector_hma_opt_in():
994
+ """
995
+ MultiConnector currently assumes HMA is opt-in: it needs
996
+ --no-disable-hybrid-kv-cache-manager to be enabled.
997
+
998
+ At runtime, _all_support_hma is True only when every sub-connector
999
+ implements SupportsHMA. Test all combinations of HMA / non-HMA
1000
+ sub-connectors.
1001
+ """
1002
+
1003
+ assert supports_hma(MultiConnector)
1004
+
1005
+ # -- All non-HMA connectors => _all_support_hma is False --
1006
+ mc_none = _make_multi_connector(["MockConnector", "MockConnector"])
1007
+ assert not supports_hma(mc_none._connectors[0])
1008
+ assert not supports_hma(mc_none._connectors[1])
1009
+ assert mc_none._all_support_hma is False
1010
+
1011
+ # -- All HMA connectors => _all_support_hma is True --
1012
+ mc_all = _make_multi_connector(["MockHMAConnector", "MockHMAConnector"])
1013
+ assert supports_hma(mc_all._connectors[0])
1014
+ assert supports_hma(mc_all._connectors[1])
1015
+ assert mc_all._all_support_hma is True
1016
+
1017
+ # -- Mixed: first HMA, second non-HMA => _all_support_hma is False --
1018
+ mc_mixed1 = _make_multi_connector(["MockHMAConnector", "MockConnector"])
1019
+ assert supports_hma(mc_mixed1._connectors[0])
1020
+ assert not supports_hma(mc_mixed1._connectors[1])
1021
+ assert mc_mixed1._all_support_hma is False
1022
+
1023
+ # -- Mixed: first non-HMA, second HMA => _all_support_hma is False --
1024
+ mc_mixed2 = _make_multi_connector(["MockConnector", "MockHMAConnector"])
1025
+ assert not supports_hma(mc_mixed2._connectors[0])
1026
+ assert supports_hma(mc_mixed2._connectors[1])
1027
+ assert mc_mixed2._all_support_hma is False
1028
+
1029
+
1030
+@pytest.mark.skipif(
1031
+ not torch.cuda.is_available(), reason="Requires GPU to instantiate LLM"
1032
+)
1033
+def test_multi_connector_mixed_hma_disables_hybrid_kv_cache(monkeypatch):
1034
+ """
1035
+ When MultiConnector wraps a mix of HMA (NixlConnector) and non-HMA
1036
+ (MockConnector) sub-connectors, verify that:
1037
+ 1. The scheduler's MultiConnector has _all_support_hma == False.
1038
+ 2. vLLM auto-disables the hybrid KV cache manager (no preference expressed by user)
1039
+ """
1040
+ from unittest.mock import patch
1041
+
1042
+ from tests.v1.kv_connector.unit.test_nixl_connector import FakeNixlWrapper
1043
+
1044
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
1045
+
1046
+ kv_transfer_config = KVTransferConfig(
1047
+ kv_connector="MultiConnector",
1048
+ kv_role="kv_both",
1049
+ kv_connector_extra_config={
1050
+ "connectors": [
1051
+ {
1052
+ "kv_connector": "NixlConnector",
1053
+ "kv_role": "kv_both",
1054
+ },
1055
+ {
1056
+ "kv_connector": "MockConnector",
1057
+ "kv_role": "kv_both",
1058
+ "kv_connector_module_path": (
1059
+ "tests.v1.kv_connector.unit.test_multi_connector"
1060
+ ),
1061
+ },
1062
+ ],
1063
+ },
1064
+ )
1065
+
1066
+ with patch(
1067
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1068
+ FakeNixlWrapper,
1069
+ ):
1070
+ llm = LLM(
1071
+ model="Qwen/Qwen3-0.6B",
1072
+ enforce_eager=True,
1073
+ gpu_memory_utilization=0.3,
1074
+ max_model_len=128,
1075
+ max_num_seqs=1,
1076
+ max_num_batched_tokens=128,
1077
+ kv_transfer_config=kv_transfer_config,
1078
+ )
1079
+ try:
1080
+ # HMA should be auto-disabled when user has not expressed a preference.
1081
+ assert (
1082
+ llm.llm_engine.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
1083
+ is True
1084
+ )
1085
+ # The scheduler-side MultiConnector should detect the mixed
1086
+ # HMA support among its sub-connectors.
1087
+ scheduler = llm.llm_engine.engine_core.engine_core.scheduler
1088
+ mc = scheduler.connector
1089
+ assert isinstance(mc, MultiConnector)
1090
+ assert mc._all_support_hma is False
1091
+ finally:
1092
+ llm.llm_engine.engine_core.shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+45
-7
@@ -1,9 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import copy
4
-from collections.abc import Iterable
4
+from collections.abc import Callable, Iterable
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, Any
6
+from typing import TYPE_CHECKING, Any, cast
7
8
import torch
9
@@ -18,6 +18,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
18
KVConnectorMetadata,
19
KVConnectorRole,
20
KVConnectorWorkerMetadata,
21
+ SupportsHMA,
22
+ supports_hma,
23
)
24
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
25
KVConnectorPromMetrics,
@@ -123,7 +125,7 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
125
self._prom_metrics[connector_id].observe(stats_data["data"], engine_idx)
126
127
126
-class MultiConnector(KVConnectorBase_V1):
128
+class MultiConnector(KVConnectorBase_V1, SupportsHMA):
129
"""
130
A wrapper for using multiple KVConnectors at the same time.
131
@@ -166,6 +168,12 @@ class MultiConnector(KVConnectorBase_V1):
168
self._connectors.append(connector_cls(temp_config, role, kv_cache_config))
169
self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
170
171
+ self._all_support_hma = all(supports_hma(c) for c in self._connectors)
172
+ assert (
173
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
174
+ or self._all_support_hma
175
+ ), "HMA should not be enabled unless all sub-connectors support it"
176
+
177
# A mapping from request id to the index of the connector chosen to
178
# load the request from (if any).
179
self._requests_to_connector: dict[str, int] = {}
@@ -436,15 +444,17 @@ class MultiConnector(KVConnectorBase_V1):
444
for c in self._connectors:
445
c.set_xfer_handshake_metadata(metadata)
446
439
- def request_finished(
447
+ def _aggregate_request_finished(
448
self,
449
request: "Request",
442
- blocks: list[int],
450
+ per_connector_fn: Callable[
451
+ [KVConnectorBase_V1], tuple[bool, dict[str, Any] | None]
452
+ ],
453
) -> tuple[bool, dict[str, Any] | None]:
454
async_saves = 0
455
kv_txfer_params = None
456
for c in self._connectors:
447
- async_save, txfer_params = c.request_finished(request, blocks)
457
+ async_save, txfer_params = per_connector_fn(c)
458
if async_save:
459
async_saves += 1
460
if txfer_params is not None:
@@ -458,11 +468,39 @@ class MultiConnector(KVConnectorBase_V1):
468
if async_saves > 1:
469
self._extra_async_saves[request.request_id] = async_saves - 1
470
461
- # Clean up other state for this request.
471
self._requests_to_connector.pop(request.request_id, None)
472
473
return async_saves > 0, kv_txfer_params
474
475
+ def request_finished(
476
+ self,
477
+ request: "Request",
478
+ blocks: list[int],
479
+ ) -> tuple[bool, dict[str, Any] | None]:
480
+ return self._aggregate_request_finished(
481
+ request,
482
+ lambda c: c.request_finished(request, blocks),
483
+ )
484
+
485
+ def request_finished_all_groups(
486
+ self,
487
+ request: "Request",
488
+ block_ids: tuple[list[int], ...],
489
+ ) -> tuple[bool, dict[str, Any] | None]:
490
+ if not self._all_support_hma:
491
+ assert len(block_ids) == 1, (
492
+ "HMA with multiple kv_cache_groups requires all "
493
+ "sub-connectors to support HMA"
494
+ )
495
+ return self.request_finished(request, block_ids[0])
496
+
497
+ return self._aggregate_request_finished(
498
+ request,
499
+ lambda c: cast(SupportsHMA, c).request_finished_all_groups(
500
+ request, block_ids
501
+ ),
502
+ )
503
+
504
def take_events(self) -> Iterable["KVCacheEvent"]:
505
for c in self._connectors:
506
yield from c.take_events()