[Bugfix][Hunyuan/Bagel]Avoid payload connector for KV-only senders (#5744)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>

tests/worker/test_omni_connector_mixin.py

@@ -105,6 +105,34 @@ class _FakeTPGroup:
 # ------------------------------------------------------------------ #
 
 
+@pytest.mark.parametrize(
+    ("role", "custom_func", "expected"),
+    [
+        ("sender", None, False),
+        ("sender", "package.build_payload", True),
+        ("receiver", None, True),
+    ],
+)
+def test_init_payload_connector_ownership(role, custom_func, expected):
+    model_config = _make_model_config(custom_func=custom_func)
+    model_config.stage_connector_config = {
+        "name": "MooncakeTransferEngineConnector",
+        "extra": {"role": role},
+    }
+
+    host = MixinHost()
+    connector = MockConnector()
+    with (
+        patch.object(host, "_create_connector", return_value=connector) as create,
+        patch.object(host, "_load_custom_func", return_value=(None, None)),
+    ):
+        host.init_omni_connectors(model_config)
+
+    assert (create.call_count == 1) is expected
+    assert (host._omni_connector is connector) is expected
+    host.shutdown_omni_connectors()
+
+
 class TestMixinAsyncChunkSendRecv(unittest.TestCase):
     """Test 2: Async chunk send/recv + bg threads."""
 

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -25,7 +25,10 @@ from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
-from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+from vllm_omni.distributed.omni_connectors.utils.config import (
+    ConnectorSpec,
+    get_stage_connector_role,
+)
 from vllm_omni.outputs import OmniConnectorOutput
 from vllm_omni.worker.payload_span import (
     get_tensor_span,
@@ -46,6 +49,19 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _should_create_payload_connector(model_config: Any) -> bool:
+    """Whether this stage owns runner payload transport for its edge.
+
+    Sender edges may instead be owned solely by KV transfer. Receivers still
+    need a connector even though they do not declare a downstream payload hook.
+    """
+    if get_stage_connector_role(model_config) != "sender":
+        return True
+
+    next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None)
+    return isinstance(next_stage_func, str) and bool(next_stage_func)
+
+
 def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool:
     """Producer-side structural gate.
 
@@ -96,7 +112,9 @@ class OmniConnectorModelRunnerMixin:
             model_config: Stage-level model config with connector settings.
             kv_transfer_manager: Existing KV transfer manager to delegate to.
         """
-        self._omni_connector: OmniConnectorBase | None = self._create_connector(model_config)
+        self._omni_connector: OmniConnectorBase | None = (
+            self._create_connector(model_config) if _should_create_payload_connector(model_config) else None
+        )
         self._kv_transfer_manager = kv_transfer_manager
 
         self._async_chunk: bool = getattr(model_config, "async_chunk", False)