[Config] Add structured VllmOmniConfig classes (#4425)

Signed-off-by: Acerak01-fy <wfy2003324@163.com> Signed-off-by: zwhzzz0821 <2831474076@qq.com> Signed-off-by: Wang fuyin <wfy2003324@163.com> Co-authored-by: zwhzzz0821 <2831474076@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-25 17:48:20 +08:00
1 parent 2e6c83c89c
commit 8ddcb3621f
@@ -0,0 +1,752 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the additive structured Omni config."""
+
+from __future__ import annotations
+
+from dataclasses import fields
+from pathlib import Path
+
+import pytest
+from pydantic import ValidationError
+
+from vllm_omni.config import omni_config as omni_config_module
+from vllm_omni.config.omni_config import (
+    BaseVllmOmniStageConfig,
+    OmniStageCacheConfig,
+    OmniStageConnectorConfig,
+    OmniStageDiffusionParallelConfig,
+    OmniStageLoadConfig,
+    OmniStageModelConfig,
+    OmniStageParallelConfig,
+    OmniStageRuntimeConfig,
+    OmniStageSchedulerConfig,
+    VllmOmniARStageConfig,
+    VllmOmniConfig,
+    VllmOmniDiffusionStageConfig,
+    VllmOmniGenerationStageConfig,
+)
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.stage_config import (
+    _STAGE_DEPLOY_FIELDS,
+    PIPELINE_WIDE_ENGINE_FIELDS,
+    DeployConfig,
+    PipelineConfig,
+    StageDeployConfig,
+    StageExecutionType,
+    load_deploy_config,
+    merge_pipeline_deploy,
+)
+from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
+
+
+@pytest.fixture(autouse=True)
+def _stable_test_platform(monkeypatch):
+    from vllm_omni import platforms
+
+    platform = platforms.current_omni_platform
+    monkeypatch.setattr(platform, "device_name", "cpu", raising=False)
+    monkeypatch.setattr(platform, "device_type", "cpu", raising=False)
+
+
+def _load_default_deploy(model_type: str) -> DeployConfig:
+    deploy_path = _DEPLOY_DIR / f"{model_type}.yaml"
+    if deploy_path.exists():
+        return load_deploy_config(deploy_path)
+    return DeployConfig()
+
+
+def _resolve_pipeline_or_skip(model_type: str) -> PipelineConfig:
+    registered = OMNI_PIPELINES[model_type]
+    pipeline = registered(None) if callable(registered) else registered
+    if pipeline is None:
+        pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
+    return pipeline
+
+
+@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
+def test_vllm_omni_config_from_registry_matches_merge_pipeline_deploy(model_type: str):
+    pipeline = _resolve_pipeline_or_skip(model_type)
+    legacy_deploy = _load_default_deploy(model_type)
+
+    legacy_stages = merge_pipeline_deploy(pipeline, legacy_deploy)
+    omni_config = VllmOmniConfig.from_registry(model_type)
+
+    assert omni_config.pipeline_config is pipeline
+    assert len(omni_config.stage_configs) == len(legacy_stages)
+
+    for legacy_stage, omni_stage in zip(legacy_stages, omni_config.stage_configs, strict=True):
+        assert omni_config.stage_by_id(legacy_stage.stage_id) is omni_stage
+
+        assert omni_stage.stage_pipeline_config is pipeline.get_stage(legacy_stage.stage_id)
+        assert omni_stage.model_config.default_sampling_params == legacy_stage.yaml_extras.get(
+            "default_sampling_params"
+        )
+        assert omni_stage.connector_config.output_connectors == legacy_stage.yaml_extras.get("output_connectors")
+        assert omni_stage.connector_config.input_connectors == legacy_stage.yaml_extras.get("input_connectors")
+        assert omni_stage.runtime_config.devices == legacy_stage.yaml_runtime.get("devices")
+        assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
+
+        engine_args = legacy_stage.yaml_engine_args
+        assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
+        assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
+        assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
+        assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization", 0.90)
+        assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs", 128)
+        assert omni_stage.scheduler_config.max_num_batched_tokens == engine_args.get("max_num_batched_tokens")
+        assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling", True)
+        legacy_parallel_config = engine_args.get("parallel_config") or {}
+        assert omni_stage.parallel_config.tensor_parallel_size == legacy_parallel_config.get(
+            "tensor_parallel_size",
+            engine_args.get("tensor_parallel_size", 1),
+        )
+
+        if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.DIFFUSION:
+            assert isinstance(omni_stage, VllmOmniDiffusionStageConfig)
+            assert omni_stage.diffusion_config is not None
+            assert omni_stage.diffusion_config.stage_id == legacy_stage.stage_id
+            assert omni_stage.diffusion_config.model_arch == engine_args.get("model_arch")
+        elif omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR:
+            assert isinstance(omni_stage, VllmOmniARStageConfig)
+            assert not hasattr(omni_stage, "diffusion_config")
+        else:
+            assert isinstance(omni_stage, VllmOmniGenerationStageConfig)
+            assert not hasattr(omni_stage, "diffusion_config")
+
+
+def test_stage_by_id_raises_for_unknown_stage():
+    omni_config = VllmOmniConfig.from_registry("qwen3_tts")
+
+    with pytest.raises(KeyError, match="no stage 99"):
+        omni_config.stage_by_id(99)
+
+
+def test_resolve_execution_mode_rejects_unknown_execution_type():
+    with pytest.raises(ValueError, match="Unsupported stage execution type"):
+        omni_config_module._resolve_execution_mode("unknown_execution_type")
+
+
+def test_from_registry_preserves_current_pipeline_config_object():
+    omni_config = VllmOmniConfig.from_registry("minicpmo_4_5")
+    pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")
+
+    assert omni_config.pipeline_config is pipeline
+    assert not hasattr(omni_config, "pipeline")
+    assert "hf_config_predicate" in {f.name for f in fields(PipelineConfig)}
+    assert omni_config.pipeline_config.hf_config_predicate is pipeline.hf_config_predicate
+
+
+def test_from_registry_normalizes_stage_engine_extras_without_expanding_stage_deploy_config():
+    assert not hasattr(StageDeployConfig, "model_config")
+    assert not hasattr(StageDeployConfig, "parallel_config")
+
+    stage = VllmOmniConfig.from_registry("dreamzero", deploy_config_path="dreamzero_tp1_cfg2").stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.tensor_parallel_size == 1
+    assert stage.parallel_config.cfg_parallel_size == 2
+    assert stage.diffusion_config.model_config["default_robot_embodiment"] == "roboarena"
+
+
+def test_from_registry_applies_cli_overrides_without_stage_config_runtime_bridge():
+    omni_config = VllmOmniConfig.from_registry(
+        "qwen3_tts",
+        cli_overrides={
+            "stage_0_max_num_seqs": 7,
+            "stage_1_tensor_parallel_size": 2,
+        },
+    )
+
+    stage0 = omni_config.stage_by_id(0)
+    stage1 = omni_config.stage_by_id(1)
+
+    assert stage0.scheduler_config.max_num_seqs == 7
+    assert stage1.parallel_config.tensor_parallel_size == 2
+    assert stage1.runtime_config.num_gpus == stage1.parallel_config.world_size
+
+
+def test_runtime_num_gpus_is_derived_from_parallel_world_size():
+    omni_config = VllmOmniConfig.from_registry("hunyuan_image3_dit")
+    stage = omni_config.stage_by_id(0)
+
+    assert stage.parallel_config.tensor_parallel_size == 4
+    assert stage.parallel_config.world_size == 4
+    assert stage.runtime_config.num_gpus == 4
+
+
+def test_runtime_num_gpus_ignores_stale_runtime_override():
+    omni_config = VllmOmniConfig.from_registry(
+        "hunyuan_image3_dit",
+        cli_overrides={
+            "stage_0_num_gpus": 1,
+        },
+    )
+    stage = omni_config.stage_by_id(0)
+
+    assert stage.parallel_config.world_size == 4
+    assert stage.runtime_config.num_gpus == 4
+
+
+def test_from_registry_does_not_route_server_cli_keys_to_diffusion_stage():
+    omni_config = VllmOmniConfig.from_registry(
+        "dreamzero",
+        deploy_config_path="dreamzero_tp1_cfg2",
+        cli_overrides={
+            "host": "0.0.0.0",
+            "port": 8000,
+            "api_key": "secret",
+            "stage_0_host": "127.0.0.1",
+            "stage_0_port": 23456,
+        },
+    )
+
+    stage = omni_config.stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.host == "127.0.0.1"
+    assert stage.diffusion_config.port == 23456
+    assert not hasattr(stage.diffusion_config, "api_key")
+
+
+def test_pipeline_deploy_cli_fields_reuse_legacy_pipeline_wide_engine_fields():
+    assert omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS is PIPELINE_WIDE_ENGINE_FIELDS
+    assert "active_stream_window" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS
+    assert "custom_voice_dir" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS
+
+
+def test_pipeline_wide_model_fields_are_retained_on_structured_stage_configs(tmp_path):
+    custom_voice_dir = tmp_path / "voices"
+    omni_config = VllmOmniConfig.from_registry(
+        "qwen3_tts",
+        cli_overrides={
+            "active_stream_window": 2,
+            "custom_voice_dir": str(custom_voice_dir),
+        },
+    )
+
+    assert {stage.model_config.active_stream_window for stage in omni_config.stage_configs} == {2}
+    assert {stage.model_config.custom_voice_dir for stage in omni_config.stage_configs} == {str(custom_voice_dir)}
+
+
+def test_stage_deploy_engine_fields_reuse_legacy_stage_deploy_fields():
+    assert omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS == tuple(_STAGE_DEPLOY_FIELDS)
+    assert "tensor_parallel_size" in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS
+    assert "stage_id" not in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS
+
+
+def test_public_config_exports_use_stage_specific_sub_config_names():
+    import vllm_omni.config as config_pkg
+
+    generic_names = {
+        "CacheConfig",
+        "ConnectorConfig",
+        "LoadConfig",
+        "ModelConfig",
+        "OrchestratorConfig",
+        "ParallelConfig",
+        "RuntimeConfig",
+        "SchedulerConfig",
+    }
+
+    assert generic_names.isdisjoint(config_pkg.__all__)
+    assert {
+        "OmniStageCacheConfig",
+        "OmniStageConnectorConfig",
+        "OmniStageDiffusionParallelConfig",
+        "OmniStageLoadConfig",
+        "OmniStageModelConfig",
+        "VllmOmniOrchestratorConfig",
+        "OmniStageParallelConfig",
+        "OmniStageRuntimeConfig",
+        "OmniStageSchedulerConfig",
+        "StageConfigType",
+    }.issubset(config_pkg.__all__)
+
+
+def test_from_registry_keeps_worker_backend_separate_from_distributed_executor_backend():
+    omni_config = VllmOmniConfig.from_registry("dreamzero", deploy_config_path="dreamzero_tp1_cfg2")
+
+    stage = omni_config.stage_by_id(0)
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.distributed_executor_backend == "mp"
+    assert omni_config.orchestrator_config.worker_backend == "multi_process"
+
+
+def test_from_registry_maps_orchestrator_cli_overrides():
+    omni_config = VllmOmniConfig.from_registry(
+        "qwen3_tts",
+        cli_overrides={
+            "stage_init_timeout": 1200,
+            "init_timeout": 1800,
+            "worker_backend": "ray",
+            "ray_address": "ray://127.0.0.1:10001",
+            "omni_master_address": "127.0.0.1",
+            "omni_master_port": 12345,
+            "omni_dp_size_local": 2,
+            "omni_lb_policy": "round_robin",
+            "omni_heartbeat_timeout": 9.5,
+            "shm_threshold_bytes": 4096,
+            "batch_timeout": 3,
+        },
+    )
+
+    orchestrator_config = omni_config.orchestrator_config
+    assert orchestrator_config.stage_init_timeout == 1200
+    assert orchestrator_config.init_timeout == 1800
+    assert orchestrator_config.worker_backend == "ray"
+    assert orchestrator_config.ray_address == "ray://127.0.0.1:10001"
+    assert orchestrator_config.omni_master_address == "127.0.0.1"
+    assert orchestrator_config.omni_master_port == 12345
+    assert orchestrator_config.omni_dp_size_local == 2
+    assert orchestrator_config.omni_lb_policy == "round_robin"
+    assert orchestrator_config.omni_heartbeat_timeout == 9.5
+    assert orchestrator_config.shm_threshold_bytes == 4096
+    assert orchestrator_config.batch_timeout == 3
+
+
+def test_from_registry_records_loaded_deploy_path_on_orchestrator_config():
+    omni_config = VllmOmniConfig.from_registry("dreamzero", deploy_config_path="dreamzero_tp1_cfg2")
+
+    assert omni_config.pipeline_config.model_type == "dreamzero"
+    assert omni_config.orchestrator_config.deploy_config_path == str(_DEPLOY_DIR / "dreamzero_tp1_cfg2.yaml")
+
+
+def test_from_registry_dispatches_async_chunk_processors_without_mutating_topology():
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+
+    async_config = VllmOmniConfig.from_registry("qwen3_tts")
+    assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk")
+    assert async_config.stage_by_id(1).custom_process_input_func is None
+
+    sync_config = VllmOmniConfig.from_registry("qwen3_tts", cli_overrides={"async_chunk": False})
+    assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
+    assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
+
+    assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
+    assert pipeline.get_stage(1).custom_process_input_func is None
+
+
+def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
+    assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config")
+    assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config")
+
+    public_fields = {f.name for f in fields(BaseVllmOmniStageConfig)}
+
+    assert public_fields == {
+        "stage_pipeline_config",
+        "model_config",
+        "load_config",
+        "cache_config",
+        "scheduler_config",
+        "connector_config",
+        "runtime_config",
+        "parallel_config",
+        "quantization_config",
+    }
+    assert "diffusion_config" not in public_fields
+    assert {f.name for f in fields(VllmOmniDiffusionStageConfig)} == public_fields | {"diffusion_config"}
+    assert {f.name for f in fields(VllmOmniARStageConfig)} == public_fields
+    assert {f.name for f in fields(VllmOmniGenerationStageConfig)} == public_fields
+
+
+def test_runtime_config_fields_match_rfc_runtime_scope():
+    assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
+        "devices",
+        "num_replicas",
+        "env",
+        "num_gpus",
+        "log_level",
+        "log_stats",
+        "profiler_config",
+    }
+
+
+def test_sub_config_fields_match_rfc_scopes():
+    assert {f.name for f in fields(OmniStageModelConfig)} == {
+        "active_stream_window",
+        "enable_sleep_mode",
+        "default_sampling_params",
+        "subtalker_sampling_params",
+        "has_sampling_extra_args",
+        "custom_voice_dir",
+        "task_type",
+        "codec_frame_rate_hz",
+        "enforce_eager",
+        "enable_flashinfer_autotune",
+        "compilation_config",
+        "enable_multithread_weight_load",
+        "num_weight_load_threads",
+        "disable_autocast",
+    }
+    assert {f.name for f in fields(OmniStageLoadConfig)} == {
+        "load_format",
+        "tokenizer_mode",
+        "config_format",
+        "skip_mm_profiling",
+    }
+    assert {f.name for f in fields(OmniStageCacheConfig)} == {
+        "gpu_memory_utilization",
+        "enable_prefix_caching",
+        "disable_hybrid_kv_cache_manager",
+        "mm_processor_cache_gb",
+    }
+    assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
+        "max_num_seqs",
+        "max_num_batched_tokens",
+        "max_model_len",
+        "enable_chunked_prefill",
+        "async_scheduling",
+    }
+    assert {f.name for f in fields(OmniStageConnectorConfig)} == {
+        "stage_connector",
+        "output_connectors",
+        "input_connectors",
+    }
+    assert {f.name for f in fields(OmniStageParallelConfig)} == {
+        "pipeline_parallel_size",
+        "data_parallel_size",
+        "tensor_parallel_size",
+        "enable_expert_parallel",
+        "world_size",
+    }
+    assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
+        "pipeline_parallel_size",
+        "data_parallel_size",
+        "tensor_parallel_size",
+        "sequence_parallel_size",
+        "ulysses_degree",
+        "ring_degree",
+        "ulysses_mode",
+        "cfg_parallel_size",
+        "vae_patch_parallel_size",
+        "use_hsdp",
+        "mask_sp_padding",
+        "hsdp_shard_size",
+        "hsdp_replicate_size",
+        "enable_expert_parallel",
+        "world_size",
+    }
+
+
+def test_diffusion_parallel_config_fields_cover_legacy_surface():
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+    legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
+    structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
+    expected_upstream_fields = {"mask_sp_padding"}
+
+    assert legacy_fields | expected_upstream_fields <= structured_fields
+    assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
+
+
+def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
+    cfg = OmniStageDiffusionParallelConfig(
+        pipeline_parallel_size=2,
+        data_parallel_size=3,
+        tensor_parallel_size=4,
+        cfg_parallel_size=3,
+        mask_sp_padding=True,
+    )
+
+    assert cfg.pipeline_parallel_size == 2
+    assert cfg.data_parallel_size == 3
+    assert cfg.cfg_parallel_size == 3
+    assert cfg.mask_sp_padding is True
+    assert cfg.world_size == 72
+
+
+def test_parallel_config_derived_fields_are_not_init_inputs():
+    with pytest.raises(ValidationError):
+        OmniStageParallelConfig(world_size=4)
+
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(world_size=4)
+
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(sequence_parallel_size=2)
+
+
+def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae_patch_parallel():
+    cfg = OmniStageDiffusionParallelConfig(
+        tensor_parallel_size=2,
+        cfg_parallel_size=2,
+        vae_patch_parallel_size=4,
+    )
+
+    assert cfg.vae_patch_parallel_size == 4
+    assert cfg.world_size == 4
+
+
+def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
+    cfg = OmniStageDiffusionParallelConfig(
+        pipeline_parallel_size=2,
+        ulysses_degree=2,
+        use_hsdp=True,
+        hsdp_shard_size=-1,
+        hsdp_replicate_size=2,
+    )
+
+    assert cfg.hsdp_shard_size == 2
+    assert cfg.world_size == 4
+
+
+def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
+    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
+
+    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
+
+
+def test_from_registry_preserves_legacy_pp_dp_for_world_size():
+    cfg = VllmOmniConfig.from_registry("hunyuan_image3_dit").stage_by_id(0).parallel_config
+
+    assert cfg.pipeline_parallel_size == 1
+    assert cfg.data_parallel_size == 1
+    assert cfg.tensor_parallel_size == 4
+    assert cfg.world_size == 4
+
+
+def test_from_registry_derives_sequence_parallel_size_from_degrees(tmp_path):
+    deploy_path = tmp_path / "dreamzero_derived_parallel.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    parallel_config:",
+                "      sequence_parallel_size: 99",
+                "      ulysses_degree: 2",
+                "      ring_degree: 3",
+            ]
+        )
+    )
+
+    stage = VllmOmniConfig.from_registry("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.sequence_parallel_size == 6
+    assert stage.parallel_config.world_size == 6
+
+
+def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound():
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
+
+
+def test_stage_realizations_use_stage_specific_parallel_config_types():
+    qwen_config = VllmOmniConfig.from_registry("qwen3_tts")
+    ar_stage = qwen_config.stage_by_id(0)
+    generation_stage = qwen_config.stage_by_id(1)
+    diffusion_stage = VllmOmniConfig.from_registry("hunyuan_image3_dit").stage_by_id(0)
+
+    assert isinstance(ar_stage, VllmOmniARStageConfig)
+    assert type(ar_stage.parallel_config) is OmniStageParallelConfig
+    assert not hasattr(ar_stage.parallel_config, "cfg_parallel_size")
+    assert not hasattr(ar_stage.parallel_config, "sequence_parallel_size")
+    assert not hasattr(ar_stage.parallel_config, "ulysses_degree")
+
+    assert isinstance(generation_stage, VllmOmniGenerationStageConfig)
+    assert type(generation_stage.parallel_config) is OmniStageParallelConfig
+    assert not hasattr(generation_stage.parallel_config, "cfg_parallel_size")
+    assert not hasattr(generation_stage.parallel_config, "sequence_parallel_size")
+    assert not hasattr(generation_stage.parallel_config, "ulysses_degree")
+
+    assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig)
+    assert isinstance(diffusion_stage.parallel_config, OmniStageDiffusionParallelConfig)
+    assert diffusion_stage.parallel_config.cfg_parallel_size == 1
+    assert diffusion_stage.parallel_config.sequence_parallel_size == 1
+    assert diffusion_stage.parallel_config.ulysses_degree == 1
+
+
+def test_from_registry_preserves_diffusion_parallel_mask_sp_padding(tmp_path):
+    deploy_path = tmp_path / "dreamzero_mask_sp_padding.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    parallel_config:",
+                "      mask_sp_padding: true",
+            ]
+        )
+    )
+
+    stage = VllmOmniConfig.from_registry("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.mask_sp_padding is True
+
+
+def test_from_registry_matches_stage_config_to_omegaconf_behavior_for_representative_stage():
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy("qwen3_tts"))[0]
+    omega_stage = legacy_stage.to_omegaconf()
+    omni_stage = VllmOmniConfig.from_registry("qwen3_tts").stage_by_id(legacy_stage.stage_id)
+
+    assert omega_stage.stage_id == omni_stage.stage_id
+    assert omega_stage.stage_type == omni_stage.stage_type.value
+    assert omega_stage.engine_input_source == omni_stage.input_sources
+    assert omega_stage.final_output == omni_stage.final_output
+    assert omega_stage.final_output_type == omni_stage.final_output_type
+    assert omega_stage.is_comprehension == omni_stage.is_comprehension
+    assert omega_stage.engine_args.model_stage == omni_stage.model_stage
+    assert omega_stage.engine_args.worker_type == omni_stage.worker_type
+    assert omega_stage.engine_args.scheduler_cls == omni_stage.scheduler_cls
+    assert omega_stage.runtime.process is True
+    assert omega_stage.runtime.requires_multimodal_data == omni_stage.requires_multimodal_data
+
+
+def test_from_registry_matches_to_omegaconf_diffusion_parallel_config():
+    pipeline = _resolve_pipeline_or_skip("hunyuan_image3_dit")
+    legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy("hunyuan_image3_dit"))[0]
+    omega_stage = legacy_stage.to_omegaconf()
+    omni_stage = VllmOmniConfig.from_registry("hunyuan_image3_dit").stage_by_id(legacy_stage.stage_id)
+
+    assert (
+        omega_stage.engine_args.parallel_config.pipeline_parallel_size
+        == omni_stage.parallel_config.pipeline_parallel_size
+    )
+    assert omega_stage.engine_args.parallel_config.data_parallel_size == omni_stage.parallel_config.data_parallel_size
+    assert (
+        omega_stage.engine_args.parallel_config.tensor_parallel_size == omni_stage.parallel_config.tensor_parallel_size
+    )
+    assert (
+        omega_stage.engine_args.parallel_config.sequence_parallel_size
+        == omni_stage.parallel_config.sequence_parallel_size
+    )
+    assert omega_stage.engine_args.parallel_config.cfg_parallel_size == omni_stage.parallel_config.cfg_parallel_size
+    assert (
+        omega_stage.engine_args.parallel_config.vae_patch_parallel_size
+        == omni_stage.parallel_config.vae_patch_parallel_size
+    )
+
+
+def test_from_registry_matches_build_engine_args_dict_behavior_for_representative_stage(monkeypatch):
+    from vllm_omni.engine import stage_init_utils
+
+    monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy("qwen3_tts"))[0]
+    omega_stage = legacy_stage.to_omegaconf()
+    legacy_engine_args = build_engine_args_dict(
+        omega_stage,
+        model="/tmp/qwen3-tts",
+        stage_connector_spec={"name": "SharedMemoryConnector", "extra": {}},
+    )
+    omni_stage = VllmOmniConfig.from_registry("qwen3_tts").stage_by_id(legacy_stage.stage_id)
+
+    assert legacy_engine_args["model"] == "/tmp/qwen3-tts"
+    assert legacy_engine_args["stage_id"] == omni_stage.stage_id
+    assert legacy_engine_args["model_stage"] == omni_stage.model_stage
+    assert legacy_engine_args["worker_type"] == omni_stage.worker_type
+    assert legacy_engine_args["scheduler_cls"] == omni_stage.scheduler_cls
+    assert legacy_engine_args["stage_connector_spec"] == {"name": "SharedMemoryConnector", "extra": {}}
+    assert legacy_engine_args["has_sampling_extra_args"] == bool(
+        (omni_stage.model_config.default_sampling_params or {}).get("extra_args")
+    )
+    assert omni_stage.model_config.has_sampling_extra_args == legacy_engine_args["has_sampling_extra_args"]
+
+
+def test_from_registry_derives_has_sampling_extra_args_from_stage_defaults():
+    stage = VllmOmniConfig.from_registry("voxtral_tts").stage_by_id(0)
+
+    assert (stage.model_config.default_sampling_params or {}).get("extra_args")
+    assert stage.model_config.has_sampling_extra_args is True
+
+
+def test_diffusion_config_preserves_existing_coercion_hooks():
+    import torch
+
+    from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
+
+    cfg = omni_config_module._DiffusionConfigProjection(
+        dtype="float32",
+        cache_config={"rel_l1_thresh": 0.3},
+        diffusion_attention_config={"default": "flash_attn"},
+        diffusion_kv_cache_skip_steps="0-2,4",
+        diffusion_kv_cache_skip_layers=[1, 3],
+    )
+
+    assert cfg.dtype is torch.float32
+    assert isinstance(cfg.cache_config, DiffusionCacheConfig)
+    assert isinstance(cfg.diffusion_attention_config, AttentionConfig)
+    assert cfg.diffusion_attention_config.default.backend == "flash_attn"
+    assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1, 2, 4}
+    assert cfg.diffusion_kv_cache_skip_layer_indices == {1, 3}
+    assert cfg.max_cpu_loras == 1
+
+
+def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
+    monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
+
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+        diffusion_attention_backend="flash_attn",
+        kv_cache_dtype="fp8",
+        kv_cache_skip_steps="0-1",
+        kv_cache_skip_layers=[2],
+        static_lora_scale=0.25,
+        diffusers_load_kwargs=None,
+        diffusers_call_kwargs=None,
+    )
+
+    assert cfg.diffusion_attention_config.default.backend == "flash_attn"
+    assert cfg.diffusion_kv_cache_dtype == "fp8"
+    assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1}
+    assert cfg.diffusion_kv_cache_skip_layer_indices == {2}
+    assert cfg.lora_scale == 0.25
+    assert cfg.cache_backend == "tea_cache"
+    assert cfg.diffusers_load_kwargs == {}
+    assert cfg.diffusers_call_kwargs == {}
+
+
+def test_from_registry_normalizes_diffusion_config_aliases_from_engine_args(tmp_path):
+    deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    diffusion_attention_backend: flash_attn",
+            ]
+        )
+    )
+
+    stage = VllmOmniConfig.from_registry(
+        "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"
+
+
+def test_diffusion_config_field_classification_covers_current_fields():
+    classified_fields = (
+        omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
+        | omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
+        | omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
+    )
+
+    assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)}
+    assert {
+        "enable_prompt_embed_cache",
+        "prompt_embed_cache_size",
+        "diffusion_kv_cache_dtype",
+    } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
+    assert {
+        "revision",
+        "trust_remote_code",
+        "distributed_executor_backend",
+    } <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
+    assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
@@ -2,9 +2,29 @@
 Configuration module for vLLM-Omni.
 """

+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.config.model import OmniModelConfig
+from vllm_omni.config.omni_config import (
+    BaseVllmOmniStageConfig,
+    OmniStageCacheConfig,
+    OmniStageConnectorConfig,
+    OmniStageDiffusionParallelConfig,
+    OmniStageLoadConfig,
+    OmniStageModelConfig,
+    OmniStageParallelConfig,
+    OmniStageRuntimeConfig,
+    OmniStageSchedulerConfig,
+    StageConfigType,
+    VllmOmniARStageConfig,
+    VllmOmniConfig,
+    VllmOmniDiffusionStageConfig,
+    VllmOmniGenerationStageConfig,
+    VllmOmniOrchestratorConfig,
+)
+from vllm_omni.config.pipeline_registry import register_pipeline
 from vllm_omni.config.stage_config import (
+    PIPELINE_WIDE_ENGINE_FIELDS,
    DeployConfig,
    PipelineConfig,
    StageConfig,
@@ -23,17 +43,40 @@ from vllm_omni.config.yaml_util import (
 )

 __all__ = [
-    "OmniModelConfig",
+    # Legacy model-level configs.
    "LoRAConfig",
+    "OmniModelConfig",
+    # Structured Omni config entry points.
+    "VllmOmniConfig",
+    "BaseVllmOmniStageConfig",
+    "VllmOmniARStageConfig",
+    "VllmOmniGenerationStageConfig",
+    "VllmOmniDiffusionStageConfig",
+    "StageConfigType",
+    # Structured Omni sub-configs.
+    "OmniStageCacheConfig",
+    "OmniStageConnectorConfig",
+    "OmniStageDiffusionParallelConfig",
+    "OmniStageLoadConfig",
+    "OmniStageModelConfig",
+    "VllmOmniOrchestratorConfig",
+    "OmniStageParallelConfig",
+    "OmniStageRuntimeConfig",
+    "OmniStageSchedulerConfig",
+    # Legacy pipeline/stage deploy config surface.
+    "PIPELINE_WIDE_ENGINE_FIELDS",
+    "DeployConfig",
+    "PipelineConfig",
    "StageConfig",
+    "StageConfigFactory",
+    "StageDeployConfig",
    "StageType",
    "StageExecutionType",
    "StagePipelineConfig",
-    "PipelineConfig",
-    "StageDeployConfig",
-    "DeployConfig",
    "load_deploy_config",
    "merge_pipeline_deploy",
+    "register_pipeline",
+    # YAML utility helpers.
    "create_config",
    "load_yaml_config",
    "merge_configs",
@@ -690,6 +690,7 @@ _PIPELINE_WIDE_ENGINE_FIELDS: tuple[str, ...] = (
    "active_stream_window",
    "custom_voice_dir",
 )
+PIPELINE_WIDE_ENGINE_FIELDS = _PIPELINE_WIDE_ENGINE_FIELDS


 def _build_engine_args(
@@ -33,6 +33,62 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)


+def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
+    """Normalize legacy diffusion kwargs before config construction."""
+    normalized = dict(kwargs)
+
+    # 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 normalized:
+        if "lora_scale" not in normalized:
+            normalized["lora_scale"] = normalized["static_lora_scale"]
+        normalized.pop("static_lora_scale", None)
+
+    # Backwards-compatibility: map "quantization" to "quantization_config"
+    # so callers using the old field name still work.
+    if "quantization" in normalized and normalized.get("quantization_config", None) is None:
+        normalized["quantization_config"] = normalized.pop("quantization")
+    else:
+        normalized.pop("quantization", None)
+
+    # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
+    if normalized.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in normalized:
+        normalized["diffusion_kv_cache_dtype"] = normalized.pop("kv_cache_dtype")
+    else:
+        normalized.pop("kv_cache_dtype", None)
+    if normalized.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in normalized:
+        normalized["diffusion_kv_cache_skip_steps"] = normalized.pop("kv_cache_skip_steps")
+    else:
+        normalized.pop("kv_cache_skip_steps", None)
+    if normalized.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in normalized:
+        normalized["diffusion_kv_cache_skip_layers"] = normalized.pop("kv_cache_skip_layers")
+    else:
+        normalized.pop("kv_cache_skip_layers", None)
+
+    # Handle "diffusion_attention_backend" shorthand: merge into
+    # diffusion_attention_config before field filtering.
+    diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
+    if diffusion_attn_backend is not None:
+        existing = normalized.get("diffusion_attention_config")
+        normalized["diffusion_attention_config"] = parse_attention_config(
+            existing,
+            attention_backend=diffusion_attn_backend,
+        )
+
+    # Check environment variable as fallback for cache_backend.
+    # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
+    if "cache_backend" not in normalized:
+        cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
+        normalized["cache_backend"] = cache_backend.lower() if cache_backend else "none"
+
+    # Convert optional YAML null values to empty containers.
+    for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
+        if key in normalized and normalized[key] is None:
+            normalized[key] = {}
+
+    return normalized
+
+
 def parse_kv_cache_skip_selector(
    selector: str | list[int] | tuple[int, ...] | set[int] | None,
 ) -> set[int] | None:
@@ -1083,55 +1139,7 @@ class OmniDiffusionConfig:

    @classmethod
    def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
-        # Backwards-compatibility: older callers may use a diffusion-specific
-        # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale"
-        # before constructing the dataclass to avoid TypeError on unknown fields.
-        if "static_lora_scale" in kwargs:
-            if "lora_scale" not in kwargs:
-                kwargs["lora_scale"] = kwargs["static_lora_scale"]
-            kwargs.pop("static_lora_scale", None)
-
-        # Backwards-compatibility: map "quantization" to "quantization_config"
-        # so callers using the old field name still work.
-        if "quantization" in kwargs and kwargs.get("quantization_config", None) is None:
-            kwargs["quantization_config"] = kwargs.pop("quantization")
-        else:
-            kwargs.pop("quantization", None)
-
-        # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
-        if kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in kwargs:
-            kwargs["diffusion_kv_cache_dtype"] = kwargs.pop("kv_cache_dtype")
-        else:
-            kwargs.pop("kv_cache_dtype", None)
-        if kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in kwargs:
-            kwargs["diffusion_kv_cache_skip_steps"] = kwargs.pop("kv_cache_skip_steps")
-        else:
-            kwargs.pop("kv_cache_skip_steps", None)
-        if kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in kwargs:
-            kwargs["diffusion_kv_cache_skip_layers"] = kwargs.pop("kv_cache_skip_layers")
-        else:
-            kwargs.pop("kv_cache_skip_layers", None)
-
-        # Handle "diffusion_attention_backend" shorthand: merge into
-        # diffusion_attention_config before field filtering.
-        diffusion_attn_backend = kwargs.pop("diffusion_attention_backend", None)
-        if diffusion_attn_backend is not None:
-            existing = kwargs.get("diffusion_attention_config")
-            kwargs["diffusion_attention_config"] = parse_attention_config(
-                existing,
-                attention_backend=diffusion_attn_backend,
-            )
-
-        # Check environment variable as fallback for cache_backend
-        # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND for backwards compatibility
-        if "cache_backend" not in kwargs:
-            cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
-            kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
-
-        # Convert optional YAML null values to empty containers.
-        for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
-            if key in kwargs and kwargs[key] is None:
-                kwargs[key] = {}
+        kwargs = normalize_omni_diffusion_kwargs(kwargs)

        # Filter kwargs to only include valid fields
        valid_fields = {f.name for f in fields(cls)}