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 2, 20:45 UTC
[KV Transfer] Enable HMA by default for connectors that support it (#41847)
Signed-off-by: Ethan Feng <ethan.fengch@gmail.com>
Ethan Feng committed
May 26, 2026 at 18:28 UTC
755043cf3cd68e6cd93bfcfde2152f28c9fb42bd
7 files changed
+161
-41
tests/v1/kv_connector/unit/test_hma_auto_config.py
new
+107
@@ -0,0 +1,107 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Regression tests for HMA auto-disable with KV transfer connectors."""
5
+
6
+import pytest
7
+
8
+from vllm.config import DeviceConfig, KVTransferConfig, SchedulerConfig, VllmConfig
9
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
10
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
11
+from vllm.platforms import current_platform
12
+from vllm.v1.kv_cache_interface import KVCacheConfig
13
+
14
+pytestmark = pytest.mark.cpu_test
15
+
16
+
17
+@pytest.fixture(autouse=True)
18
+def mock_hybrid_kv_cache_supported(monkeypatch):
19
+ monkeypatch.setattr(current_platform, "support_hybrid_kv_cache", lambda: True)
20
+
21
+
22
+@pytest.mark.parametrize(
23
+ "kv_transfer_config,expect_disabled",
24
+ [
25
+ ( # HMA-supporting connector → HMA stays enabled
26
+ KVTransferConfig(
27
+ kv_connector="SimpleCPUOffloadConnector",
28
+ kv_role="kv_both",
29
+ kv_connector_extra_config={"cpu_bytes_to_use": 1 << 30},
30
+ ),
31
+ False,
32
+ ),
33
+ ( # Non-HMA connector → HMA is auto-disabled
34
+ KVTransferConfig(kv_connector="ExampleConnector", kv_role="kv_both"),
35
+ True,
36
+ ),
37
+ ( # MultiConnector: all HMA children → HMA stays enabled
38
+ KVTransferConfig(
39
+ kv_connector="MultiConnector",
40
+ kv_role="kv_both",
41
+ kv_connector_extra_config={
42
+ "connectors": [
43
+ {
44
+ "kv_connector": "SimpleCPUOffloadConnector",
45
+ "kv_role": "kv_both",
46
+ "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
47
+ },
48
+ {
49
+ "kv_connector": "OffloadingConnector",
50
+ "kv_role": "kv_both",
51
+ "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
52
+ },
53
+ ]
54
+ },
55
+ ),
56
+ False,
57
+ ),
58
+ ( # MultiConnector: mixed children → HMA is auto-disabled
59
+ KVTransferConfig(
60
+ kv_connector="MultiConnector",
61
+ kv_role="kv_both",
62
+ kv_connector_extra_config={
63
+ "connectors": [
64
+ {
65
+ "kv_connector": "SimpleCPUOffloadConnector",
66
+ "kv_role": "kv_both",
67
+ "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
68
+ },
69
+ {"kv_connector": "ExampleConnector", "kv_role": "kv_both"},
70
+ ]
71
+ },
72
+ ),
73
+ True,
74
+ ),
75
+ ],
76
+ ids=["hma_connector", "non_hma_connector", "multi_all_hma", "multi_mixed"],
77
+)
78
+def test_hma_auto_config(kv_transfer_config, expect_disabled):
79
+ vllm_config = VllmConfig(
80
+ device_config=DeviceConfig("cpu"),
81
+ kv_transfer_config=kv_transfer_config,
82
+ )
83
+ assert (
84
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is expect_disabled
85
+ )
86
+
87
+
88
+def test_explicit_hma_with_non_hma_connector_errors_at_factory():
89
+ vllm_config = VllmConfig(
90
+ device_config=DeviceConfig("cpu"),
91
+ scheduler_config=SchedulerConfig(
92
+ max_model_len=16,
93
+ is_encoder_decoder=False,
94
+ disable_hybrid_kv_cache_manager=False,
95
+ ),
96
+ kv_transfer_config=KVTransferConfig(
97
+ kv_connector="ExampleConnector",
98
+ kv_role="kv_both",
99
+ ),
100
+ )
101
+ kv_cache_config = KVCacheConfig(
102
+ num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
103
+ )
104
+ with pytest.raises(ValueError, match="does not support HMA but HMA is enabled"):
105
+ KVConnectorFactory.create_connector(
106
+ vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
107
+ )
tests/v1/kv_connector/unit/test_multi_connector.py
+1
-4
@@ -1000,11 +1000,8 @@ def _make_multi_connector(connector_names: list[str]) -> MultiConnector:
1000
)
1001
1002
1003
-def test_multi_connector_hma_opt_in():
1003
+def test_multi_connector_hma_support_detection():
1004
"""
1005
- MultiConnector currently assumes HMA is opt-in: it needs
1006
- --no-disable-hybrid-kv-cache-manager to be enabled.
1007
-
1005
At runtime, _all_support_hma is True only when every sub-connector
1006
implements SupportsHMA. Test all combinations of HMA / non-HMA
1007
sub-connectors.
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+1
-2
@@ -723,8 +723,7 @@ def test_has_mamba_init(
723
724
block_size = 16
725
vllm_config = create_vllm_config(block_size=block_size)
726
- # VllmConfig.__post_init__ auto-disables HMA when kv_transfer_config
727
- # is set; override so we can test the scheduler's own derivation.
726
+ # Explicitly enable HMA so we can test the scheduler's own derivation.
727
vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
728
kv_cache_config = make_kv_cache_config(
729
block_size=block_size,
tests/v1/kv_connector/unit/test_offloading_connector.py
+1
-1
@@ -280,7 +280,7 @@ def test_cpu_offloading(
280
kv_events_config=kv_events_config,
281
kv_transfer_config=kv_transfer_config,
282
**({"attention_config": {"backend": attn_backend}} if attn_backend else {}),
283
- # HMA models need explicit opt-in when kv_transfer_config is set
283
+ # Keep HMA explicitly enabled for HMA model coverage.
284
**({"disable_hybrid_kv_cache_manager": False} if uses_hma else {}),
285
**({"enable_prefix_caching": True} if force_prefix_caching else {}),
286
# ROCm: batch size 1 to reduce variability
vllm/config/vllm.py
+12
-30
@@ -1406,7 +1406,7 @@ class VllmConfig:
1406
# Hybrid KV cache manager (HMA) runtime rules:
1407
# - Explicit enable (--no-disable-kv-cache-manager): error if runtime
1408
# disables it
1409
- # - No preference: auto-disable for unsupported features (e.g. kv connector)
1409
+ # - No preference: auto-disable for unsupported features or connector configs
1410
# - Explicit disable (--disable-kv-cache-manager): always respect it
1411
need_disable_hybrid_kv_cache_manager = False
1412
# logger should only print warning message for hybrid models. As we
@@ -1438,43 +1438,25 @@ class VllmConfig:
1438
need_disable_hybrid_kv_cache_manager = True
1439
1440
if self.scheduler_config.disable_hybrid_kv_cache_manager is None:
1441
- # Default to disable HMA, but only if the user didn't express a preference.
1441
+ # Auto-disable HMA only when the connector config does not support it.
1442
if self.kv_transfer_config is not None:
1443
- from vllm.config.kv_transfer import KVTransferConfig
1443
from vllm.distributed.kv_transfer.kv_connector.factory import (
1444
KVConnectorFactory,
1445
)
1447
- from vllm.distributed.kv_transfer.kv_connector.v1.base import (
1448
- supports_hma,
1449
- )
1446
1451
- connector_cls = KVConnectorFactory.get_connector_class(
1452
- self.kv_transfer_config
1453
- )
1454
- all_support_hma = supports_hma(connector_cls)
1455
- # MultiConnector subclasses SupportsHMA; only effectively
1456
- # supports HMA when every sub-connector does.
1457
- if all_support_hma and connector_cls.__name__ == "MultiConnector":
1458
- sub_ktcs = self.kv_transfer_config.kv_connector_extra_config.get(
1459
- "connectors", []
1460
- )
1461
- all_support_hma = all(
1462
- supports_hma(
1463
- KVConnectorFactory.get_connector_class(
1464
- KVTransferConfig(**sub)
1465
- )
1466
- )
1467
- for sub in sub_ktcs
1468
- )
1469
- if not all_support_hma:
1447
+ if not KVConnectorFactory.supports_hma_config(self.kv_transfer_config):
1448
need_disable_hybrid_kv_cache_manager = True
1449
logger.warning(
1450
"Turning off hybrid kv cache manager because "
1473
- "connector %s does not subclass `SupportsHMA`. "
1474
- "This will reduce performance on models with "
1475
- "sliding window or Mamba attention. See "
1476
- "kv_connector/v1/base.py for details.",
1477
- connector_cls.__name__,
1451
+ "`--kv-transfer-config` selects a KV connector that "
1452
+ "does not support it. Impact: hybrid SSM models "
1453
+ "(e.g. Jamba, Bamba) require HMA and will fail at "
1454
+ "startup without it; models with sliding window "
1455
+ "attention will run with reduced performance. "
1456
+ "To add HMA support to a KV connector, subclass "
1457
+ "`SupportsHMA` defined in kv_connector/v1/base.py "
1458
+ "(for MultiConnector, all child connectors must "
1459
+ "support HMA)."
1460
)
1461
self.scheduler_config.disable_hybrid_kv_cache_manager = (
1462
need_disable_hybrid_kv_cache_manager
vllm/distributed/kv_transfer/kv_connector/factory.py
+19
-2
@@ -5,6 +5,7 @@ import importlib
5
from collections.abc import Callable
6
from typing import TYPE_CHECKING, cast
7
8
+from vllm.config.kv_transfer import KVTransferConfig
9
from vllm.distributed.kv_transfer.kv_connector.base import (
10
KVConnectorBase,
11
KVConnectorBaseType,
@@ -18,7 +19,6 @@ from vllm.utils.func_utils import supports_kw
19
20
if TYPE_CHECKING:
21
from vllm.config import VllmConfig
21
- from vllm.config.kv_transfer import KVTransferConfig
22
from vllm.v1.kv_cache_interface import KVCacheConfig
23
24
logger = init_logger(__name__)
@@ -53,7 +53,7 @@ class KVConnectorFactory:
53
54
# check if the connector supports HMA
55
hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager
56
- if hma_enabled and not supports_hma(connector_cls):
56
+ if hma_enabled and not cls.supports_hma_config(kv_transfer_config):
57
raise ValueError(
58
f"Connector {connector_cls.__name__} does not support HMA but "
59
f"HMA is enabled. Please set `--disable-hybrid-kv-cache-manager`."
@@ -127,6 +127,23 @@ class KVConnectorFactory:
127
raise ValueError(f"Unsupported connector type: {connector_name}")
128
return connector_cls
129
130
+ @classmethod
131
+ def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool:
132
+ """Return whether this KV transfer config supports HMA.
133
+
134
+ MultiConnector is a special case: the wrapper class implements
135
+ SupportsHMA, but effective support depends on every configured child.
136
+ """
137
+ connector_cls = cls.get_connector_class(kv_transfer_config)
138
+ if kv_transfer_config.kv_connector != "MultiConnector":
139
+ return supports_hma(connector_cls)
140
+
141
+ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
142
+ MultiConnector,
143
+ )
144
+
145
+ return MultiConnector.all_children_support_hma(kv_transfer_config)
146
+
147
148
# Register various connectors here.
149
# The registration should not be done in each individual file, as we want to
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+20
-2
@@ -19,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19
KVConnectorRole,
20
KVConnectorWorkerMetadata,
21
SupportsHMA,
22
- supports_hma,
22
)
23
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
24
KVConnectorPromMetrics,
@@ -151,6 +150,22 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
150
return True
151
return False
152
153
+ @classmethod
154
+ def all_children_support_hma(cls, kv_transfer_config: "KVTransferConfig") -> bool:
155
+ """Return True only if every configured child connector supports HMA."""
156
+ connectors_config = kv_transfer_config.kv_connector_extra_config.get(
157
+ "connectors", []
158
+ )
159
+ if not connectors_config:
160
+ return False
161
+ for conn_config in connectors_config:
162
+ child_config = KVTransferConfig(
163
+ **{"engine_id": kv_transfer_config.engine_id, **conn_config}
164
+ )
165
+ if not KVConnectorFactory.supports_hma_config(child_config):
166
+ return False
167
+ return True
168
+
169
def __init__(
170
self,
171
vllm_config: "VllmConfig",
@@ -169,7 +184,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
184
self._connectors.append(connector_cls(temp_config, role, kv_cache_config))
185
self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
186
172
- self._all_support_hma = all(supports_hma(c) for c in self._connectors)
187
+ assert vllm_config.kv_transfer_config is not None
188
+ self._all_support_hma = MultiConnector.all_children_support_hma(
189
+ vllm_config.kv_transfer_config
190
+ )
191
assert (
192
vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
193
or self._all_support_hma