diff --git a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py index 9502ed5e6..6537b1cac 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py +++ b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py @@ -65,7 +65,12 @@ def _stage_config(): # gains enough VRAM headroom for a no-eager run, add a companion test. return modify_stage_config( _CI_DEPLOY, - updates={"stages": {0: {"enforce_eager": True}, 1: {"enforce_eager": True}}}, + updates={ + "stages": { + 0: {"enforce_eager": True, "moe_backend": "auto"}, + 1: {"enforce_eager": True}, + } + }, ) diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index e2196dfd1..250aff884 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -11,6 +11,7 @@ from pathlib import Path import pytest from pydantic.fields import FieldInfo +from transformers import Qwen3OmniMoeConfig from vllm.config import CacheConfig as VllmCacheConfig from vllm.config import CompilationConfig as VllmCompilationConfig from vllm.config import LoadConfig as VllmLoadConfig @@ -19,6 +20,7 @@ from vllm.config import ProfilerConfig as VllmProfilerConfig from vllm.config import SchedulerConfig as VllmSchedulerConfig from vllm.engine.arg_utils import EngineArgs +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.config.omni_config import ( _LLM_STAGE_ENGINE_FIELDS, OmniStageCacheConfig, @@ -155,6 +157,44 @@ def test_qwen3_omni_preserves_explicit_moe_backend(): assert engine_args["moe_backend"] == "flashinfer" +def test_qwen3_omni_preserves_explicit_auto_moe_backend(): + engine_args = { + "model_arch": "Qwen3OmniMoeForConditionalGeneration", + "moe_backend": "auto", + } + + stage_init_utils._maybe_set_qwen3_omni_moe_backend(engine_args) + + assert engine_args["moe_backend"] == "auto" + + +def test_qwen3_omni_nvfp4_ci_config_preserves_auto_moe_backend(tmp_path): + """The NVFP4 CI deploy opts stage 0 into vLLM's auto backend selection.""" + hf_config = Qwen3OmniMoeConfig(enable_audio_output=True) + pipeline = resolve_pipeline_config("qwen3_omni_moe", hf_config) + assert pipeline is not None + nvfp4_deploy = modify_stage_config( + get_deploy_config_path("ci/qwen3_omni_moe.yaml"), + updates={"stages": {0: {"moe_backend": "auto"}}}, + ) + deploy = load_deploy_config(nvfp4_deploy) + legacy_stages = merge_pipeline_deploy(pipeline, copy.deepcopy(deploy)) + omni_config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=copy.deepcopy(deploy), + cli_overrides={"model": str(tmp_path)}, + ) + + for stage_id, expected_backend in ((0, "auto"), (1, "triton")): + legacy_stage = next(stage for stage in legacy_stages if stage.stage_id == stage_id) + typed_stage = omni_config.stage_by_id(stage_id) + for engine_args in ( + build_engine_args_dict(legacy_stage.to_omegaconf(), str(tmp_path)), + build_engine_args_dict_from_omni_stage_config(typed_stage, str(tmp_path)), + ): + assert engine_args["moe_backend"] == expected_backend + + def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]: parent_model = tmp_path / "parent-model" stage_model = tmp_path / "stage-model" diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index aaa9c3512..81eb76f0b 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -400,7 +400,7 @@ def _get_deploy_config( @config -class OmniStageModelConfig: +class OmniStageModelConfig(_TrackExplicitConfigFields): """Per-stage model behavior and resolved model-engine inputs.""" model: str | None = None diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 0562d2b08..765549bcb 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -764,12 +764,15 @@ def prepare_engine_environment() -> None: pass -def _maybe_set_qwen3_omni_moe_backend(engine_args_dict: dict[str, Any]) -> None: +def _maybe_set_qwen3_omni_moe_backend( + engine_args_dict: dict[str, Any], + *, + moe_backend_is_explicit: bool | None = None, +) -> None: """Choose the stable MoE backend when Qwen3-Omni has no explicit choice.""" - if ( - engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration" - and engine_args_dict.get("moe_backend", "auto") == "auto" - ): + if moe_backend_is_explicit is None: + moe_backend_is_explicit = "moe_backend" in engine_args_dict + if engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration" and not moe_backend_is_explicit: engine_args_dict["moe_backend"] = "triton" logger.info("[stage_init] Set moe_backend=triton for Qwen3-Omni stage") @@ -1121,6 +1124,10 @@ def _project_omni_stage_engine_args( # The legacy builder always emits this key, including for pipelines such # as Audex that intentionally defer architecture discovery to HF config. engine_args["model_arch"] = copy.deepcopy(stage_config.model_config.model_arch) + _maybe_set_qwen3_omni_moe_backend( + engine_args, + moe_backend_is_explicit="moe_backend" in getattr(stage_config.model_config, "_omni_explicit_fields", ()), + ) topology = stage_config.stage_pipeline_config topology_engine_args = { @@ -1283,7 +1290,6 @@ def _finalize_engine_args_dict( engine_args_dict["has_sampling_extra_args"] = has_sampling_extra_args engine_args_dict["sampling_extra_args_keys"] = sampling_extra_args_keys - # Select the typed backend option during stage argument finalization. _maybe_set_qwen3_omni_moe_backend(engine_args_dict) return engine_args_dict