[Refactor][2/N Scheduler]Replace _FULL_PAYLOAD_INPUT_STAGES with a resolved stage transport capability. (#6149)

Signed-off-by: wy17003 <yaow17003@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/config/test_config_factory.py

@@ -2753,22 +2753,23 @@ class TestSentinelDefaultPrecedence:
         )
 
     def test_ming_flash_omni_topology(self):
-        """Guard ming_flash_omni's SIP cleanup: stage 0 has no full-payload
-        producer hook (arch is not in ``_FULL_PAYLOAD_INPUT_STAGES``), and
-        stage 1 uses only ``thinker2talker_token_only`` (sync_process_input_func).
+        """Guard ming_flash_omni's SIP cleanup: stage 0 has no
+        ``custom_process_next_stage_input_func``, and stage 1 does not set
+        ``requires_full_payload_input=True``. Stage 1 uses only
+        ``thinker2talker_token_only`` (sync_process_input_func).
         The dead ``thinker2talker`` (custom_process_input_func) was removed
         because ming_flash_omni has no async_chunk support and both functions
         called ``_build_talker_inputs`` identically.
-        Merge under either async_chunk mode must not re-introduce a
-        stage-0 full-payload hook."""
+        Merge under either async_chunk mode must not re-introduce either
+        full-payload transport gate."""
         pipeline = resolve_pipeline_config("ming_flash_omni")
         assert isinstance(pipeline, PipelineConfig)
 
         stage0, stage1 = pipeline.stages
         assert stage0.custom_process_next_stage_input_func is None, (
-            "ming_flash_omni stage 0 must not declare a full-payload producer "
-            "(connector path is not active for this arch)."
+            "ming_flash_omni stage 0 must not declare a full-payload producer hook."
         )
+        assert stage1.requires_full_payload_input is False
         assert stage1.custom_process_input_func is None
         assert stage1.sync_process_input_func is not None
         assert stage1.sync_process_input_func.endswith("thinker2talker_token_only")

tests/config/test_omni_config.py

@@ -244,6 +244,36 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
         VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
 
 
+@pytest.mark.parametrize(
+    ("engine_extras", "cli_overrides"),
+    [
+        ({"requires_full_payload_input": False}, {}),
+        ({}, {"stage_1_requires_full_payload_input": False}),
+    ],
+    ids=["deploy", "stage-cli"],
+)
+def test_from_pipeline_config_rejects_full_payload_input_capability_overrides(
+    engine_extras,
+    cli_overrides,
+):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    deploy = DeployConfig(
+        stages=[
+            StageDeployConfig(
+                stage_id=1,
+                engine_extras=engine_extras,
+            )
+        ]
+    )
+
+    with pytest.raises(ValueError, match=r"no structured config owner: requires_full_payload_input"):
+        VllmOmniConfig.from_pipeline_config(
+            pipeline,
+            user_deploy_config=deploy,
+            cli_overrides=cli_overrides,
+        )
+
+
 @pytest.mark.parametrize(
     ("cli_overrides", "stage_id"),
     [
@@ -443,10 +473,12 @@ def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating
     async_config = _from_pipeline_key("qwen3_tts")
     assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk")
     assert async_config.stage_by_id(1).custom_process_input_func is None
+    assert async_config.stage_by_id(1).model_config.requires_full_payload_input is True
 
     sync_config = _from_pipeline_key("qwen3_tts", cli_overrides={"async_chunk": False})
     assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
     assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
+    assert sync_config.stage_by_id(1).model_config.requires_full_payload_input is True
 
     assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
     assert pipeline.get_stage(1).custom_process_input_func is None
@@ -567,6 +599,7 @@ def test_sub_config_fields_match_structured_scopes():
         # legacy engine-args path.
         "model_subdir",
         "tokenizer_subdir",
+        "requires_full_payload_input",
     }
     vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
     assert issubclass(OmniStageLoadConfig, VllmLoadConfig)

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -4,10 +4,8 @@ from types import SimpleNamespace
 import pytest
 from vllm.v1.engine import FinishReason
 
-from vllm_omni.core.sched.omni_scheduler_mixin import (
-    DEFAULT_INPUT_WAIT_TIMEOUT_S,
-    OmniSchedulerMixin,
-)
+from vllm_omni.core.sched import omni_scheduler_mixin
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniChunkRecvHandle
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -17,6 +15,77 @@ class _Scheduler(OmniSchedulerMixin):
     pass
 
 
+def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
+    created_adapters = []
+
+    def adapter_factory(config):
+        adapter = SimpleNamespace(vllm_config=config)
+        created_adapters.append(adapter)
+        return adapter
+
+    monkeypatch.setattr(omni_scheduler_mixin, "OmniChunkTransferAdapter", adapter_factory)
+
+    def init_scheduler(**model_config):
+        scheduler = _Scheduler()
+        scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(**model_config))
+        scheduler._init_omni_io_scheduling_state()
+        return scheduler
+
+    producer = init_scheduler(
+        stage_id=0,
+        async_chunk=True,
+        requires_full_payload_input=False,
+        custom_process_next_stage_input_func="test.pipeline.produce_async_chunk",
+    )
+    receiver = init_scheduler(
+        stage_id=1,
+        async_chunk=True,
+        requires_full_payload_input=True,
+        custom_process_next_stage_input_func=None,
+    )
+    terminal_stage = init_scheduler(
+        stage_id=2,
+        async_chunk=True,
+        requires_full_payload_input=False,
+        custom_process_next_stage_input_func=None,
+    )
+
+    assert producer.chunk_transfer_adapter is not None
+    assert receiver.chunk_transfer_adapter is not None
+    assert terminal_stage.chunk_transfer_adapter is not None
+    assert receiver.input_coordinator is None
+    assert len(created_adapters) == 3
+
+
+@pytest.mark.parametrize(
+    ("stage_id", "async_chunk", "required", "enabled"),
+    [
+        (0, False, True, False),
+        (1, True, True, False),
+        (1, False, False, False),
+        (1, False, True, True),
+    ],
+)
+def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, async_chunk, required, enabled):
+    monkeypatch.setattr(
+        omni_scheduler_mixin,
+        "OmniChunkTransferAdapter",
+        lambda _config: SimpleNamespace(),
+    )
+    scheduler = _Scheduler()
+    scheduler.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            stage_id=stage_id,
+            async_chunk=async_chunk,
+            requires_full_payload_input=required,
+        )
+    )
+
+    scheduler._init_omni_io_scheduling_state()
+
+    assert (scheduler.input_coordinator is not None) is enabled
+
+
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     calls = []
     scheduler = _Scheduler()
@@ -58,7 +127,7 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
         ("process", scheduler.waiting, scheduler.running, scheduler.requests),
         # The chunk deadline runs after chunks are applied, so a chunk that
         # arrived this cycle resets the clock before it is measured (R1.1).
-        ("chunk-timeouts", DEFAULT_INPUT_WAIT_TIMEOUT_S),
+        ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
         ("failed-sends",),
         ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
         ("restore-full", scheduler.waiting),

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -20,7 +20,6 @@ import torch
 import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
-    uses_full_payload_input_coordinator,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -96,109 +95,6 @@ class MockQueue:
 # ------------------------------------------------------------------ #
 
 
-class TestFullPayloadCoordinatorSelection(unittest.TestCase):
-    """Tests for the (model_arch, model_stage) whitelist gate.
-
-    The init_omni_connectors arch allowlist is keyed by ``model_arch`` and
-    is a superset of the stages registered here -- consumer-wait stages
-    must be registered explicitly in ``_FULL_PAYLOAD_INPUT_STAGES``, while
-    the init allowlist covers both producer- and consumer-side runners.
-    These tests pin which ``(arch, stage)`` pairs the gate fires for today.
-    """
-
-    # Expected whitelist (model_arch, model_stage).  Hardcoded to avoid the
-    # tautology of importing _FULL_PAYLOAD_INPUT_STAGES and asserting it
-    # against itself; any drift between this matrix and the whitelist will
-    # fail loudly here.
-    EXPECTED_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
-        {
-            ("Qwen3OmniMoeForConditionalGeneration", "talker"),
-            ("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
-            ("Qwen2_5OmniForConditionalGeneration", "talker"),
-            ("Qwen2_5OmniForConditionalGeneration", "code2wav"),
-            ("CovoAudioForConditionalGeneration", "code2wav"),
-            ("MiMoAudioModel", "code2wav"),
-            ("Qwen3TTSCode2Wav", "code2wav"),
-            ("MiniCPMO45Code2Wav", "code2wav"),
-            # nemotron_voicechat: sync-path RVQ-VAE decode waits for the
-            # talker's full [frames, 31] code-stack payload.
-            ("NemotronVoiceChatCode2Wav", "code2wav"),
-            ("CosyVoice3Model", "cosyvoice3_code2wav"),
-            # audex: sync-path TTS decoder (async_chunk: false deploys) and
-            # the TTA XCodec1 stage (always sync full-payload).
-            ("AudexCode2Wav", "audex_code2wav"),
-            ("AudexXCodec1", "audex_xcodec"),
-            ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
-            ("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"),
-            ("DyninOmniForConditionalGeneration", "token2image"),
-            ("DyninOmniForConditionalGeneration", "token2audio"),
-        }
-    )
-
-    def test_whitelist_matches_expected_matrix(self):
-        """_FULL_PAYLOAD_INPUT_STAGES must equal the hardcoded expected matrix.
-
-        Catches both accidental additions (which would silently enable the
-        consumer-wait gate for a new arch) and accidental removals (which
-        would silently disable an enabled arch).
-        """
-        from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
-
-        self.assertEqual(
-            frozenset(_FULL_PAYLOAD_INPUT_STAGES),
-            self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES,
-            msg="_FULL_PAYLOAD_INPUT_STAGES drifted from the expected matrix; "
-            "update EXPECTED_FULL_PAYLOAD_INPUT_STAGES if intentional.",
-        )
-
-    def test_all_whitelisted_arch_stage_pairs_fire_gate(self):
-        """Every (arch, stage) pair in the expected matrix must fire
-        the gate when stage_id > 0 and async_chunk=False.
-        """
-        for arch, stage in self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES:
-            model_config = SimpleNamespace(
-                stage_id=1,
-                async_chunk=False,
-                model_arch=arch,
-                model_stage=stage,
-            )
-            self.assertTrue(
-                uses_full_payload_input_coordinator(model_config),
-                msg=f"expected gate to fire for {arch}/{stage}",
-            )
-
-    def test_other_arch_or_stage_or_mode_does_not_fire(self):
-        cases = [
-            SimpleNamespace(
-                stage_id=1, async_chunk=True, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="talker"
-            ),
-            SimpleNamespace(
-                stage_id=0, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="thinker"
-            ),
-            SimpleNamespace(
-                stage_id=1,
-                async_chunk=False,
-                model_arch="Qwen3OmniMoeForConditionalGeneration",
-                model_stage="some_future_stage",
-            ),
-            SimpleNamespace(
-                stage_id=1, async_chunk=False, model_arch="Qwen3TTSForConditionalGeneration", model_stage="code2wav"
-            ),
-            SimpleNamespace(
-                stage_id=1, async_chunk=False, model_arch="MingFlashOmniForConditionalGeneration", model_stage="talker"
-            ),
-            SimpleNamespace(stage_id=1, async_chunk=False, model_arch=None, model_stage="talker"),
-            SimpleNamespace(
-                stage_id=1, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage=None
-            ),
-        ]
-        for model_config in cases:
-            self.assertFalse(
-                uses_full_payload_input_coordinator(model_config),
-                msg=f"expected gate OFF for {model_config}",
-            )
-
-
 class TestCoordinatorUpdateRequestMetadata(unittest.TestCase):
     """Test update_request_metadata applies scheduling metadata to requests."""
 

tests/engine/test_arg_utils.py

@@ -8,6 +8,8 @@ import inspect
 import json
 import os
 import shutil
+import sys
+import types
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -20,6 +22,8 @@ from vllm.engine.arg_utils import EngineArgs
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -42,7 +46,102 @@ def test_default_stage_id_is_concrete_int():
     assert cfg.stage_id == 0
 
 
-def test_multimodal_kwarg_overrides(mocker):
+def test_full_payload_capability_reaches_omni_model_config(monkeypatch):
+    """Keep the pipeline capability intact through the engine adapter."""
+    captured: dict[str, object] = {}
+    baseline_config = Mock()
+    worker_module = types.ModuleType("test_connector_worker")
+
+    class ConnectorRunner(OmniConnectorModelRunnerMixin):
+        pass
+
+    class ConnectorWorker:
+        model_runner_cls = ConnectorRunner
+
+    worker_module.ConnectorWorker = ConnectorWorker
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+    monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda _self: None)
+    monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config)
+
+    def capture_omni_config(_cls, model_config, **omni_kwargs):
+        assert model_config is baseline_config
+        captured.update(omni_kwargs)
+        return baseline_config
+
+    monkeypatch.setattr(
+        OmniModelConfig,
+        "from_vllm_model_config",
+        classmethod(capture_omni_config),
+    )
+
+    OmniEngineArgs(
+        stage_id=1,
+        requires_full_payload_input=True,
+        worker_cls="test_connector_worker.ConnectorWorker",
+    ).create_model_config()
+
+    assert captured["requires_full_payload_input"] is True
+
+
+def test_stage_without_connector_configuration_accepts_plain_runner(monkeypatch):
+    worker_module = types.ModuleType("test_async_scheduler_worker")
+
+    class WorkerWithoutConnector:
+        model_runner_cls = object
+
+    worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+    args = OmniEngineArgs(
+        stage_id=1,
+        async_chunk=True,
+        worker_cls="test_async_scheduler_worker.WorkerWithoutConnector",
+    )
+
+    assert args.requires_full_payload_input is False
+
+
+def test_full_payload_capability_requires_selected_worker_connector(monkeypatch):
+    worker_module = types.ModuleType("test_worker_without_connector")
+
+    class WorkerWithoutConnector:
+        model_runner_cls = object
+
+    worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+    with pytest.raises(ValueError, match="does not provide an Omni connector model runner"):
+        OmniEngineArgs(
+            stage_id=1,
+            requires_full_payload_input=True,
+            worker_cls="test_worker_without_connector.WorkerWithoutConnector",
+        )
+
+
+def test_full_payload_capability_validates_platform_selected_worker(monkeypatch):
+    worker_module = types.ModuleType("test_platform_worker_without_connector")
+
+    class WorkerWithoutConnector:
+        model_runner_cls = object
+
+    worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+    monkeypatch.setattr(
+        current_omni_platform,
+        "get_omni_ar_worker_cls",
+        lambda: "test_platform_worker_without_connector.WorkerWithoutConnector",
+    )
+
+    with pytest.raises(ValueError, match="does not provide an Omni connector model runner"):
+        OmniEngineArgs(
+            stage_id=1,
+            requires_full_payload_input=True,
+            worker_type="ar",
+        )
+
+
+def test_multimodal_kwarg_overrides(monkeypatch):
     """Ensure that overrides in the multimodal config are preserved."""
     sig = inspect.signature(OmniEngineArgs)
     default_mm_cache = sig.parameters["mm_processor_cache_gb"].default
@@ -56,8 +155,8 @@ def test_multimodal_kwarg_overrides(mocker):
         assert self.mm_processor_cache_gb == override_val
         return fake_model_config
 
-    mocker.patch.object(EngineArgs, "create_model_config", _fake_parent_create_model_config)
-    mocker.patch.object(OmniModelConfig, "from_vllm_model_config", side_effect=lambda model_config, **_: model_config)
+    monkeypatch.setattr(EngineArgs, "create_model_config", _fake_parent_create_model_config)
+    monkeypatch.setattr(OmniModelConfig, "from_vllm_model_config", lambda model_config, **_: model_config)
 
     cfg = OmniEngineArgs(
         model="Qwen/Qwen2-VL-2B-Instruct",

tests/engine/test_stage_engine_args.py

@@ -48,6 +48,7 @@ from vllm_omni.engine.stage_init_utils import (
     build_engine_args_dict_from_omni_stage_config,
     build_legacy_engine_args_dict,
 )
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -123,6 +124,17 @@ def _stable_engine_arg_environment(monkeypatch, tmp_path):
 
     monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", _resolve_test_worker)
 
+    class _TestConnectorRunner(OmniConnectorModelRunnerMixin):
+        pass
+
+    class _TestWorker:
+        model_runner_cls = _TestConnectorRunner
+
+    for worker_type in ("ar", "generation"):
+        worker_module = types.ModuleType(f"test.{worker_type}")
+        worker_module.Worker = _TestWorker
+        monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
 
 def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]:
     parent_model = tmp_path / "parent-model"

tests/model_executor/models/test_nemotron_voicechat_registration.py

@@ -67,15 +67,6 @@ def test_architectures_registered() -> None:
         assert arch in _OMNI_MODELS, f"{arch} missing from _OMNI_MODELS"
 
 
-def test_full_payload_whitelists_cover_code2wav_only() -> None:
-    from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
-
-    assert ("NemotronVoiceChatCode2Wav", "code2wav") in _FULL_PAYLOAD_INPUT_STAGES
-    # The thinker->talker hop is token-path only: the talker must NOT wait on a
-    # connector payload or stage 1 hangs.
-    assert ("NemotronVoiceChatTalkerForConditionalGeneration", "talker") not in _FULL_PAYLOAD_INPUT_STAGES
-
-
 def test_config_parses_nemo_layout_without_hub_access(monkeypatch) -> None:
     # Provide the thinker text config explicitly so the test never touches HF.
     cfg = NemotronVoiceChatConfig(

vllm_omni/config/model.py

@@ -152,6 +152,7 @@ class OmniModelConfig(ModelConfig):
     # place it can learn which request-shaping conventions a stage uses (e.g.
     # ``cfg_role`` for classifier-free-guidance request pairs).
     sampling_extra_args_keys: tuple[str, ...] = ()
+    requires_full_payload_input: bool = False
 
     @property
     def registry(self):

vllm_omni/config/omni_config.py

@@ -404,6 +404,7 @@ class OmniStageModelConfig:
     # StagePipelineConfig.model_subdir/tokenizer_subdir on the legacy path.
     model_subdir: str | None = None
     tokenizer_subdir: str | None = None
+    requires_full_payload_input: bool = False
 
 
 @_enforce_keyword_only_init
@@ -1603,6 +1604,7 @@ def _build_model_config(
 ) -> OmniStageModelConfig:
     default_sampling_params = _stage_sampling_params(stage_deploy, topology)
     kwargs = _config_kwargs(engine)
+    kwargs["requires_full_payload_input"] = topology.requires_full_payload_input
     kwargs["model"] = _first_defined(kwargs.get("model"), model)
     if "model_arch" not in kwargs:
         kwargs["model_arch"] = topology.model_arch or pipeline.model_arch or None

vllm_omni/config/stage_config.py

@@ -201,6 +201,9 @@ class StagePipelineConfig:
     # by ``stage_init_utils._resolve_model_tokenizer_paths``.
     model_subdir: str | None = None
     tokenizer_subdir: str | None = None
+    # Whether the non-async path waits for a complete upstream payload from
+    # the model-runner connector before scheduling this stage.
+    requires_full_payload_input: bool = False
     extras: dict[str, Any] = field(default_factory=dict)
 
 
@@ -854,6 +857,7 @@ def _build_engine_args(
         engine_args["duplex_max_sessions"] = deploy.duplex_session.max_sessions
     if ps.omni_kv_config:
         engine_args["omni_kv_config"] = dict(ps.omni_kv_config)
+    engine_args["requires_full_payload_input"] = ps.requires_full_payload_input
     return engine_args
 
 

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -22,75 +22,19 @@ from vllm_omni.core.sched.output import OmniChunkRecvHandle
 logger = init_logger(__name__)
 
 
-# (arch, model_stage) pairs that route their full_payload stage input via
-# the worker connector and therefore need the scheduler-side coordinator to
-# park requests in WAITING_FOR_INPUT until the recv side delivers.  This set
-# must stay aligned with the arch scope of `init_omni_connectors` in
-# gpu_ar_model_runner.py and gpu_generation_model_runner.py.  Adding a stage
-# here without also wiring its worker connector init produces a permanent
-# Stage 1 hang (gate parks the request, no transport ever releases it).
-#
-_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
-    {
-        ("Qwen3OmniMoeForConditionalGeneration", "talker"),
-        ("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
-        # qwen2_5_omni thinker->talker uses the real full-payload
-        # producer builder (text_hidden_states routed via
-        # pooler_output["hidden"] -> accumulator -> connector).  Both
-        # stages of qwen2_5_omni are enabled.
-        ("Qwen2_5OmniForConditionalGeneration", "talker"),
-        ("Qwen2_5OmniForConditionalGeneration", "code2wav"),
-        # covo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1).
-        ("CovoAudioForConditionalGeneration", "code2wav"),
-        # mimo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1).
-        ("MiMoAudioModel", "code2wav"),
-        # qwen3_tts: Qwen3TTSTalkerForConditionalGeneration (Stage 0)
-        # -> Qwen3TTSCode2Wav (Stage 1).  Stage 1 is the consumer.
-        ("Qwen3TTSCode2Wav", "code2wav"),
-        # minicpmo_4_5: Talker (Stage 1) -> Code2Wav (Stage 2).
-        ("MiniCPMO45Code2Wav", "code2wav"),
-        # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
-        ("CosyVoice3Model", "cosyvoice3_code2wav"),
-        # nemotron_voicechat: talker (Stage 1) -> code2wav (Stage 2). Stage 2
-        # waits for the talker's full-payload code stacks; the thinker (Stage 0)
-        # -> talker (Stage 1) hop is token-path only and must NOT be listed.
-        ("NemotronVoiceChatCode2Wav", "code2wav"),
-        # audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1).
-        # The default deploy is async_chunk; this covers async_chunk: false.
-        ("AudexCode2Wav", "audex_code2wav"),
-        # audex TTA: tta thinker (Stage 0) -> XCodec1 (Stage 1, always sync
-        # full-payload — CNN codec decoded over the full sequence).
-        ("AudexXCodec1", "audex_xcodec"),
-        # indextts2 / indextts2_5: talker (Stage 0) -> s2mel decoder
-        # (Stage 1). Stage 1 consumes the complete mel/optional-latent payload.
-        ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
-        ("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"),
-        # dynin: token2text (Stage 0) -> token2image (Stage 1) ->
-        # token2audio (Stage 2).  Producer wires via
-        # custom_process_next_stage_input_func: *_full_payload in deploy yaml.
-        ("DyninOmniForConditionalGeneration", "token2image"),
-        ("DyninOmniForConditionalGeneration", "token2audio"),
-    }
-)
-
-
 def uses_full_payload_input_coordinator(model_config: Any) -> bool:
     """Returns True if this stage parks pending requests in
     WAITING_FOR_INPUT awaiting a full_payload delivery on the worker connector.
 
-    Gated by (model_arch, model_stage) — see _FULL_PAYLOAD_INPUT_STAGES for the
-    rationale on why this is a whitelist instead of a marker-driven structural
-    gate.
+    Gated by the topology-declared ``requires_full_payload_input`` capability on
+    downstream (stage_id > 0) stages, and only on the non-async-chunk path
+    (async-chunk stages are fed through the streamed connector instead).
     """
     if getattr(model_config, "stage_id", 0) <= 0:
         return False
     if getattr(model_config, "async_chunk", False):
         return False
-    key = (
-        getattr(model_config, "model_arch", None),
-        getattr(model_config, "model_stage", None),
-    )
-    return key in _FULL_PAYLOAD_INPUT_STAGES
+    return bool(getattr(model_config, "requires_full_payload_input", False))
 
 
 class OmniSchedulingCoordinator:

vllm_omni/engine/arg_utils.py

@@ -16,6 +16,7 @@ from vllm_omni.config import OmniModelConfig
 from vllm_omni.outputs.output_modality import OutputModality
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.plugins import load_omni_general_plugins
+from vllm_omni.worker.omni_connector_validation import validate_worker_omni_connector
 
 logger = init_logger(__name__)
 
@@ -186,6 +187,7 @@ class OmniEngineArgs(EngineArgs):
     engine_output_type: str | None = None
     hf_config_name: str | None = None
     custom_process_next_stage_input_func: str | None = None
+    requires_full_payload_input: bool = False
     stage_connector_spec: dict[str, Any] = field(default_factory=dict)
     subtalker_sampling_params: dict[str, Any] | None = None
     silence_ban_frames: int = 0
@@ -240,12 +242,19 @@ class OmniEngineArgs(EngineArgs):
     sampling_extra_args_keys: tuple[str, ...] = ()
 
     def __post_init__(self) -> None:
-        if self.worker_cls is None:
+        if self.worker_cls in (None, "auto"):
             if self.worker_type == "ar":
                 self.worker_cls = current_omni_platform.get_omni_ar_worker_cls()
             elif self.worker_type == "generation":
                 self.worker_cls = current_omni_platform.get_omni_generation_worker_cls()
         load_omni_general_plugins()
+
+        connector_extra = self.stage_connector_spec.get("extra")
+        connector_role = connector_extra.get("role") if isinstance(connector_extra, dict) else None
+        needs_connector = bool(
+            self.requires_full_payload_input or self.custom_process_next_stage_input_func or connector_role is not None
+        )
+        validate_worker_omni_connector(self.worker_cls, needs_connector)
         super().__post_init__()
 
     def _ensure_omni_models_registered(self):
@@ -416,6 +425,7 @@ class OmniEngineArgs(EngineArgs):
             engine_output_type=self.engine_output_type,
             hf_config_name=self.hf_config_name,
             custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
+            requires_full_payload_input=self.requires_full_payload_input,
             stage_connector_config=stage_connector_config,
             subtalker_sampling_params=self.subtalker_sampling_params,
             silence_ban_frames=self.silence_ban_frames,

vllm_omni/model_executor/models/audex/pipeline.py

@@ -66,6 +66,7 @@ AUDEX_TTS_PIPELINE = PipelineConfig(
             tokenizer_subdir="checkpoint_folder_audiogen",
             sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )
@@ -127,6 +128,7 @@ AUDEX_S2S_PIPELINE = PipelineConfig(
             tokenizer_subdir="checkpoint_folder_audiogen",
             sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )
@@ -190,6 +192,7 @@ AUDEX_TTA_PIPELINE = PipelineConfig(
             # payload ships via thinker2xcodec_full_payload.
             sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/aura_omni/pipeline.py

@@ -80,6 +80,7 @@ AURA_OMNI_PIPELINE = PipelineConfig(
             sync_process_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
             extras={"tts_args": {"max_instructions_length": 500}},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/cosyvoice3/pipeline.py

@@ -48,6 +48,7 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="latent",
             sync_process_input_func=f"{_PROC}.text2flow_token_only",
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/covo_audio/pipeline.py

@@ -47,6 +47,7 @@ COVO_AUDIO_PIPELINE = PipelineConfig(
             engine_output_type="audio",
             sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/dynin_omni/pipeline.py

@@ -50,6 +50,7 @@ DYNIN_OMNI_PIPELINE = PipelineConfig(
             engine_output_type="latent",
             custom_process_input_func=f"{_PROC}.token2text_to_token2image",
             custom_process_next_stage_input_func=f"{_PROC}.token2image_to_token2audio_full_payload",
+            requires_full_payload_input=True,
         ),
         StagePipelineConfig(
             stage_id=2,
@@ -60,6 +61,7 @@ DYNIN_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="latent",
             custom_process_input_func=f"{_PROC}.token2image_to_token2audio",
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/indextts2/pipeline.py

@@ -45,6 +45,7 @@ INDEXTTS2_PIPELINE = PipelineConfig(
             sync_process_input_func=f"{_PROC}.talker2s2mel_token_only",
             extras={"skip_tokenizer_init": True},
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )
@@ -81,6 +82,7 @@ INDEXTTS25_PIPELINE = PipelineConfig(
             sync_process_input_func=f"{_PROC}.talker2s2mel_token_only",
             extras={"skip_tokenizer_init": True},
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/mimo_audio/pipeline.py

@@ -59,6 +59,7 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             engine_output_type="audio",
             sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py

@@ -80,6 +80,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig(
             model_arch="MiniCPMO45Code2Wav",
             sync_process_input_func=f"{_PROC}.tts2code2wav_token_only",
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py

@@ -80,6 +80,7 @@ NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig(
             engine_output_type="audio",
             sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
             sampling_constraints={"detokenize": False},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py

@@ -47,6 +47,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
                 "detokenize": True,
                 "stop_token_ids": [8294],
             },
+            requires_full_payload_input=True,
         ),
         StagePipelineConfig(
             stage_id=2,
@@ -59,6 +60,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
             engine_output_type="audio",
             sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -59,6 +59,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
                 "detokenize": False,
                 "stop_token_ids": [2150],
             },
+            requires_full_payload_input=True,
         ),
         StagePipelineConfig(
             stage_id=2,
@@ -70,6 +71,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             hf_config_name="thinker_config",
             engine_output_type="audio",
             sampling_constraints={"detokenize": True},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/models/qwen3_tts/pipeline.py

@@ -52,6 +52,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
             sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
             extras={"tts_args": {"max_instructions_length": 500}},
+            requires_full_payload_input=True,
         ),
     ),
 )

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -23,8 +23,7 @@ TALKER_CODEC_END_TOKEN_ID = 8294
 # ============================================================================
 # Worker-connector data plane (non-async-chunk path).
 # Both transitions ship payloads via the worker connector
-# (registered in ``_FULL_PAYLOAD_INPUT_STAGES`` in
-# omni_scheduling_coordinator):
+# (requires_full_payload_input=True in the pipeline topology):
 # - thinker->talker reads accumulated ``pooling_output["hidden"]`` and
 #   packs an OmniPayload-shaped dict (embed.prefill /
 #   hidden_states.output / ids.prompt / ids.output) for the talker, which
@@ -155,9 +154,9 @@ def thinker2talker_token_only(
     ``thinker2talker_full_payload`` via the worker connector and lands
     in ``model_intermediate_buffer`` before the talker's forward() runs.
 
-    Consumer-wait gating is whitelist-driven via
-    ``_FULL_PAYLOAD_INPUT_STAGES`` (see the mixin
-    ``should_accumulate_full_payload_output`` docstring).
+    Consumer-wait gating is driven by
+    ``requires_full_payload_input=True`` in the pipeline topology
+    (see the mixin ``should_accumulate_full_payload_output`` docstring).
     """
     thinker_outputs = source_outputs
     talker_inputs = []

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -44,15 +44,15 @@ from vllm_ascend.worker.model_runner_v1 import graph_capture
 
 from vllm_omni.data_entry_keys import flatten_payload
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
-from vllm_omni.distributed.omni_connectors.utils.config import (
-    get_stage_connector_role,
-    stage_sends_async_output,
-)
+from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output
 from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
-from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+from vllm_omni.worker.omni_connector_model_runner_mixin import (
+    OmniConnectorModelRunnerMixin,
+    needs_omni_connector,
+)
 from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
 
 
@@ -119,23 +119,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         # Initialize KV cache manager (preserve vllm_config fallback behavior)
         self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
         self._async_chunk = getattr(self.model_config, "async_chunk", False)
-        #  -------------------------------------- Omni-new -------------------------------------------------
-        _OMNI_CONNECTOR_INIT_ARCHS = {
-            "Qwen3OmniMoeForConditionalGeneration",
-            "Qwen2_5OmniForConditionalGeneration",
-            "CovoAudioForConditionalGeneration",
-            "MiMoAudioModel",
-            "Qwen3TTSTalkerForConditionalGeneration",
-            "Qwen3TTSCode2Wav",
-            "CosyVoice3Model",
-            "DyninOmniForConditionalGeneration",
-            "IndexTTS2TalkerForConditionalGeneration",
-        }
-        stage_archs = set(getattr(self.model_config, "architectures", None) or ())
-        model_arch_override = getattr(self.model_config, "model_arch", None)
-        if model_arch_override:
-            stage_archs.add(model_arch_override)
-        if stage_archs & _OMNI_CONNECTOR_INIT_ARCHS or get_stage_connector_role(self.model_config) is not None:
+        if needs_omni_connector(self.model_config):
             self.init_omni_connectors(
                 model_config=self.model_config,
                 kv_transfer_manager=self.kv_transfer_manager,

vllm_omni/platforms/npu/worker/npu_ar_worker.py

@@ -11,9 +11,11 @@ from vllm_omni.worker.mixins import OmniWorkerMixin
 class NPUARWorker(OmniWorkerMixin, OmniNPUWorkerBase):
     """NPU AR worker for thinker/talker stages in Omni model."""
 
+    model_runner_cls = NPUARModelRunner
+
     def init_device(self):
         self.device = self._init_device()
         num_ubatches = 1
         init_workspace_manager(self.device, num_ubatches)
 
-        self.model_runner = NPUARModelRunner(self.vllm_config, self.device)
+        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -33,12 +33,14 @@ from vllm_ascend.ops.rotary_embedding import update_cos_sin
 from vllm_ascend.utils import enable_sp, lmhead_tp_enable
 from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE
 
-from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_role
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState, _ensure_tensor_values
 from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 from vllm_omni.utils.mm_outputs import partition_payload_list
-from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+from vllm_omni.worker.omni_connector_model_runner_mixin import (
+    OmniConnectorModelRunnerMixin,
+    needs_omni_connector,
+)
 
 
 class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
@@ -47,31 +49,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._async_chunk = getattr(self.model_config, "async_chunk", False)
-        #  -------------------------------------- Omni-new -------------------------------------------------
-        # Mirrors the init allowlist in gpu_generation_model_runner.py. The
-        # connector-role fallback keeps consumer stages whose arch is not on the
-        # allowlist (e.g. MiniCPM-o 4.5 stage 2) from starving on input.
-        _OMNI_CONNECTOR_INIT_ARCHS = {
-            "Qwen3OmniMoeForConditionalGeneration",
-            "Qwen2_5OmniForConditionalGeneration",
-            "CovoAudioForConditionalGeneration",
-            "MiMoAudioModel",
-            "Qwen3TTSTalkerForConditionalGeneration",
-            "Qwen3TTSCode2Wav",
-            "AudexCode2Wav",
-            "AudexXCodec1",
-            "CosyVoice3Model",
-            "DyninOmniForConditionalGeneration",
-            "IndexTTS2S2MelDecoder",
-        }
-        if (
-            getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS
-            or get_stage_connector_role(self.model_config) is not None
-        ):
+        if needs_omni_connector(self.model_config):
             self.init_omni_connectors(
                 model_config=self.model_config,
             )
-        #  -------------------------------------- Omni-new -------------------------------------------------
 
     def _update_request_states(self, scheduler_output: SchedulerOutput):
         # remove requests

vllm_omni/platforms/npu/worker/npu_generation_worker.py

@@ -11,9 +11,11 @@ from vllm_omni.worker.mixins import OmniWorkerMixin
 class NPUGenerationWorker(OmniWorkerMixin, OmniNPUWorkerBase):
     """NPU generation worker for code2wav stage in Omni model."""
 
+    model_runner_cls = NPUGenerationModelRunner
+
     def init_device(self):
         self.device = self._init_device()
         num_ubatches = 1
         init_workspace_manager(self.device, num_ubatches)
 
-        self.model_runner = NPUGenerationModelRunner(self.vllm_config, self.device)
+        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)

vllm_omni/platforms/xpu/worker/xpu_ar_worker.py

@@ -13,10 +13,12 @@ logger = init_logger(__name__)
 class XPUARWorker(OmniWorkerMixin, XPUWorker):
     """XPU AR worker for thinker/talker stages in Omni model."""
 
+    model_runner_cls = XPUARModelRunner
+
     def init_device(self):
         super().init_device()
         if self.use_v2_model_runner:
             # OMNI: v2 model runner does not yet include omni hooks.
             logger.warning("OMNI XPUARWorker forces v1 model runner for omni hooks.")
             self.use_v2_model_runner = False
-        self.model_runner: XPUARModelRunner = XPUARModelRunner(self.vllm_config, self.device)
+        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)

vllm_omni/platforms/xpu/worker/xpu_generation_worker.py

@@ -13,10 +13,12 @@ logger = init_logger(__name__)
 class XPUGenerationWorker(OmniWorkerMixin, XPUWorker):
     """XPU generation worker for the code2wav (non-AR waveform generation) stage in the Omni model."""
 
+    model_runner_cls = XPUGenerationModelRunner
+
     def init_device(self):
         super().init_device()
         if self.use_v2_model_runner:
             # OMNI: v2 model runner does not yet include omni hooks.
             logger.warning("OMNI XPUGenerationWorker forces v1 model runner for omni hooks.")
             self.use_v2_model_runner = False
-        self.model_runner: XPUGenerationModelRunner = XPUGenerationModelRunner(self.vllm_config, self.device)
+        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)

vllm_omni/worker/gpu_ar_model_runner.py

@@ -41,10 +41,7 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
 from vllm_omni.data_entry_keys import flatten_payload
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
-from vllm_omni.distributed.omni_connectors.utils.config import (
-    get_stage_connector_role,
-    stage_sends_async_output,
-)
+from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output
 from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.utils.mm_outputs import (
@@ -53,7 +50,10 @@ from vllm_omni.utils.mm_outputs import (
     snapshot_mm_payload,
 )
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+from vllm_omni.worker.omni_connector_model_runner_mixin import (
+    OmniConnectorModelRunnerMixin,
+    needs_omni_connector,
+)
 from vllm_omni.worker.output.payload_build import build_omni_mm_payload
 from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding, sanitize_min_tokens_stop_ids
@@ -323,42 +323,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         # Initialize KV cache manager (preserve vllm_config fallback behavior)
         self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
         self._async_chunk = getattr(self.model_config, "async_chunk", False)
-        # Worker-connector init is gated by a per-`model_arch` allowlist
-        # (covers both producer-side and consumer-side runners for the
-        # arches below).  Consumer-wait stages must be registered
-        # separately as `(model_arch, model_stage)` tuples in
-        # `omni_scheduling_coordinator._FULL_PAYLOAD_INPUT_STAGES`;
-        # forgetting that produces a Stage-1 hang on the consumer.
-        _OMNI_CONNECTOR_INIT_ARCHS = {
-            "Qwen3OmniMoeForConditionalGeneration",
-            "Qwen2_5OmniForConditionalGeneration",
-            "CovoAudioForConditionalGeneration",
-            "MiMoAudioModel",
-            "Qwen3TTSTalkerForConditionalGeneration",
-            "Qwen3TTSCode2Wav",
-            "CosyVoice3Model",
-            "NemotronDenseForCausalLM",
-            "NemotronHForCausalLM",
-            "DyninOmniForConditionalGeneration",
-            "IndexTTS2TalkerForConditionalGeneration",
-            # nemotron_voicechat: the talker (stage 1) is the full-payload
-            # producer for code2wav (stage 2); the thinker (stage 0) only
-            # produces over the connector in async-chunk (streaming) mode.
-            "NemotronVoiceChatTalkerForConditionalGeneration",
-            "NemotronVoiceChatThinkerForConditionalGeneration",
-        }
-        # The stage-level ``model_arch`` override may be blank so the class
-        # resolves from the checkpoint's own ``architectures`` (e.g. the Audex
-        # TTA/TTS thinker stages, which bind dense on the 2B and NemotronH on
-        # the 30B-A3B). Gate on the RESOLVED architectures, not only the raw
-        # override — otherwise a blank-override producer stage never creates
-        # its worker connector and the sync full-payload flush silently never
-        # runs (downstream stage starves on connector input).
-        stage_archs = set(getattr(self.model_config, "architectures", None) or ())
-        model_arch_override = getattr(self.model_config, "model_arch", None)
-        if model_arch_override:
-            stage_archs.add(model_arch_override)
-        if stage_archs & _OMNI_CONNECTOR_INIT_ARCHS or get_stage_connector_role(self.model_config) is not None:
+        # Connector init is derived from pipeline capabilities and active
+        # producer/KV/connector configuration. A stage needs it when it is:
+        # - a full-payload consumer, or
+        # - a producer (custom_process_next_stage_input_func is set, meaning
+        #   it sends inter-stage data to a downstream consumer via connector), or
+        # - explicitly configured with a connector role via stage_connector_config.
+        if needs_omni_connector(self.model_config):
             self.init_omni_connectors(
                 model_config=self.model_config,
                 kv_transfer_manager=self.kv_transfer_manager,

vllm_omni/worker/gpu_ar_worker.py

@@ -28,6 +28,8 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
     model runners for text generation stages (e.g., thinker stages).
     """
 
+    model_runner_cls = GPUARModelRunner
+
     @instrument(span_name="Init device")
     def init_device(self):
         if self.device_config.device_type in ("cuda", "musa"):
@@ -117,7 +119,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
             self.use_v2_model_runner = False
 
         # Construct the model runner
-        self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
+        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm_omni/worker/gpu_generation_model_runner.py

@@ -35,12 +35,14 @@ from vllm.v1.worker.gpu_model_runner import (
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
-from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_role
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.utils.mm_outputs import partition_payload_list
 from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+from vllm_omni.worker.omni_connector_model_runner_mixin import (
+    OmniConnectorModelRunnerMixin,
+    needs_omni_connector,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -56,27 +58,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._async_chunk = getattr(self.model_config, "async_chunk", False)
-        # Mirrors the init allowlist in gpu_ar_model_runner.py.
-        _OMNI_CONNECTOR_INIT_ARCHS = {
-            "Qwen3OmniMoeForConditionalGeneration",
-            "Qwen2_5OmniForConditionalGeneration",
-            "CovoAudioForConditionalGeneration",
-            "MiMoAudioModel",
-            "Qwen3TTSTalkerForConditionalGeneration",
-            "Qwen3TTSCode2Wav",
-            "AudexCode2Wav",
-            "AudexXCodec1",
-            "CosyVoice3Model",
-            "DyninOmniForConditionalGeneration",
-            "IndexTTS2S2MelDecoder",
-            # nemotron_voicechat: code2wav (stage 2) consumes the talker's
-            # full-payload code stacks.
-            "NemotronVoiceChatCode2Wav",
-        }
-        if (
-            getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS
-            or get_stage_connector_role(self.model_config) is not None
-        ):
+        if needs_omni_connector(self.model_config):
             self.init_omni_connectors(
                 model_config=self.model_config,
             )

vllm_omni/worker/gpu_generation_worker.py

@@ -27,6 +27,8 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
     ``execution_type=StageExecutionType.LLM_GENERATION``.
     """
 
+    model_runner_cls = GPUGenerationModelRunner
+
     @instrument(span_name="Init device")
     def init_device(self):
         if self.device_config.device_type in ("cuda", "musa"):
@@ -99,7 +101,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
             logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.")
             self.use_v2_model_runner = False
 
-        self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device)
+        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -49,6 +49,15 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def needs_omni_connector(model_config: Any) -> bool:
+    """Whether a runner owns an input, output, or explicitly routed connector."""
+    return (
+        bool(getattr(model_config, "requires_full_payload_input", False))
+        or bool(getattr(model_config, "custom_process_next_stage_input_func", None))
+        or get_stage_connector_role(model_config) is not None
+    )
+
+
 def _should_create_payload_connector(model_config: Any) -> bool:
     """Whether this stage owns runner payload transport for its edge.
 

vllm_omni/worker/omni_connector_validation.py

@@ -0,0 +1,42 @@
+"""Validation for Omni connector support on selected LLM workers."""
+
+import importlib
+from typing import Any
+
+
+def validate_worker_omni_connector(
+    worker_cls_path: str | type[Any] | None,
+    required: bool,
+) -> None:
+    """Require the selected worker's runner to implement the connector mixin."""
+    if not required:
+        return
+    if worker_cls_path is None:
+        raise ValueError("Omni connector support requires a resolved worker_cls.")
+
+    worker_cls = _resolve_worker_cls(worker_cls_path)
+    model_runner_cls = getattr(worker_cls, "model_runner_cls", None)
+
+    from vllm_omni.worker.omni_connector_model_runner_mixin import (
+        OmniConnectorModelRunnerMixin,
+    )
+
+    if not isinstance(model_runner_cls, type) or not issubclass(
+        model_runner_cls,
+        OmniConnectorModelRunnerMixin,
+    ):
+        worker_name = getattr(worker_cls, "__name__", str(worker_cls_path))
+        raise ValueError(f"Worker {worker_name!r} does not provide an Omni connector model runner.")
+
+
+def _resolve_worker_cls(worker_cls_path: str | type[Any]) -> type[Any]:
+    if not isinstance(worker_cls_path, str):
+        return worker_cls_path
+    module_path, separator, class_name = worker_cls_path.rpartition(".")
+    if not separator:
+        raise ValueError(f"worker_cls must be a fully qualified class path, got {worker_cls_path!r}.")
+    try:
+        worker_module = importlib.import_module(module_path)
+        return getattr(worker_module, class_name)
+    except (ImportError, AttributeError) as exc:
+        raise ValueError(f"Unable to resolve worker_cls {worker_cls_path!r} for connector validation.") from exc