From b454e34966f3d14d1605f29aca19c962d3a16152 Mon Sep 17 00:00:00 2001 From: wy17003 Date: Tue, 25 Aug 2026 15:24:59 +0800 Subject: [PATCH] [Refactor][2/N Scheduler]Replace _FULL_PAYLOAD_INPUT_STAGES with a resolved stage transport capability. (#6149) Signed-off-by: wy17003 Co-authored-by: Hongsheng Liu --- tests/config/test_config_factory.py | 15 +-- tests/config/test_omni_config.py | 33 ++++++ .../sched/test_omni_scheduler_mixin_shared.py | 79 ++++++++++++- .../sched/test_omni_scheduling_coordinator.py | 104 ----------------- tests/engine/test_arg_utils.py | 105 +++++++++++++++++- tests/engine/test_stage_engine_args.py | 12 ++ .../test_nemotron_voicechat_registration.py | 9 -- vllm_omni/config/model.py | 1 + vllm_omni/config/omni_config.py | 2 + vllm_omni/config/stage_config.py | 4 + .../core/sched/omni_scheduling_coordinator.py | 64 +---------- vllm_omni/engine/arg_utils.py | 12 +- .../model_executor/models/audex/pipeline.py | 3 + .../models/aura_omni/pipeline.py | 1 + .../models/cosyvoice3/pipeline.py | 1 + .../models/covo_audio/pipeline.py | 1 + .../models/dynin_omni/pipeline.py | 2 + .../models/indextts2/pipeline.py | 2 + .../models/mimo_audio/pipeline.py | 1 + .../models/minicpmo_4_5/pipeline.py | 1 + .../models/nemotron_voicechat/pipeline.py | 1 + .../models/qwen2_5_omni/pipeline.py | 2 + .../models/qwen3_omni/pipeline.py | 2 + .../models/qwen3_tts/pipeline.py | 1 + .../stage_input_processors/qwen2_5_omni.py | 9 +- .../npu/worker/npu_ar_model_runner.py | 28 +---- .../platforms/npu/worker/npu_ar_worker.py | 4 +- .../npu/worker/npu_generation_model_runner.py | 29 +---- .../npu/worker/npu_generation_worker.py | 4 +- .../platforms/xpu/worker/xpu_ar_worker.py | 4 +- .../xpu/worker/xpu_generation_worker.py | 4 +- vllm_omni/worker/gpu_ar_model_runner.py | 53 ++------- vllm_omni/worker/gpu_ar_worker.py | 4 +- .../worker/gpu_generation_model_runner.py | 28 +---- vllm_omni/worker/gpu_generation_worker.py | 4 +- .../omni_connector_model_runner_mixin.py | 9 ++ vllm_omni/worker/omni_connector_validation.py | 42 +++++++ 37 files changed, 370 insertions(+), 310 deletions(-) create mode 100644 vllm_omni/worker/omni_connector_validation.py diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 374a6d212..e5b6d607a 100644 --- a/tests/config/test_config_factory.py +++ b/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") diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 0db9803a6..1292bf0a0 100644 --- a/tests/config/test_omni_config.py +++ b/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) diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py index 4f98fd092..bb9fd8859 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_shared.py +++ b/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), diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index 444d21c34..ae1b3924e 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/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.""" diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index ed1031b64..7e202932f 100644 --- a/tests/engine/test_arg_utils.py +++ b/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", diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 9ec7e520a..632b3891e 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/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" diff --git a/tests/model_executor/models/test_nemotron_voicechat_registration.py b/tests/model_executor/models/test_nemotron_voicechat_registration.py index 09d85b271..62e006aa3 100644 --- a/tests/model_executor/models/test_nemotron_voicechat_registration.py +++ b/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( diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 736a504bd..70cdea18a 100644 --- a/vllm_omni/config/model.py +++ b/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): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 280b3c86b..e5c87dc88 100644 --- a/vllm_omni/config/omni_config.py +++ b/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 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 1825fbdfe..33636c561 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 75f234d8c..93f40daa0 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/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: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index e1c289357..c618d60d0 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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, diff --git a/vllm_omni/model_executor/models/audex/pipeline.py b/vllm_omni/model_executor/models/audex/pipeline.py index 8723a2138..a7a4d7960 100644 --- a/vllm_omni/model_executor/models/audex/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py index 14aa5ff30..de163778d 100644 --- a/vllm_omni/model_executor/models/aura_omni/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py index 413dc1dc3..d6bb0e5bb 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/covo_audio/pipeline.py b/vllm_omni/model_executor/models/covo_audio/pipeline.py index c2e1911ef..b3bd9c0c6 100644 --- a/vllm_omni/model_executor/models/covo_audio/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/dynin_omni/pipeline.py b/vllm_omni/model_executor/models/dynin_omni/pipeline.py index 584b6ab5c..6f59ebb36 100644 --- a/vllm_omni/model_executor/models/dynin_omni/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/indextts2/pipeline.py b/vllm_omni/model_executor/models/indextts2/pipeline.py index e29d61563..6e8754daa 100644 --- a/vllm_omni/model_executor/models/indextts2/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/mimo_audio/pipeline.py b/vllm_omni/model_executor/models/mimo_audio/pipeline.py index 9e54979de..684004b7a 100644 --- a/vllm_omni/model_executor/models/mimo_audio/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index f4e8e4c32..d3d1da05b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py index 4f4834740..683228c21 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index de79e1aae..2036c42aa 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index c93e764ef..1e515d0dd 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 8a68e63cd..6dfe8736d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/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, ), ), ) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 42377faa8..42a9b19e3 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/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 = [] diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index c43a4f8d0..9cbbc4fb5 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/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, diff --git a/vllm_omni/platforms/npu/worker/npu_ar_worker.py b/vllm_omni/platforms/npu/worker/npu_ar_worker.py index d34d83ef9..19f3f89c8 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_worker.py +++ b/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) diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 0f9c155ee..b033dab3f 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/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 diff --git a/vllm_omni/platforms/npu/worker/npu_generation_worker.py b/vllm_omni/platforms/npu/worker/npu_generation_worker.py index f77975a5d..ac66c6a54 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_worker.py +++ b/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) diff --git a/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py b/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py index 9acf869f7..e25ea4cc6 100644 --- a/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py +++ b/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) diff --git a/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py b/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py index d9bb790d3..3036f5a5e 100644 --- a/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py +++ b/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) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 1197f5b44..be084f8f5 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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, diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 2afa59576..654cd4dc5 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/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. diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 19a6bdb93..73fe406e6 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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, ) diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 4a7aaf711..d6a8d6c8e 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/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. diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 70e9dafd2..0f233b1da 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/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. diff --git a/vllm_omni/worker/omni_connector_validation.py b/vllm_omni/worker/omni_connector_validation.py new file mode 100644 index 000000000..867a6d95f --- /dev/null +++ b/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