[Refactor] Reject unknown diffusion config fields (#5172)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Signed-off-by: zuiho <2324465096@qq.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -9,7 +9,7 @@ import torch
 from pydantic import ValidationError
 
 from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs
 from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
 from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
@@ -21,7 +21,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _terminal_config(stage_cfg: dict) -> OmniDiffusionConfig:
-    return OmniDiffusionConfig.from_kwargs(**stage_cfg["engine_args"])
+    kwargs = extract_diffusion_stage_config_kwargs(stage_cfg["engine_args"], stage_id=stage_cfg["stage_id"])
+    return OmniDiffusionConfig.from_kwargs(**kwargs)
 
 
 def test_default_stage_config_includes_cache_backend():
@@ -32,6 +33,9 @@ def test_default_stage_config_includes_cache_backend():
             "cache_config": '{"Fn_compute_blocks": 2}',
             "vae_use_slicing": True,
             "ulysses_degree": 2,
+            "seed": 7,
+            "kv_cache_dtype": "fp8",
+            "diffusion_kv_cache_dtype": "fp8_e4m3",
         }
     )[0]
 
@@ -42,6 +46,9 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["vae_use_slicing"] is True
     assert engine_args["parallel_config"]["ulysses_degree"] == 2
     assert engine_args["model_stage"] == "diffusion"
+    assert "seed" not in engine_args
+    assert "kv_cache_dtype" not in engine_args
+    assert engine_args["diffusion_kv_cache_dtype"] == "fp8_e4m3"
 
 
 def test_default_stage_config_preserves_ulysses_a2a_permute() -> None:
@@ -130,6 +137,67 @@ def test_default_stage_rejects_unknown_nested_parallel_config_key():
         )
 
 
+def test_default_stage_routes_ar_profiler_away_before_diffusion_build(mocker):
+    stage_dict = StageConfigFactory.create_default_diffusion({"enable_ar_profiler": True})[0]
+    assert "enable_ar_profiler" not in stage_dict["engine_args"]
+
+
+@pytest.mark.parametrize(
+    ("field_name", "value"),
+    [
+        ("enable_sleep_mod", None),
+        ("enable_lora", True),
+        ("kv_cache_dtype", "fp8"),
+        ("seed", 7),
+    ],
+)
+def test_legacy_diffusion_stage_rejects_unowned_field(field_name, value):
+    from vllm_omni.config.config_factory import StageConfigFactory
+    from vllm_omni.config.yaml_util import create_config
+    from vllm_omni.engine.stage_init_utils import build_diffusion_config
+
+    stage_dict = StageConfigFactory.create_default_diffusion({"model": "unused"})[0]
+    stage_dict["engine_args"][field_name] = value
+    stage_cfg = create_config(stage_dict)
+    metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None)
+
+    with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"):
+        build_diffusion_config("unused", stage_cfg, metadata)
+
+
+@pytest.mark.parametrize(
+    ("field_name", "value"),
+    [
+        ("enable_sleep_mod", None),
+        ("enable_lora", True),
+    ],
+)
+def test_default_diffusion_factory_rejects_unowned_field(field_name, value):
+    with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"):
+        StageConfigFactory.create_default_diffusion({field_name: value})
+
+
+def test_legacy_default_stage_build_accepts_engine_adapter_metadata(monkeypatch):
+    from vllm_omni.config.config_factory import StageConfigFactory
+    from vllm_omni.config.yaml_util import create_config
+    from vllm_omni.engine import stage_init_utils
+
+    stage_dict = StageConfigFactory.create_default_diffusion(
+        {
+            "model": "unused",
+            "api_key": "frontend-owned",
+        }
+    )[0]
+    assert "api_key" not in stage_dict["engine_args"]
+    stage_cfg = create_config(stage_dict)
+    metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None)
+    monkeypatch.setattr(stage_init_utils.current_omni_platform, "get_device_count", lambda: 1)
+
+    config = stage_init_utils.build_diffusion_config("unused", stage_cfg, metadata)
+
+    assert config.model == "unused"
+
+
 def test_default_cache_config_used_when_missing():
     """Ensure default cache_config is synthesized when only backend is given."""
     stage_cfg = StageConfigFactory.create_default_diffusion(
@@ -671,10 +739,11 @@ def test_serve_cli_rejects_invalid_request_batch_max_wait_ms(bad_wait: str):
         )
 
 
-def test_serve_cli_accepts_additional_config():
+@pytest.mark.parametrize("subcommand_dest", ["command", "subparser"])
+def test_serve_cli_accepts_additional_config(subcommand_dest):
     """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config."""
     parser = TrackingArgumentParser()
-    subparsers = parser.add_subparsers(dest="command")
+    subparsers = parser.add_subparsers(dest=subcommand_dest)
     OmniServeCommand().subparser_init(subparsers)
 
     args = parser.parse_args(
@@ -687,7 +756,7 @@ def test_serve_cli_accepts_additional_config():
         ]
     )
 
-    stage_cfg = StageConfigFactory.create_default_diffusion(vars(args))[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(args.get_explicit_kwargs_dict())[0]
 
     engine_args = stage_cfg["engine_args"]
 
@@ -757,6 +826,24 @@ def test_default_stage_config_includes_quantization_config():
     assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
 
 
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+def test_default_diffusion_factory_preserves_engine_quantization(typed, monkeypatch):
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+    monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
+    kwargs = {"quantization": "fp8"}
+
+    if typed:
+        stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0]
+        config = stage.diffusion_config
+        config.enrich_config()
+    else:
+        config = _terminal_config(StageConfigFactory.create_default_diffusion(kwargs)[0])
+
+    assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "fp8"
+    assert config.quantization_config_is_auto_detected is False
+
+
 @pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"])
 def test_generic_diffusion_uses_canonical_video_output_type(model_class_name):
     config = StageConfigFactory.create_typed_default_diffusion(