[Config] Remove invalid LLM-only engine_args from diffusion stage configs (#2622)

Signed-off-by: Yiyang Liu <yiyangliu@microsoft.com>
Co-authored-by: Yiyang Liu <yiyangliu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

tests/test_diffusion_config_fields.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ensure diffusion stage YAML configs only use valid OmniDiffusionConfig fields.
+
+Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
+"""
+
+from dataclasses import fields
+from pathlib import Path
+
+import pytest
+import yaml
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+try:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+except Exception:
+    OmniDiffusionConfig = None
+
+
+@pytest.mark.skipif(
+    OmniDiffusionConfig is None,
+    reason="OmniDiffusionConfig could not be imported (missing torch?)",
+)
+def test_diffusion_stage_configs_only_contain_valid_fields():
+    """Diffusion stage engine_args must only contain OmniDiffusionConfig fields.
+
+    Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
+    """
+    # Scan both main configs and test configs
+    repo_root = Path(__file__).parent.parent
+    config_dirs = [
+        repo_root / "vllm_omni" / "model_executor" / "stage_configs",
+    ]
+    # Also scan test directories recursively
+    test_dir = repo_root / "tests"
+
+    yaml_paths: list[Path] = []
+    for config_dir in config_dirs:
+        yaml_paths.extend(sorted(config_dir.glob("*.yaml")))
+    yaml_paths.extend(sorted(test_dir.rglob("*.yaml")))
+
+    valid_fields = {f.name for f in fields(OmniDiffusionConfig)}
+    # model_stage is consumed by the stage init layer, not OmniDiffusionConfig
+    valid_fields.add("model_stage")
+    # model_arch is consumed by the stage init layer for diffusion model class resolution
+    valid_fields.add("model_arch")
+    # "quantization" is mapped to "quantization_config" by from_kwargs() backwards-compat
+    valid_fields.add("quantization")
+
+    invalid_entries: list[tuple[str, set[str]]] = []
+    for yaml_path in yaml_paths:
+        with open(yaml_path) as fh:
+            config = yaml.safe_load(fh)
+
+        stages = config.get("stage_args", config.get("stages", []))
+        for stage in stages:
+            if stage.get("stage_type") != "diffusion":
+                continue
+            engine_args = stage.get("engine_args", {})
+            invalid = set(engine_args.keys()) - valid_fields
+            if invalid:
+                invalid_entries.append((yaml_path.relative_to(repo_root), invalid))
+
+    assert not invalid_entries, "Diffusion stage configs contain fields not in OmniDiffusionConfig:\n" + "\n".join(
+        f"  {name}: {sorted(bad)}" for name, bad in invalid_entries
+    )