From e1c5f59aab091b5aed864808d118903ceb77d6ac Mon Sep 17 00:00:00 2001 From: Wang fuyin Date: Tue, 18 Aug 2026 21:27:58 +0800 Subject: [PATCH] [Config] Reuse vLLM configs in VllmOmniConfig (#6050) Signed-off-by: Acerak01-fy Co-authored-by: Hongsheng Liu --- docs/design/module/vllm_omni_config.md | 61 ++++- tests/config/test_omni_config.py | 203 +++++++++++--- tests/engine/test_stage_engine_args.py | 142 +++++++++- vllm_omni/config/omni_config.py | 356 +++++++++++++++++++++---- vllm_omni/engine/stage_init_utils.py | 60 ++++- 5 files changed, 734 insertions(+), 88 deletions(-) diff --git a/docs/design/module/vllm_omni_config.md b/docs/design/module/vllm_omni_config.md index 085e4cee6..eabb9908f 100644 --- a/docs/design/module/vllm_omni_config.md +++ b/docs/design/module/vllm_omni_config.md @@ -18,8 +18,8 @@ validation_paths: - tests/config/** upstream_refs: - vllm.config -last_reviewed: 2026-08-07 -last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc +last_reviewed: 2026-08-11 +last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f --- # vLLM-Omni configuration @@ -39,11 +39,66 @@ The PR review proposal to capture environment-derived values in configuration objects at initialization is an open design question, not a current invariant. +## Current structured-stage reuse behavior + +The current RFC #4021 implementation uses upstream vLLM config classes as the +schema source for AR and generation stages: + +- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`. +- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`. +- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`. +- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`. + +This inheritance preserves the existing load, cache, scheduler, and parallel +concern boundaries while reducing duplicate downstream declarations. It does +not, by itself, make every inherited field an effective vLLM-Omni runtime +option. Construction and runtime consumption are separate surfaces: + +| Surface | Current behavior | +| --- | --- | +| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. | +| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. | +| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. | +| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. | + +The generated projection maps include known upstream naming differences such +as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and +`data_parallel_master_ip` to `data_parallel_address`. A field added to both an +upstream concern config and `EngineArgs` therefore enters the AR/generation +projection without requiring a second downstream allowlist. Explicit inherited +fields that have no projection raise an error rather than being silently +dropped. + +Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`, +cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns +`distributed_executor_backend` and `worker_cls`; vLLM's private API-process +fields are terminal internals. These exclusions prevent one input from being +constructed or projected through two config concerns. Effective-engine-argument +tests cover both the dynamic mapping and this exclusion boundary. + +`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields: +`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML +inputs and complete upstream value objects. Mapping inputs are validated and +materialized as their concrete upstream types when the structured config is +constructed, so downstream consumers see one resolved representation. +Prebuilt upstream objects retain their type across the structured projection +boundary. + +Quantization remains on its existing Omni-owned transport contract. This +reuse step does not adopt upstream `QuantizationConfigArgs`, because the +current engine-specific quantization split needs to be resolved separately. + +Diffusion keeps its existing engine-field ownership and projection surface. +Shared Python inheritance must not be treated as evidence that an LLM-only +parallel input is effective for a diffusion stage. + ## Safe-change guide while deferred Changes should trace every affected producer to its runtime consumer and test the structured, legacy, CLI, and deployment paths that are actually supported. -Do not infer a stable contract from this placeholder page. +For inherited vLLM fields, tests should separately cover constructor behavior +and effective `EngineArgs` projection. Do not infer a stable contract from +this placeholder page. ## Promotion gate diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index cc08ed397..108d997a5 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -5,11 +5,20 @@ from __future__ import annotations from dataclasses import fields +from inspect import Parameter, signature from pathlib import Path +import msgspec import pytest from pydantic import ValidationError +from pydantic.fields import FieldInfo from transformers import Qwen3OmniMoeConfig +from vllm.config import CacheConfig as VllmCacheConfig +from vllm.config import CompilationConfig as VllmCompilationConfig +from vllm.config import LoadConfig as VllmLoadConfig +from vllm.config import ParallelConfig as VllmParallelConfig +from vllm.config import ProfilerConfig as VllmProfilerConfig +from vllm.config import SchedulerConfig as VllmSchedulerConfig from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.config import omni_config as omni_config_module @@ -40,6 +49,7 @@ from vllm_omni.config.stage_config import ( merge_pipeline_deploy, ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.engine.stage_engine_startup import _serialize_stage_config from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -243,9 +253,29 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras ], ids=["global", "ar-stage", "generation-stage"], ) -def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id): - with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"): - _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides) +def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id): + stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id) + + assert stage.cache_config.cache_dtype == "fp8" + assert "cache_dtype" in stage.cache_config._omni_explicit_fields + + +@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"]) +def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + deploy = DeployConfig( + stages=[ + StageDeployConfig( + stage_id=0, + engine_extras={"parallel_config": {field_name: "10.0.0.1"}}, + ) + ] + ) + + stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0) + + assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1" + assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage(): @@ -437,6 +467,8 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): "connector_config", "runtime_config", "parallel_config", + "compilation_config", + "profiler_config", "quantization_config", } assert "diffusion_config" not in public_fields @@ -455,10 +487,47 @@ def test_runtime_config_fields_match_structured_runtime_scope(): "num_gpus", "log_level", "log_stats", - "profiler_config", } +def test_upstream_config_compatible_fields_materialize_mapping_inputs(): + compilation_config = {"backend": "eager"} + profiler_config = {"profiler": "cuda"} + + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0], + compilation_config=compilation_config, + profiler_config=profiler_config, + ) + + assert isinstance(stage_config.compilation_config, VllmCompilationConfig) + assert stage_config.compilation_config.backend == "eager" + assert isinstance(stage_config.profiler_config, VllmProfilerConfig) + assert stage_config.profiler_config.profiler == "cuda" + + +def test_upstream_config_compatible_fields_validate_mapping_inputs(): + with pytest.raises(ValidationError): + VllmOmniARStageConfig( + stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0], + profiler_config={"delay_iterations": -1}, + ) + + +def test_upstream_config_compatible_fields_keep_prebuilt_config_types(): + compilation_config = VllmCompilationConfig(backend="eager") + profiler_config = VllmProfilerConfig(profiler="cuda") + + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0], + compilation_config=compilation_config, + profiler_config=profiler_config, + ) + + assert stage_config.compilation_config is compilation_config + assert stage_config.profiler_config is profiler_config + + def test_sub_config_fields_match_structured_scopes(): assert {f.name for f in fields(OmniStageModelConfig)} == { "model", @@ -489,7 +558,6 @@ def test_sub_config_fields_match_structured_scopes(): "enforce_eager", "max_cudagraph_capture_size", "enable_flashinfer_autotune", - "compilation_config", "enable_multithread_weight_load", "num_weight_load_threads", "disable_autocast", @@ -499,29 +567,25 @@ def test_sub_config_fields_match_structured_scopes(): "model_subdir", "tokenizer_subdir", } - assert {f.name for f in fields(OmniStageLoadConfig)} == { + vllm_load_fields = {f.name for f in fields(VllmLoadConfig)} + assert issubclass(OmniStageLoadConfig, VllmLoadConfig) + assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | { "tokenizer", - "download_dir", "skip_tokenizer_init", - "load_format", "tokenizer_mode", "config_format", "skip_mm_profiling", } - assert {f.name for f in fields(OmniStageCacheConfig)} == { - "kv_cache_memory_bytes", - "gpu_memory_utilization", - "enable_prefix_caching", + assert OmniStageLoadConfig(load_format="PT").load_format == "pt" + assert issubclass(OmniStageCacheConfig, VllmCacheConfig) + assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | { "disable_hybrid_kv_cache_manager", "mm_processor_cache_gb", "mamba_ssm_cache_dtype", } - assert {f.name for f in fields(OmniStageSchedulerConfig)} == { - "max_num_seqs", - "max_num_batched_tokens", + assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig) + assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | { "max_model_len", - "enable_chunked_prefill", - "async_scheduling", } assert {f.name for f in fields(OmniStageConnectorConfig)} == { "async_chunk", @@ -530,17 +594,10 @@ def test_sub_config_fields_match_structured_scopes(): "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", + vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)} + assert issubclass(OmniStageParallelConfig, VllmParallelConfig) + assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields + assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | { "sequence_parallel_size", "ulysses_degree", "ring_degree", @@ -554,20 +611,86 @@ def test_sub_config_fields_match_structured_scopes(): "mask_sp_padding", "hsdp_shard_size", "hsdp_replicate_size", - "enable_expert_parallel", - "world_size", } +def test_inherited_sub_configs_initialize_transport_safe_derived_fields(): + scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096) + assert scheduler_config.max_num_encoder_input_tokens == 4096 + assert scheduler_config.encoder_cache_size == 4096 + + unresolved_scheduler_config = OmniStageSchedulerConfig() + assert unresolved_scheduler_config.max_num_encoder_input_tokens is None + assert unresolved_scheduler_config.encoder_cache_size is None + + parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2) + assert parallel_config.data_parallel_index == 2 + + diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2) + assert diffusion_parallel_config.data_parallel_index == 2 + + +@pytest.mark.parametrize( + "config_cls", + [ + OmniStageLoadConfig, + OmniStageCacheConfig, + OmniStageSchedulerConfig, + OmniStageParallelConfig, + OmniStageDiffusionParallelConfig, + ], +) +def test_inherited_sub_configs_are_keyword_only(config_cls): + assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values()) + + with pytest.raises(TypeError): + config_cls(1) + + +@pytest.mark.parametrize( + "config_cls", + [ + OmniStageLoadConfig, + OmniStageCacheConfig, + OmniStageSchedulerConfig, + OmniStageParallelConfig, + OmniStageDiffusionParallelConfig, + ], +) +def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls): + config = config_cls() + + uninitialized_fields = [ + config_field.name + for config_field in fields(config) + if isinstance(getattr(config, config_field.name), FieldInfo) + ] + assert uninitialized_fields == [] + + msgspec.msgpack.encode(_serialize_stage_config(config)) + + +def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe(): + omni_config = _from_pipeline_key("qwen3_tts") + + for stage_config in omni_config.stage_configs: + payload = { + "stage_id": stage_config.stage_id, + "stage_config": _serialize_stage_config(stage_config), + } + msgspec.msgpack.encode(payload) + + 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"} + vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)} assert legacy_fields | expected_upstream_fields <= structured_fields - assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"} + assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set() def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface(): @@ -1078,3 +1201,21 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config) assert cfg.quantization_config == quantization_config + + +def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + quantization_config = {"method": "int8", "activation_scheme": "dynamic"} + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( + quantization_config=quantization_config, + ) + + # Exercise the terminal construction performed by the future typed startup + # path without probing ports or loading remote model metadata. + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) + monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None) + cfg.enrich_config() + + assert cfg.quantization_config is not None + assert cfg.quantization_config.get_name() == "int8" diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 55ae43603..e7c9b2311 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -8,8 +8,23 @@ from pathlib import Path import pytest from pydantic.fields import FieldInfo +from vllm.config import CacheConfig as VllmCacheConfig +from vllm.config import CompilationConfig as VllmCompilationConfig +from vllm.config import LoadConfig as VllmLoadConfig +from vllm.config import ParallelConfig as VllmParallelConfig +from vllm.config import ProfilerConfig as VllmProfilerConfig +from vllm.config import SchedulerConfig as VllmSchedulerConfig +from vllm.engine.arg_utils import EngineArgs -from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.omni_config import ( + _LLM_STAGE_ENGINE_FIELDS, + OmniStageCacheConfig, + OmniStageLoadConfig, + OmniStageParallelConfig, + OmniStageSchedulerConfig, + VllmOmniARStageConfig, + VllmOmniConfig, +) from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config from vllm_omni.config.stage_config import ( DeployConfig, @@ -53,6 +68,30 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, _LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs)) _DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig)) +_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( + { + "active_stream_window", + "codec_frame_rate_hz", + "custom_voice_dir", + "devices", + "disable_autocast", + "enable_multithread_weight_load", + "env", + "has_sampling_extra_args", + "log_level", + "log_stats", + "model_arch", + "model_subdir", + "num_gpus", + "num_replicas", + "num_weight_load_threads", + "omni_kv_config", + "parallel_config", + "subtalker_sampling_params", + "task_type", + "tokenizer_subdir", + } +) @pytest.fixture(autouse=True) @@ -197,6 +236,100 @@ def _legacy_and_typed_stages( return legacy_stages, omni_config +def test_llm_stage_engine_field_schema_tracks_upstream_engine_args(): + upstream_engine_fields = frozenset(field.name for field in fields(EngineArgs)) + + assert _LLM_STAGE_ENGINE_FIELDS - upstream_engine_fields == _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS + + +def test_typed_llm_projection_discovers_explicit_upstream_config_fields(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + load_config=OmniStageLoadConfig(ignore_patterns=["*.bin"]), + cache_config=OmniStageCacheConfig(block_size=32, cache_dtype="fp8"), + scheduler_config=OmniStageSchedulerConfig( + prefill_schedule_interval=2, + policy="priority", + ), + parallel_config=OmniStageParallelConfig( + data_parallel_master_ip="10.0.0.1", + enable_dbo=True, + ), + ) + + engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config) + + assert engine_args["ignore_patterns"] == ["*.bin"] + assert engine_args["block_size"] == 32 + assert engine_args["kv_cache_dtype"] == "fp8" + assert engine_args["prefill_schedule_interval"] == 2 + assert engine_args["scheduling_policy"] == "priority" + assert engine_args["data_parallel_address"] == "10.0.0.1" + assert engine_args["enable_dbo"] is True + + +def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + ) + + engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config) + + inherited_defaults = { + "ignore_patterns", + "block_size", + "kv_cache_dtype", + "prefill_schedule_interval", + "scheduling_policy", + "data_parallel_address", + "data_parallel_rank", + "data_parallel_size_local", + "data_parallel_rpc_port", + "enable_dbo", + } + assert inherited_defaults.isdisjoint(engine_args) + + +def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + scheduler_config=OmniStageSchedulerConfig(scheduler_cls="test.CustomScheduler"), + ) + + with pytest.raises(ValueError, match=r"no EngineArgs projection: scheduler_cls"): + stage_init_utils._project_omni_stage_engine_args(stage_config) + + +@pytest.mark.parametrize("stage_id", [0, 1], ids=["ar", "generation"]) +def test_typed_llm_engine_args_preserve_upstream_config_objects(tmp_path, stage_id): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + compilation_config = VllmCompilationConfig(backend="eager") + profiler_config = VllmProfilerConfig(profiler="cuda") + stage_prefix = f"stage_{stage_id}_" + omni_config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=copy.deepcopy(deploy), + cli_overrides={ + "model": model, + f"{stage_prefix}compilation_config": compilation_config, + f"{stage_prefix}profiler_config": profiler_config, + }, + ) + + stage_config = omni_config.stage_by_id(stage_id) + typed_args = build_engine_args_dict_from_omni_stage_config(stage_config, model) + + assert isinstance(stage_config.compilation_config, VllmCompilationConfig) + assert isinstance(typed_args["compilation_config"], VllmCompilationConfig) + assert typed_args["compilation_config"] is not stage_config.compilation_config + assert typed_args["compilation_config"].backend == "eager" + + assert isinstance(stage_config.profiler_config, VllmProfilerConfig) + assert isinstance(typed_args["profiler_config"], VllmProfilerConfig) + assert typed_args["profiler_config"] is not stage_config.profiler_config + assert typed_args["profiler_config"].profiler == "cuda" + + def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path): pipeline, deploy, model = _engine_arg_inputs(tmp_path) legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model) @@ -222,6 +355,13 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path): assert {name: typed_args[name] for name in legacy_args} == legacy_args thinker_args = typed_args_by_stage[0] + inherited_vllm_fields = { + field.name + for config_cls in (VllmLoadConfig, VllmCacheConfig, VllmSchedulerConfig, VllmParallelConfig) + for field in fields(config_cls) + } + topology_projected_fields = {"scheduler_cls"} + assert (inherited_vllm_fields - _LLM_STAGE_ENGINE_FIELDS - topology_projected_fields).isdisjoint(thinker_args) assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model") assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer") assert thinker_args["stage_id"] == 0 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 4a8a54eca..9cd7ed518 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -12,12 +12,22 @@ from __future__ import annotations import copy from collections.abc import Mapping -from dataclasses import dataclass, field, fields +from dataclasses import InitVar, dataclass, field, fields +from functools import wraps +from inspect import Parameter, signature from pathlib import Path from typing import Any, Literal, TypeAlias, TypedDict, cast -from pydantic import ConfigDict, Field +from pydantic import ConfigDict, Field, model_validator +from typing_extensions import Self +from vllm.config import CacheConfig as VllmCacheConfig +from vllm.config import CompilationConfig as VllmCompilationConfig +from vllm.config import LoadConfig as VllmLoadConfig +from vllm.config import ParallelConfig as VllmParallelConfig +from vllm.config import ProfilerConfig as VllmProfilerConfig +from vllm.config import SchedulerConfig as VllmSchedulerConfig from vllm.config.utils import config +from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.config.stage_config import ( @@ -77,6 +87,49 @@ class _QuantizationEngineOverrides(TypedDict, total=False): quantization: str +class _TrackExplicitConfigFields: + """Record constructor inputs without adding a serialized config field.""" + + @model_validator(mode="wrap") + @classmethod + def _record_explicit_config_fields(cls, value: Any, handler: Any) -> Any: + result = handler(value) + kwargs = getattr(value, "kwargs", None) + if kwargs is not None: + explicit_fields = frozenset(kwargs) + elif isinstance(value, cls): + explicit_fields = getattr(value, "_omni_explicit_fields", frozenset()) + else: + explicit_fields = frozenset() + object.__setattr__(result, "_omni_explicit_fields", explicit_fields) + return result + + +def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]: + """Make inherited Pydantic dataclass fields keyword-only as well.""" + generated_init = cls.__init__ + generated_signature = signature(cls) + keyword_only_signature = generated_signature.replace( + parameters=[ + parameter.replace(kind=Parameter.KEYWORD_ONLY) + if parameter.kind in {Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD} + else parameter + for parameter in generated_signature.parameters.values() + ] + ) + + @wraps(generated_init) + def keyword_only_init(self: Any, *args: Any, **kwargs: Any) -> None: + if args: + raise TypeError(f"{cls.__name__}() accepts keyword arguments only; got {len(args)} positional argument(s)") + generated_init(self, **kwargs) + + cls.__init__ = keyword_only_init + cls.__signature__ = keyword_only_signature + cls.__match_args__ = () + return cls + + class _ModelEngineOverrides(TypedDict, total=False): model: str model_arch: str @@ -106,7 +159,6 @@ class _ModelEngineOverrides(TypedDict, total=False): enforce_eager: bool max_cudagraph_capture_size: int enable_flashinfer_autotune: bool - compilation_config: dict[str, Any] enable_multithread_weight_load: bool num_weight_load_threads: int disable_autocast: bool @@ -148,7 +200,6 @@ class _RuntimeEngineOverrides(TypedDict, total=False): num_gpus: int log_level: str log_stats: bool - profiler_config: dict[str, Any] class _ParallelConfigEngineOverrides(TypedDict, total=False): @@ -192,6 +243,8 @@ class _StageEngineValues: runtime: _RuntimeEngineOverrides parallel: _ParallelEngineOverrides diffusion: _DiffusionEngineOverrides + compilation_config: Mapping[str, Any] | VllmCompilationConfig | None + profiler_config: Mapping[str, Any] | VllmProfilerConfig | None @dataclass(frozen=True) @@ -341,7 +394,6 @@ class OmniStageModelConfig: enforce_eager: bool = False max_cudagraph_capture_size: int | None = Field(default=None, ge=0) 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 @@ -352,21 +404,21 @@ class OmniStageModelConfig: tokenizer_subdir: str | None = None -@config -class OmniStageLoadConfig: - """Per-stage loading behavior and resolved tokenizer input.""" +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig): + """vLLM loading behavior plus Omni stage-specific tokenizer inputs.""" tokenizer: str | None = None - download_dir: str | None = None skip_tokenizer_init: bool = False - load_format: str = "auto" tokenizer_mode: str = "auto" config_format: str | None = None skip_mm_profiling: bool | None = None -@config -class OmniStageCacheConfig: +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageCacheConfig(_TrackExplicitConfigFields, VllmCacheConfig): """Per-stage engine cache and memory behavior. This is separate from ``_DiffusionConfigProjection.cache_config``, which configures @@ -383,17 +435,26 @@ class OmniStageCacheConfig: mamba_ssm_cache_dtype: str | None = None -@config -class OmniStageSchedulerConfig: +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig): """Per-stage request scheduling behavior.""" + # Upstream receives max_model_len only while materializing SchedulerConfig. + # Omni retains it as unresolved stage input until the owning engine process. max_num_seqs: int | None = Field(default=None, ge=1) max_num_batched_tokens: int | None = Field(default=None, ge=1) max_model_len: int | None = Field(default=None, ge=-1) + is_encoder_decoder: InitVar[bool] = False enable_chunked_prefill: bool | None = None async_scheduling: bool | None = None - def __post_init__(self) -> None: + def __post_init__(self, is_encoder_decoder: bool = False) -> None: + # Upstream initializes these derived fields in its terminal post-init. + # Keep them serializable here without running model-dependent checks. + self.max_num_encoder_input_tokens = self.max_num_batched_tokens + self.encoder_cache_size = self.max_num_batched_tokens + if ( self.max_num_batched_tokens is not None and self.max_num_seqs is not None @@ -432,24 +493,73 @@ class OmniStageRuntimeConfig: 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: +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig): """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) + # EngineArgs intentionally leaves these unresolved. The upstream terminal + # ParallelConfig defaults (rank 0, local size 1, port 29550, worker + # ``"auto"``) must not turn into explicit engine inputs in the head + # process merely because this transport class inherits ParallelConfig. + data_parallel_size_local: int | None = Field(default=None, ge=0) + data_parallel_rank: int | None = Field(default=None, ge=0) + data_parallel_rpc_port: int | None = None + worker_cls: str | None = None + + @model_validator(mode="after") + def _validate_parallel_config(self) -> Self: + """Run upstream validation without resolving deferred engine inputs.""" + deferred_values = ( + self.data_parallel_size_local, + self.data_parallel_rank, + self.data_parallel_rpc_port, + self.worker_cls, + self.all2all_backend, + self.disable_custom_all_reduce, + ) + if self.data_parallel_size_local is not None and self.data_parallel_size_local > self.data_parallel_size: + raise ValueError( + f"data_parallel_size_local ({self.data_parallel_size_local}) " + f"must be <= data_parallel_size ({self.data_parallel_size})" + ) + if self.data_parallel_rank is not None and not 0 <= self.data_parallel_rank < self.data_parallel_size: + raise ValueError( + f"data_parallel_rank ({self.data_parallel_rank}) must be in the range [0, {self.data_parallel_size})" + ) + self.data_parallel_size_local = self.data_parallel_size + self.data_parallel_rank = 0 + self.data_parallel_rpc_port = VllmParallelConfig.data_parallel_rpc_port + self.worker_cls = VllmParallelConfig.worker_cls + try: + VllmParallelConfig._validate_parallel_config(self) + finally: + ( + self.data_parallel_size_local, + self.data_parallel_rank, + self.data_parallel_rpc_port, + self.worker_cls, + self.all2all_backend, + self.disable_custom_all_reduce, + ) = deferred_values + return self def __post_init__(self) -> None: + # Keep config construction transport-safe. Upstream runtime backend, + # rank, port, and platform resolution happens in the owning process. + self.data_parallel_index = self.data_parallel_rank self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size + @property + def world_size_across_dp(self) -> int: + # Omni's public world_size has historically included DP. + return self.world_size -@config + +@_enforce_keyword_only_init +@config(kw_only=True) class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): """Diffusion-stage distributed parallelism behavior.""" @@ -468,6 +578,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): hsdp_replicate_size: int = Field(default=1, ge=1) def __post_init__(self) -> None: + self.data_parallel_index = self.data_parallel_rank self.sequence_parallel_size = ( self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree ) @@ -843,6 +954,58 @@ _DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMP "stage_id", } + +def _upstream_engine_field_map( + config_cls: type[Any], + *, + aliases: Mapping[str, str] = {}, + exclude: frozenset[str] = frozenset(), +) -> dict[str, str]: + """Map reusable upstream config fields to their EngineArgs inputs.""" + engine_fields = frozenset(config_field.name for config_field in fields(VllmEngineArgs)) + return { + config_field.name: engine_name + for config_field in fields(config_cls) + if config_field.init + and config_field.name not in exclude + and (engine_name := aliases.get(config_field.name, config_field.name)) in engine_fields + } + + +_LOAD_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(VllmLoadConfig) +_CACHE_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map( + VllmCacheConfig, + aliases={"cache_dtype": "kv_cache_dtype"}, +) +_SCHEDULER_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map( + VllmSchedulerConfig, + aliases={"policy": "scheduling_policy"}, + # ``scheduler_cls`` is selected by the immutable stage topology, while + # ``disable_hybrid_kv_cache_manager`` belongs to the cache concern in the + # Omni schema. Keep both out of the dynamic SchedulerConfig projection so + # one input cannot acquire two owners. + exclude=frozenset( + { + "scheduler_cls", + "disable_hybrid_kv_cache_manager", + } + ), +) +_PARALLEL_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map( + VllmParallelConfig, + aliases={"data_parallel_master_ip": "data_parallel_address"}, + # Runtime owns worker selection. The private API-process fields are + # terminal vLLM internals rather than per-stage user inputs. + exclude=frozenset( + { + "distributed_executor_backend", + "worker_cls", + "_api_process_count", + "_api_process_rank", + } + ), +) + _QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__) _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__) _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__) @@ -850,7 +1013,17 @@ _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__) _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__) _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__) _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) -_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__) +_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"}) +_LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values()) +_LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values()) +_LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values()) +_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP.values()) +_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset( + f.name for f in fields(OmniStageDiffusionParallelConfig) +) & frozenset(_ParallelConfigEngineOverrides.__annotations__) +_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP = {name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS} +_LLM_PARALLEL_CONFIG_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP) +_PARALLEL_CONFIG_ENGINE_FIELDS = _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS _PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"} _COMMON_STAGE_ENGINE_FIELDS = ( _QUANTIZATION_ENGINE_FIELDS @@ -860,12 +1033,16 @@ _COMMON_STAGE_ENGINE_FIELDS = ( | _SCHEDULER_ENGINE_FIELDS | _CONNECTOR_ENGINE_FIELDS | _RUNTIME_ENGINE_FIELDS + | _DIRECT_VLLM_CONFIG_ENGINE_FIELDS +) +_LLM_STAGE_ENGINE_FIELDS = ( + _COMMON_STAGE_ENGINE_FIELDS + | _LLM_LOAD_ENGINE_FIELDS + | _LLM_CACHE_ENGINE_FIELDS + | _LLM_SCHEDULER_ENGINE_FIELDS + | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS + | {"parallel_config"} ) -_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageParallelConfig)) - {"world_size"} -_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageDiffusionParallelConfig)) - { - "world_size" -} -_LLM_STAGE_ENGINE_FIELDS = _COMMON_STAGE_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"} _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = ( _COMMON_STAGE_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS @@ -886,6 +1063,27 @@ _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = { StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS, StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS, } +_PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = { + StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS, + StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS, + StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP), +} + +_LOAD_STAGE_ENGINE_FIELD_MAP = { + **{name: name for name in _LOAD_ENGINE_FIELDS}, + **_LOAD_CONFIG_ENGINE_FIELD_MAP, +} +_CACHE_STAGE_ENGINE_FIELD_MAP = { + **{name: name for name in _CACHE_ENGINE_FIELDS}, + **_CACHE_CONFIG_ENGINE_FIELD_MAP, +} +_SCHEDULER_STAGE_ENGINE_FIELD_MAP = { + **{name: name for name in _SCHEDULER_ENGINE_FIELDS}, + **_SCHEDULER_CONFIG_ENGINE_FIELD_MAP, +} +_DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP = {name: name for name in _LOAD_ENGINE_FIELDS} +_DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP = {name: name for name in _CACHE_ENGINE_FIELDS} +_DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {name: name for name in _SCHEDULER_ENGINE_FIELDS} def _validate_stage_engine_override_ownership( @@ -896,13 +1094,19 @@ def _validate_stage_engine_override_ownership( try: owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] + parallel_config_fields = _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE[execution_type] except KeyError as exc: raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc unowned_fields = set(overrides) - owner_fields parallel_config = overrides.get("parallel_config") if isinstance(parallel_config, Mapping): - unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - parallel_owner_fields) + # Nested values traditionally use upstream config names, while flat + # EngineArgs uses aliases such as ``data_parallel_address``. Accept + # either spelling at the boundary; the builder canonicalizes aliases + # before instantiating the inherited config. + accepted_parallel_names = parallel_config_fields | parallel_owner_fields + unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - accepted_parallel_names) if unowned_fields: names = ", ".join(sorted(unowned_fields)) raise ValueError( @@ -973,17 +1177,28 @@ def _stage_engine_values( topology.execution_type, engine, ) + if topology.execution_type in { + StageExecutionType.LLM_AR, + StageExecutionType.LLM_GENERATION, + }: + load_engine_fields = _LLM_LOAD_ENGINE_FIELDS + cache_engine_fields = _LLM_CACHE_ENGINE_FIELDS + scheduler_engine_fields = _LLM_SCHEDULER_ENGINE_FIELDS + else: + load_engine_fields = _LOAD_ENGINE_FIELDS + cache_engine_fields = _CACHE_ENGINE_FIELDS + scheduler_engine_fields = _SCHEDULER_ENGINE_FIELDS 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)), + 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), + _select_engine_overrides(engine, scheduler_engine_fields), ), connector=cast( _ConnectorEngineOverrides, @@ -992,6 +1207,8 @@ def _stage_engine_values( runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)), parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)), diffusion=_DiffusionEngineOverrides.from_engine(engine), + compilation_config=_copy_value(engine.get("compilation_config")), + profiler_config=_copy_value(engine.get("profiler_config")), ) @@ -1046,6 +1263,8 @@ class BaseVllmOmniStageConfig: connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig) runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig) parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) + compilation_config: VllmCompilationConfig | None = None + profiler_config: VllmProfilerConfig | None = None quantization_config: _QuantizationConfigType = None @property @@ -1197,6 +1416,8 @@ def _build_common_stage_config_kwargs( parallel_config, ), "parallel_config": parallel_config, + "compilation_config": _copy_value(engine.compilation_config), + "profiler_config": _copy_value(engine.profiler_config), "quantization_config": _copy_value(quantization_config), }, input_proc, @@ -1391,12 +1612,51 @@ def _build_load_config( return OmniStageLoadConfig(**kwargs) +def _config_kwargs_from_engine_args( + engine: Mapping[str, Any], + field_map: Mapping[str, str], +) -> dict[str, Any]: + """Convert EngineArgs names back to their upstream config field names.""" + return { + config_name: _copy_value(engine[engine_name]) + for config_name, engine_name in field_map.items() + if engine_name in engine and engine[engine_name] is not None + } + + +def _normalize_config_mapping( + values: Mapping[str, Any], + field_map: Mapping[str, str], +) -> dict[str, Any]: + """Canonicalize config and EngineArgs spellings to config field names.""" + engine_to_config = {engine_name: config_name for config_name, engine_name in field_map.items()} + normalized: dict[str, Any] = {} + source_names: dict[str, str] = {} + for name, value in values.items(): + if value is None: + continue + config_name = name if name in field_map else engine_to_config.get(name, name) + if config_name in normalized: + raise ValueError( + f"Config mapping specifies both {source_names[config_name]!r} and {name!r} " + f"for upstream field {config_name!r}" + ) + normalized[config_name] = _copy_value(value) + source_names[config_name] = name + return normalized + + +def _omni_config_kwargs(engine: Mapping[str, Any], fields_: frozenset[str]) -> dict[str, Any]: + return {name: _copy_value(engine[name]) for name in fields_ if name in engine and engine[name] is not None} + + def _build_cache_config( deploy: DeployConfig, engine: _CacheEngineOverrides, execution_type: StageExecutionType, ) -> OmniStageCacheConfig: - kwargs = _config_kwargs(engine) + kwargs = _config_kwargs_from_engine_args(engine, _CACHE_CONFIG_ENGINE_FIELD_MAP) + kwargs.update(_omni_config_kwargs(engine, _CACHE_ENGINE_FIELDS)) 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) if "disable_hybrid_kv_cache_manager" not in kwargs and execution_type == StageExecutionType.LLM_GENERATION: @@ -1410,7 +1670,8 @@ def _build_scheduler_config( engine: _SchedulerEngineOverrides, execution_type: StageExecutionType, ) -> OmniStageSchedulerConfig: - kwargs = _config_kwargs(engine) + kwargs = _config_kwargs_from_engine_args(engine, _SCHEDULER_CONFIG_ENGINE_FIELD_MAP) + kwargs.update(_omni_config_kwargs(engine, _SCHEDULER_ENGINE_FIELDS)) 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) if "async_scheduling" not in kwargs and execution_type == StageExecutionType.LLM_AR: @@ -1458,16 +1719,21 @@ def _build_parallel_config( 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 - } + config_fields = frozenset(signature(config_cls).parameters) + is_diffusion = issubclass(config_cls, OmniStageDiffusionParallelConfig) + field_map = _DIFFUSION_PARALLEL_CONFIG_FIELD_MAP if is_diffusion else _PARALLEL_CONFIG_ENGINE_FIELD_MAP + owned_parallel_fields = frozenset(field_map) + kwargs = _config_kwargs(engine) if is_diffusion else _config_kwargs_from_engine_args(engine, field_map) + kwargs = {name: value for name, value in kwargs.items() if name in owned_parallel_fields} 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}) + nested_kwargs = _normalize_config_mapping(parallel_config, field_map) + kwargs.update( + { + name: value + for name, value in nested_kwargs.items() + if name in owned_parallel_fields and 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: diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 9da6f2cfa..683726d28 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -14,7 +14,7 @@ import importlib import multiprocessing as mp import os import time -from collections.abc import Callable, Generator, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass, fields, replace from typing import Any, Literal, cast @@ -29,6 +29,14 @@ from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.executor import Executor from vllm_omni.config.omni_config import ( + _CACHE_STAGE_ENGINE_FIELD_MAP, + _DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP, + _DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP, + _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS, + _DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP, + _LOAD_STAGE_ENGINE_FIELD_MAP, + _PARALLEL_CONFIG_ENGINE_FIELD_MAP, + _SCHEDULER_STAGE_ENGINE_FIELD_MAP, BaseVllmOmniStageConfig, VllmOmniDiffusionStageConfig, ) @@ -781,20 +789,38 @@ def stage_runtime_env(stage_id: int, runtime_cfg: Any) -> Generator[None, None, def _project_omni_config_fields( config: Any, *, + field_map: Mapping[str, str] | None = None, exclude: frozenset[str] = frozenset(), ) -> dict[str, Any]: """Copy defined typed config fields into backend adapter kwargs.""" projected: dict[str, Any] = {} for config_field in fields(config): name = config_field.name - if name in exclude: + if name in exclude or (field_map is not None and name not in field_map): continue value = getattr(config, name) if value is not None: - projected[name] = copy.deepcopy(value) + projected[field_map.get(name, name) if field_map is not None else name] = copy.deepcopy(value) return projected +def _project_upstream_config_fields( + config: Any, + field_map: Mapping[str, str], +) -> dict[str, Any]: + """Project every explicit upstream input, including newly added fields.""" + explicit_fields = getattr(config, "_omni_explicit_fields", frozenset()) + unprojected_fields = explicit_fields - frozenset(field_map) + if unprojected_fields: + names = ", ".join(sorted(unprojected_fields)) + raise ValueError(f"{type(config).__name__} has explicit field(s) with no EngineArgs projection: {names}") + return _project_omni_config_fields( + config, + field_map=field_map, + exclude=frozenset(field_map) - explicit_fields, + ) + + def _project_omni_stage_engine_args( stage_config: BaseVllmOmniStageConfig, ) -> dict[str, Any]: @@ -810,9 +836,6 @@ def _project_omni_stage_engine_args( stage_config.model_config, frozenset({"default_sampling_params", "has_sampling_extra_args"}), ), - (stage_config.load_config, frozenset()), - (stage_config.cache_config, frozenset()), - (stage_config.scheduler_config, frozenset()), ( stage_config.runtime_config, frozenset({"devices", "num_replicas", "env", "num_gpus"}), @@ -825,6 +848,26 @@ def _project_omni_stage_engine_args( ) ) + if is_diffusion: + for config, field_map in ( + (stage_config.load_config, _DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP), + (stage_config.cache_config, _DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP), + (stage_config.scheduler_config, _DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP), + ): + engine_args.update(_project_omni_config_fields(config, field_map=field_map)) + else: + for config, field_map in ( + (stage_config.load_config, _LOAD_STAGE_ENGINE_FIELD_MAP), + (stage_config.cache_config, _CACHE_STAGE_ENGINE_FIELD_MAP), + (stage_config.scheduler_config, _SCHEDULER_STAGE_ENGINE_FIELD_MAP), + ): + engine_args.update(_project_upstream_config_fields(config, field_map)) + + for name in ("compilation_config", "profiler_config"): + value = getattr(stage_config, name) + if value is not None: + engine_args[name] = copy.deepcopy(value) + # The legacy builder always emits this key, including for pipelines such # as Audex that intentionally defer architecture discovery to HF config. engine_args["model_arch"] = copy.deepcopy(stage_config.model_config.model_arch) @@ -851,13 +894,14 @@ def _project_omni_stage_engine_args( if is_diffusion: engine_args["parallel_config"] = _project_omni_config_fields( stage_config.parallel_config, + field_map={name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS}, exclude=frozenset({"world_size"}), ) else: engine_args.update( - _project_omni_config_fields( + _project_upstream_config_fields( stage_config.parallel_config, - exclude=frozenset({"world_size"}), + _PARALLEL_CONFIG_ENGINE_FIELD_MAP, ) )