fix: reserve Qwen3-Omni NVFP4 backend fix (#7200)

Signed-off-by: kunkunblueberry <1833921874@qq.com>

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},
+            }
+        },
     )
 
 

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"

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

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