[Bugfix] Narrow AuxOutput KV restrictions to known PD connectors (#58150)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Sep 22, 2026 at 23:09 UTC
feb87b1934c10046a19cfef610c44f1710b825b8
2 files changed
+56
-11
tests/config/test_aux_output_config.py
+43
-3
@@ -7,6 +7,7 @@ from typing import Any
7
import pytest
8
9
from vllm.config import AuxOutputConfig, VllmConfig
10
+from vllm.config.kv_transfer import KVRole, KVTransferConfig
11
from vllm.engine.arg_utils import EngineArgs
12
13
pytestmark = pytest.mark.cpu_test
@@ -19,6 +20,7 @@ def _config(
20
dcp: int = 1,
21
pcp: int = 1,
22
connector: str | None = None,
23
+ kv_role: KVRole = "kv_both",
24
runner_type: str = "generate",
25
is_moe: bool = True,
26
sliding_window: int | None = None,
@@ -51,9 +53,9 @@ def _config(
53
kv_transfer_config=(
54
None
55
if connector is None
54
- else SimpleNamespace(
55
- is_kv_transfer_instance=True,
56
+ else KVTransferConfig(
57
kv_connector=connector,
58
+ kv_role=kv_role,
59
)
60
),
61
)
@@ -95,7 +97,6 @@ def test_legacy_routed_experts_flag_updates_aux_output_config():
97
({"pp": 2}, "pipeline parallelism"),
98
({"dcp": 2}, "context parallelism"),
99
({"pcp": 2}, "context parallelism"),
98
- ({"connector": "MooncakeConnector"}, "incompatible with KV connectors"),
100
],
101
)
102
def test_aux_output_connector_rejects_unsupported_configuration(kwargs, error):
@@ -103,6 +104,45 @@ def test_aux_output_connector_rejects_unsupported_configuration(kwargs, error):
104
VllmConfig._verify_aux_output_compatibility(_config(**kwargs))
105
106
107
+@pytest.mark.parametrize(
108
+ ("connector", "blocked"),
109
+ [
110
+ ("NixlConnector", True),
111
+ ("NixlPullConnector", True),
112
+ ("NixlPushConnector", True),
113
+ ("MoRIIOConnector", True),
114
+ ("MooncakeConnector", True),
115
+ ("MooncakeStoreConnector", False),
116
+ ("OffloadingConnector", False),
117
+ ("LMCacheConnectorV1", False),
118
+ ("LMCacheMPConnector", False),
119
+ ("SimpleCPUOffloadConnector", False),
120
+ ],
121
+)
122
+@pytest.mark.parametrize("kv_role", ["kv_both", "kv_producer", "kv_consumer"])
123
+@pytest.mark.parametrize("multi", [False, True])
124
+def test_aux_output_connector_policy_is_independent_of_role(
125
+ connector, blocked, kv_role, multi
126
+):
127
+ config = _config(connector=connector, kv_role=kv_role)
128
+ if multi:
129
+ config.kv_transfer_config = KVTransferConfig(
130
+ kv_connector="MultiConnector",
131
+ kv_role=kv_role,
132
+ kv_connector_extra_config={
133
+ "connectors": [
134
+ {"kv_connector": "OffloadingConnector", "kv_role": "kv_both"},
135
+ {"kv_connector": connector, "kv_role": kv_role},
136
+ ]
137
+ },
138
+ )
139
+ if blocked:
140
+ with pytest.raises(ValueError, match=f"incompatible with {connector}"):
141
+ VllmConfig._verify_aux_output_compatibility(config)
142
+ else:
143
+ VllmConfig._verify_aux_output_compatibility(config)
144
+
145
+
146
@pytest.mark.parametrize(
147
"kwargs",
148
[{"sliding_window": 4096}, {"attention_chunk_size": 4096}],
vllm/config/vllm.py
+13
-8
@@ -1126,14 +1126,19 @@ class VllmConfig:
1126
)
1127
1128
kv_transfer_config = self.kv_transfer_config
1129
- if (
1130
- kv_transfer_config is not None
1131
- and kv_transfer_config.is_kv_transfer_instance
1132
- ):
1133
- raise ValueError(
1134
- "--enable-return-routed-experts is incompatible with KV "
1135
- "connectors (PD disaggregation and KV cache offload)."
1136
- )
1129
+ if kv_transfer_config is not None:
1130
+ for connector_name in (
1131
+ "NixlConnector",
1132
+ "NixlPullConnector",
1133
+ "NixlPushConnector",
1134
+ "MoRIIOConnector",
1135
+ "MooncakeConnector",
1136
+ ):
1137
+ if kv_transfer_config.has_connector(connector_name):
1138
+ raise ValueError(
1139
+ "--enable-return-routed-experts is incompatible with "
1140
+ f"{connector_name}; PD auxiliary output is not supported."
1141
+ )
1142
1143
def _verify_kv_transfer_compat(self) -> None:
1144
"""Reject configurations that silently corrupt KV transfers."""