[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/config/test_omni_config.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import warnings
 from dataclasses import fields
 from inspect import Parameter, signature
 from multiprocessing.reduction import ForkingPickler
@@ -24,6 +25,7 @@ from vllm.config import SchedulerConfig as VllmSchedulerConfig
 
 from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.config import omni_config as omni_config_module
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.omni_config import (
     BaseVllmOmniStageConfig,
     OmniStageCacheConfig,
@@ -150,6 +152,30 @@ def test_nested_stage_override_deep_merges_structured_model_config() -> None:
     assert model_config["policy_server_config"]["action_space"] == "joint_position"
 
 
+def _build_single_diffusion_config(
+    *,
+    cli_overrides: dict | None = None,
+    engine_extras: dict | None = None,
+) -> VllmOmniConfig:
+    pipeline = PipelineConfig(
+        model_type="diffusion-owner-test",
+        stages=(
+            StagePipelineConfig(
+                stage_id=0,
+                model_stage="diffusion",
+                execution_type=StageExecutionType.DIFFUSION,
+                final_output=True,
+            ),
+        ),
+    )
+    deploy_stages = [StageDeployConfig(stage_id=0, engine_extras=engine_extras)] if engine_extras is not None else []
+    return VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=DeployConfig(async_chunk=False, stages=deploy_stages),
+        cli_overrides=cli_overrides,
+    )
+
+
 @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
 def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
     pipeline = _resolve_pipeline_or_skip(model_type)
@@ -475,7 +501,7 @@ def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_na
     assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
 
 
-def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
+def test_from_pipeline_config_keeps_global_kv_cache_dtype_outside_diffusion_stage():
     stage = _from_pipeline_key(
         "dreamzero",
         deploy_config_path="dreamzero_tp1_cfg2",
@@ -483,11 +509,57 @@ def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_st
     ).stage_by_id(0)
 
     assert isinstance(stage, VllmOmniDiffusionStageConfig)
-    assert stage.diffusion_config.diffusion_kv_cache_dtype == "fp8"
+    assert stage.diffusion_config.diffusion_kv_cache_dtype is None
 
 
 def test_stage_cli_field_selection_defers_ownership_validation_until_sources_are_merged():
-    assert omni_config_module._stage_cli_overrides(0, {"enable_lora": True}) == {"enable_lora": True}
+    assert omni_config_module._stage_cli_overrides(
+        0,
+        {"enable_lora": True},
+    ) == {"enable_lora": True}
+
+
+@pytest.mark.parametrize(
+    ("cli_overrides", "expected_dtype"),
+    [
+        ({}, torch.float32),
+        ({"dtype": None}, torch.float32),
+        ({"dtype": "float16"}, torch.float16),
+        ({"dtype": "auto"}, torch.bfloat16),
+        ({"stage_0_dtype": "bfloat16"}, torch.bfloat16),
+    ],
+)
+def test_diffusion_deploy_dtype_survives_unset_cli_overrides(cli_overrides, expected_dtype):
+    stage = _build_single_diffusion_config(
+        engine_extras={"dtype": "float32", "cache_backend": "tea_cache"},
+        cli_overrides=cli_overrides,
+    ).stage_by_id(0)
+
+    assert stage.diffusion_config.dtype == expected_dtype
+    assert stage.diffusion_config.cache_backend == "tea_cache"
+
+
+@pytest.mark.parametrize(
+    "coordination_kwargs",
+    [
+        {"omni_master_address": "127.0.0.1", "omni_master_port": 30000},
+        {"omni_dp_size_local": 2, "omni_heartbeat_timeout": 60.0, "omni_lb_policy": "random"},
+    ],
+)
+def test_diffusion_ingress_routes_coordination_fields_away(coordination_kwargs):
+    normalize = omni_config_module.normalize_and_validate_diffusion_engine_ingress_kwargs
+    assert normalize({**coordination_kwargs, "enable_sleep_mode": True}, stage_id=0) == {"enable_sleep_mode": True}
+    with pytest.raises(ValueError, match="omni_master_unknown_field"):
+        normalize({**coordination_kwargs, "omni_master_unknown_field": "unknown"}, stage_id=0)
+
+
+def test_diffusion_ingress_defers_defaults(monkeypatch):
+    monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache")
+
+    assert omni_config_module.normalize_and_validate_diffusion_engine_ingress_kwargs({}, stage_id=0) == {}
+    assert omni_config_module.normalize_and_validate_diffusion_engine_ingress_kwargs(
+        {"dtype": None, "cache_backend": None}, stage_id=0
+    ) == {"dtype": None, "cache_backend": None}
 
 
 @pytest.mark.parametrize(
@@ -1415,18 +1487,19 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
     monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
     monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
 
-    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
-        diffusion_attention_backend="flash_attn",
-        fa_deterministic=True,
-        kv_cache_dtype="fp8",
-        kv_cache_skip_steps="0-1",
-        kv_cache_skip_layers=[2],
-        static_lora_scale=0.25,
-        diffusion_kv_mode="paged_scheduler",
-        diffusion_kv_max_rows_per_request=2,
-        diffusers_load_kwargs=None,
-        diffusers_call_kwargs=None,
-    )
+    with pytest.warns(FutureWarning):
+        cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+            diffusion_attention_backend="flash_attn",
+            fa_deterministic=True,
+            kv_cache_dtype="fp8",
+            kv_cache_skip_steps="0-1",
+            kv_cache_skip_layers=[2],
+            static_lora_scale=0.25,
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+            diffusers_load_kwargs=None,
+            diffusers_call_kwargs=None,
+        )
 
     assert cfg.diffusion_attention_config.default.backend == "flash_attn"
     assert cfg.fa_deterministic is True
@@ -1463,47 +1536,6 @@ def test_diffusion_config_none_values_preserve_dataclass_defaults():
     assert config.enable_sleep_mode is False
 
 
-@pytest.mark.parametrize(
-    ("canonical_key", "alias_key", "canonical_value", "alias_value"),
-    [
-        ("lora_scale", "static_lora_scale", 0.75, 0.25),
-        (
-            "quantization_config",
-            "diffusion_quantization_config",
-            {"method": "canonical"},
-            {"method": "diffusion-alias"},
-        ),
-        (
-            "quantization_config",
-            "quantization",
-            {"method": "canonical"},
-            "legacy-alias",
-        ),
-        ("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"),
-        ("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"),
-        ("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"),
-        ("streaming_output", "diffusion_streaming_output", False, True),
-    ],
-)
-def test_diffusion_alias_conflicts_prefer_canonical_key(
-    canonical_key,
-    alias_key,
-    canonical_value,
-    alias_value,
-):
-    from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
-
-    normalized = normalize_omni_diffusion_kwargs(
-        {
-            canonical_key: canonical_value,
-            alias_key: alias_value,
-        }
-    )
-
-    assert normalized[canonical_key] == canonical_value
-    assert alias_key not in normalized
-
-
 def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
     from vllm_omni.platforms import current_omni_platform
 
@@ -1528,10 +1560,11 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
         )
     )
 
-    stage = _from_pipeline_key(
-        "dreamzero",
-        deploy_config_path=str(deploy_path),
-    ).stage_by_id(0)
+    with pytest.warns(FutureWarning):
+        stage = _from_pipeline_key(
+            "dreamzero",
+            deploy_config_path=str(deploy_path),
+        ).stage_by_id(0)
 
     assert isinstance(stage, VllmOmniDiffusionStageConfig)
     assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
@@ -1544,7 +1577,10 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
 
     engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path))
-    od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+    diffusion_kwargs = omni_config_module.extract_diffusion_stage_config_kwargs(
+        engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True
+    )
+    od_config = OmniDiffusionConfig.from_kwargs(**diffusion_kwargs)
     assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
 
 
@@ -1697,10 +1733,8 @@ def test_video_output_transport_survives_stage_override_filtering() -> None:
 
 
 def test_video_output_transport_reaches_default_diffusion_stage() -> None:
-    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-
     transport = {"enable_device_postprocess": True}
-    stages = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stages = StageConfigFactory.create_default_diffusion(
         {
             "model": "unused",
             "model_class_name": "UnknownPipeline",
@@ -1830,3 +1864,95 @@ def test_diffusion_quantization_origin_survives_projection_and_transport(monkeyp
     restored = ForkingPickler.loads(ForkingPickler.dumps(cfg))
     assert restored.quantization_config_is_auto_detected is not explicit
     assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37])
+
+
+@pytest.mark.parametrize(
+    "factory",
+    [
+        omni_config_module._DiffusionConfigProjection.from_kwargs,
+        pytest.param(
+            lambda **kwargs: __import__(
+                "vllm_omni.diffusion.data",
+                fromlist=["OmniDiffusionConfig"],
+            ).OmniDiffusionConfig.from_kwargs(**kwargs),
+            id="runtime-config",
+        ),
+    ],
+)
+def test_diffusion_config_entrypoints_reject_unknown_none(factory):
+    with pytest.raises((ValueError, ValidationError), match="enable_sleep_mod"):
+        factory(enable_sleep_mod=None)
+
+
+def test_diffusion_config_projection_rejects_unknown_none():
+    with pytest.raises(ValueError, match="enable_sleep_mod"):
+        omni_config_module._DiffusionConfigProjection.from_kwargs(enable_sleep_mod=None)
+
+
+def test_runtime_diffusion_config_rejects_unknown_none():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with pytest.raises(ValueError, match="enable_sleep_mod"):
+        OmniDiffusionConfig.from_kwargs(enable_sleep_mod=None)
+
+
+def test_direct_diffusion_aliases_promote_none_and_reject_real_conflicts(monkeypatch):
+    from vllm_omni.diffusion import data as diffusion_data
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    monkeypatch.setattr(diffusion_data, "build_quant_config", lambda config: config)
+    with pytest.warns(FutureWarning) as warnings:
+        config = OmniDiffusionConfig.from_kwargs(
+            quantization={"method": "example"},
+            quantization_config=None,
+            kv_cache_dtype="fp8",
+            diffusion_kv_cache_dtype=None,
+            enable_sleep_mode=False,
+            request_batch_max_wait_ms=0.0,
+        )
+
+    warning_messages = [str(warning.message) for warning in warnings]
+    assert any("quantization" in message for message in warning_messages)
+    assert any("kv_cache_dtype" in message for message in warning_messages)
+    assert config.quantization_config == {"method": "example"}
+    assert config.diffusion_kv_cache_dtype == "fp8"
+    assert config.enable_sleep_mode is False
+    assert config.request_batch_max_wait_ms == 0.0
+    with pytest.raises(ValueError, match=r"quantization.*quantization_config"):
+        OmniDiffusionConfig.from_kwargs(quantization="fp8", quantization_config={"method": "example"})
+
+
+def test_structured_diffusion_stage_rejects_unknown_engine_extra():
+    with pytest.raises(ValueError, match=r"stage 0.*enable_sleep_mod"):
+        _build_single_diffusion_config(engine_extras={"enable_sleep_mod": None})
+
+    with pytest.raises(ValueError, match=r"stage 0.*enable_sleep_mod"):
+        _build_single_diffusion_config(cli_overrides={"enable_sleep_mod": None})
+
+    with pytest.raises(ValueError, match=r"stage 0.*enable_sleep_mod"):
+        _build_single_diffusion_config(cli_overrides={"stage_0_enable_sleep_mod": None})
+
+    with pytest.raises(ValueError, match=r"no structured config owner: parallel_config\.unknown_field"):
+        _build_single_diffusion_config(engine_extras={"parallel_config": {"unknown_field": 1}})
+
+
+def test_structured_diffusion_stage_keeps_shared_globals_outside_diffusion():
+    with warnings.catch_warnings(record=True) as caught:
+        stage = _build_single_diffusion_config(
+            cli_overrides={"kv_cache_dtype": "fp8", "seed": 7},
+        ).stage_by_id(0)
+
+    assert stage.diffusion_config.diffusion_kv_cache_dtype is None
+    assert not any(issubclass(warning.category, FutureWarning) for warning in caught)
+
+
+@pytest.mark.parametrize(
+    ("field_name", "config_kwargs"),
+    [
+        ("seed", {"cli_overrides": {"stage_0_seed": None}}),
+        ("kv_cache_dtype", {"engine_extras": {"kv_cache_dtype": "fp8"}}),
+    ],
+)
+def test_structured_diffusion_stage_rejects_explicit_shared_engine_field(field_name, config_kwargs):
+    with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"):
+        _build_single_diffusion_config(**config_kwargs)

tests/config/test_sleep_mode_deploy.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Check the sleep/wake E2E topology before loading large BAGEL weights."""
+
+import pytest
+
+from tests.entrypoints.test_omni_sleep_mode import _sleep_deploy_config
+from vllm_omni.config.omni_config import (
+    VllmOmniConfig,
+    VllmOmniDiffusionStageConfig,
+    extract_diffusion_stage_config_kwargs,
+)
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+from vllm_omni.config.stage_config import load_deploy_config
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("diffusion_only,tp_size", [(False, 1), (False, 2), (True, 2)])
+def test_sleep_deploy_resolves_expected_stages(diffusion_only, tp_size):
+    deploy = load_deploy_config(_sleep_deploy_config(diffusion_only=diffusion_only, tp_size=tp_size))
+    pipeline = resolve_pipeline_config(deploy.pipeline or "bagel")
+    config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
+
+    assert len(config.stage_configs) == (1 if diffusion_only else 2)
+    for stage in config.stage_configs:
+        is_diffusion = diffusion_only or stage.stage_id == 1
+        assert isinstance(stage, VllmOmniDiffusionStageConfig) == is_diffusion
+        # Keep BAGEL's batch budget: a single multimodal item needs 8625 tokens.
+        assert stage.scheduler_config.max_num_batched_tokens == 32768
+        assert stage.model_config.enable_sleep_mode is True
+        assert stage.parallel_config.tensor_parallel_size == tp_size
+        assert stage.runtime_config.devices == ",".join(str(stage.stage_id * tp_size + rank) for rank in range(tp_size))
+        if not is_diffusion:
+            assert stage.cache_config.gpu_memory_utilization == 0.8
+        if is_diffusion:
+            engine_args = build_engine_args_dict_from_omni_stage_config(stage, "test-bagel")
+            terminal = OmniDiffusionConfig.from_kwargs(
+                **extract_diffusion_stage_config_kwargs(
+                    engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True
+                )
+            )
+            assert terminal.enable_sleep_mode is True
+            assert terminal.parallel_config.tensor_parallel_size == tp_size
+
+    if not diffusion_only:
+        assert config.stage_by_id(1).stage_pipeline_config.input_sources == (0,)
+        assert deploy.stages[1].input_connectors["from_stage_0"] == "shared_memory_connector"

tests/diffusion/quantization/test_quantization_fp8.py

@@ -414,18 +414,16 @@ def test_quantization_key_maps_to_quantization_config():
 
 
 @hardware_test(res={"cuda": ["L4", "B200"]})
-def test_quantization_config_key_takes_priority():
-    """When both 'quantization' and 'quantization_config' are set,
-    'quantization_config' takes priority."""
+def test_quantization_and_quantization_config_conflict():
+    """Two explicit quantization sources are rejected instead of silently merged."""
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    config = OmniDiffusionConfig.from_kwargs(
-        model="test",
-        quantization="fp8",
-        quantization_config={"method": "fp8", "activation_scheme": "static"},
-    )
-    assert config.quantization_config is not None
-    assert config.quantization_config.activation_scheme == "static"
+    with pytest.raises(ValueError, match=r"quantization.*quantization_config"):
+        OmniDiffusionConfig.from_kwargs(
+            model="test",
+            quantization="fp8",
+            quantization_config={"method": "fp8", "activation_scheme": "static"},
+        )
 
 
 @hardware_test(res={"cuda": ["L4", "B200"]})

tests/diffusion/test_diffusion_config_propagation.py

@@ -11,6 +11,7 @@ import pytest
 import torch
 
 from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.omni_config import extract_diffusion_stage_config_kwargs
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
@@ -40,7 +41,8 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
     """
     stages = StageConfigFactory.create_default_diffusion(kwargs)
     engine_args = dict(stages[0]["engine_args"])
-    return OmniDiffusionConfig.from_kwargs(**engine_args)
+    diffusion_kwargs = extract_diffusion_stage_config_kwargs(engine_args, stage_id=0)
+    return OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None})
 
 
 class TestParallelConfigPropagation:
@@ -57,9 +59,7 @@ class TestParallelConfigPropagation:
         stages = StageConfigFactory.create_default_diffusion({"parallel_config": pc, "model": "x"})
         assert stages[0]["runtime"]["devices"] == "0,1,2,3"
 
-        # Let __post_init__ reconstruct from dict (real code path)
-        ea = dict(stages[0]["engine_args"])
-        od = OmniDiffusionConfig.from_kwargs(**ea)
+        od = _roundtrip_diffusion_config(parallel_config=pc, model="x")
         assert od.parallel_config.tensor_parallel_size == 4
         assert od.parallel_config.world_size == 4
 
@@ -193,6 +193,17 @@ class TestCreateDefaultDiffusion:
         assert od.max_num_batched_tokens == 2048
         assert od.max_model_len == 4096
 
+    @pytest.mark.parametrize(
+        ("field_name", "value"),
+        [
+            ("enable_sleep_mod", None),
+            ("enable_lora", True),
+        ],
+    )
+    def test_unowned_raw_field_is_rejected(self, field_name, value):
+        with pytest.raises(ValueError, match=field_name):
+            StageConfigFactory.create_default_diffusion({"model": "x", field_name: value})
+
 
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")

tests/distributed/omni_connectors/test_nixl_connector.py

@@ -317,6 +317,7 @@ def copying_native_agent(consumer, monkeypatch):
 
 
 @pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"])
+@pytest.mark.usefixtures("reliable_claim_queries")
 def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consumer, copying_native_agent, case):
     empty = torch.empty((2, 0, 3), dtype=torch.float64)
     other = torch.empty((0,), dtype=torch.int64)

tests/e2e/features/helpers/verl_omni_server.py

@@ -10,11 +10,9 @@ Ray workers can unpickle ``vLLMOmniHttpServerLocal`` via
 
 from __future__ import annotations
 
-import argparse
 import json
 import logging
 import os
-from dataclasses import asdict
 from enum import Enum
 from typing import Any
 
@@ -22,14 +20,13 @@ import torchvision.transforms as T
 from omegaconf import DictConfig, OmegaConf
 from pydantic import BaseModel, ConfigDict
 from vllm.entrypoints.openai.api_server import build_app
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 import vllm_omni.entrypoints.cli.serve
-from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.api_server import omni_init_app_state
 from vllm_omni.inputs.data import OmniCustomPrompt, OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 logger = logging.getLogger(__name__)
 
@@ -183,7 +180,6 @@ class vLLMOmniHttpServerLocal:
             "tensor_parallel_size": _cfg_get(self.config, "tensor_model_parallel_size", 1),
             "seed": self.replica_rank + int(_cfg_get(self.config, "seed", 0) or 0),
             "override_generation_config": json.dumps({}),
-            "scheduling_policy": _cfg_get(self.config, "scheduling_policy", "fcfs"),
             "compilation_config": compilation_config,
             **engine_kwargs,
         }
@@ -191,7 +187,7 @@ class vLLMOmniHttpServerLocal:
         model_path = str(_cfg_get(self.model_config, "local_path", None) or _cfg_get(self.model_config, "path"))
         server_args_list = ["serve", model_path] + build_cli_args_from_config(args)
 
-        parser = FlexibleArgumentParser(description="vLLM-Omni CLI")
+        parser = TrackingArgumentParser(description="vLLM-Omni CLI")
         subparsers = parser.add_subparsers(required=False, dest="subparser")
         cmds: dict[str, Any] = {}
         for cmd in vllm_omni.entrypoints.cli.serve.cmd_init():
@@ -204,11 +200,10 @@ class vLLMOmniHttpServerLocal:
 
         await self.run_server(parsed)
 
-    async def run_server(self, args: argparse.Namespace) -> None:
-        engine_args = OmniEngineArgs.from_cli_args(args)
-        engine_args = asdict(engine_args)
-
-        engine_args["enable_dummy_pipeline"] = True
+    async def run_server(self, args: TrackingNamespace) -> None:
+        # Match the production entrypoint: parser defaults are not user overrides.
+        engine_args = args.get_explicit_kwargs_dict()
+        engine_args["model"] = args.model
         engine_args["custom_pipeline_args"] = {"pipeline_class": CUSTOM_PIPELINE_CLASS}
 
         engine_client = AsyncOmni(**engine_args)

tests/engine/test_stage_engine_args.py

@@ -29,6 +29,7 @@ from vllm_omni.config.omni_config import (
     OmniStageSchedulerConfig,
     VllmOmniARStageConfig,
     VllmOmniConfig,
+    extract_diffusion_stage_config_kwargs,
 )
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
 from vllm_omni.config.stage_config import (
@@ -567,7 +568,9 @@ def test_engine_args_consume_stage_diffusion_attention_shorthand(tmp_path):
         assert engine_args.get("diffusion_attention_backend") is None
         assert isinstance(engine_args["diffusion_attention_config"], AttentionConfig)
         assert engine_args["diffusion_attention_config"].default.backend == "TORCH_SDPA"
-        od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+        od_config = OmniDiffusionConfig.from_kwargs(
+            **extract_diffusion_stage_config_kwargs(engine_args, stage_id=2, include_engine_adapter_metadata=True)
+        )
         assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
 
 
@@ -600,7 +603,9 @@ def test_engine_args_apply_cli_attention_shorthand_over_yaml_config(tmp_path, ya
         attention_config = engine_args["diffusion_attention_config"]
         assert attention_config.default.backend == "TORCH_SDPA"
         assert attention_config.per_role["cross"].backend == "SAGE_ATTN"
-        od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+        od_config = OmniDiffusionConfig.from_kwargs(
+            **extract_diffusion_stage_config_kwargs(engine_args, stage_id=2, include_engine_adapter_metadata=True)
+        )
         assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
 
 

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(

tests/entrypoints/test_omni_sleep_mode.py

@@ -18,6 +18,7 @@ import torch
 from vllm import SamplingParams
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -98,16 +99,6 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
 MODEL_DIFF = "riverclouds/qwen_image_random"
 
 
-def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str:
-    total_gpus = torch.accelerator.device_count()
-    gpus_per_stage = tp_size
-    start_idx = stage_idx * gpus_per_stage
-    if start_idx + gpus_per_stage > total_gpus:
-        start_idx = start_idx % total_gpus
-    device_ids = [str(start_idx + i) for i in range(gpus_per_stage)]
-    return ",".join(device_ids)
-
-
 async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
     """Best-effort full wake + resume so the next shared-engine test starts active."""
     try:
@@ -120,33 +111,28 @@ async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
         logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e)
 
 
-def _build_llm_stages() -> tuple[list[dict], list[dict]]:
-    common_args = {
-        "worker_type": "ar",
-        "enable_sleep_mode": True,
-        "dtype": "bfloat16",
-        "trust_remote_code": True,
-        "max_model_len": 2048,
-        "max_num_batched_tokens": 8192,
-        "enforce_eager": True,
-    }
-    stages = [
-        {
-            "stage_id": 0,
-            "stage_type": "llm",
-            "runtime": {"process": True, "devices": "0", "max_batch_size": 1},
-            "engine_args": {**common_args, "model_stage": "thinker", "gpu_memory_utilization": 0.1},
+def _sleep_deploy_config(*, diffusion_only: bool = False, tp_size: int = 1) -> str:
+    """Use registered BAGEL topology and override only deployment settings."""
+    filename = "bagel_single_stage.yaml" if diffusion_only else "bagel.yaml"
+    stage_ids = (0,) if diffusion_only else (0, 1)
+    return modify_stage_config(
+        get_deploy_config_path(filename),
+        updates={
+            "stages": {
+                stage_id: {
+                    "devices": ",".join(str(stage_id * tp_size + rank) for rank in range(tp_size)),
+                    "tensor_parallel_size": tp_size,
+                    "enable_sleep_mode": True,
+                    "enforce_eager": True,
+                    "dtype": "bfloat16",
+                    # Stages use disjoint GPUs. On an 80 GiB H100, 40% cannot
+                    # fit TP=1 Thinker weights, activation peaks and KV cache.
+                    "gpu_memory_utilization": 0.8 if not diffusion_only and stage_id == 0 else 0.4,
+                }
+                for stage_id in stage_ids
+            }
         },
-        {
-            "stage_id": 1,
-            "stage_type": "llm",
-            "engine_input_source": [0],
-            "runtime": {"process": True, "devices": "1", "max_batch_size": 1, "connector_type": "queue"},
-            "engine_args": {**common_args, "model_stage": "talker", "gpu_memory_utilization": 0.1},
-        },
-    ]
-    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
-    return stages, connectors
+    )
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -163,11 +149,10 @@ def _module_device_cleanup():
 
 @pytest_asyncio.fixture(scope="class", loop_scope="class")
 async def llm_engine():
-    """Shared 2-stage BAGEL LLM engine for sleep/wake protocol tests."""
+    """Shared BAGEL Thinker + DiT engine for LLM sleep/wake protocol tests."""
     if current_omni_platform.is_rocm():
         clean_device_envs()
-    stages, connectors = _build_llm_stages()
-    engine = AsyncOmni(model=MODEL, stages=stages, connectors=connectors, init_timeout=600, enable_sleep_mode=True)
+    engine = AsyncOmni(model=MODEL, deploy_config=_sleep_deploy_config(), init_timeout=600, enable_sleep_mode=True)
     yield engine
     engine.shutdown()
     await asyncio.sleep(1.5)
@@ -178,27 +163,12 @@ async def diffusion_engine():
     """Shared BAGEL diffusion TP=2 engine for sleep/wake + generate checks."""
     if current_omni_platform.is_rocm():
         clean_device_envs()
-    stages = [
-        {
-            "stage_id": 0,
-            "stage_type": "diffusion",
-            "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1},
-            "engine_args": {
-                "model_stage": "base",
-                "gpu_memory_utilization": 0.1,
-                "model_class_name": "BagelPipeline",
-                "enable_sleep_mode": True,
-                "enforce_eager": True,
-                "max_num_batched_tokens": 8192,
-                "parallel_config": {
-                    "tensor_parallel_size": 2,
-                },
-            },
-            "final_output": True,
-            "final_output_type": "image",
-        }
-    ]
-    engine = AsyncOmni(model=MODEL, stages=stages, init_timeout=600, enable_sleep_mode=True)
+    engine = AsyncOmni(
+        model=MODEL,
+        deploy_config=_sleep_deploy_config(diffusion_only=True, tp_size=2),
+        init_timeout=600,
+        enable_sleep_mode=True,
+    )
     yield engine
     engine.shutdown()
     await asyncio.sleep(1.5)
@@ -208,7 +178,7 @@ class TestOmniLlmSleepMode:
     """LLM sleep/wake protocol tests sharing one class-scoped engine."""
 
     @pytest.mark.asyncio(loop_scope="class")
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
     async def test_llm_sleep_ack(self, llm_engine: AsyncOmni):
         """LLM Thinker (GPU0) Signal and Physical Recycling Audit"""
         device_id = 0
@@ -333,11 +303,10 @@ class TestOmniDiffusionSleepMode:
             # --- integrity: baseline generate at original resolution/steps ---
             prompt = "A huge swimming pool, with many people swimming."
             sp = OmniDiffusionSamplingParams(num_inference_steps=4, height=512, width=512, seed=42)
-            llm_sp = SamplingParams()
 
             logger.info("Running Baseline Generation...")
             base_output = None
-            async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params_list=[llm_sp, sp]):
+            async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params=sp):
                 base_output = output
             assert base_output is not None and len(base_output.images) > 0
             logger.info("Baseline Generation successful.")
@@ -384,7 +353,7 @@ class TestOmniDiffusionSleepMode:
             # --- integrity + lifecycle: post-wake generate ---
             logger.info("Running Post-Wakeup Generation...")
             post_output = None
-            async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params_list=[llm_sp, sp]):
+            async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params=sp):
                 post_output = output
             assert post_output is not None
             assert len(base_output.images) == len(post_output.images)
@@ -458,41 +427,22 @@ class TestOmniCoordinatedSleepMode:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.asyncio
 async def test_multistage_llm_diffusion_sleep_wake(tp_size: int):
-    """Explicit 2-stage (llm + diffusion) + connectors; sleep/wake both stages."""
+    """Registered BAGEL Thinker + DiT topology; sleep/wake both stages."""
     if current_omni_platform.is_rocm():
         clean_device_envs()
     num_gpus = torch.accelerator.device_count()
     if num_gpus < tp_size * 2:
         pytest.skip("Not enough GPUs")
 
-    stages = []
-    for i in range(2):
-        devs = get_dynamic_devices(i, 2, tp_size)
-        stages.append(
-            {
-                "stage_id": i,
-                "stage_type": "llm" if i == 0 else "diffusion",
-                "runtime": {"process": True, "devices": devs},
-                "engine_args": {
-                    "model": MODEL,
-                    "model_stage": "thinker" if i == 0 else "base",
-                    "tensor_parallel_size": tp_size,
-                    "gpu_memory_utilization": 0.4,
-                    "dtype": "bfloat16",
-                    "enable_sleep_mode": True,
-                    "trust_remote_code": True,
-                },
-            }
-        )
-
-    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
-
     engine = AsyncOmni(
-        model=MODEL, stages=stages, connectors=connectors, enable_sleep_mode=True, stage_init_timeout=1200
+        model=MODEL,
+        deploy_config=_sleep_deploy_config(tp_size=tp_size),
+        enable_sleep_mode=True,
+        stage_init_timeout=1200,
     )
     try:
         sp = OmniDiffusionSamplingParams(num_inference_steps=2)
-        async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]):
+        async for _ in engine.generate("warmup", sampling_params_list=[engine.default_sampling_params_list[0], sp]):
             pass
 
         acks = await engine.sleep(stage_ids=[0, 1], level=1)
@@ -500,7 +450,7 @@ async def test_multistage_llm_diffusion_sleep_wake(tp_size: int):
 
         await engine.wake_up(stage_ids=[0, 1])
         await engine.resume_generation(stage_ids=[0, 1])
-        async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]):
+        async for _ in engine.generate("verify", sampling_params_list=[engine.default_sampling_params_list[0], sp]):
             pass
     finally:
         engine.shutdown()

vllm_omni/config/config_factory.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import functools
 import json
 from collections.abc import Mapping
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, fields
 from pathlib import Path
 from typing import Any
 
@@ -18,7 +18,10 @@ from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
 from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
 
 from vllm_omni.config.endpoint_policy import EndpointRestriction
-from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.omni_config import (
+    VllmOmniConfig,
+    normalize_and_validate_diffusion_engine_ingress_kwargs,
+)
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
 from vllm_omni.config.stage_config import (
     _DEPLOY_DIR,
@@ -645,6 +648,14 @@ class StageConfigFactory:
         kwargs: dict[str, Any],
     ) -> tuple[dict[str, Any], DiffusionParallelConfig, dict[str, Any], str]:
         """Normalize inputs shared by typed and compatibility diffusion builders."""
+        # Preserve the default builder's explicit top-level override of model
+        # extras before the strict same-source alias validator runs.
+        if kwargs.get("auxiliary_text_encoder") is not None:
+            kwargs = dict(kwargs)
+            extras = dict(kwargs.get("extras") or {})
+            extras.pop("auxiliary_text_encoder", None)
+            kwargs["extras"] = extras
+        kwargs = normalize_and_validate_diffusion_engine_ingress_kwargs(kwargs, stage_id=0)
         raw_sampling_params = kwargs.get("default_sampling_params")
         if isinstance(raw_sampling_params, str):
             try:
@@ -659,7 +670,15 @@ class StageConfigFactory:
         parallel_config = DiffusionParallelConfig.from_stage_overrides(kwargs)
         if kwargs.get("num_gpus") is not None:
             parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"]))
-        engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs)
+        # Ingress keeps the engine spelling; canonicalize it before filtering
+        # to terminal diffusion fields. Conflicts were validated at ingress.
+        quantization = kwargs.pop("quantization", None)
+        if quantization is not None:
+            kwargs["quantization_config"] = quantization
+        diffusion_config_fields = frozenset(config_field.name for config_field in fields(OmniDiffusionConfig))
+        engine_args = OmniDiffusionConfig.normalize_init_kwargs(
+            {name: value for name, value in kwargs.items() if name in diffusion_config_fields}
+        )
 
         extras = dict(engine_args.get("extras") or {})
         for key, default in (

vllm_omni/config/omni_config.py

@@ -18,6 +18,7 @@ from inspect import Parameter, signature
 from pathlib import Path
 from typing import Any, Literal, TypeAlias, TypedDict, cast
 
+import regex as re
 from pydantic import ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 from vllm.config import CacheConfig as VllmCacheConfig
@@ -74,10 +75,14 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
     {
         "async_chunk",
         "disable_log_stats",
+        "command",
+        "headless",
         "model",
+        "model_tag",
         "omni",
         "output_modalities",
         "stage_id",
+        "subparser",
         "tokenizer",
     }
 )
@@ -95,6 +100,8 @@ _LEGACY_STAGE_METADATA_EXTRA_FIELDS = frozenset(
 
 _QuantizationConfigType: TypeAlias = QuantizationConfig | str | Mapping[str, Any] | None
 
+_DIFFUSION_SHARED_ONLY_ENGINE_FIELDS = frozenset({"kv_cache_dtype", "seed"})
+
 
 class _QuantizationEngineOverrides(TypedDict, total=False):
     quantization_config: _QuantizationConfigType
@@ -290,10 +297,6 @@ class _DiffusionEngineOverrides:
 
     _values: dict[str, Any]
 
-    @classmethod
-    def from_engine(cls, engine: Mapping[str, Any]) -> _DiffusionEngineOverrides:
-        return cls(_select_engine_overrides(engine, _DIFFUSION_STAGE_ENGINE_FIELDS))
-
     def to_kwargs(self) -> dict[str, Any]:
         return {name: _copy_value(value) for name, value in self._values.items()}
 
@@ -359,12 +362,32 @@ def _stage_cli_overrides(
     *,
     execution_type: StageExecutionType | None = None,
 ) -> dict[str, Any]:
+    if execution_type == StageExecutionType.DIFFUSION:
+        stage_override_pattern = re.compile(r"^stage_\d+_")
+        stage_scoped = {
+            name: _copy_value(value) for name, value in cli_overrides.items() if stage_override_pattern.match(name)
+        }
+        global_inputs = {name: _copy_value(value) for name, value in cli_overrides.items() if name not in stage_scoped}
+        cli_overrides = {
+            **normalize_and_validate_diffusion_engine_ingress_kwargs(
+                global_inputs,
+                stage_id=stage_id,
+            ),
+            **stage_scoped,
+        }
     runtime_overrides = build_stage_runtime_overrides(stage_id, dict(cli_overrides))
     global_stage_fields = _global_stage_cli_fields()
     owned_fields = None if execution_type is None else _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
     result: dict[str, Any] = {}
     for key, value in runtime_overrides.items():
-        stage_specific = f"stage_{stage_id}_{key}" in cli_overrides
+        stage_key = f"stage_{stage_id}_{key}"
+        stage_specific = stage_key in cli_overrides
+        if (
+            execution_type == StageExecutionType.DIFFUSION
+            and key in _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS
+            and not stage_specific
+        ):
+            continue
         if stage_specific or (key in global_stage_fields and (owned_fields is None or key in owned_fields)):
             result[key] = _copy_value(value)
 
@@ -372,6 +395,19 @@ def _stage_cli_overrides(
     # argument. Keep global and stage-scoped CLI values off AR/generation stages.
     if execution_type is not None and execution_type is not StageExecutionType.DIFFUSION:
         result.pop("step_execution", None)
+    if execution_type == StageExecutionType.DIFFUSION:
+        prefix = f"stage_{stage_id}_"
+        for key, value in cli_overrides.items():
+            if key.startswith(prefix):
+                field_name = key.removeprefix(prefix)
+                if field_name in {"model", "model_arch", "stage_id"}:
+                    raise ValueError(f"Diffusion stage {stage_id} cannot override identity field {field_name!r}.")
+                if field_name in _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS:
+                    raise ValueError(
+                        f"Diffusion stage {stage_id} cannot override shared engine field {field_name!r}; "
+                        "it has no diffusion stage-config consumer."
+                    )
+                result.setdefault(field_name, _copy_value(value))
     return result
 
 
@@ -387,6 +423,10 @@ def _validate_global_stage_cli_ownership(
         field for stage in pipeline.stages for field in _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[stage.execution_type]
     }
     unowned_fields = explicit_global_fields - owned_fields
+    if any(stage.execution_type == StageExecutionType.DIFFUSION for stage in pipeline.stages):
+        # Mixed engine ingress accepts these shared globals, but diffusion
+        # stages deliberately leave them outside their terminal config.
+        unowned_fields -= _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS
     if unowned_fields:
         names = ", ".join(sorted(unowned_fields))
         raise ValueError(
@@ -736,7 +776,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             self.world_size = other_parallel_world_size
 
 
-@config(config=ConfigDict(arbitrary_types_allowed=True))
+@config(config=ConfigDict(arbitrary_types_allowed=True, extra="forbid"))
 class _DiffusionConfigProjection:
     """Diffusion-specific per-stage settings.
 
@@ -852,15 +892,19 @@ class _DiffusionConfigProjection:
 
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> _DiffusionConfigProjection:
-        from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
+        from vllm_omni.diffusion.data import (
+            normalize_omni_diffusion_kwargs,
+            validate_omni_diffusion_kwargs,
+        )
         from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config
 
-        config_kwargs = normalize_omni_diffusion_kwargs(kwargs)
+        valid_fields = frozenset(f.name for f in fields(cast(Any, cls)))
+        normalized = normalize_omni_diffusion_kwargs(kwargs)
+        validate_omni_diffusion_kwargs(normalized, valid_fields)
         # Validate before stage construction while retaining the raw mapping
         # needed by dataclass/config serialization across process boundaries.
-        parse_diffusion_offload_config(config_kwargs.get("diffusion_offload_config"))
-        valid_fields = {f.name for f in fields(cast(Any, cls))}
-        return cls(**{k: v for k, v in config_kwargs.items() if k in valid_fields})
+        parse_diffusion_offload_config(normalized.get("diffusion_offload_config"))
+        return cls(**{name: value for name, value in normalized.items() if value is not None})
 
     def __post_init__(self) -> None:
         # Keep diffusion imports lazy so importing vllm_omni.config does not
@@ -1095,20 +1139,7 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
 
 _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS)
 
-_DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset(
-    {
-        "diffusion_attention_backend",
-        "fastvideo_vsa_topk",
-        "kv_cache_dtype",
-        "kv_cache_skip_layers",
-        "kv_cache_skip_steps",
-        "static_lora_scale",
-    }
-)
-_DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMPAT_ENGINE_FIELDS) - {
-    "model",
-    "stage_id",
-}
+_DIFFUSION_STAGE_ENGINE_FIELDS = _DIFFUSION_CONFIG_FIELDS - {"model", "stage_id"}
 
 
 def _upstream_engine_field_map(
@@ -1248,6 +1279,8 @@ def _validate_stage_engine_override_ownership(
     stage_id: int,
     execution_type: StageExecutionType,
     overrides: Mapping[str, Any],
+    *,
+    validate_top_level: bool = True,
 ) -> None:
     try:
         owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
@@ -1256,7 +1289,7 @@ def _validate_stage_engine_override_ownership(
     except KeyError as exc:
         raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc
 
-    unowned_fields = set(overrides) - owner_fields
+    unowned_fields = set(overrides) - owner_fields if validate_top_level else set()
     parallel_config = overrides.get("parallel_config")
     if isinstance(parallel_config, Mapping):
         # Nested values traditionally use upstream config names, while flat
@@ -1273,6 +1306,144 @@ def _validate_stage_engine_override_ownership(
         )
 
 
+_DIFFUSION_STAGE_METADATA_FIELDS = frozenset(
+    {
+        "async_chunk",
+        "custom_process_next_stage_input_func",
+        "engine_output_type",
+        "has_sampling_extra_args",
+        "hf_config_name",
+        "model_arch",
+        "model_stage",
+        "retains_state_across_chunks",
+        "scheduler_cls",
+        "stage_connector_spec",
+        "worker_type",
+    }
+)
+
+_DIFFUSION_ENGINE_ADAPTER_METADATA_FIELDS = frozenset(
+    {
+        # Serialized shared model settings are consumed outside the terminal
+        # diffusion config; accepting them here does not widen raw ingress.
+        "duplex_max_sessions",
+        "has_sampling_extra_args",
+        "inline_diffusion",
+        "requires_full_payload_input",
+        "sampling_extra_args_keys",
+        "session_mode",
+    }
+)
+_DIFFUSION_DEFAULT_FACTORY_FIELDS = frozenset(
+    {"default_llama_model_id", "default_sampling_params", "devices", "stage_0_devices"}
+)
+
+
+def _frontend_cli_fields() -> frozenset[str]:
+    """Return vLLM server fields that are consumed before stage startup."""
+    from vllm.entrypoints.launchers.cli_args import FrontendArgs
+
+    return frozenset(config_field.name for config_field in fields(FrontendArgs))
+
+
+def normalize_and_validate_diffusion_engine_ingress_kwargs(
+    kwargs: Mapping[str, Any],
+    *,
+    stage_id: int | str,
+) -> dict[str, Any]:
+    """Normalize and validate raw diffusion input without inserting defaults."""
+    from vllm_omni.diffusion.data import (
+        OmniDiffusionConfig,
+        normalize_omni_diffusion_kwargs,
+        validate_omni_diffusion_kwargs,
+    )
+    from vllm_omni.engine.arg_utils import orchestrator_field_names
+
+    mixed_kwargs = {name: _copy_value(value) for name, value in kwargs.items()}
+    engine_owned = {
+        name: mixed_kwargs.pop(name)
+        for name in _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS | {"quantization"}
+        if name in mixed_kwargs
+    }
+    normalized = normalize_omni_diffusion_kwargs(mixed_kwargs, apply_defaults=False)
+    if engine_owned.get("quantization") is not None and normalized.get("quantization_config") is not None:
+        raise ValueError("Diffusion config fields 'quantization' and 'quantization_config' cannot both be provided.")
+    normalized.update(engine_owned)
+
+    diffusion_fields = frozenset(config_field.name for config_field in fields(OmniDiffusionConfig))
+    stage_consumed_fields = (
+        diffusion_fields
+        | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS
+        | frozenset(_STAGE_DEPLOY_ENGINE_FIELDS)
+        | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS)
+        | _DIFFUSION_STAGE_METADATA_FIELDS
+        | _DIFFUSION_DEFAULT_FACTORY_FIELDS
+    ) - _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS
+    externally_consumed_fields = (
+        _DIFFUSION_SHARED_ONLY_ENGINE_FIELDS
+        | _NON_STAGE_ENGINE_CLI_FIELDS
+        | _frontend_cli_fields()
+        | orchestrator_field_names()
+        # Coordination fields also live on the typed orchestrator config, not
+        # all of them are present on the CLI-only OrchestratorArgs dataclass.
+        | frozenset(config_field.name for config_field in fields(VllmOmniOrchestratorConfig))
+    )
+    allowed_fields = stage_consumed_fields | externally_consumed_fields
+    validate_omni_diffusion_kwargs(normalized, allowed_fields, stage_id=stage_id)
+    return {name: value for name, value in normalized.items() if name in stage_consumed_fields}
+
+
+def extract_diffusion_stage_config_kwargs(
+    kwargs: Mapping[str, Any],
+    *,
+    stage_id: int | str,
+    include_engine_adapter_metadata: bool = False,
+) -> dict[str, Any]:
+    """Take the diffusion-owned payload from resolved mixed stage arguments."""
+    from vllm_omni.diffusion.data import (
+        OmniDiffusionConfig,
+        normalize_omni_diffusion_kwargs,
+        validate_omni_diffusion_kwargs,
+    )
+
+    shared_only_fields = sorted(_DIFFUSION_SHARED_ONLY_ENGINE_FIELDS.intersection(kwargs))
+    if shared_only_fields:
+        field_names = ", ".join(repr(name) for name in shared_only_fields)
+        raise ValueError(
+            f"Diffusion stage {stage_id} cannot consume shared engine field(s) {field_names}; "
+            "use diffusion-owned fields or request sampling parameters instead."
+        )
+
+    diffusion_fields = frozenset(config_field.name for config_field in fields(OmniDiffusionConfig))
+    fields_owned_elsewhere = (
+        frozenset(_STAGE_DEPLOY_ENGINE_FIELDS)
+        | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS)
+        | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS
+        | _DIFFUSION_STAGE_METADATA_FIELDS
+    )
+    if include_engine_adapter_metadata:
+        fields_owned_elsewhere |= _DIFFUSION_ENGINE_ADAPTER_METADATA_FIELDS
+
+    # ``quantization`` is canonical in the mixed engine namespace. Hand it to
+    # the compatibility adapter without treating it as a deprecated alias.
+    mixed_kwargs = {name: _copy_value(value) for name, value in kwargs.items()}
+    engine_quantization = mixed_kwargs.pop("quantization", None)
+    normalized = normalize_omni_diffusion_kwargs(mixed_kwargs)
+    if engine_quantization is not None:
+        if normalized.get("quantization_config") is not None:
+            raise ValueError(
+                "Diffusion config fields 'quantization' and 'quantization_config' cannot both be provided."
+            )
+        normalized["quantization_config"] = engine_quantization
+
+    validate_omni_diffusion_kwargs(
+        normalized,
+        diffusion_fields | fields_owned_elsewhere,
+        stage_id=stage_id,
+    )
+    return {name: _copy_value(value) for name, value in normalized.items() if name in diffusion_fields}
+
+
 def _global_stage_cli_fields() -> frozenset[str]:
     # Lazy import avoids vllm_omni.config -> omni_config -> engine.arg_utils ->
     # vllm_omni.config during package-level config imports.
@@ -1285,6 +1456,7 @@ def _global_stage_cli_fields() -> frozenset[str]:
     )
     externally_consumed = (
         _NON_STAGE_ENGINE_CLI_FIELDS
+        | _frontend_cli_fields()
         | frozenset(f.name for f in fields(cast(Any, VllmOmniOrchestratorConfig)))
         | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS)
     )
@@ -1343,10 +1515,15 @@ def _stage_engine_values(
                 engine[key] = _get_recursively_merged_dict(existing, dict(value))
             else:
                 engine[key] = _copy_value(value)
+    if topology.execution_type == StageExecutionType.DIFFUSION:
+        diffusion_kwargs = extract_diffusion_stage_config_kwargs(engine, stage_id=topology.stage_id)
+    else:
+        diffusion_kwargs = {}
     _validate_stage_engine_override_ownership(
         topology.stage_id,
         topology.execution_type,
         engine,
+        validate_top_level=topology.execution_type != StageExecutionType.DIFFUSION,
     )
     if topology.execution_type in {
         StageExecutionType.LLM_AR,
@@ -1383,7 +1560,7 @@ def _stage_engine_values(
         ),
         runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, runtime_engine_fields)),
         parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)),
-        diffusion=_DiffusionEngineOverrides.from_engine(engine),
+        diffusion=_DiffusionEngineOverrides(_select_engine_overrides(diffusion_kwargs, _DIFFUSION_STAGE_ENGINE_FIELDS)),
         compilation_config=_copy_value(engine.get("compilation_config")),
         profiler_config=_copy_value(engine.get("profiler_config")),
     )

vllm_omni/diffusion/data.py

@@ -6,6 +6,7 @@ import json
 import math
 import os
 import random
+import warnings
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field, fields
 from enum import Enum
@@ -44,55 +45,63 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
-def normalize_omni_diffusion_kwargs(raw_kwargs: Mapping[str, Any]) -> dict[str, Any]:
-    """Normalize legacy diffusion kwargs before config construction."""
+def _move_diffusion_alias(
+    normalized: dict[str, Any],
+    legacy_name: str,
+    canonical_name: str,
+) -> None:
+    legacy_value = normalized.pop(legacy_name, None)
+    if legacy_value is None:
+        return
+    if normalized.get(canonical_name) is not None:
+        raise ValueError(f"Diffusion config fields {legacy_name!r} and {canonical_name!r} cannot both be provided.")
+    warnings.warn(
+        f"Diffusion config field {legacy_name!r} is deprecated; use {canonical_name!r}.",
+        FutureWarning,
+        stacklevel=3,
+    )
+    normalized[canonical_name] = legacy_value
+
+
+def normalize_omni_diffusion_kwargs(
+    raw_kwargs: Mapping[str, Any],
+    *,
+    apply_defaults: bool = True,
+) -> dict[str, Any]:
+    """Normalize diffusion kwargs, deferring defaults until sources are merged."""
     config_kwargs = dict(raw_kwargs)
 
     dtype = config_kwargs.get("dtype")
     if dtype is None:
-        config_kwargs["dtype"] = "auto"
+        if apply_defaults:
+            config_kwargs["dtype"] = "auto"
     elif isinstance(dtype, torch.dtype):
         config_kwargs["dtype"] = str(dtype).removeprefix("torch.")
     elif not isinstance(dtype, str):
         raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(dtype).__name__}")
 
-    # Backwards-compatibility: older callers may use a diffusion-specific
-    # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale".
-    if "static_lora_scale" in config_kwargs:
-        if "lora_scale" not in config_kwargs:
-            config_kwargs["lora_scale"] = config_kwargs["static_lora_scale"]
-        config_kwargs.pop("static_lora_scale", None)
-
-    diffusion_quantization = config_kwargs.pop("diffusion_quantization_config", None)
-    if config_kwargs.get("quantization_config") is None and diffusion_quantization is not None:
-        config_kwargs["quantization_config"] = diffusion_quantization
-
-    # Backwards-compatibility: map "quantization" to "quantization_config"
-    # so callers using the old field name still work.
-    if "quantization" in config_kwargs and config_kwargs.get("quantization_config", None) is None:
-        config_kwargs["quantization_config"] = config_kwargs.pop("quantization")
-    else:
-        config_kwargs.pop("quantization", None)
-
-    # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
-    if config_kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in config_kwargs:
-        config_kwargs["diffusion_kv_cache_dtype"] = config_kwargs.pop("kv_cache_dtype")
-    else:
-        config_kwargs.pop("kv_cache_dtype", None)
-    if config_kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in config_kwargs:
-        config_kwargs["diffusion_kv_cache_skip_steps"] = config_kwargs.pop("kv_cache_skip_steps")
-    else:
-        config_kwargs.pop("kv_cache_skip_steps", None)
-    if config_kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in config_kwargs:
-        config_kwargs["diffusion_kv_cache_skip_layers"] = config_kwargs.pop("kv_cache_skip_layers")
-    else:
-        config_kwargs.pop("kv_cache_skip_layers", None)
+    for legacy, canonical in (
+        ("static_lora_scale", "lora_scale"),
+        ("quantization", "quantization_config"),
+        ("diffusion_quantization_config", "quantization_config"),
+        ("max_batch_size", "max_num_seqs"),
+        ("kv_cache_skip_steps", "diffusion_kv_cache_skip_steps"),
+        ("kv_cache_skip_layers", "diffusion_kv_cache_skip_layers"),
+    ):
+        _move_diffusion_alias(config_kwargs, legacy, canonical)
+    _move_diffusion_alias(config_kwargs, "kv_cache_dtype", "diffusion_kv_cache_dtype")
 
     # Handle "diffusion_attention_backend" shorthand: merge into
     # diffusion_attention_config before field filtering.
     diffusion_attn_backend = config_kwargs.pop("diffusion_attention_backend", None)
     fastvideo_vsa_topk = config_kwargs.pop("fastvideo_vsa_topk", None)
     if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None:
+        if diffusion_attn_backend is not None:
+            warnings.warn(
+                "Diffusion config field 'diffusion_attention_backend' is deprecated; use 'diffusion_attention_config'.",
+                FutureWarning,
+                stacklevel=2,
+            )
         existing = config_kwargs.get("diffusion_attention_config")
         config_kwargs["diffusion_attention_config"] = parse_attention_config(
             existing,
@@ -100,12 +109,23 @@ def normalize_omni_diffusion_kwargs(raw_kwargs: Mapping[str, Any]) -> dict[str,
             fastvideo_vsa_topk=fastvideo_vsa_topk,
         )
 
+    auxiliary_text_encoder = config_kwargs.pop("auxiliary_text_encoder", None)
+    if auxiliary_text_encoder is not None:
+        extras = dict(config_kwargs.get("extras") or {})
+        if extras.get("auxiliary_text_encoder") is not None:
+            raise ValueError(
+                "Diffusion engine field 'auxiliary_text_encoder' cannot be provided both at the top level and in "
+                "'extras'."
+            )
+        extras["auxiliary_text_encoder"] = auxiliary_text_encoder
+        config_kwargs["extras"] = extras
+
     # Check environment variable as fallback for cache_backend.
     # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
-    if "cache_backend" not in config_kwargs:
+    if "cache_backend" not in config_kwargs and apply_defaults:
         cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
         config_kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
-    elif config_kwargs["cache_backend"] is None:
+    elif "cache_backend" in config_kwargs and config_kwargs["cache_backend"] is None and apply_defaults:
         # Callers (e.g. example CLIs with `default=None`) pass an explicit
         # None for "no cache"; canonicalize it so every consumer sees the
         # declared `str` value instead of relying on per-model None handling.
@@ -162,6 +182,20 @@ def validate_dlo_host_registration_options(
     return value
 
 
+def validate_omni_diffusion_kwargs(
+    kwargs: Mapping[str, Any],
+    allowed_fields: set[str] | frozenset[str],
+    *,
+    stage_id: int | str | None = None,
+) -> None:
+    """Reject every field without an owner."""
+    unknown = set(kwargs) - allowed_fields
+    if unknown:
+        names = ", ".join(repr(name) for name in sorted(unknown))
+        suffix = "" if stage_id is None else f" for stage {stage_id}"
+        raise ValueError(f"Unknown diffusion config field(s){suffix}: {names}")
+
+
 def parse_kv_cache_skip_selector(
     selector: str | list[int] | tuple[int, ...] | set[int] | None,
 ) -> set[int] | None:
@@ -1631,8 +1665,9 @@ class OmniDiffusionConfig:
 
     @classmethod
     def normalize_init_kwargs(cls, raw_kwargs: Mapping[str, Any]) -> dict[str, Any]:
+        valid_fields = frozenset(f.name for f in fields(cls))
         config_kwargs = normalize_omni_diffusion_kwargs(raw_kwargs)
-        valid_fields = {f.name for f in fields(cls)}
+        validate_omni_diffusion_kwargs(config_kwargs, valid_fields)
         # Remaining ``None`` values mean "unset" at the CLI/deploy boundary.
         # Drop them so non-optional dataclass defaults are not overwritten.
         # Fields where ``None`` has normalization semantics (for example dtype

vllm_omni/engine/stage_init_utils.py

@@ -1965,13 +1965,20 @@ def build_diffusion_config(
     metadata: StageMetadata,
 ) -> Any:
     """Build diffusion config for a stage."""
+    from vllm_omni.config.omni_config import extract_diffusion_stage_config_kwargs
 
     engine_args_dict = (
         build_engine_args_dict_from_omni_stage_config(stage_cfg, model)
         if isinstance(stage_cfg, BaseVllmOmniStageConfig)
         else build_engine_args_dict(stage_cfg, model)
     )
-    od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict)
+    stage_id = engine_args_dict.get("stage_id", metadata.stage_id)
+    diffusion_kwargs = extract_diffusion_stage_config_kwargs(
+        engine_args_dict,
+        stage_id=stage_id,
+        include_engine_adapter_metadata=True,
+    )
+    od_config = OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None})
 
     num_devices_per_stage = od_config.parallel_config.world_size
     device_control_env = current_omni_platform.device_control_env_var