diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py new file mode 100644 index 000000000..c8cc8a3f4 --- /dev/null +++ b/tests/config/test_omni_config.py @@ -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 diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index a43cc41d9..3069025be 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -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", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py new file mode 100644 index 000000000..4ab8858fa --- /dev/null +++ b/vllm_omni/config/omni_config.py @@ -0,0 +1,1303 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structured vLLM-Omni configuration classes. + +This module is additive for Phase 2 of RFC #4021. ``VllmOmniConfig.from_registry`` +builds the structured view directly from the pipeline registry and deploy config +so parity can be proven before later PRs cut consumers over to these classes. +""" + +from __future__ import annotations + +import copy +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from pathlib import Path +from typing import Any, TypeAlias, TypedDict, cast + +from pydantic import ConfigDict, Field +from vllm.config.utils import config +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +from vllm_omni.config.stage_config import ( + _DEPLOY_DIR, + _STAGE_DEPLOY_FIELDS, + PIPELINE_WIDE_ENGINE_FIELDS, + DeployConfig, + PipelineConfig, + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, + StageType, + _apply_platform_overrides, + _resolve_scheduler, + _scheduler_path, + _select_processor_funcs, + build_stage_runtime_overrides, + load_deploy_config, +) + +_EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = { + StageExecutionType.LLM_AR: (StageType.LLM, "ar"), + StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"), + StageExecutionType.DIFFUSION: (StageType.DIFFUSION, None), +} + +_PIPELINE_DEPLOY_CLI_FIELDS = PIPELINE_WIDE_ENGINE_FIELDS + +_QuantizationConfigType: TypeAlias = QuantizationConfig | str | Mapping[str, Any] | None + + +class _QuantizationEngineOverrides(TypedDict, total=False): + quantization_config: _QuantizationConfigType + quantization: str + + +class _ModelEngineOverrides(TypedDict, total=False): + active_stream_window: int + enable_sleep_mode: bool + subtalker_sampling_params: dict[str, Any] + has_sampling_extra_args: bool + custom_voice_dir: str + task_type: str + codec_frame_rate_hz: float + enforce_eager: bool + enable_flashinfer_autotune: bool + compilation_config: dict[str, Any] + enable_multithread_weight_load: bool + num_weight_load_threads: int + disable_autocast: bool + + +class _LoadEngineOverrides(TypedDict, total=False): + load_format: str + tokenizer_mode: str + config_format: str + skip_mm_profiling: bool + + +class _CacheEngineOverrides(TypedDict, total=False): + gpu_memory_utilization: float + enable_prefix_caching: bool + disable_hybrid_kv_cache_manager: bool + mm_processor_cache_gb: float + + +class _SchedulerEngineOverrides(TypedDict, total=False): + max_num_seqs: int + max_num_batched_tokens: int + max_model_len: int + enable_chunked_prefill: bool + async_scheduling: bool + + +class _RuntimeEngineOverrides(TypedDict, total=False): + devices: str + num_replicas: int + env: dict[str, Any] + num_gpus: int + log_level: str + log_stats: bool + profiler_config: dict[str, Any] + + +class _ParallelConfigEngineOverrides(TypedDict, total=False): + pipeline_parallel_size: int + data_parallel_size: int + tensor_parallel_size: int + sequence_parallel_size: int + ulysses_degree: int + ring_degree: int + ulysses_mode: str + cfg_parallel_size: int + vae_patch_parallel_size: int + use_hsdp: bool + mask_sp_padding: bool + hsdp_shard_size: int + hsdp_replicate_size: int + enable_expert_parallel: bool + + +class _ParallelEngineOverrides(_ParallelConfigEngineOverrides, total=False): + parallel_config: _ParallelConfigEngineOverrides | Mapping[str, Any] + + +@dataclass(frozen=True) +class _StageEngineValues: + """Typed projections of legacy flat per-stage ``yaml_engine_args``.""" + + quantization: _QuantizationEngineOverrides + model: _ModelEngineOverrides + load: _LoadEngineOverrides + cache: _CacheEngineOverrides + scheduler: _SchedulerEngineOverrides + runtime: _RuntimeEngineOverrides + parallel: _ParallelEngineOverrides + diffusion: _DiffusionEngineOverrides + + +@dataclass(frozen=True) +class _DiffusionEngineOverrides: + """Validated diffusion projection of legacy flat per-stage engine args.""" + + _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()} + + +_IMMUTABLE_CONFIG_VALUE_TYPES = (str, int, float, bool, bytes, type(None)) + + +def _can_share_config_value(value: Any) -> bool: + if isinstance(value, _IMMUTABLE_CONFIG_VALUE_TYPES): + return True + if isinstance(value, tuple): + return all(_can_share_config_value(item) for item in value) + return False + + +def _copy_value(value: Any) -> Any: + """Copy nested config values so the structured view owns its data.""" + if _can_share_config_value(value): + return value + return copy.deepcopy(value) + + +def _config_kwargs(overrides: Mapping[str, Any]) -> dict[str, Any]: + return {name: _copy_value(value) for name, value in overrides.items() if value is not None} + + +def _first_defined(*values: Any) -> Any: + for value in values: + if value is not None: + return _copy_value(value) + return None + + +def _validate_async_chunk_support(pipeline: PipelineConfig, deploy: DeployConfig) -> None: + if deploy.async_chunk and not any( + stage.async_chunk_process_next_stage_input_func or stage.custom_process_next_stage_input_func + for stage in pipeline.stages + ): + raise ValueError( + f"Pipeline {pipeline.model_type!r} has async_chunk=True in deploy but no stage " + "declares a next-stage input processor " + "(``async_chunk_process_next_stage_input_func`` or ``custom_process_next_stage_input_func``). " + "Either set async_chunk=False or implement an async-chunk processor on the pipeline." + ) + + +def _resolve_execution_mode(execution_type: StageExecutionType) -> tuple[StageType, str | None]: + try: + return _EXECUTION_TYPE_TO_STAGE_WORKER[execution_type] + except KeyError as exc: + raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc + + +def _resolve_scheduler_path(execution_type: StageExecutionType, async_scheduling: bool = True) -> str | None: + return _scheduler_path(_resolve_scheduler(execution_type, async_scheduling)) + + +def _stage_cli_overrides(stage_id: int, cli_overrides: Mapping[str, Any]) -> dict[str, Any]: + runtime_overrides = build_stage_runtime_overrides(stage_id, dict(cli_overrides)) + global_stage_fields = _global_stage_cli_fields() + result: dict[str, Any] = {} + for key, value in runtime_overrides.items(): + if key in global_stage_fields or f"stage_{stage_id}_{key}" in cli_overrides: + result[key] = _copy_value(value) + return result + + +def _resolve_deploy_path(model_type: str, deploy_config_path: str | None = None) -> Path: + if deploy_config_path is None: + return _DEPLOY_DIR / f"{model_type}.yaml" + + deploy_path = Path(deploy_config_path) + if not deploy_path.exists() and deploy_path.parent == Path("."): + bare_name = deploy_path.name + if not bare_name.endswith(".yaml"): + bare_name = f"{bare_name}.yaml" + candidate = _DEPLOY_DIR / bare_name + if candidate.exists(): + return candidate + return deploy_path + + +def _registered_pipeline_keys() -> list[str]: + from vllm_omni.config.pipeline_registry import OMNI_PIPELINES + + return sorted(OMNI_PIPELINES) + + +def _resolve_registered_pipeline(model_type: str) -> PipelineConfig | None: + from vllm_omni.config.pipeline_registry import OMNI_PIPELINES + + registered = OMNI_PIPELINES[model_type] + return registered(None) if callable(registered) else registered + + +@config +class OmniStageModelConfig: + """Per-stage model behavior.""" + + active_stream_window: int = Field(default=0, ge=0) + enable_sleep_mode: bool = False + default_sampling_params: dict[str, Any] | None = None + subtalker_sampling_params: dict[str, Any] | None = None + has_sampling_extra_args: bool = False + custom_voice_dir: str | None = None + task_type: str | None = None + codec_frame_rate_hz: float | None = None + enforce_eager: bool = False + enable_flashinfer_autotune: bool | None = None + compilation_config: dict[str, Any] | None = None + enable_multithread_weight_load: bool = True + num_weight_load_threads: int = Field(default=4, ge=1) + disable_autocast: bool = False + + +@config +class OmniStageLoadConfig: + """Per-stage loading behavior.""" + + load_format: str = "auto" + tokenizer_mode: str = "auto" + config_format: str | None = None + skip_mm_profiling: bool | None = None + + +@config +class OmniStageCacheConfig: + """Per-stage engine cache and memory behavior. + + This is separate from ``_DiffusionConfigProjection.cache_config``, which configures + vLLM-Omni diffusion-specific cache backends such as TeaCache and Cache-DiT. + """ + + gpu_memory_utilization: float = Field(default=0.90, gt=0.0, le=1.0) + enable_prefix_caching: bool = False + disable_hybrid_kv_cache_manager: bool = False + mm_processor_cache_gb: float | None = Field(default=None, ge=0.0) + + +@config +class OmniStageSchedulerConfig: + """Per-stage request scheduling behavior.""" + + max_num_seqs: int = Field(default=128, ge=1) + max_num_batched_tokens: int | None = Field(default=None, ge=1) + max_model_len: int | None = Field(default=None, ge=-1) + enable_chunked_prefill: bool = False + async_scheduling: bool = True + + def __post_init__(self) -> None: + if self.max_num_batched_tokens is not None and self.max_num_batched_tokens < self.max_num_seqs: + raise ValueError( + f"max_num_batched_tokens ({self.max_num_batched_tokens}) must be >= max_num_seqs ({self.max_num_seqs})" + ) + + +@config +class OmniStageConnectorConfig: + """Per-stage inter-stage connector wiring.""" + + stage_connector: dict[str, Any] = field( + default_factory=lambda: { + "name": "SharedMemoryConnector", + "extra": {}, + } + ) + output_connectors: dict[str, Any] | None = None + input_connectors: dict[str, Any] | None = None + + +@config +class OmniStageRuntimeConfig: + """Per-stage process placement and runtime behavior.""" + + devices: str | None = None + num_replicas: int = Field(default=1, ge=1) + env: dict[str, Any] | None = None + num_gpus: int = Field(default=1, ge=1) + log_level: str = "info" + log_stats: bool = False + profiler_config: dict[str, Any] | None = None + + +@config +class OmniStageParallelConfig: + """Common per-stage distributed parallelism behavior.""" + + pipeline_parallel_size: int = Field(default=1, ge=1) + data_parallel_size: int = Field(default=1, ge=1) + tensor_parallel_size: int = Field(default=1, ge=1) + enable_expert_parallel: bool = False + world_size: int = Field(default=1, ge=1, init=False) + + def __post_init__(self) -> None: + self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size + + +@config +class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): + """Diffusion-stage distributed parallelism behavior.""" + + sequence_parallel_size: int = Field(default=1, ge=1, init=False) + ulysses_degree: int = Field(default=1, ge=1) + ring_degree: int = Field(default=1, ge=1) + ulysses_mode: str = "strict" + cfg_parallel_size: int = Field(default=1, ge=1, le=3) + vae_patch_parallel_size: int = Field(default=1, ge=1) + use_hsdp: bool = False + mask_sp_padding: bool = False + hsdp_shard_size: int = -1 + hsdp_replicate_size: int = Field(default=1, ge=1) + + def __post_init__(self) -> None: + self.sequence_parallel_size = self.ulysses_degree * self.ring_degree + if self.ulysses_mode not in {"strict", "advanced_uaa"}: + raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") + + other_parallel_world_size = ( + self.pipeline_parallel_size + * self.data_parallel_size + * self.tensor_parallel_size + * self.sequence_parallel_size + * self.cfg_parallel_size + ) + if self.use_hsdp: + if self.tensor_parallel_size > 1 or self.data_parallel_size > 1: + raise ValueError( + "HSDP (use_hsdp=True) cannot be used with TP or DP " + f"(tensor_parallel_size={self.tensor_parallel_size}, " + f"data_parallel_size={self.data_parallel_size})" + ) + if self.hsdp_shard_size == -1: + if other_parallel_world_size == 1: + raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1") + if other_parallel_world_size % self.hsdp_replicate_size != 0: + raise ValueError( + f"hsdp_replicate_size ({self.hsdp_replicate_size}) must evenly divide " + f"world_size ({other_parallel_world_size}) when hsdp_shard_size is -1" + ) + self.hsdp_shard_size = other_parallel_world_size // self.hsdp_replicate_size + self.world_size = other_parallel_world_size + else: + if self.hsdp_shard_size <= 0: + raise ValueError("hsdp_shard_size must be > 0 when use_hsdp=True") + hsdp_world_size = self.hsdp_replicate_size * self.hsdp_shard_size + if other_parallel_world_size == 1: + self.world_size = hsdp_world_size + else: + if hsdp_world_size != other_parallel_world_size: + raise ValueError( + f"HSDP dimensions ({self.hsdp_replicate_size} x {self.hsdp_shard_size} = " + f"{hsdp_world_size}) must equal world_size from other parallelism " + f"({other_parallel_world_size})" + ) + self.world_size = other_parallel_world_size + else: + self.world_size = other_parallel_world_size + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class _DiffusionConfigProjection: + """Diffusion-specific per-stage settings. + + Shared AR/diffusion fields are projected into the other sub-configs. This + class keeps the diffusion-only knobs from ``OmniDiffusionConfig`` without + running its startup-time side effects such as port probing or HF metadata + loading. + """ + + stage_id: int = 0 + model: str | None = None + model_class_name: str | None = None + model_arch: str | None = None + dtype: Any = "auto" + trust_remote_code: bool = False + revision: str | None = None + distributed_executor_backend: str = "mp" + dist_timeout: int | None = None + nccl_port: int | None = None + master_port: int | None = None + host: str | None = None + port: int | None = None + model_config: dict[str, Any] = field(default_factory=dict) + tf_model_config: Any = None + diffusion_attention_config: Any = None + cache_strategy: str = "none" + cache_backend: str = "none" + cache_config: Any = field(default_factory=dict) + enable_cache_dit_summary: bool = False + enable_prompt_embed_cache: bool = False + prompt_embed_cache_size: int = Field(default=32, ge=1) + diffusion_load_format: str = "default" + diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict) + diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict) + diffusers_pipeline_cls: Any = None + lora_path: str | None = None + lora_scale: float = 1.0 + max_cpu_loras: int | None = None + output_type: str = "pil" + enable_cpu_offload: bool = False + enable_layerwise_offload: bool = False + pin_cpu_memory: bool = True + vae_use_slicing: bool = False + vae_use_tiling: bool = False + mask_strategy_file_path: str | None = None + skip_time_steps: int = 15 + VSA_sparsity: float = 0.0 + moba_config_path: str | None = None + boundary_ratio: float | None = None + flow_shift: float | None = None + diffusion_kv_cache_dtype: str | None = None + diffusion_kv_cache_skip_steps: str | list[int] | tuple[int, ...] | set[int] | None = None + diffusion_kv_cache_skip_layers: str | list[int] | tuple[int, ...] | set[int] | None = None + diffusion_kv_cache_skip_step_indices: set[int] | None = None + diffusion_kv_cache_skip_layer_indices: set[int] | None = None + force_cutlass_fp8: bool = False + enable_diffusion_pipeline_profiler: bool = False + step_execution: bool = False + supports_multimodal_inputs: bool = False + max_multimodal_image_inputs: int | None = None + model_paths: dict[str, str] = field(default_factory=dict) + model_loaded: dict[str, bool] = field( + default_factory=lambda: { + "transformer": True, + "vae": True, + } + ) + override_transformer_cls_name: str | None = None + worker_extension_cls: str | None = None + custom_pipeline_args: dict[str, Any] | None = None + additional_config: dict[str, Any] = field(default_factory=dict) + enable_stage_verification: bool = True + prompt_file_path: str | None = None + quantization_config: _QuantizationConfigType = None + extras: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_kwargs(cls, **kwargs: Any) -> _DiffusionConfigProjection: + from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs + + normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs) + valid_fields = {f.name for f in fields(cls)} + return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields}) + + def __post_init__(self) -> None: + # Keep diffusion imports lazy so importing vllm_omni.config does not + # pull in the full diffusion stack unless a diffusion stage is built. + from vllm_omni.diffusion.data import ( + AttentionConfig, + DiffusionCacheConfig, + TransformerConfig, + build_attention_config, + parse_kv_cache_skip_selector, + ) + from vllm_omni.quantization import build_quant_config + + if self.tf_model_config is None: + self.tf_model_config = TransformerConfig() + elif isinstance(self.tf_model_config, Mapping): + self.tf_model_config = TransformerConfig.from_dict(dict(self.tf_model_config)) + + if self.additional_config is None: + self.additional_config = {} + elif isinstance(self.additional_config, Mapping): + self.additional_config = dict(self.additional_config) + else: + raise TypeError(f"additional_config must be a mapping or None, got {type(self.additional_config)!r}") + + if isinstance(self.dtype, str): + # Import torch only when string dtype normalization is needed. + import torch + + dtype_map = { + "auto": torch.bfloat16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float16": torch.float16, + "fp16": torch.float16, + "half": torch.float16, + "float32": torch.float32, + "fp32": torch.float32, + "float": torch.float32, + } + self.dtype = dtype_map.get(self.dtype.lower(), torch.bfloat16) + + if isinstance(self.cache_config, Mapping): + self.cache_config = DiffusionCacheConfig.from_dict(dict(self.cache_config)) + elif not isinstance(self.cache_config, DiffusionCacheConfig): + self.cache_config = DiffusionCacheConfig() + + self._propagate_quantization_from_tf_config(self.tf_model_config) + if self.quantization_config is not None: + if isinstance(self.quantization_config, QuantizationConfig): + pass + elif isinstance(self.quantization_config, str): + self.quantization_config = build_quant_config(self.quantization_config) + elif isinstance(self.quantization_config, Mapping): + self.quantization_config = build_quant_config(dict(self.quantization_config)) + else: + raise TypeError( + "quantization_config must be str, dict, QuantizationConfig, or None, " + f"got {type(self.quantization_config)!r}" + ) + + if self.diffusion_attention_config is None or isinstance( + self.diffusion_attention_config, + (AttentionConfig, Mapping), + ): + self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config) + else: + raise TypeError( + "diffusion_attention_config must be an AttentionConfig, mapping, or None, " + f"got {type(self.diffusion_attention_config)!r}" + ) + + self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps) + self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers) + + if self.max_cpu_loras is None: + self.max_cpu_loras = 1 + elif self.max_cpu_loras < 1: + raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): + raise ValueError( + "diffusers_load_kwargs and diffusers_call_kwargs are only " + "valid together with diffusion_load_format=diffusers" + ) + + def _propagate_quantization_from_tf_config(self, tf_config: Any) -> None: + quant_config = getattr(tf_config, "quant_config", None) + if quant_config is None: + return + quant_method = getattr(tf_config, "quant_method", None) + is_checkpoint_fp8 = bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False)) + is_checkpoint_nvfp4 = bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False)) + should_use_checkpoint_config = ( + self.quantization_config is None + or (is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config)) + or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config)) + ) + if should_use_checkpoint_config: + self.quantization_config = quant_config + if quant_method is not None: + self.additional_config.setdefault("auto_detected_quant_method", quant_method) + + @staticmethod + def _is_generic_fp8_quant_config(quant_config: object) -> bool: + if isinstance(quant_config, str): + return quant_config.lower() == "fp8" + if isinstance(quant_config, Mapping): + method = quant_config.get("method", quant_config.get("quant_method")) + return isinstance(method, str) and method.lower() == "fp8" + if hasattr(quant_config, "get_name"): + return quant_config.get_name() == "fp8" + return False + + @staticmethod + def _is_generic_nvfp4_quant_config(quant_config: object) -> bool: + if isinstance(quant_config, str): + return quant_config.lower() in {"fp4", "nvfp4", "modelopt_fp4"} + if isinstance(quant_config, Mapping): + method = quant_config.get("method", quant_config.get("quant_method")) + return isinstance(method, str) and method.lower() in {"fp4", "nvfp4", "modelopt_fp4"} + if hasattr(quant_config, "get_name"): + return quant_config.get_name() == "modelopt_fp4" + return False + + def set_tf_model_config(self, tf_config: Any) -> None: + self.tf_model_config = tf_config + self._propagate_quantization_from_tf_config(tf_config) + + def enrich_config(self) -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig + + omni_diffusion_fields = frozenset(f.name for f in fields(OmniDiffusionConfig)) + kwargs = { + name: _copy_value(getattr(self, name)) for name in _DIFFUSION_CONFIG_FIELDS if name in omni_diffusion_fields + } + omni_diffusion_config = OmniDiffusionConfig(**kwargs) + omni_diffusion_config.enrich_config() + for name in _DIFFUSION_CONFIG_FIELDS: + if hasattr(omni_diffusion_config, name): + setattr(self, name, _copy_value(getattr(omni_diffusion_config, name))) + + +_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(_DiffusionConfigProjection)) + +# Current OmniDiffusionConfig still contains a flat mix of shared engine, +# runtime, parallel, and diffusion-specific knobs. Keep this classification +# explicit while Phase 2 is additive; later cutover PRs can move or remove +# fields without rediscovering the current boundary. +_DIFFUSION_SHARED_CONFIG_FIELDS = frozenset( + { + "stage_id", + "model", + "model_arch", + "dtype", + "trust_remote_code", + "revision", + "distributed_executor_backend", + "dist_timeout", + "model_config", + "quantization_config", + } +) +_DIFFUSION_RUNTIME_CONFIG_FIELDS = frozenset( + { + "host", + "port", + "nccl_port", + "master_port", + "worker_extension_cls", + "enable_stage_verification", + "prompt_file_path", + } +) +_DIFFUSION_ONLY_CONFIG_FIELDS = ( + _DIFFUSION_CONFIG_FIELDS - _DIFFUSION_SHARED_CONFIG_FIELDS - _DIFFUSION_RUNTIME_CONFIG_FIELDS +) +_DIFFUSION_MOVED_SHARED_FIELDS = frozenset( + { + "parallel_config", + "num_gpus", + "log_level", + "profiler_config", + "omni_kv_config", + "cfg_kv_collect_func", + "max_num_seqs", + "enable_sleep_mode", + "enforce_eager", + "enable_multithread_weight_load", + "num_weight_load_threads", + "disable_autocast", + } +) + + +_STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS) + +_DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset( + { + "diffusion_attention_backend", + "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", +} + +_QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__) +_MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__) +_LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__) +_CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__) +_SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__) +_RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) +_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__) +_PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"} + + +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. + from vllm_omni.engine.arg_utils import OmniEngineArgs + + return ( + frozenset(f.name for f in fields(OmniEngineArgs)) + | frozenset(_STAGE_DEPLOY_ENGINE_FIELDS) + | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS) + ) - {"model", "stage_id", "stage_configs_path", "async_chunk"} + + +def _mapping_or_empty(value: Any) -> Mapping[str, Any]: + return value if isinstance(value, Mapping) else {} + + +def _select_engine_overrides(engine: Mapping[str, Any], keys: set[str] | frozenset[str]) -> dict[str, Any]: + return {name: _copy_value(engine[name]) for name in keys if name in engine and engine[name] is not None} + + +def _stage_engine_overrides(stage_deploy: StageDeployConfig | None) -> dict[str, Any]: + if stage_deploy is None: + return {} + + overrides: dict[str, Any] = {} + for name in _STAGE_DEPLOY_ENGINE_FIELDS: + value = getattr(stage_deploy, name) + if value is not None: + overrides[name] = _copy_value(value) + overrides.update(_copy_value(stage_deploy.engine_extras)) + return overrides + + +def _stage_engine_values( + stage_deploy: StageDeployConfig | None, + stage_cli_overrides: Mapping[str, Any] | None = None, +) -> _StageEngineValues: + engine = _stage_engine_overrides(stage_deploy) + if stage_cli_overrides: + engine.update(_copy_value(stage_cli_overrides)) + return _StageEngineValues( + quantization=cast( + _QuantizationEngineOverrides, + _select_engine_overrides(engine, _QUANTIZATION_ENGINE_FIELDS), + ), + model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)), + load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, _LOAD_ENGINE_FIELDS)), + cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, _CACHE_ENGINE_FIELDS)), + scheduler=cast( + _SchedulerEngineOverrides, + _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS), + ), + 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), + ) + + +def _stage_sampling_params( + stage_deploy: StageDeployConfig | None, + topology: StagePipelineConfig, +) -> dict[str, Any] | None: + sampling: dict[str, Any] = {} + if stage_deploy is not None and stage_deploy.default_sampling_params: + sampling.update(_copy_value(stage_deploy.default_sampling_params)) + sampling.update(_copy_value(topology.sampling_constraints)) + return sampling or None + + +def _orchestrator_cli_overrides(cli_overrides: Mapping[str, Any]) -> dict[str, Any]: + overrides: dict[str, Any] = {} + for config_field in fields(VllmOmniOrchestratorConfig): + name = config_field.name + if name == "deploy_config_path": + continue + if cli_overrides.get(name) is not None: + overrides[name] = _copy_value(cli_overrides[name]) + return overrides + + +@config +class VllmOmniOrchestratorConfig: + """Configuration consumed by the orchestrator process only.""" + + stage_init_timeout: int = Field(default=300, ge=1) + init_timeout: int = Field(default=600, ge=1) + worker_backend: str = "multi_process" + ray_address: str | None = None + deploy_config_path: str | None = None + omni_master_address: str | None = None + omni_master_port: int | None = None + omni_dp_size_local: int = Field(default=1, ge=1) + omni_lb_policy: str = "random" + omni_heartbeat_timeout: float = Field(default=30.0, gt=0.0) + shm_threshold_bytes: int = Field(default=65536, ge=0) + batch_timeout: int = Field(default=10, ge=0) + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class BaseVllmOmniStageConfig: + """Common structured config contract shared by all Omni stage realizations.""" + + stage_pipeline_config: StagePipelineConfig + model_config: OmniStageModelConfig = field(default_factory=OmniStageModelConfig) + load_config: OmniStageLoadConfig = field(default_factory=OmniStageLoadConfig) + cache_config: OmniStageCacheConfig = field(default_factory=OmniStageCacheConfig) + scheduler_config: OmniStageSchedulerConfig = field(default_factory=OmniStageSchedulerConfig) + connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig) + runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig) + parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) + quantization_config: _QuantizationConfigType = None + + @property + def stage_id(self) -> int: + return self.stage_pipeline_config.stage_id + + @property + def model_stage(self) -> str: + return self.stage_pipeline_config.model_stage + + @property + def input_sources(self) -> list[int]: + return list(self.stage_pipeline_config.input_sources) + + @property + def final_output(self) -> bool: + return self.stage_pipeline_config.final_output + + @property + def final_output_type(self) -> str | None: + return self.stage_pipeline_config.final_output_type + + @property + def hf_config_name(self) -> str | None: + return self.stage_pipeline_config.hf_config_name + + @property + def stage_type(self) -> StageType: + stage_type, _ = _resolve_execution_mode(self.stage_pipeline_config.execution_type) + return stage_type + + @property + def worker_type(self) -> str | None: + _, worker_type = _resolve_execution_mode(self.stage_pipeline_config.execution_type) + return worker_type + + @property + def scheduler_cls(self) -> str | None: + return _resolve_scheduler_path( + self.stage_pipeline_config.execution_type, + self.scheduler_config.async_scheduling, + ) + + @property + def custom_process_input_func(self) -> str | None: + return getattr( + self, + "_resolved_custom_process_input_func", + self.stage_pipeline_config.custom_process_input_func, + ) + + @property + def custom_process_next_stage_input_func(self) -> str | None: + return getattr( + self, + "_resolved_custom_process_next_stage_input_func", + self.stage_pipeline_config.custom_process_next_stage_input_func, + ) + + @property + def is_comprehension(self) -> bool: + return self.stage_pipeline_config.owns_tokenizer + + @property + def engine_output_type(self) -> str | None: + return self.stage_pipeline_config.engine_output_type + + @property + def requires_multimodal_data(self) -> bool: + return self.stage_pipeline_config.requires_multimodal_data + + @property + def prompt_expand_func(self) -> str | None: + return self.stage_pipeline_config.prompt_expand_func + + @property + def cfg_kv_collect_func(self) -> str | None: + return self.stage_pipeline_config.cfg_kv_collect_func + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniARStageConfig(BaseVllmOmniStageConfig): + """Structured config for autoregressive LLM stages.""" + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniGenerationStageConfig(BaseVllmOmniStageConfig): + """Structured config for generation LLM stages.""" + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniDiffusionStageConfig(BaseVllmOmniStageConfig): + """Structured config for diffusion stages.""" + + parallel_config: OmniStageDiffusionParallelConfig = field(default_factory=OmniStageDiffusionParallelConfig) + diffusion_config: _DiffusionConfigProjection = field(default_factory=_DiffusionConfigProjection) + + +StageConfigType: TypeAlias = VllmOmniARStageConfig | VllmOmniGenerationStageConfig | VllmOmniDiffusionStageConfig + + +def _build_common_stage_config_kwargs( + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + parallel_config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig, +) -> tuple[dict[str, Any], str | None, str | None]: + input_proc, next_stage_proc = _select_processor_funcs(topology, bool(deploy.async_chunk)) + quantization_config = _build_quantization_config(deploy, engine.quantization) + parallel_config = _build_parallel_config(deploy, engine.parallel, parallel_config_cls) + + return ( + { + "stage_pipeline_config": topology, + "model_config": _build_model_config(topology, stage_deploy, engine.model), + "load_config": _build_load_config(engine.load), + "cache_config": _build_cache_config(deploy, engine.cache), + "scheduler_config": _build_scheduler_config(deploy, engine.scheduler), + "connector_config": _build_connector_config(stage_deploy), + "runtime_config": _build_runtime_config(stage_deploy, engine.runtime, parallel_config), + "parallel_config": parallel_config, + "quantization_config": _copy_value(quantization_config), + }, + input_proc, + next_stage_proc, + ) + + +def _with_resolved_processors( + stage_config: StageConfigType, + input_proc: str | None, + next_stage_proc: str | None, +) -> StageConfigType: + stage_config._resolved_custom_process_input_func = input_proc + stage_config._resolved_custom_process_next_stage_input_func = next_stage_proc + return stage_config + + +def _build_ar_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> VllmOmniARStageConfig: + common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + deploy, + topology, + stage_deploy, + engine, + ) + return cast( + VllmOmniARStageConfig, + _with_resolved_processors( + VllmOmniARStageConfig(**common_kwargs), + input_proc, + next_stage_proc, + ), + ) + + +def _build_generation_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> VllmOmniGenerationStageConfig: + common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + deploy, + topology, + stage_deploy, + engine, + ) + return cast( + VllmOmniGenerationStageConfig, + _with_resolved_processors( + VllmOmniGenerationStageConfig(**common_kwargs), + input_proc, + next_stage_proc, + ), + ) + + +def _build_diffusion_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> VllmOmniDiffusionStageConfig: + common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + deploy, + topology, + stage_deploy, + engine, + OmniStageDiffusionParallelConfig, + ) + common_kwargs["diffusion_config"] = _build_diffusion_config_projection( + pipeline, + deploy, + topology, + engine.diffusion, + model=model, + quantization_config=common_kwargs["quantization_config"], + ) + return cast( + VllmOmniDiffusionStageConfig, + _with_resolved_processors( + VllmOmniDiffusionStageConfig(**common_kwargs), + input_proc, + next_stage_proc, + ), + ) + + +_STAGE_CONFIG_BUILDERS = { + StageExecutionType.LLM_AR: _build_ar_stage_config, + StageExecutionType.LLM_GENERATION: _build_generation_stage_config, + StageExecutionType.DIFFUSION: _build_diffusion_stage_config, +} + + +def _build_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> StageConfigType: + try: + builder = _STAGE_CONFIG_BUILDERS[topology.execution_type] + except KeyError as exc: + raise ValueError(f"Unsupported stage execution type: {topology.execution_type!r}") from exc + return builder( + pipeline, + deploy, + topology, + stage_deploy, + engine, + model=model, + ) + + +def _build_quantization_config( + deploy: DeployConfig, + engine: _QuantizationEngineOverrides, +) -> _QuantizationConfigType: + return _first_defined( + engine.get("quantization_config"), + engine.get("quantization"), + deploy.quantization, + ) + + +def _build_model_config( + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _ModelEngineOverrides, +) -> OmniStageModelConfig: + default_sampling_params = _stage_sampling_params(stage_deploy, topology) + kwargs = _config_kwargs(engine) + if "has_sampling_extra_args" not in kwargs: + kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args")) + return OmniStageModelConfig( + default_sampling_params=default_sampling_params, + **kwargs, + ) + + +def _build_load_config(engine: _LoadEngineOverrides) -> OmniStageLoadConfig: + return OmniStageLoadConfig(**_config_kwargs(engine)) + + +def _build_cache_config( + deploy: DeployConfig, + engine: _CacheEngineOverrides, +) -> OmniStageCacheConfig: + kwargs = _config_kwargs(engine) + if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None: + kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching) + return OmniStageCacheConfig(**kwargs) + + +def _build_scheduler_config( + deploy: DeployConfig, + engine: _SchedulerEngineOverrides, +) -> OmniStageSchedulerConfig: + kwargs = _config_kwargs(engine) + if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None: + kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill) + return OmniStageSchedulerConfig(**kwargs) + + +def _build_connector_config(stage_deploy: StageDeployConfig | None) -> OmniStageConnectorConfig: + output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None + input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None + return OmniStageConnectorConfig( + output_connectors=_copy_value(output_connectors) if output_connectors else None, + input_connectors=_copy_value(input_connectors) if input_connectors else None, + ) + + +def _build_runtime_config( + stage_deploy: StageDeployConfig | None, + engine: _RuntimeEngineOverrides, + parallel_config: OmniStageParallelConfig, +) -> OmniStageRuntimeConfig: + kwargs = _config_kwargs(engine) + if "devices" not in kwargs and stage_deploy is not None and stage_deploy.devices is not None: + kwargs["devices"] = _copy_value(stage_deploy.devices) + if "num_replicas" not in kwargs and stage_deploy is not None: + kwargs["num_replicas"] = stage_deploy.num_replicas + if "env" not in kwargs and stage_deploy is not None and stage_deploy.env is not None: + kwargs["env"] = _copy_value(stage_deploy.env) + kwargs["num_gpus"] = parallel_config.world_size + return OmniStageRuntimeConfig(**kwargs) + + +def _build_parallel_config( + deploy: DeployConfig, + engine: _ParallelEngineOverrides, + config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig, +) -> OmniStageParallelConfig: + parallel_config = _mapping_or_empty(engine.get("parallel_config")) + config_fields = { + config_field.name + for config_field in fields(config_cls) + if getattr(config_field.default, "init", None) is not False + } + kwargs = { + name: _copy_value(value) for name in _PARALLEL_CONFIG_ENGINE_FIELDS if (value := engine.get(name)) is not None + } + kwargs = {name: value for name, value in kwargs.items() if name in config_fields} + kwargs.update({name: value for name, value in _config_kwargs(parallel_config).items() if name in config_fields}) + if "pipeline_parallel_size" not in kwargs and deploy.pipeline_parallel_size is not None: + kwargs["pipeline_parallel_size"] = _copy_value(deploy.pipeline_parallel_size) + if "data_parallel_size" not in kwargs and deploy.data_parallel_size is not None: + kwargs["data_parallel_size"] = _copy_value(deploy.data_parallel_size) + return config_cls(**kwargs) + + +def _build_diffusion_config_projection( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + engine: _DiffusionEngineOverrides, + *, + model: str | None, + quantization_config: _QuantizationConfigType, +) -> _DiffusionConfigProjection: + diffusion_kwargs = engine.to_kwargs() + diffusion_kwargs["stage_id"] = topology.stage_id + diffusion_kwargs["model_arch"] = _first_defined( + diffusion_kwargs.get("model_arch"), + topology.model_arch, + pipeline.model_arch, + ) + if "dtype" not in diffusion_kwargs and deploy.dtype is not None: + diffusion_kwargs["dtype"] = _copy_value(deploy.dtype) + if "trust_remote_code" not in diffusion_kwargs and deploy.trust_remote_code is not None: + diffusion_kwargs["trust_remote_code"] = _copy_value(deploy.trust_remote_code) + if "distributed_executor_backend" not in diffusion_kwargs and deploy.distributed_executor_backend is not None: + diffusion_kwargs["distributed_executor_backend"] = _copy_value(deploy.distributed_executor_backend) + if model is not None: + diffusion_kwargs["model"] = model + if quantization_config is not None: + diffusion_kwargs["quantization_config"] = _copy_value(quantization_config) + + return _DiffusionConfigProjection.from_kwargs(**{k: v for k, v in diffusion_kwargs.items() if v is not None}) + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniConfig: + """Top-level structured Omni config built once from registry inputs.""" + + pipeline_config: PipelineConfig + stage_configs: tuple[StageConfigType, ...] + orchestrator_config: VllmOmniOrchestratorConfig = field(default_factory=VllmOmniOrchestratorConfig) + + def stage_by_id(self, stage_id: int) -> StageConfigType: + for stage in self.stage_configs: + if stage.stage_id == stage_id: + return stage + raise KeyError(f"no stage {stage_id}") + + @classmethod + def from_registry( + cls, + model_type: str, + deploy_config_path: str | None = None, + cli_overrides: dict[str, Any] | None = None, + ) -> VllmOmniConfig: + """Create a structured config from a registered pipeline and deploy YAML.""" + if cli_overrides is None: + cli_overrides = {} + + deploy_path = _resolve_deploy_path(model_type, deploy_config_path) + loaded_deploy_config_path = str(deploy_path) if deploy_path.exists() else None + if loaded_deploy_config_path is not None: + deploy = load_deploy_config(deploy_path) + else: + deploy = DeployConfig() + + pipeline_key = deploy.pipeline or model_type + if pipeline_key not in _registered_pipeline_keys(): + raise KeyError( + f"Pipeline {pipeline_key!r} not in registry " + f"(resolved from {deploy_path.name!r}). Available: " + f"{_registered_pipeline_keys()}" + ) + pipeline = _resolve_registered_pipeline(pipeline_key) + if pipeline is None: + raise ValueError( + f"Pipeline {pipeline_key!r} did not resolve to a concrete PipelineConfig without an HF config" + ) + + deploy_for_registry = copy.deepcopy(deploy) + if cli_overrides.get("async_chunk") is not None: + deploy_for_registry.async_chunk = bool(cli_overrides["async_chunk"]) + for name in _PIPELINE_DEPLOY_CLI_FIELDS: + if cli_overrides.get(name) is not None: + setattr(deploy_for_registry, name, _copy_value(cli_overrides[name])) + + deploy_for_registry = _apply_platform_overrides(deploy_for_registry) + _validate_async_chunk_support(pipeline, deploy_for_registry) + deploy_by_id = {stage.stage_id: stage for stage in deploy_for_registry.stages} + model = cli_overrides.get("model") + + stage_configs = tuple( + _build_stage_config( + pipeline, + deploy_for_registry, + topology, + deploy_by_id.get(topology.stage_id), + _stage_engine_values( + deploy_by_id.get(topology.stage_id), + _stage_cli_overrides(topology.stage_id, cli_overrides), + ), + model=model, + ) + for topology in pipeline.stages + ) + + orchestrator_config = VllmOmniOrchestratorConfig( + deploy_config_path=loaded_deploy_config_path, + **_orchestrator_cli_overrides(cli_overrides), + ) + return cls( + pipeline_config=pipeline, + stage_configs=stage_configs, + orchestrator_config=orchestrator_config, + ) + + +__all__ = [ + "OmniStageCacheConfig", + "OmniStageConnectorConfig", + "BaseVllmOmniStageConfig", + "OmniStageLoadConfig", + "OmniStageModelConfig", + "VllmOmniOrchestratorConfig", + "OmniStageDiffusionParallelConfig", + "OmniStageParallelConfig", + "OmniStageRuntimeConfig", + "OmniStageSchedulerConfig", + "StageConfigType", + "VllmOmniARStageConfig", + "VllmOmniConfig", + "VllmOmniDiffusionStageConfig", + "VllmOmniGenerationStageConfig", +] diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 45bd80c7a..30d9081fd 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -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( diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 421a42ed0..bac5f90da 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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)}