[Config] Reuse vLLM configs in VllmOmniConfig (#6050)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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
 

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"

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

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:

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,
             )
         )