Merge main and adapt multi-API startup to structured configs

Signed-off-by: Sy03 <1370724210@qq.com>

docs/design/architecture_overview.md

@@ -323,7 +323,7 @@ inputs, fields, and ownership rules are described below.
 flowchart TB
     layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
     layer2["Layer 2 · Production resolution boundary<br/>resolve_omni_config()<br/>(StageConfigFactory + VllmOmniConfig internally)"]
-    layer3["Layer 3 · Startup hand-off<br/>OmniConfigResolution<br/>PipelineConfig + temporary OmegaConf stage bridge"]
+    layer3["Layer 3 · Startup hand-off<br/>OmniConfigResolution<br/>PipelineConfig + typed stage configs"]
     layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
     layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
 
@@ -335,10 +335,10 @@ The production resolution boundary is `resolve_omni_config()` in
 It delegates typed construction to `StageConfigFactory.create_from_model()` and
 `VllmOmniConfig.from_pipeline_config()`, then returns an
 `OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup.
-Until `StageRuntime` consumes typed stage configs directly, this envelope carries
-the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as
-a temporary runtime bridge; both views describe the same resolved topology,
-including injected stages.
+This envelope carries the effective `PipelineConfig` alongside typed
+`stage_configs`, which `StageRuntime` consumes directly. Both views describe the
+same resolved topology, including injected stages. Backend arguments are
+projected from these typed stages when each engine is initialized.
 
 The legacy `stage_args` YAML path has been removed. Model topology now resolves
 through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`.
@@ -360,7 +360,7 @@ The important ownership rules are:
 3. CLI and Python overrides are applied at the resolution boundary, with
    per-stage overrides taking precedence over global values where supported.
 4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config`
-   and temporary `stage_configs` compatibility view must describe the same
+   and typed `stage_configs` must describe the same
    topology.
 5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
    is runtime-private state, not a user configuration object.

tests/config/composable_parallel/test_factory_integration.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Integration tests: strategy overlay on the real merge_pipeline_deploy seam."""
 
 from __future__ import annotations
@@ -20,6 +20,7 @@ from vllm_omni.config.composable_parallel import (
     apply_strategy_specs,
 )
 from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.omni_config import VllmOmniConfig
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
 
@@ -137,3 +138,31 @@ def test_cli_overrides_strategy_with_warning():
     assert _resolved(stages, "talker").runtime.num_replicas == 3
     # ...and the override was warned about, naming the conflicting field.
     assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages)
+
+
+def test_typed_cli_overrides_strategy():
+    pipeline = OMNI_PIPELINES["qwen2_5_omni"]
+    config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=load_deploy_config(_DEPLOY),
+        cli_overrides={"stage_1_num_replicas": 3},
+        strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
+    )
+
+    talker = next(stage for stage in config.stage_configs if stage.model_stage == "talker")
+    assert talker.runtime_config.num_replicas == 3
+    assert config.orchestrator_config.omni_lb_policy == "round-robin"
+    assert config.strategy_omni_lb_policy == "round-robin"
+
+
+def test_typed_tp_only_strategy_has_no_derived_lb_policy():
+    pipeline = OMNI_PIPELINES["qwen2_5_omni"]
+    config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=load_deploy_config(_DEPLOY),
+        cli_overrides={"omni_lb_policy": "round-robin"},
+        strategy_specs={"thinker": [_tp(1)]},
+    )
+
+    assert config.orchestrator_config.omni_lb_policy == "round-robin"
+    assert config.strategy_omni_lb_policy is None

tests/config/test_forced_aligner_injection.py

@@ -2,13 +2,19 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import pytest
+from vllm.pooling_params import PoolingParams
 
 import vllm_omni.config.pipeline_registry  # noqa: F401  (populate registry)
+from vllm_omni.config.omni_config import VllmOmniConfig
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
 from vllm_omni.config.stage_config import (
     DeployConfig,
     StageExecutionType,
 )
+from vllm_omni.engine.stage_init_utils import (
+    build_engine_args_dict_from_omni_stage_config,
+    extract_stage_metadata_from_omni_stage_config,
+)
 from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
     POOLING_OUTPUT_DECODER_PATH,
     TIMESTAMPS_MODALITY,
@@ -94,3 +100,33 @@ def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_ch
     assert stages[2].engine_args.runner == "pooling"
     assert len(pipeline.stages) == 2
     assert deploy.stages == []
+
+
+def test_injected_aligner_survives_typed_config_and_runtime_projections():
+    pipeline, deploy = inject_forced_aligner_stage(
+        _PIPELINE_REGISTRY["qwen3_tts"],
+        DeployConfig(),
+        {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"},
+    )
+    config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=deploy,
+        cli_overrides={"model": "/models/Qwen3-TTS"},
+    )
+    stage = config.stage_configs[-1]
+
+    assert stage.pooling_config.runner == "pooling"
+    assert stage.pooling_config.pooling_output_decoder == POOLING_OUTPUT_DECODER_PATH
+    assert isinstance(stage.pooling_config.default_pooling_params, PoolingParams)
+    assert stage.pooling_config.default_pooling_params.task == deploy.stages[-1].default_pooling_params["task"]
+
+    metadata = extract_stage_metadata_from_omni_stage_config(stage)
+    assert isinstance(metadata.default_sampling_params, PoolingParams)
+    assert metadata.default_sampling_params.task == deploy.stages[-1].default_pooling_params["task"]
+
+    engine_args = build_engine_args_dict_from_omni_stage_config(
+        stage,
+        stage.model_config.model or "/models/Qwen3-ForcedAligner-0.6B",
+    )
+    assert engine_args["runner"] == "pooling"
+    assert engine_args["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH

tests/config/test_omni_config.py

@@ -11,6 +11,7 @@ from pathlib import Path
 
 import msgspec
 import pytest
+import torch
 from pydantic import ValidationError
 from pydantic.fields import FieldInfo
 from transformers import Qwen3OmniMoeConfig
@@ -123,6 +124,18 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
     assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
 
 
+def test_nested_stage_override_deep_merges_structured_model_config() -> None:
+    config = _from_pipeline_key(
+        "cosmos3_policy",
+        deploy_config_path=get_deploy_config_path("cosmos3_policy_droid.yaml"),
+        cli_overrides={"stage_0_model_config": {"guardrails": False}},
+    )
+
+    model_config = config.stage_by_id(0).diffusion_config.model_config
+    assert model_config["guardrails"] is False
+    assert model_config["policy_server_config"]["action_space"] == "joint_position"
+
+
 @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
 def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
     pipeline = _resolve_pipeline_or_skip(model_type)
@@ -654,6 +667,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
         "load_config",
         "cache_config",
         "scheduler_config",
+        "pooling_config",
         "connector_config",
         "runtime_config",
         "parallel_config",
@@ -889,6 +903,17 @@ def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_saf
         msgspec.msgpack.encode(payload)
 
 
+def test_structured_diffusion_torch_dtype_is_msgpack_transport_safe():
+    stage_config = _from_pipeline_key(
+        "hunyuan_image3_dit",
+        cli_overrides={"dtype": torch.bfloat16},
+    ).stage_by_id(0)
+
+    serialized = _serialize_stage_config(stage_config)
+    assert serialized["diffusion_config"]["dtype"] == "bfloat16"
+    msgspec.msgpack.encode(serialized)
+
+
 def test_diffusion_parallel_config_fields_cover_legacy_surface():
     from vllm_omni.diffusion.data import DiffusionParallelConfig
 

tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -21,6 +21,7 @@ import torch.distributed as dist
 from transformers import AutoTokenizer
 
 from tests.helpers.mark import hardware_test
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
 from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
 from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
 from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -136,6 +137,15 @@ async def _generate_once(
 
 def _assert_live_uniproc_worker(engine: AsyncOmni) -> None:
     """Default single-GPU path must use the in-process executor, not mp IPC."""
+    stage_configs = engine.engine.stage_configs
+    assert len(stage_configs) == 1
+    stage_config = stage_configs[0]
+    assert isinstance(stage_config, VllmOmniDiffusionStageConfig)
+    assert stage_config.model_config.model == MODEL
+    assert stage_config.model_config.enforce_eager is True
+    assert stage_config.scheduler_config.max_num_seqs == 1
+    assert stage_config.diffusion_config.custom_pipeline_args == {"pipeline_class": CUSTOM_PIPELINE_CLASS}
+    assert stage_config.diffusion_config.worker_extension_cls == WORKER_EXTENSION_CLASS
     clients = engine.engine.stage_clients
     assert len(clients) == 1
     client = clients[0]

tests/e2e/offline_inference/test_hunyuanimage3.py

@@ -1,5 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: E501
+import copy
 import importlib.util
+import os
 import sys
 from collections.abc import Generator
 from pathlib import Path
@@ -10,20 +15,37 @@ import torch.nn.functional as F
 from PIL import Image
 from transformers import CLIPModel, CLIPProcessor
 
+from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni import Omni
+from vllm_omni.config.omni_config import (
+    VllmOmniARStageConfig,
+    VllmOmniDiffusionStageConfig,
+)
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
 from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name
 from vllm_omni.platforms import current_omni_platform
 
 PROMPT = "A brown and white dog is running on the grass"
-MODEL_NAME = "tencent/HunyuanImage-3.0"
+MODEL_NAME = os.environ.get("HUNYUAN_IMAGE3_MODEL", "tencent/HunyuanImage-3.0-Instruct")
 LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32"
-DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml")
+_PRODUCTION_DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml")
+# Keep the production Mooncake connector, but use its TCP transport for this
+# single-host test.  The shared server's RDMA device is not usable reliably;
+# TCP still exercises the Mooncake AR→DiT data plane without changing deploy
+# semantics in production.
+DEPLOY_CONFIG_PATH = modify_stage_config(
+    _PRODUCTION_DEPLOY_CONFIG_PATH,
+    updates={"connectors.rdma_connector.extra.protocol": "tcp"},
+)
 
-pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion]
+pytestmark = [
+    pytest.mark.advanced_model,
+    pytest.mark.diffusion,
+]
 
 # System prompt type. Options: None, dynamic, en_vanilla, en_recaption, en_think_recaption, en_unified
 # Below are the CLIP embedding tensors from the official HunyuanImage model (seed=1234, prompt: "A brown and white dog is running on the grass").
@@ -285,25 +307,97 @@ def omni() -> Generator[Omni, None, None]:
         yield runner.omni
 
 
+@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+def test_structured_mixed_pipeline_config_reaches_runtime(omni: Omni) -> None:
+    """Mixed-pipeline deploy settings reach both terminal engine configs."""
+    engine = omni.engine
+
+    # Topology, connector selection, and placement remain Omni-owned.
+    resolved_stages = engine.stage_configs
+    assert len(resolved_stages) == 2
+    ar_stage, diffusion_stage = resolved_stages
+    assert isinstance(ar_stage, VllmOmniARStageConfig)
+    assert ar_stage.stage_id == 0
+    assert ar_stage.model_stage == "AR"
+    assert ar_stage.final_output_type == "text"
+    assert ar_stage.runtime_config.devices == "0,1"
+    assert ar_stage.connector_config.output_connectors == {"to_stage_1": "rdma_connector"}
+    assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig)
+    assert diffusion_stage.stage_id == 1
+    assert diffusion_stage.model_stage == "dit"
+    assert diffusion_stage.final_output_type == "image"
+    assert diffusion_stage.input_sources == [0]
+    assert diffusion_stage.runtime_config.devices == "2,3"
+    assert diffusion_stage.connector_config.input_connectors == {"from_stage_0": "rdma_connector"}
+
+    # Stage 0 is materialized as the VllmConfig actually passed to its engine.
+    ar_vllm_config, diffusion_vllm_config = engine.stage_vllm_configs
+    assert ar_vllm_config is not None
+    assert diffusion_vllm_config is None
+    assert ar_vllm_config.model_config.enforce_eager is True
+    assert ar_vllm_config.model_config.trust_remote_code is True
+    assert ar_vllm_config.parallel_config.tensor_parallel_size == 2
+    assert ar_vllm_config.scheduler_config.max_num_seqs == 1
+    assert ar_vllm_config.scheduler_config.max_num_batched_tokens == 32768
+    assert ar_vllm_config.cache_config.gpu_memory_utilization == 0.95
+    ar_omni_config = ar_vllm_config.model_config.omni_config
+    assert ar_omni_config.omni_kv_config["need_send_cache"] is True
+
+    expected_ar_sampling = {
+        "temperature": 0.0,
+        "top_p": 1,
+        "top_k": -1,
+        "max_tokens": 8192,
+        "detokenize": True,
+        "skip_special_tokens": False,
+        "include_stop_str_in_output": True,
+    }
+    ar_sampling, diffusion_sampling = engine.default_sampling_params_list
+    assert all(getattr(ar_sampling, name) == value for name, value in expected_ar_sampling.items())
+    assert diffusion_sampling.num_inference_steps == 50
+    assert diffusion_sampling.guidance_scale == 0
+
+    # Stage 1 is a diffusion engine, so its terminal owner is the launched
+    # OmniDiffusionConfig rather than a VllmConfig.
+    od_config = omni.get_diffusion_od_config()
+    assert od_config is not None
+    assert od_config.model == MODEL_NAME
+    assert od_config.enforce_eager is True
+    assert od_config.distributed_executor_backend == "mp"
+    assert od_config.parallel_config.tensor_parallel_size == 2
+    assert od_config.parallel_config.enable_expert_parallel is True
+    assert od_config.max_num_seqs == 1
+    diffusion_kv = od_config.omni_kv_config
+    assert diffusion_kv["need_recv_cache"] is True
+    assert diffusion_kv["enable_kv_async_prefetch"] is True
+    assert diffusion_kv["kv_prefetch_min_free_mem_ratio"] == 0.2
+
+
 def _extract_generated_image(outputs: list[object]) -> Image.Image:
     if not outputs:
         raise AssertionError("No outputs were returned from Omni.generate()")
 
-    first_output = outputs[0]
-    if images := getattr(first_output, "images", None):
-        return images[0]
-
-    request_output = first_output
-    if request_output is not None and (images := getattr(request_output, "images", None)):
-        return images[0]
+    for output in outputs:
+        if images := getattr(output, "images", None):
+            return images[0]
+        multimodal_output = getattr(output, "multimodal_output", None)
+        if isinstance(multimodal_output, dict) and (images := multimodal_output.get("images")):
+            return images[0]
 
     raise AssertionError("No generated image found in Omni output")
 
 
 def extract_embedding(image: Image.Image, clip_model: CLIPModel, clip_processor: CLIPProcessor) -> torch.Tensor:
-    inputs = clip_processor(images=image.convert("RGB"), return_tensors="pt")
+    # The full CLIP forward API is stable across supported Transformers
+    # versions, but it expects text inputs as well as pixels.  An empty text
+    # prompt is sufficient when only the projected image embedding is needed.
+    inputs = clip_processor(text=[""], images=[image.convert("RGB")], return_tensors="pt", padding=True)
     with torch.inference_mode():
-        features = clip_model.get_image_features(**inputs)
+        # Use the public CLIP forward output, matching the shared CLIPScorer
+        # helper.  ``get_image_features`` changed from Tensor to
+        # BaseModelOutputWithPooling in Transformers 5.x.
+        features = clip_model(**inputs).image_embeds
         features = F.normalize(features, p=2, dim=-1)
     return features.squeeze(0)
 
@@ -320,16 +414,36 @@ def compare_semantic(
 
 
 def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
-    generator_device = current_omni_platform.device_type or "cuda"
-    sampling_params = OmniDiffusionSamplingParams(
-        seed=1234,
-        generator=torch.Generator(device=generator_device).manual_seed(1234),
-        num_outputs_per_prompt=1,
-    )
-    if use_system_prompt is not None:
-        sampling_params.extra_args = {"use_system_prompt": use_system_prompt}
+    from transformers import AutoTokenizer
 
-    outputs = omni.generate({"prompt": PROMPT}, sampling_params)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    prompt_tokens = build_prompt_tokens(
+        PROMPT,
+        tokenizer,
+        task="t2i",
+        bot_task="think",
+        sys_type=use_system_prompt or "None",
+    )
+    stop_token_ids = resolve_stop_token_ids(task="t2i", bot_task="think", tokenizer=tokenizer, image_size="1024x1024")
+    params_list = copy.deepcopy(omni.default_sampling_params_list)
+    generator_device = current_omni_platform.device_type or "cuda"
+    for params in params_list:
+        if isinstance(params, OmniDiffusionSamplingParams):
+            params.seed = 1234
+            params.generator = torch.Generator(device=generator_device).manual_seed(1234)
+            params.num_outputs_per_prompt = 1
+        elif hasattr(params, "stop_token_ids"):
+            params.stop_token_ids = stop_token_ids
+
+    prompt = {
+        "prompt": PROMPT,
+        "prompt_token_ids": prompt_tokens.token_ids,
+        "use_system_prompt": prompt_tokens.system_prompt_type,
+        "modalities": ["image"],
+        "height": 1024,
+        "width": 1024,
+    }
+    outputs = omni.generate(prompt, params_list)
     return _extract_generated_image(outputs)
 
 
@@ -342,7 +456,7 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
 MIN_SEMANTIC_SIMILARITY = 0.5
 
 
-@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
+@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
 @pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES)
 def test_system_prompt_scores(
     omni: Omni,
@@ -378,7 +492,7 @@ def _load_text_to_image_module():
     return module
 
 
-@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
+@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
 def test_shared_script_ar_path_reaches_generation(
     omni: Omni,
     clip_bundle: tuple[CLIPModel, CLIPProcessor],

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -17,6 +17,10 @@ import pytest
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_video
 from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.config.omni_config import (
+    VllmOmniARStageConfig,
+    VllmOmniGenerationStageConfig,
+)
 from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
@@ -25,6 +29,7 @@ thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"]
 # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
 # platforms: section. Only CUDA needs an extra enforce_eager tweak.
 _CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
+_PRODUCTION_DEPLOY = get_deploy_config_path("qwen3_omni_moe.yaml")
 
 
 def get_cuda_graph_config():
@@ -46,6 +51,7 @@ else:
 
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+production_test_params = [(model, _PRODUCTION_DEPLOY) for model in models]
 # we can use the same config for a model that only has thinker (i.e., does not
 # enable audio output) because the resolver should figure out that it doesn't
 # need the full pipeline based on the HF config.
@@ -59,6 +65,85 @@ def get_question(prompt_type="video"):
     return prompts.get(prompt_type, prompts["video"])
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_runner", production_test_params, indirect=True)
+def test_structured_multistage_config_reaches_runtime(omni_runner, offline_client) -> None:
+    """Deploy settings reach the materialized configs used by live stages."""
+    engine = omni_runner.omni.engine
+
+    # The resolved typed stages remain the Omni-owned source for topology and
+    # placement, which are not fields of the terminal vLLM config.
+    resolved_stages = engine.stage_configs
+    assert len(resolved_stages) == 3
+    thinker, talker, code2wav = resolved_stages
+    assert isinstance(thinker, VllmOmniARStageConfig)
+    assert isinstance(talker, VllmOmniARStageConfig)
+    assert isinstance(code2wav, VllmOmniGenerationStageConfig)
+    assert [stage.stage_id for stage in resolved_stages] == [0, 1, 2]
+    assert [stage.model_stage for stage in resolved_stages] == ["thinker", "talker", "code2wav"]
+    assert code2wav.final_output_type == "audio"
+    assert thinker.runtime_config.devices == "0"
+    assert talker.runtime_config.devices == code2wav.runtime_config.devices == "1"
+    assert talker.connector_config.input_connectors == {"from_stage_0": "connector_of_shared_memory"}
+    assert code2wav.connector_config.input_connectors == {"from_stage_1": "connector_of_shared_memory"}
+
+    # Engine-owned deploy settings must survive the startup projection into
+    # the actual VllmConfig objects retained by the stage pools.
+    assert len(engine.stage_vllm_configs) == 3
+    thinker_vllm, talker_vllm, code2wav_vllm = engine.stage_vllm_configs
+    assert thinker_vllm is not None
+    assert talker_vllm is not None
+    assert code2wav_vllm is not None
+    assert thinker_vllm.scheduler_config.max_num_seqs == 64
+    assert talker_vllm.scheduler_config.max_num_seqs == 64
+    assert code2wav_vllm.scheduler_config.max_num_seqs == 64
+    assert thinker_vllm.scheduler_config.max_num_batched_tokens == 32768
+    assert talker_vllm.scheduler_config.max_num_batched_tokens == 32768
+    assert code2wav_vllm.scheduler_config.max_num_batched_tokens == 65536
+    assert thinker_vllm.cache_config.gpu_memory_utilization == 0.9
+    assert talker_vllm.cache_config.gpu_memory_utilization == 0.6
+    assert code2wav_vllm.cache_config.gpu_memory_utilization == 0.1
+    assert all(config.cache_config.enable_prefix_caching is False for config in engine.stage_vllm_configs)
+    assert all(config.model_config.trust_remote_code is True for config in engine.stage_vllm_configs)
+    assert code2wav_vllm.scheduler_config.enable_chunked_prefill is False
+    assert code2wav_vllm.scheduler_config.async_scheduling is False
+
+    # Sampling defaults are consumed from StageClient metadata, rather than
+    # from either the resolver output or VllmConfig.
+    expected_sampling = (
+        {"temperature": 0.0, "max_tokens": 2048},
+        {"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "repetition_penalty": 1.05},
+        {
+            "temperature": 0.0,
+            "top_p": 1.0,
+            "top_k": -1,
+            "max_tokens": 65536,
+            "repetition_penalty": 1.1,
+        },
+    )
+    assert len(engine.default_sampling_params_list) == len(expected_sampling)
+    for runtime_params, expected in zip(engine.default_sampling_params_list, expected_sampling, strict=True):
+        # SamplingParams normalization may add backend defaults (for example
+        # ``detokenize=True``). Verify deploy-owned values without rejecting
+        # those non-lossy normalized fields.
+        assert all(getattr(runtime_params, name) == value for name, value in expected.items())
+
+    if current_omni_platform.is_cuda():
+        assert thinker_vllm.model_config.enforce_eager is False
+        assert talker_vllm.model_config.enforce_eager is False
+
+    # Exercise the materialized stages and their connector path with a real
+    # request; the assertions above are not merely resolver-shape checks.
+    offline_client.send_omni_request(
+        {
+            "prompts": "Answer with one short sentence: what is the capital of China?",
+            "modalities": ["text"],
+        }
+    )
+
+
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)

tests/e2e/offline_inference/test_wan22_t2v.py

@@ -5,6 +5,7 @@ import pytest
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OfflineOmniClient
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -12,6 +13,31 @@ PROMPT = "Two anthropomorphic cats in boxing gear on a spotlighted stage."
 NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
 
 
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1})
+@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
+def test_structured_diffusion_config_reaches_runtime(omni_runner) -> None:
+    """The default WAN resolver launches the typed diffusion config itself."""
+    omni = omni_runner.omni
+    resolved_stages = omni.engine.stage_configs
+    assert len(resolved_stages) == 1
+    stage_config = resolved_stages[0]
+
+    assert isinstance(stage_config, VllmOmniDiffusionStageConfig)
+    assert stage_config.stage_id == 0
+    assert stage_config.model_stage == "dit"
+    assert stage_config.final_output_type == "video"
+
+    assert omni.engine.stage_vllm_configs == [None]
+    od_config = omni.get_diffusion_od_config()
+    assert od_config is not None
+    assert od_config.stage_id == 0
+    assert od_config.model == MODEL
+    assert od_config.model_class_name == "WanPipeline"
+    assert od_config.output_type == "pil"
+
+
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1})

tests/engine/test_async_omni_engine_abort_ack.py

@@ -66,9 +66,10 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture):
         assert msg.rpc_id == "abort-rpc-1"
         assert not task.done()
 
-        rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
+        partial_output = SimpleNamespace(request_id="req-a", finished=True)
+        rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True, abort_outputs=[partial_output]))
         result = await task
-        assert result == []
+        assert result == [partial_output]
 
     try:
         asyncio.run(_run())

tests/engine/test_async_omni_engine_stage_init.py

@@ -14,6 +14,7 @@ import pytest
 from omegaconf import OmegaConf
 from vllm.v1.engine.utils import EngineZmqAddresses
 
+from vllm_omni.config.omni_config import OmniStageRuntimeConfig
 from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine import omni_engine_base as async_omni_engine_module
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -24,6 +25,7 @@ from vllm_omni.engine.stage_init_utils import (
     build_stage0_input_processor,
     compute_replica_layout,
     split_devices_for_replicas,
+    stage_runtime_env,
 )
 from vllm_omni.engine.stage_runtime import StageRuntime
 
@@ -56,6 +58,16 @@ class _FakeStageConfig:
     runtime: _FakeRuntimeConfig
 
 
+def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
+    env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV"
+    monkeypatch.delenv(env_key, raising=False)
+
+    with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "typed-value"})):
+        assert os.environ[env_key] == "typed-value"
+
+    assert env_key not in os.environ
+
+
 def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch):
     engine = object.__new__(AsyncOmniEngine)
     engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True)
@@ -254,6 +266,50 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_
     assert captured["client_config"] is engine._client_config
 
 
+def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch):
+    import vllm_omni.engine.omni_engine_base as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.stage_configs = [types.SimpleNamespace()]
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.single_stage_mode = False
+    engine.async_chunk = False
+    engine.tokenizer = None
+    engine._single_stage_id_filter = None
+    engine._omni_master_address = None
+    engine._omni_master_port = None
+    engine._omni_dp_size_local = 1
+    engine._omni_heartbeat_timeout = 30.0
+    engine._omni_lb_policy = "random"
+    engine.request_queue = types.SimpleNamespace()
+    engine._log_stats = False
+    engine._parallel_stage_init = False
+
+    prompt_transform = object()
+    client = types.SimpleNamespace(
+        prompt_transform_func=prompt_transform,
+        prompt_expand_func=None,
+        default_sampling_params=types.SimpleNamespace(),
+        final_output=True,
+        final_output_type="text",
+        stage_type="llm",
+        model_stage="text_encoder",
+        is_comprehension=True,
+    )
+    pool = types.SimpleNamespace(
+        stage_client=client,
+        stage_vllm_config=None,
+        output_processor=None,
+    )
+    runtime = types.SimpleNamespace(stage_pools=[pool], initialize=lambda: None)
+    monkeypatch.setattr(engine_mod, "create_stage_runtime", lambda **_kwargs: runtime)
+
+    engine._initialize_stages(stage_init_timeout=7)
+
+    assert engine.prompt_transform_func is prompt_transform
+
+
 def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -2029,6 +2085,51 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology():
     assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
 
 
+def test_inject_kv_stage_info_updates_typed_connector_config():
+    from vllm_omni.config.omni_config import (
+        OmniStageConnectorConfig,
+        OmniStageDiffusionParallelConfig,
+        VllmOmniDiffusionStageConfig,
+    )
+    from vllm_omni.config.stage_config import StageExecutionType, StagePipelineConfig
+    from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
+    from vllm_omni.entrypoints.utils import inject_omni_kv_config
+
+    stage0 = VllmOmniDiffusionStageConfig(
+        stage_pipeline_config=StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+        ),
+        connector_config=OmniStageConnectorConfig(
+            omni_kv_config={
+                "need_send_cache": True,
+                "omni_from_stage": "0",
+                "omni_to_stage": "1",
+            }
+        ),
+        parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=4),
+    )
+    stage1 = VllmOmniDiffusionStageConfig(
+        stage_pipeline_config=StagePipelineConfig(
+            stage_id=1,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(0,),
+        ),
+        connector_config=OmniStageConnectorConfig(omni_kv_config={"need_recv_cache": True}),
+        parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=2),
+    )
+
+    inject_omni_kv_config(stage0, {"kv_connector": "P2pNcclConnector"}, "0", "1")
+    inject_kv_stage_info(stage0, 0, [stage0, stage1])
+
+    assert stage0.connector_config.omni_kv_config["stage_id"] == 0
+    assert stage0.connector_config.omni_kv_config["connector_config"] == {"kv_connector": "P2pNcclConnector"}
+    assert stage0.connector_config.omni_kv_config["engine_input_source"] == []
+    assert stage0.connector_config.omni_kv_config["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
+
+
 def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
     """ROCm default attention logic only applies to LLM stages, not diffusion."""
     from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata
@@ -2262,3 +2363,80 @@ def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes,
     with pytest.raises(ConfigResolvedError):
         engine.__init__("test-model")
     assert engine.async_chunk is expected
+
+
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+def test_dist_stage_runtime_applies_local_dp_to_stage_config(typed):
+    from vllm_omni.config.omni_config import VllmOmniARStageConfig
+    from vllm_omni.config.stage_config import StagePipelineConfig
+    from vllm_omni.engine.stage_runtime import DistStageRuntime
+
+    if typed:
+        stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar"))
+        runtime_cfg = stage.runtime_config
+    else:
+        runtime_cfg = types.SimpleNamespace(num_replicas=1)
+        stage = types.SimpleNamespace(stage_id=0, runtime=runtime_cfg)
+    runtime = DistStageRuntime(
+        stage_configs=[stage],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+        single_stage_id_filter=0,
+        omni_master_address="127.0.0.1",
+        omni_master_port=12345,
+        omni_dp_size_local=2,
+    )
+
+    runtime._validate_single_stage_mode_replica_constraints()
+
+    assert runtime_cfg.num_replicas == 2
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2, 3])
+def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas):
+    from vllm_omni.config.config_factory import StageConfigFactory
+    from vllm_omni.engine import stage_runtime as runtime_module
+
+    stage = StageConfigFactory.create_typed_default_diffusion(
+        "generic-diffusion", {"model_class_name": "QwenImagePipeline"}
+    ).stage_configs[0]
+    devices = ",".join(str(i) for i in range(num_replicas))
+    stage.runtime_config.devices = devices
+    stage.runtime_config.num_replicas = num_replicas
+    runtime = StageRuntime(
+        stage_configs=[stage],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+    )
+    mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
+    mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
+    mocker.patch.object(runtime, "_stage_device_scope", side_effect=lambda *_: contextlib.nullcontext())
+    client = mocker.Mock()
+    launch = mocker.patch.object(runtime_module, "launch_diffusion_stage_replica", return_value=(client, None))
+
+    # Replanning must not inherit the first replica's narrowed device slice.
+    for _ in range(2):
+        counts, device_map = compute_replica_layout([stage])
+        plans = runtime._build_logical_stage_init_plans(
+            omni_transfer_config=None,
+            replicas_per_stage=counts,
+            replica_devices_map=device_map,
+        )
+        replicas = plans[0].replicas
+        assert stage.runtime_config.devices == devices
+        assert len({id(plan.stage_cfg) for plan in replicas}) == num_replicas
+        for i, plan in enumerate(replicas):
+            assert plan.metadata.runtime_cfg is plan.stage_cfg.runtime_config
+            assert plan.stage_cfg.runtime_config.devices == str(i)
+            if num_replicas > 1:
+                assert plan.stage_cfg is not stage
+            assert runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) is client
+            assert launch.call_args.kwargs["stage_config"] is plan.stage_cfg
+            assert launch.call_args.kwargs["metadata"] is plan.metadata
+            assert launch.call_args.kwargs["use_inline"] is (num_replicas == 1)
+
+    assert launch.call_count == 2 * num_replicas

tests/engine/test_parallel_stage_init.py

@@ -306,6 +306,23 @@ def test_check_admission_admits_diffusion_with_explicit_util():
     assert ledgers[0].kv_budget_bytes == int(cap * 0.3)
 
 
+def test_check_admission_reads_typed_diffusion_cache_config():
+    cap = 80 * 1024**3
+    diff = _local_diffusion_replica(0, 0, "0", util=None)
+    diff.stage_cfg.cache_config = types.SimpleNamespace(gpu_memory_utilization=0.35)
+    plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff])
+
+    ledgers = check_admission(
+        [plan],
+        resolve_physical_devices=lambda _r: [0],
+        device_total_memory=lambda _d: cap,
+        external_reserve_bytes=0,
+        safety_margin_bytes=0,
+    )
+
+    assert ledgers[0].kv_budget_bytes == int(cap * 0.35)
+
+
 def test_check_admission_unresolved_local_raises():
     """A local replica invisible to the ledger must fail admission, not skip it:
     it would otherwise initialize concurrently with unbounded demand."""

tests/engine/test_single_stage_mode.py

@@ -16,7 +16,8 @@ from pytest_mock import MockerFixture
 from vllm.v1.engine.utils import EngineZmqAddresses
 
 from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
+from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig
+from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, StageExecutionType, StagePipelineConfig
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
 from vllm_omni.engine.stage_engine_startup import (
@@ -57,6 +58,10 @@ def _make_llm_plan(
     vllm_config: Any | None = None,
 ) -> LogicalStageInitPlan:
     stage_cfg = _make_stage_cfg(stage_id)
+    if launch_mode == "remote":
+        stage_cfg = VllmOmniARStageConfig(
+            stage_pipeline_config=StagePipelineConfig(stage_id=stage_id, model_stage="thinker"),
+        )
     metadata = SimpleNamespace(
         stage_id=stage_id,
         stage_type="llm",
@@ -98,6 +103,16 @@ def _make_diffusion_plan(
     launch_mode: str,
 ) -> LogicalStageInitPlan:
     stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion")
+    if launch_mode == "remote":
+        stage_cfg = VllmOmniDiffusionStageConfig(
+            stage_pipeline_config=StagePipelineConfig(
+                stage_id=stage_id,
+                model_stage="diffusion",
+                execution_type=StageExecutionType.DIFFUSION,
+                final_output=True,
+                final_output_type="image",
+            ),
+        )
     metadata = SimpleNamespace(
         stage_id=stage_id,
         stage_type="diffusion",
@@ -960,6 +975,7 @@ class TestSingleStageReplicaInitialization:
         assert mock_connect.call_args.kwargs["stage_id"] == 1
         assert mock_connect.call_args.kwargs["replica_id"] == 0
         assert client_kwargs["log_stats"] is True
+        assert client_kwargs["metadata"].model_stage == "thinker"
         assert events == ["enter", "exit", "attach"]
 
     def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture):
@@ -1125,11 +1141,9 @@ class TestSingleStageReplicaInitialization:
 
             return _ctx()
 
-        remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata
         plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0]
         sentinel_client = SimpleNamespace()
 
-        mocker.patch.object(runtime_mod, "extract_legacy_stage_metadata", return_value=remote_metadata)
         mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect)
         mock_from_addresses = mocker.patch(
             "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses",
@@ -1146,6 +1160,7 @@ class TestSingleStageReplicaInitialization:
             replica_id=0,
         )
         mock_from_addresses.assert_called_once()
+        assert mock_from_addresses.call_args.args[0].final_output_type == "image"
 
     def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture):
         import vllm_omni.engine.stage_runtime as runtime_mod

tests/engine/test_stage_device_layout.py

@@ -157,6 +157,43 @@ def test_replica_guard_and_splitter_share_local_world_size(devices, expected):
     assert replica_devices_map == {0: expected}
 
 
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize("local_dp", [None, 0, 1, 2])
+@pytest.mark.parametrize("pool_devices", [False, True], ids=["template", "pool"])
+def test_replica_layout_preserves_local_dp_semantics(typed, local_dp, pool_devices):
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig
+    from vllm_omni.engine.stage_engine_startup import get_headless_replica_devices
+
+    engine_args = {
+        "tensor_parallel_size": 2,
+        "data_parallel_size": 4,
+        "data_parallel_size_local": local_dp,
+        "pipeline_parallel_size": 2,
+    }
+    width = 2 * max(1, local_dp if local_dp is not None else 4) * 2
+    devices = ",".join(str(i) for i in range(width * (2 if pool_devices else 1)))
+    if typed:
+        pipeline = PipelineConfig(
+            model_type="test",
+            model_arch="TestModel",
+            stages=(StagePipelineConfig(stage_id=0, model_stage="ar", final_output=True),),
+        )
+        stage = VllmOmniConfig.from_pipeline_config(
+            pipeline,
+            cli_overrides={
+                f"stage_0_{key}": value for key, value in {**engine_args, "devices": devices, "num_replicas": 2}.items()
+            },
+        ).stage_configs[0]
+    else:
+        stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args)
+    expected = [",".join(str(i) for i in range(r * width, (r + 1) * width)) for r in range(2)]
+
+    assert get_stage_devices_per_replica(stage) == width
+    assert compute_replica_layout([stage]) == ([2], {0: expected})
+    assert get_headless_replica_devices(stage, stage_id=0, omni_dp_size_local=2) == expected
+
+
 def test_non_tp_layout_error_uses_generic_guidance():
     """A PP mismatch must not be reported as a top-level TP broadcast."""
     stage = _stage(0, devices="0")

tests/engine/test_stage_engine_args.py

@@ -95,6 +95,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
         "num_weight_load_threads",
         "omni_kv_config",
         "parallel_config",
+        "pooling_output_decoder",
         "silence_ban_frames",
         "subtalker_sampling_params",
         "task_type",
@@ -366,6 +367,21 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
     assert inherited_defaults.isdisjoint(engine_args)
 
 
+def test_typed_llm_projection_omits_diffusion_only_and_process_only_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)
+
+    assert {
+        "enable_multithread_weight_load",
+        "num_weight_load_threads",
+        "disable_autocast",
+        "log_level",
+    }.isdisjoint(engine_args)
+
+
 def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
     monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
 

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -127,6 +127,13 @@ async def test_omni_run_server_worker_forwards_client_config(monkeypatch):
 # TODO: extend this set and the expected census if @router.head / @router.options appear.
 _HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
 
+
+def test_profiler_endpoints_detect_typed_stage_config():
+    stage = SimpleNamespace(profiler_config=SimpleNamespace(profiler="torch"))
+
+    assert api_server._should_enable_profiler_endpoints([stage])
+
+
 # Exact Omni router census on main. Update intentionally when routes change.
 _EXPECTED_ROUTER_ROUTES = {
     ("POST", "/v1/chat/completions"),

tests/entrypoints/openai_api/test_audex_serving_guards.py

@@ -15,6 +15,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm_omni.config.stage_config import StagePipelineConfig
 from vllm_omni.entrypoints.openai.serving_speech import (
     _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES,
     OmniOpenAIServingSpeech,
@@ -251,9 +252,28 @@ class TestAudexStageDetection:
         stage = serving._find_tts_stage()
         assert stage is not None and stage.engine_args.model_stage == "audex_thinker"
 
+    def test_typed_pipeline_detected_without_legacy_engine_args(self):
+        serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+        stages = [
+            SimpleNamespace(
+                stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage),
+                model_config=SimpleNamespace(model_arch=None),
+                worker_type="generation",
+            )
+            for index, model_stage in enumerate(("audex_thinker", "audex_code2wav"))
+        ]
+        serving.engine_client = SimpleNamespace(stage_configs=stages)
+
+        stage = serving._find_tts_stage()
+        serving._tts_stage = stage
+
+        assert stage is stages[0]
+        assert serving._detect_tts_model_type() == "audex"
+
 
 def test_guided_injection_sets_guided_temperature():
     """Guided requests decode at the measured CFG temperature (0.05)."""
+
     serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
     cond = TestAudexCfgInjection()._cond_prompt()
     params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
@@ -265,6 +285,31 @@ def test_guided_injection_sets_guided_temperature():
     assert out.temperature == 0.1
 
 
+def test_typed_full_pipeline_uses_full_checkpoint_tokenizer(mocker):
+    from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+
+    serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+    serving._tts_stage = SimpleNamespace(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="audex_omni"),
+        model_config=SimpleNamespace(model_arch=None),
+        worker_type="ar",
+    )
+    serving.engine_client = SimpleNamespace(
+        model_config=SimpleNamespace(model="/models/audex"),
+    )
+    ensure_snapshot = mocker.patch(
+        "vllm_omni.model_executor.models.audex.checkpoint.ensure_audex_snapshot",
+        return_value="/snapshots/audex",
+    )
+    from_pretrained = mocker.patch("transformers.AutoTokenizer.from_pretrained", return_value=object())
+
+    adapter = AudexAdapter(SpeechServingContext(server=serving, engine_client=serving.engine_client))
+    adapter._get_tokenizer()
+
+    ensure_snapshot.assert_called_once_with("/models/audex", profile="full")
+    from_pretrained.assert_called_once_with("/snapshots/audex/checkpoint_folder_full")
+
+
 # ---------------------------------------------------------------- TTA serving surface
 
 from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter  # noqa: E402

tests/entrypoints/openai_api/test_openpi_serving.py

@@ -138,6 +138,36 @@ def test_policy_server_config_reads_stage_config_model_config():
     assert serving.policy_server_config.to_dict() == policy_config
 
 
+@pytest.mark.parametrize("inline", [False, True], ids=["out_of_process", "inline"])
+def test_policy_server_config_reads_typed_runtime(inline):
+    from tests.helpers.stage_config import get_deploy_config_path
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+    typed_stage = VllmOmniConfig.from_pipeline_config(
+        PI0_PIPELINE, deploy_config_path=get_deploy_config_path("pi0.yaml")
+    ).stage_configs[0]
+    policy_config = typed_stage.diffusion_config.model_config["policy_server_config"]
+    stage_client = SimpleNamespace(stage_type="diffusion")
+    if inline:
+        # Worker-enriched inline metadata remains authoritative over the
+        # startup stage config; only the process client needs the fallback.
+        policy_config = {**policy_config, "worker_metadata": "loaded"}
+        stage_client.od_config = SimpleNamespace(model_config={"policy_server_config": policy_config})
+    engine_client = AsyncOmni.__new__(AsyncOmni)
+    engine_client.engine = SimpleNamespace(
+        stage_clients=[stage_client],
+        stage_configs=[typed_stage],
+        stage_vllm_configs=[None],
+        get_diffusion_od_config=lambda: SimpleNamespace(model_class_name="Pi0Pipeline"),
+    )
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    assert serving.policy_server_config.to_dict() == policy_config
+
+
 def test_policy_server_config_reads_omegaconf_stage_config():
     engine_client = SimpleNamespace(
         get_diffusion_od_config=lambda: None,

tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py

@@ -27,7 +27,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 @dataclass
 class _EngineClientStub:
-    stage_configs: list[dict[str, object]]
+    stage_configs: list[object]
 
 
 @dataclass
@@ -149,6 +149,37 @@ async def test_chat_reference_audio_reaches_code2wav_transport(
     assert code2wav_payload.meta.ref_audio_sr == 16000
 
 
+@pytest.mark.asyncio
+async def test_typed_stage_enables_reference_audio(
+    serving_chat: OmniOpenAIServingChat,
+    mocker: MockerFixture,
+) -> None:
+    reference_audio_url = "data:audio/wav;base64,AAAA"
+    reference_waveform = np.array([0.25, -0.5], dtype=np.float32)
+    serving_chat.engine_client = _EngineClientStub(
+        [
+            SimpleNamespace(
+                model_config=SimpleNamespace(model_arch="MiniCPMO45OmniForConditionalGeneration"),
+                model_stage="llm",
+            )
+        ]
+    )
+    serving_chat.model_config = _ModelConfigStub()
+    engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+    media_connector = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+    ).return_value
+    media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000))
+
+    await serving_chat._attach_minicpmo45_reference_audio(
+        engine_prompt,
+        _chat_request(reference_audio_url),
+    )
+
+    media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url)
+    assert MINICPMO45_REFERENCE_AUDIO_KEY in cast(dict[str, object], engine_prompt)
+
+
 @pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"])
 @pytest.mark.asyncio
 async def test_reference_audio_accepts_request_extra_mappings(

tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for MiniCPM-o 4.5 audio-streaming output-kind coercion.
 
 Chat ``stream=true`` coerces every AR stage to DELTA. For MiniCPM-o 4.5 that
@@ -33,6 +34,13 @@ def _stage(model_arch: str, model_stage: str) -> SimpleNamespace:
     return SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage=model_stage))
 
 
+def _typed_stage(model_arch: str, model_stage: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        model_config=SimpleNamespace(model_arch=model_arch),
+        model_stage=model_stage,
+    )
+
+
 def _minicpmo45_stages():
     return [_stage(MINICPMO45_ARCH, "llm"), _stage(MINICPMO45_ARCH, "tts")]
 
@@ -50,6 +58,20 @@ def test_thinker_forced_final_only_talker_stays_delta(serving_chat):
     assert out[1].output_kind is RequestOutputKind.DELTA
 
 
+def test_typed_stages_force_thinker_final_only(serving_chat):
+    serving_chat.engine_client = SimpleNamespace(
+        stage_configs=[
+            _typed_stage(MINICPMO45_ARCH, "llm"),
+            _typed_stage(MINICPMO45_ARCH, "tts"),
+        ]
+    )
+
+    out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(), ["text", "audio"])
+
+    assert out[0].output_kind is RequestOutputKind.FINAL_ONLY
+    assert out[1].output_kind is RequestOutputKind.DELTA
+
+
 @pytest.mark.parametrize("modalities", [None, [], ["text"]])
 def test_non_audio_request_is_untouched(serving_chat, modalities):
     serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages())

tests/entrypoints/openai_api/test_serving_speech.py

@@ -27,6 +27,7 @@ from pydantic import ValidationError
 from pytest_mock import MockerFixture
 from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse
 
+from vllm_omni.config.stage_config import StagePipelineConfig
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
 from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
@@ -1039,6 +1040,38 @@ class TestTTSMethods:
             asyncio.run(server._prepare_speech_generation(request))
         server.shutdown()
 
+    def test_prepare_speech_rejects_typed_non_tts_omni_model(self, mocker: MockerFixture):
+        """Typed Qwen3-Omni stages must hit the same talker guard as legacy stages."""
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.tts_max_instructions_length = None
+        mock_engine_client.stage_configs = [
+            SimpleNamespace(
+                stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage),
+                model_config=SimpleNamespace(model_arch=None),
+                worker_type=worker_type,
+            )
+            for index, (model_stage, worker_type) in enumerate(
+                (
+                    ("thinker", "ar"),
+                    ("talker", "ar"),
+                    ("code2wav", "generation"),
+                )
+            )
+        ]
+
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        request = OpenAICreateSpeechRequest(input="Hello world")
+        with pytest.raises(ValueError, match="only supported for dedicated TTS models"):
+            asyncio.run(server._prepare_speech_generation(request))
+        server.shutdown()
+
     def test_estimate_prompt_len_fallback(self, speech_server):
         """Test prompt length estimation falls back to 2048 when model is unavailable."""
         tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]}

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -22,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import (
     detect_tts_model_type,
     resolve_adapter,
 )
+from vllm_omni.entrypoints.openai.tts_adapters.base import resolve_stage_model_path
 from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter
 from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter
 from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
@@ -99,6 +100,31 @@ def test_resolve_unknown_returns_none():
     assert resolve_adapter(None) is None
 
 
+def test_stage_model_path_prefers_typed_override():
+    engine_client = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model")),
+            SimpleNamespace(
+                model_config=SimpleNamespace(model="typed-stage-model"),
+            ),
+        ],
+        model="served-model",
+    )
+
+    assert resolve_stage_model_path(engine_client) == "typed-stage-model"
+
+
+def test_stage_model_path_falls_back_to_legacy_then_served_model():
+    legacy_client = SimpleNamespace(
+        stage_configs=[SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model"))],
+        model="served-model",
+    )
+    served_client = SimpleNamespace(stage_configs=[SimpleNamespace()], model="served-model")
+
+    assert resolve_stage_model_path(legacy_client) == "legacy-stage-model"
+    assert resolve_stage_model_path(served_client) == "served-model"
+
+
 def test_voxcpm2_resolves():
     """VoxCPM2 (the served ``latent_generator`` model) resolves cleanly.
 

tests/entrypoints/openai_api/test_video_server.py

@@ -1135,6 +1135,32 @@ def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_fal
     assert handler.supports_mixed_reference_inputs
 
 
+def test_typed_stage_drives_video_capability_checks():
+    from vllm_omni.config.config_factory import StageConfigFactory
+
+    minimax_stage = StageConfigFactory.create_typed_default_diffusion(
+        "minimax-h3",
+        {"model_class_name": "MiniMaxH3Pipeline"},
+    ).stage_configs[0]
+    minimax_handler = OmniOpenAIServingVideo.for_diffusion(
+        SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)),
+        model_name="minimax-h3",
+        stage_configs=[minimax_stage],
+    )
+    assert minimax_handler.supports_mixed_reference_inputs
+
+    cosmos_stage = StageConfigFactory.create_typed_default_diffusion(
+        "cosmos3",
+        {"model_class_name": "Cosmos3OmniDiffusersPipeline"},
+    ).stage_configs[0]
+    cosmos_handler = OmniOpenAIServingVideo.for_diffusion(
+        SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)),
+        model_name="cosmos3",
+        stage_configs=[cosmos_stage],
+    )
+    assert cosmos_handler.supported_control_upload_types == frozenset({"edge", "blur", "depth", "seg", "wsm"})
+
+
 @pytest.mark.parametrize("model_class_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
 def test_control_upload_capability_is_declared_only_by_cosmos3(test_client, model_class_name):
     handler = test_client.app.state.openai_serving_video
@@ -1192,6 +1218,35 @@ def test_cosmos3_reference_video_limit_uses_v2v_condition_frames():
     assert spec.keep == "first"
 
 
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize(
+    ("num_frames", "extra_params", "expected"),
+    [
+        (189, {"condition_frame_indexes_vision": [0, 2], "condition_video_keep": "last"}, (9, "last")),
+        (189, {"condition_frame_indexes_vision": [0, 2]}, (9, "first")),
+        (5, {"condition_frame_indexes_vision": [0, 20]}, (5, "first")),
+        (None, {"action_mode": "inverse_dynamics", "action_chunk_size": 16}, (17, "first")),
+    ],
+)
+def test_cosmos3_reference_video_decode_policy_with_runtime_configs(typed, num_frames, extra_params, expected):
+    from vllm_omni.config.config_factory import StageConfigFactory
+
+    if typed:
+        stages = list(
+            StageConfigFactory.create_typed_default_diffusion(
+                "cosmos3",
+                {"model_class_name": "Cosmos3OmniDiffusersPipeline"},
+            ).stage_configs
+        )
+    else:
+        stages = _cosmos3_stage_configs()
+    request = VideoGenerationRequest(prompt="Continue this motion.", num_frames=num_frames, extra_params=extra_params)
+
+    spec = _reference_video_decode_spec(request, stages)
+
+    assert (spec.max_frames, spec.keep) == expected
+
+
 def test_cosmos3_reference_video_limit_preserves_action_frames():
     request = VideoGenerationRequest(
         prompt="Predict the action.",

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -2,13 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
 from pydantic import ValidationError
 
 from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
+from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
 from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -117,7 +119,7 @@ def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocke
         trust_remote_code=False,
     )
 
-    assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
+    assert stage_configs[0].diffusion_config.extras["ltx2_use_conv_vae"] is True
 
 
 def test_default_stage_rejects_unknown_nested_parallel_config_key():
@@ -239,6 +241,33 @@ def test_default_stage_config_includes_default_sampling_params():
     }
 
 
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize("sampling_defaults", [{"0": {"guidance_scale": 7.5}}, '{"0":{"guidance_scale":7.5}}'])
+def test_generic_diffusion_sampling_defaults_remain_overridable(typed, sampling_defaults):
+    from vllm_omni.config.yaml_util import create_config
+    from vllm_omni.entrypoints.omni_base import OmniBase
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    kwargs = {"model_class_name": "QwenImagePipeline", "default_sampling_params": sampling_defaults}
+    if typed:
+        stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0]
+        metadata = stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage)
+    else:
+        stage = create_config(StageConfigFactory.create_default_diffusion(kwargs))[0]
+        metadata = stage_init_utils.extract_legacy_stage_metadata(stage)
+    base = OmniBase.__new__(OmniBase)
+    base.engine = SimpleNamespace(num_stages=1, stage_configs=[stage])
+    base.default_sampling_params_list = [metadata.default_sampling_params]
+    base.sampling_constraints_list = base._get_sampling_constraints_list([stage])
+
+    assert base.resolve_sampling_params_list(None)[0].guidance_scale == 7.5
+    requested = OmniDiffusionSamplingParams(guidance_scale=2.0)
+    assert base.resolve_sampling_params_list(requested)[0].guidance_scale == 2.0
+    assert requested.guidance_scale == 2.0
+    assert base.default_sampling_params_list[0].guidance_scale == 7.5
+    assert base.sampling_constraints_list == [{}]
+
+
 def test_default_stage_config_includes_diffusion_attention_backend():
     """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config."""
     stage_cfg = StageConfigFactory.create_default_diffusion(
@@ -726,3 +755,106 @@ def test_default_stage_config_includes_quantization_config():
     stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0]
 
     assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
+
+
+@pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"])
+def test_generic_diffusion_uses_canonical_video_output_type(model_class_name):
+    config = StageConfigFactory.create_typed_default_diffusion(
+        "generic-video",
+        {"model_class_name": model_class_name},
+    )
+
+    assert config.stage_configs[0].final_output_type == "video"
+
+
+def test_generic_diffusion_resolves_structured_stage_without_legacy_conversion(mocker):
+    """Generic diffusion reaches runtime as the structured stage itself."""
+    mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+    mocker.patch(
+        "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+        return_value=(True, "FakeDiffusionPipeline"),
+    )
+
+    resolved = resolve_omni_config(
+        "generic-diffusion",
+        trust_remote_code=False,
+        deploy_config_path=None,
+        cli_overrides={
+            "num_gpus": 4,
+            "tensor_parallel_size": 2,
+            "default_sampling_params": '{"0": {"guidance_scale": 7.5}}',
+        },
+        stage_overrides=None,
+        strategy_config_path=None,
+    )
+
+    assert resolved.pipeline_config is not None
+    assert resolved.pipeline_config.model_type == "generic_diffusion"
+    assert len(resolved.stage_configs) == 1
+    stage = resolved.stage_configs[0]
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.model_config.model == "generic-diffusion"
+    assert stage.model_config.default_sampling_params == {"guidance_scale": 7.5}
+    assert stage.parallel_config.tensor_parallel_size == 2
+    assert stage.parallel_config.data_parallel_size == 2
+    assert stage.parallel_config.world_size == 4
+    assert stage.runtime_config.devices == "0,1,2,3"
+
+
+def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
+    """Standard runtime resolves and starts the typed stage through the real launcher."""
+    from vllm_omni.engine import stage_engine_startup as startup_module
+    from vllm_omni.engine import stage_runtime as runtime_module
+    from vllm_omni.engine.stage_runtime import StageRuntime
+
+    mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+    mocker.patch(
+        "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+        return_value=(True, "FakeDiffusionPipeline"),
+    )
+    resolved = resolve_omni_config(
+        "generic-diffusion",
+        trust_remote_code=False,
+        deploy_config_path=None,
+        cli_overrides={"num_gpus": 1},
+        stage_overrides=None,
+        strategy_config_path=None,
+    )
+    stage = resolved.stage_configs[0]
+    launched: dict[str, Any] = {}
+    client = SimpleNamespace(input_address=None, shutdown=mocker.Mock())
+
+    mocker.patch.object(runtime_module, "prepare_engine_environment")
+    mocker.patch.object(runtime_module, "load_omni_transfer_config_for_model", return_value=None)
+    mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
+    mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
+
+    def _initialize_typed_stage(stage_id, model, stage_config, metadata, **kwargs):
+        launched.update(
+            stage_id=stage_id,
+            model=model,
+            stage_config=stage_config,
+            metadata=metadata,
+            **kwargs,
+        )
+        return client
+
+    mocker.patch.object(startup_module, "initialize_diffusion_stage", side_effect=_initialize_typed_stage)
+
+    runtime = StageRuntime(
+        stage_configs=list(resolved.stage_configs),
+        model="generic-diffusion",
+        config_path="",
+        stage_init_timeout=10,
+        async_chunk=False,
+    )
+    runtime.initialize()
+
+    assert launched["stage_config"] is stage
+    assert isinstance(launched["stage_config"], VllmOmniDiffusionStageConfig)
+    assert launched["metadata"].stage_type == "diffusion"
+    assert launched["metadata"].model_stage == "diffusion"
+    assert launched["use_inline"] is True
+    assert launched["model"] == "generic-diffusion"
+    assert launched["stage_id"] == 0
+    assert runtime.stage_pools[0].clients == [client]

tests/entrypoints/test_omni_entrypoints.py

@@ -17,6 +17,8 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
+from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import StageConfig
 from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
 from vllm_omni.engine.messages import ErrorMessage, OutputMessage
@@ -202,6 +204,49 @@ def _make_base():
     return obj
 
 
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+def test_resolve_sampling_params_list_preserves_stage_constraints(typed):
+    from vllm_omni.config.omni_config import VllmOmniARStageConfig
+    from vllm_omni.config.stage_config import StagePipelineConfig
+
+    base = _make_base()
+    base.engine.num_stages = 1
+    base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])]
+    if typed:
+        stage = VllmOmniARStageConfig(
+            stage_pipeline_config=StagePipelineConfig(
+                stage_id=0,
+                model_stage="dummy-model",
+                sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
+            )
+        )
+    else:
+        stage = StageConfig(
+            stage_id=0,
+            model_stage="dummy-model",
+            sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
+        ).to_omegaconf()
+    base.engine.stage_configs = [stage]
+    base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
+    assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}]
+
+    resolved_defaults = base.resolve_sampling_params_list(None)
+    assert resolved_defaults[0].stop_token_ids == [42]
+
+    caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7])
+    resolved = base.resolve_sampling_params_list(caller_params)
+
+    assert resolved[0] is not caller_params
+    assert resolved[0].seed == 1234
+    assert resolved[0].max_tokens == 7
+    assert resolved[0].detokenize is False
+    assert resolved[0].stop_token_ids == [7, 42]
+    assert 7 in resolved[0]._all_stop_token_ids
+    assert 42 in resolved[0]._all_stop_token_ids
+    assert caller_params.detokenize is True
+    assert caller_params.stop_token_ids == [7]
+
+
 def test_resolve_sampling_params_list_merges_required_stop_tokens():
     base = _make_base()
     base.engine.num_stages = 1
@@ -219,22 +264,13 @@ def test_resolve_sampling_params_list_merges_required_stop_tokens():
     base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
     assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": required_stop_ids}]
 
-    resolved_defaults = base.resolve_sampling_params_list(None)
 
-    assert resolved_defaults[0].stop_token_ids == [99, 151704, 151645]
+def test_sampling_constraints_are_forwarded_by_typed_stage_configs():
+    config = VllmOmniConfig.from_pipeline_config(OMNI_PIPELINES["qwen3_tts"])
 
-    caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[100, 151704])
+    constraints = OmniBase._get_sampling_constraints_list(config.stage_configs)
 
-    resolved = base.resolve_sampling_params_list(caller_params)
-
-    assert resolved[0] is not caller_params
-    assert resolved[0].seed == 1234
-    assert resolved[0].max_tokens == 7
-    assert resolved[0].detokenize is False
-    assert resolved[0].stop_token_ids == [100, 151704, 151645]
-    assert {100, 151704, 151645}.issubset(resolved[0]._all_stop_token_ids)
-    assert caller_params.detokenize is True
-    assert caller_params.stop_token_ids == [100, 151704]
+    assert constraints == [dict(stage.stage_pipeline_config.sampling_constraints) for stage in config.stage_configs]
 
 
 @pytest.mark.parametrize("use_defaults", [False, True])

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -32,15 +32,18 @@ def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch):
         stage_overrides=None,
         strategy_config_path=None,
     )
-    engine_args = resolved.stage_configs[0].engine_args
+    stage = resolved.stage_configs[0]
 
     assert resolved.config_path is not None
     assert resolved.config_path.endswith("vllm_omni/deploy/dreamzero.yaml")
-    assert engine_args.model_class_name == "DreamZeroPipeline"
-    assert engine_args.model_config.policy_server_config.action_space == "joint_position"
+    assert stage.diffusion_config.model_class_name == "DreamZeroPipeline"
+    assert stage.diffusion_config.model_config["policy_server_config"]["action_space"] == "joint_position"
 
 
 def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
+    def _get_hf_config(path, _model):
+        return None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}
+
     monkeypatch.setattr(
         "vllm.transformers_utils.config.get_hf_file_to_dict",
         lambda path, _model, **_kwargs: (

tests/entrypoints/test_serve.py

@@ -9,11 +9,13 @@ import argparse
 from dataclasses import dataclass
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 from vllm.v1.engine.utils import EngineZmqAddresses
 
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
 from vllm_omni.config.resolver import OmniConfigResolution
 from vllm_omni.engine.stage_engine_startup import StageReplicaResources
 from vllm_omni.engine.stage_runtime import StageEngineLaunch
@@ -243,7 +245,11 @@ def test_serve_validate_rejects_sleep_mode_with_multiple_api_servers(mocker: Moc
         cmd.validate(args)
 
 
-def test_build_multi_api_stage_runtime_matches_current_constructor(mocker: MockerFixture) -> None:
+@pytest.mark.parametrize("typed", [False, True])
+@pytest.mark.parametrize("downstream_async_chunk", [False, True])
+def test_build_multi_api_stage_runtime_matches_current_constructor(
+    mocker: MockerFixture, typed: bool, downstream_async_chunk: bool
+) -> None:
     from vllm_omni.entrypoints.cli import serve as serve_module
 
     args = TrackingNamespace(
@@ -256,10 +262,25 @@ def test_build_multi_api_stage_runtime_matches_current_constructor(mocker: Mocke
         ),
         frozenset({"model"}),
     )
-    stage_config = _FakeStageConfig(
-        stage_id=0,
-        engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=False),
-    )
+    if typed:
+        from vllm_omni.config.omni_config import VllmOmniARStageConfig
+        from vllm_omni.config.stage_config import StagePipelineConfig
+
+        stages = [
+            VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=i, model_stage="ar"))
+            for i in range(2)
+        ]
+        stages[1].connector_config.async_chunk = downstream_async_chunk
+    else:
+        stages = [
+            _FakeStageConfig(
+                stage_id=i,
+                engine_args=_FakeStageEngineArgs(
+                    async_chunk=bool(i and downstream_async_chunk), enable_sleep_mode=False
+                ),
+            )
+            for i in range(2)
+        ]
     mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
     mocker.patch(
         "vllm_omni.config.config_factory.with_trust_remote_code_override",
@@ -268,16 +289,20 @@ def test_build_multi_api_stage_runtime_matches_current_constructor(mocker: Mocke
     mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
     mocker.patch(
         "vllm_omni.config.resolver.resolve_omni_config",
-        return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=(stage_config,)),
+        return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=tuple(stages)),
     )
 
     runtime = serve_module._build_multi_api_stage_runtime(args, 2)
 
     assert runtime._client_count == 1
     assert runtime._client_index == 0
+    assert runtime._async_chunk is downstream_async_chunk
 
 
-def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config(mocker: MockerFixture) -> None:
+@pytest.mark.parametrize("typed", [False, True])
+def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config(
+    mocker: MockerFixture, typed: bool
+) -> None:
     from vllm_omni.entrypoints.cli import serve as serve_module
 
     args = TrackingNamespace(
@@ -294,6 +319,13 @@ def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config(moc
         stage_id=3,
         engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=True),
     )
+    if typed:
+        from vllm_omni.config.omni_config import VllmOmniARStageConfig
+        from vllm_omni.config.stage_config import StagePipelineConfig
+
+        typed_stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=3, model_stage="ar"))
+        typed_stage.model_config.enable_sleep_mode = True
+        stage_config = typed_stage
     mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
     mocker.patch(
         "vllm_omni.config.config_factory.with_trust_remote_code_override",
@@ -574,6 +606,26 @@ def test_parse_stage_overrides_invalid_json_raises() -> None:
     assert f"Got: {bad!r}" in message
 
 
+@pytest.mark.parametrize(
+    "payload",
+    ['{"abc": {}}', '{"-1": {}}', '{"1.5": {}}', '{"\uff10": {}}'],
+)
+def test_parse_stage_overrides_rejects_invalid_stage_ids(payload: str) -> None:
+    with pytest.raises(argparse.ArgumentTypeError, match="non-negative integer stage ids"):
+        _parse_stage_overrides(payload)
+
+
+def test_parse_stage_overrides_preserves_arbitrary_override_fields() -> None:
+    parsed = _parse_stage_overrides(
+        '{"0": {"extras": {"ltx2_use_conv_vae": true}, "kv_cache_dtype": "fp8", "typo_field_xyz": 1}}'
+    )
+    assert parsed["0"] == {
+        "extras": {"ltx2_use_conv_vae": True},
+        "kv_cache_dtype": "fp8",
+        "typo_field_xyz": 1,
+    }
+
+
 def test_parse_stage_overrides_rejects_non_dict_top_level() -> None:
     """Top-level must be a JSON object (dict). A list, scalar, or non-dict
     mapping is rejected with a ValueError naming ``--stage-overrides`` and
@@ -946,6 +998,77 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
     assert manager_kwargs["omni_replica_id"] == 0
 
 
+def test_run_headless_generic_diffusion_launches_structured_stage(mocker: MockerFixture) -> None:
+    """Headless resolution starts the typed stage through the real group launcher."""
+    from vllm_omni.engine import stage_engine_startup as startup_module
+
+    mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+    mocker.patch(
+        "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+        return_value=(True, "FakeDiffusionPipeline"),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+    mocker.patch.object(startup_module.stage_init_utils, "load_omni_transfer_config_for_model", return_value=None)
+    mocker.patch.object(
+        startup_module.initialization,
+        "resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    captured: dict[str, Any] = {}
+    od_config = SimpleNamespace()
+
+    def _build_diffusion_config(model, stage_config, metadata):
+        captured.update(model=model, stage_config=stage_config, metadata=metadata)
+        return od_config
+
+    def _launch_replica_group(**kwargs):
+        captured.update(group_kwargs=kwargs)
+        captured["manager"] = kwargs["launch_one"](0)
+
+    def _launch_replica(**kwargs):
+        captured.update(replica_kwargs=kwargs)
+        return SimpleNamespace(exitcode=None)
+
+    mocker.patch.object(startup_module.stage_init_utils, "build_diffusion_config", side_effect=_build_diffusion_config)
+    mocker.patch.object(startup_module, "launch_headless_replica_group", side_effect=_launch_replica_group)
+    mocker.patch.object(startup_module, "launch_headless_diffusion_replica", side_effect=_launch_replica)
+
+    explicit_keys = frozenset(
+        {
+            "model",
+            "stage_id",
+            "omni_master_address",
+            "omni_master_port",
+            "worker_backend",
+            "model_class_name",
+            "num_gpus",
+        }
+    )
+    args = _make_headless_args(
+        explicit_keys=explicit_keys,
+        model="generic-diffusion",
+        model_class_name="FakeDiffusionPipeline",
+        num_gpus=1,
+    )
+
+    run_headless(args)
+
+    stage = captured["stage_config"]
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert captured["metadata"].stage_type == "diffusion"
+    assert captured["metadata"].model_stage == "diffusion"
+    assert captured["model"] == "generic-diffusion"
+    assert captured["group_kwargs"]["stage_id"] == 0
+    assert captured["group_kwargs"]["omni_dp_size_local"] == 1
+    assert captured["group_kwargs"]["per_replica_devices"] == ["0"]
+    assert captured["replica_kwargs"]["stage_config"] is stage
+    assert captured["replica_kwargs"]["stage_id"] == 0
+    assert captured["replica_kwargs"]["omni_master_address"] == "127.0.0.1"
+    assert captured["replica_kwargs"]["omni_master_port"] == 26000
+    assert captured["replica_kwargs"]["od_config"] is od_config
+    assert stage.runtime_config.devices == "0"
+
+
 def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None:
     """A diffusion replica that exits with a non-zero code must surface as a
     RuntimeError from ``run_headless`` (the head needs the signal to roll

tests/entrypoints/test_utils.py

@@ -13,6 +13,14 @@ import torch
 from pytest_mock import MockerFixture
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
+from vllm_omni.config.composable_parallel import (
+    Broadcast,
+    FanInByStage,
+    MeshAxisSpec,
+    RouteByStage,
+    StrategySpec,
+    TakeRank,
+)
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.resolver import (
     OmniConfigResolution,
@@ -21,7 +29,6 @@ from vllm_omni.config.resolver import (
     resolve_omni_config,
 )
 from vllm_omni.config.stage_config import PipelineConfig
-from vllm_omni.config.yaml_util import create_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.utils import (
@@ -356,16 +363,12 @@ class TestResolveOmniConfig:
         )
         assert resolved.config_path is None
         assert len(resolved.stage_configs) == 1
-        engine_args = resolved.stage_configs[0]["engine_args"]
-        assert "dtype" in engine_args
-        assert engine_args["engine_backend"] == engine_backend
-        assert engine_args["revision"] == "pinned-revision"
+        assert resolved.stage_configs[0].diffusion_config.dtype is torch.float32
+        assert resolved.stage_configs[0].diffusion_config.engine_backend == engine_backend
+        assert resolved.stage_configs[0].diffusion_config.revision == "pinned-revision"
 
     def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture):
-        mocker.patch(
-            "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
-            return_value=None,
-        )
+        mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
         resolve_model_class = mocker.patch(
             "vllm_omni.config.resolver.resolve_model_class_name",
             return_value="WanImageToVideoPipeline",
@@ -374,59 +377,98 @@ class TestResolveOmniConfig:
             "vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs",
             return_value={"WanImageToVideoPipeline"},
         )
+        load_model_class = mocker.patch("vllm_omni.config.resolver.DiffusionModelRegistry._try_load_model_cls")
 
         resolved = resolve_omni_config(
             "/models/Wan2.2-I2V",
             trust_remote_code=False,
             deploy_config_path=None,
-            cli_overrides={
-                "diffusion_load_format": "diffusers",
-                "revision": "pinned-revision",
-            },
+            cli_overrides={"diffusion_load_format": "diffusers", "revision": "pinned-revision"},
             stage_overrides=None,
             strategy_config_path=None,
         )
 
+        resolve_model_class.assert_called_once_with("/models/Wan2.2-I2V", "diffusers", "pinned-revision")
+        load_model_class.assert_not_called()
         stage = resolved.stage_configs[0]
-        resolve_model_class.assert_called_once_with(
-            "/models/Wan2.2-I2V",
-            "diffusers",
-            "pinned-revision",
-        )
-        assert stage.engine_args.model_class_name == "WanImageToVideoPipeline"
+        assert stage.diffusion_config.model_class_name == "WanImageToVideoPipeline"
+        assert stage.diffusion_config.revision == "pinned-revision"
         assert stage.final_output_type == "video"
 
+    def test_generic_diffusion_stage_overrides_reach_typed_backend(self, mocker: MockerFixture):
+        from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+        mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+        mocker.patch(
+            "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+            return_value=(True, "LTX2Pipeline"),
+        )
+        resolved = resolve_omni_config(
+            "/models/LTX-2.5-Diffusers",
+            trust_remote_code=False,
+            deploy_config_path=None,
+            cli_overrides={
+                "extras": {"ltx2_use_conv_vae": False, "keep": "global"},
+                "tensor_parallel_size": 1,
+                "diffusion_streaming_output": True,
+            },
+            stage_overrides={
+                "0": {
+                    "extras": {"ltx2_use_conv_vae": True},
+                    "tensor_parallel_size": 2,
+                    "devices": "2,3",
+                    "engine_backend": "custom.backend",
+                },
+            },
+            strategy_config_path=None,
+        )
+        stage = resolved.stage_by_id(0)
+        engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="/models/LTX-2.5-Diffusers")
+
+        assert stage.runtime_config.devices == "2,3"
+        assert stage.parallel_config.world_size == 2
+        assert engine_args["engine_backend"] == "custom.backend"
+        assert engine_args["extras"]["ltx2_use_conv_vae"] is True
+        assert engine_args["extras"]["keep"] == "global"
+        assert engine_args["streaming_output"] is True
+
     def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture):
         endpoint_restriction = SimpleNamespace(name="chat")
+        typed_stage = SimpleNamespace(stage_id=1)
         structured_config = SimpleNamespace(
-            orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"),
+            orchestrator_config=SimpleNamespace(
+                deploy_config_path="/resolved/deploy.yaml",
+                omni_lb_policy="round_robin",
+            ),
             pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)),
+            stage_configs=(typed_stage,),
+            strategy_omni_lb_policy="round-robin",
         )
-        runtime_stage = create_config(
-            {
-                "stage_id": 1,
-                "runtime": {"devices": "1,2,3", "num_replicas": 3},
-                "engine_args": {"model": "dummy-model"},
-            }
-        )
-        legacy_stage = SimpleNamespace(to_omegaconf=lambda: runtime_stage)
         create_structured = mocker.patch(
             "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
             return_value=structured_config,
         )
         create_legacy = mocker.patch(
-            "vllm_omni.config.resolver.StageConfigFactory._resolve_legacy_from_registry",
-            return_value=SimpleNamespace(
-                stage_configs=[legacy_stage],
-                pipeline_config=structured_config.pipeline_config,
-                omni_lb_policy="round_robin",
-            ),
+            "vllm_omni.config.resolver.StageConfigFactory.create_legacy_stage_configs_from_model",
         )
-        strategy_specs = {"stage_1": {"dp": 3}}
+        strategy_specs = {
+            "stage_1": [
+                StrategySpec(
+                    "stage_replica",
+                    MeshAxisSpec("stage_replica", 2),
+                    RouteByStage("round_robin"),
+                    FanInByStage(),
+                )
+            ]
+        }
         load_strategy = mocker.patch(
             "vllm_omni.config.resolver._load_strategy_specs",
             return_value=strategy_specs,
         )
+        apply_strategy = mocker.patch(
+            "vllm_omni.config.composable_parallel.apply_strategy_specs",
+            return_value=SimpleNamespace(omni_lb_policy="round_robin"),
+        )
 
         resolved = resolve_omni_config(
             "dummy-model",
@@ -447,31 +489,63 @@ class TestResolveOmniConfig:
             trust_remote_code=None,
             cli_overrides=expected_overrides,
             deploy_config_path="deploy.yaml",
-        )
-        create_legacy.assert_called_once_with(
-            structured_config.pipeline_config,
-            expected_overrides,
-            "/resolved/deploy.yaml",
             strategy_specs=strategy_specs,
         )
+        create_legacy.assert_not_called()
         load_strategy.assert_called_once_with("strategy.yaml")
+        apply_strategy.assert_not_called()
         assert resolved.config_path == "/resolved/deploy.yaml"
         assert resolved.pipeline_config is structured_config.pipeline_config
-        assert resolved.omni_lb_policy == "round_robin"
+        assert resolved.omni_lb_policy == "round-robin"
         assert resolved.endpoint_restrictions == (endpoint_restriction,)
-        assert resolved.stage_configs == (runtime_stage,)
+        assert resolved.stage_configs == (typed_stage,)
 
-    def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
-        pipeline = OMNI_PIPELINES["qwen3_tts"]
-        assert isinstance(pipeline, PipelineConfig)
+    def test_tp_only_strategy_does_not_report_default_lb_policy_as_derived(self, mocker: MockerFixture):
         structured_config = SimpleNamespace(
-            orchestrator_config=SimpleNamespace(deploy_config_path=None),
-            pipeline_config=pipeline,
+            orchestrator_config=SimpleNamespace(
+                deploy_config_path="/resolved/deploy.yaml",
+                omni_lb_policy="round-robin",
+            ),
+            pipeline_config=SimpleNamespace(endpoint_restrictions=()),
+            stage_configs=(SimpleNamespace(stage_id=0),),
+            strategy_omni_lb_policy=None,
         )
         mocker.patch(
             "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
             return_value=structured_config,
         )
+        mocker.patch(
+            "vllm_omni.config.resolver._load_strategy_specs",
+            return_value={
+                "thinker": [
+                    StrategySpec(
+                        "tp",
+                        MeshAxisSpec("tp", 2),
+                        Broadcast(),
+                        TakeRank(),
+                    )
+                ]
+            },
+        )
+
+        resolved = resolve_omni_config(
+            "dummy-model",
+            trust_remote_code=True,
+            deploy_config_path="deploy.yaml",
+            cli_overrides={"omni_lb_policy": "round-robin"},
+            stage_overrides=None,
+            strategy_config_path="strategy.yaml",
+        )
+
+        assert resolved.omni_lb_policy is None
+
+    def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
+        pipeline = OMNI_PIPELINES["qwen3_tts"]
+        assert isinstance(pipeline, PipelineConfig)
+        mocker.patch(
+            "vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config",
+            return_value=pipeline,
+        )
 
         resolved = resolve_omni_config(
             "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
@@ -483,9 +557,9 @@ class TestResolveOmniConfig:
         )
 
         assert resolved.pipeline_config is not None
-        pipeline_stage_ids = [stage.stage_id for stage in resolved.pipeline_config.stages]
-        runtime_stage_ids = [stage.stage_id for stage in resolved.stage_configs]
-        assert pipeline_stage_ids == runtime_stage_ids
+        assert [stage.stage_id for stage in resolved.pipeline_config.stages] == [
+            stage.stage_id for stage in resolved.stage_configs
+        ]
         assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner"
         assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1
 

vllm_omni/config/composable_parallel/apply.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Apply declarative parallel strategies onto merged stage configs.
 
 This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses
@@ -198,10 +198,41 @@ def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: Ro
 
 
 def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None:
+    declared = set(cfg.l1_owners.keys())
+    if hasattr(stage, "parallel_config") and hasattr(stage, "runtime_config"):
+        parallel = stage.parallel_config
+        for kind, field_name in _ENGINE_FIELD_BY_KIND.items():
+            if kind in declared:
+                existing = getattr(parallel, field_name, None)
+                value = getattr(cfg, field_name)
+                if existing not in (None, 1, value):
+                    _fail(
+                        f"role {role!r}: strategy derives {field_name}={value} "
+                        f"but structured config already set {existing}"
+                    )
+                setattr(parallel, field_name, value)
+        if "ep" in declared and cfg.enable_expert_parallel:
+            setattr(parallel, "enable_expert_parallel", True)
+        if "stage_replica" in declared:
+            existing = stage.runtime_config.num_replicas
+            if existing not in (None, 1, cfg.stage_replica_size):
+                _fail(
+                    f"role {role!r}: strategy derives "
+                    f"num_replicas={cfg.stage_replica_size} but structured config already set {existing}"
+                )
+            stage.runtime_config.num_replicas = cfg.stage_replica_size
+        check_device_layout(
+            stage.runtime_config.devices,
+            tensor_parallel_size=cfg.tensor_parallel_size,
+            data_parallel_size=cfg.data_parallel_size,
+            pipeline_parallel_size=cfg.pipeline_parallel_size,
+            num_replicas=int(stage.runtime_config.num_replicas or 1),
+            role=role,
+        )
+        return
+
     engine_args = stage.yaml_engine_args
     runtime = stage.yaml_runtime
-
-    declared = set(cfg.l1_owners.keys())
     for kind, field_name in _ENGINE_FIELD_BY_KIND.items():
         if kind in declared:
             _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role)
@@ -209,7 +240,6 @@ def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> No
         _set_explicit(engine_args, "enable_expert_parallel", True, role=role)
     if "stage_replica" in declared:
         _set_num_replicas(runtime, cfg.stage_replica_size, role=role)
-
     _check_devices(runtime, cfg, role=role)
 
 

vllm_omni/config/config_factory.py

@@ -25,6 +25,9 @@ from vllm_omni.config.stage_config import (
     DeployConfig,
     PipelineConfig,
     StageConfig,
+    StageDeployConfig,
+    StageExecutionType,
+    StagePipelineConfig,
     StageType,
     build_stage_runtime_overrides,
     load_deploy_config,
@@ -373,6 +376,7 @@ class StageConfigFactory:
         trust_remote_code: bool | None,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None,
+        strategy_specs: Mapping[Any, Any] | None = None,
     ) -> VllmOmniConfig | None:
         """Build the structured Omni config for a model/deploy pair."""
         user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
@@ -387,6 +391,20 @@ class StageConfigFactory:
         if pipeline_cfg is None:
             return None
 
+        # Keep the typed representation's topology in lockstep with the
+        # compatibility launch representation. The aligner is a real extra
+        # stage, not an engine-only setting, so it must be injected before
+        # either representation is built.
+        if user_deploy_config is None:
+            if pipeline_cfg.default_deploy_config_name is not None:
+                default_deploy_path = _DEPLOY_DIR / pipeline_cfg.default_deploy_config_name
+                user_deploy_config = load_deploy_config(default_deploy_path)
+                deploy_config_path = str(default_deploy_path)
+            else:
+                user_deploy_config = DeployConfig()
+        from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage
+
+        pipeline_cfg, user_deploy_config = inject_forced_aligner_stage(pipeline_cfg, user_deploy_config, cli_overrides)
         registry_cli_overrides = with_trust_remote_code_override(
             {**cli_overrides, "model": model},
             trust_remote_code,
@@ -396,6 +414,7 @@ class StageConfigFactory:
             user_deploy_config=user_deploy_config,
             deploy_config_path=deploy_config_path,
             cli_overrides=registry_cli_overrides,
+            strategy_specs=strategy_specs,
         )
 
     @classmethod
@@ -621,13 +640,11 @@ class StageConfigFactory:
             )
 
     @classmethod
-    def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
-        """Build the temporary runtime ABI for a generic diffusion stage.
-
-        The terminal diffusion config owns engine defaults and normalization;
-        this compatibility builder only adds Omni stage topology, request
-        defaults, and device placement.
-        """
+    def _normalize_default_diffusion(
+        cls,
+        kwargs: dict[str, Any],
+    ) -> tuple[dict[str, Any], DiffusionParallelConfig, dict[str, Any], str]:
+        """Normalize inputs shared by typed and compatibility diffusion builders."""
         raw_sampling_params = kwargs.get("default_sampling_params")
         if isinstance(raw_sampling_params, str):
             try:
@@ -643,8 +660,6 @@ class StageConfigFactory:
         if kwargs.get("num_gpus") is not None:
             parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"]))
         engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs)
-        engine_args["parallel_config"] = asdict(parallel_config)
-        engine_args["model_stage"] = "diffusion"
 
         extras = dict(engine_args.get("extras") or {})
         for key, default in (
@@ -665,6 +680,21 @@ class StageConfigFactory:
             final_output_type,
             model_class_name,
         )
+        return engine_args, parallel_config, default_sampling_params, final_output_type
+
+    @classmethod
+    def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
+        """Build the temporary runtime ABI for a generic diffusion stage.
+
+        The terminal diffusion config owns engine defaults and normalization;
+        this compatibility builder only adds Omni stage topology, request
+        defaults, and device placement.
+        """
+        engine_args, parallel_config, default_sampling_params, final_output_type = cls._normalize_default_diffusion(
+            kwargs
+        )
+        engine_args["parallel_config"] = asdict(parallel_config)
+        engine_args["model_stage"] = "diffusion"
 
         return [
             {
@@ -681,6 +711,50 @@ class StageConfigFactory:
             }
         ]
 
+    @classmethod
+    def create_typed_default_diffusion(
+        cls,
+        model: str,
+        kwargs: dict[str, Any],
+    ) -> VllmOmniConfig:
+        """Build generic diffusion directly into the structured runtime config."""
+        engine_overrides, parallel_config, default_sampling_params, final_output_type = (
+            cls._normalize_default_diffusion(kwargs)
+        )
+        engine_overrides.update(asdict(parallel_config))
+        model_class_name = engine_overrides.get("model_class_name")
+        pipeline = PipelineConfig(
+            model_type="generic_diffusion",
+            model_arch=str(model_class_name or ""),
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="diffusion",
+                    execution_type=StageExecutionType.DIFFUSION,
+                    final_output=True,
+                    final_output_type=final_output_type,
+                ),
+            ),
+        )
+
+        # These values have already been normalized for this diffusion stage.
+        # Scope them explicitly so diffusion-only fields such as engine_backend
+        # and extras are not filtered by the global LLM CLI argument surface.
+        stage_overrides = {
+            f"stage_0_{key}": value for key, value in engine_overrides.items() if key not in {"model", "stage_id"}
+        }
+        stage_overrides["stage_0_devices"] = (
+            kwargs.get("stage_0_devices")
+            or kwargs.get("devices")
+            or ",".join(str(i) for i in range(parallel_config.world_size))
+        )
+        # Caller defaults belong to deployment, not immutable pipeline constraints:
+        # explicit request sampling parameters must still be able to override them.
+        deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, default_sampling_params=default_sampling_params)])
+        return VllmOmniConfig.from_pipeline_config(
+            pipeline, user_deploy_config=deploy, cli_overrides={"model": model, **stage_overrides}
+        )
+
     @classmethod
     def _merge_cli_overrides(
         cls,

vllm_omni/config/omni_config.py

@@ -29,7 +29,9 @@ 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.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.pooling_params import PoolingParams
 
 from vllm_omni.config.stage_config import (
     _DEPLOY_DIR,
@@ -42,11 +44,13 @@ from vllm_omni.config.stage_config import (
     StagePipelineConfig,
     StageType,
     _apply_platform_overrides,
+    _get_recursively_merged_dict,
     _resolve_scheduler,
     _scheduler_path,
     _select_processor_funcs,
     build_stage_runtime_overrides,
     load_deploy_config,
+    merge_pipeline_deploy,
     merge_sampling_constraints,
     normalize_pipeline_cli_overrides,
     reconcile_diffusion_attention_overrides,
@@ -55,6 +59,8 @@ from vllm_omni.config.stage_config import (
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 
+logger = init_logger(__name__)
+
 _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = {
     StageExecutionType.LLM_AR: (StageType.LLM, "ar"),
     StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"),
@@ -186,6 +192,11 @@ class _ModelEngineOverrides(TypedDict, total=False):
     enable_prompt_embeds: bool
 
 
+class _PoolingEngineOverrides(TypedDict, total=False):
+    runner: str
+    pooling_output_decoder: str
+
+
 class _LoadEngineOverrides(TypedDict, total=False):
     tokenizer: str
     download_dir: str
@@ -261,6 +272,7 @@ class _StageEngineValues:
     quantization: _QuantizationEngineOverrides
     model: _ModelEngineOverrides
     load: _LoadEngineOverrides
+    pooling: _PoolingEngineOverrides
     cache: _CacheEngineOverrides
     scheduler: _SchedulerEngineOverrides
     connector: _ConnectorEngineOverrides
@@ -476,6 +488,15 @@ class OmniStageModelConfig(_TrackExplicitConfigFields):
     enable_prompt_embeds: bool | None = None
 
 
+@config(config=ConfigDict(arbitrary_types_allowed=True))
+class OmniStagePoolingConfig:
+    """Typed inputs owned by vLLM pooling stages."""
+
+    runner: str | None = None
+    pooling_output_decoder: str | None = None
+    default_pooling_params: PoolingParams | None = None
+
+
 @_enforce_keyword_only_init
 @config(kw_only=True)
 class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig):
@@ -833,12 +854,12 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
         from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config
 
-        normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs)
+        config_kwargs = normalize_omni_diffusion_kwargs(kwargs)
         # Validate before stage construction while retaining the raw mapping
         # needed by dataclass/config serialization across process boundaries.
-        parse_diffusion_offload_config(normalized_kwargs.get("diffusion_offload_config"))
+        parse_diffusion_offload_config(config_kwargs.get("diffusion_offload_config"))
         valid_fields = {f.name for f in fields(cast(Any, cls))}
-        return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields})
+        return cls(**{k: v for k, v in config_kwargs.items() if k in valid_fields})
 
     def __post_init__(self) -> None:
         # Keep diffusion imports lazy so importing vllm_omni.config does not
@@ -1145,6 +1166,7 @@ _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__)
 _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__)
 _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__)
 _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__)
+_POOLING_ENGINE_FIELDS = frozenset(_PoolingEngineOverrides.__annotations__)
 _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__)
 _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__)
 _DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"})
@@ -1175,6 +1197,7 @@ _LLM_STAGE_ENGINE_FIELDS = (
     | _LLM_CACHE_ENGINE_FIELDS
     | _LLM_SCHEDULER_ENGINE_FIELDS
     | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS
+    | _POOLING_ENGINE_FIELDS
     | {"parallel_config"}
 )
 _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
@@ -1308,7 +1331,12 @@ def _stage_engine_values(
         if topology.execution_type == StageExecutionType.DIFFUSION:
             # Mirror StageConfig.to_omegaconf so both projections resolve alike.
             reconcile_diffusion_attention_overrides(engine, stage_cli_overrides)
-        engine.update(_copy_value(stage_cli_overrides))
+        for key, value in stage_cli_overrides.items():
+            existing = engine.get(key)
+            if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, Mapping):
+                engine[key] = _get_recursively_merged_dict(existing, dict(value))
+            else:
+                engine[key] = _copy_value(value)
     _validate_stage_engine_override_ownership(
         topology.stage_id,
         topology.execution_type,
@@ -1334,6 +1362,10 @@ def _stage_engine_values(
         ),
         model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)),
         load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)),
+        pooling=cast(
+            _PoolingEngineOverrides,
+            _select_engine_overrides(engine, _POOLING_ENGINE_FIELDS),
+        ),
         cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)),
         scheduler=cast(
             _SchedulerEngineOverrides,
@@ -1406,6 +1438,7 @@ class BaseVllmOmniStageConfig:
     cache_config: OmniStageCacheConfig = field(default_factory=OmniStageCacheConfig)
     scheduler_config: OmniStageSchedulerConfig = field(default_factory=OmniStageSchedulerConfig)
     connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig)
+    pooling_config: OmniStagePoolingConfig = field(default_factory=OmniStagePoolingConfig)
     runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig)
     parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig)
     compilation_config: VllmCompilationConfig | None = None
@@ -1490,6 +1523,10 @@ class BaseVllmOmniStageConfig:
     def prompt_transform_func(self) -> str | None:
         return self.stage_pipeline_config.prompt_transform_func
 
+    @property
+    def sampling_constraints(self) -> dict[str, Any]:
+        return dict(self.stage_pipeline_config.sampling_constraints)
+
     @property
     def cfg_kv_collect_func(self) -> str | None:
         return self.stage_pipeline_config.cfg_kv_collect_func
@@ -1543,6 +1580,7 @@ def _build_common_stage_config_kwargs(
                 model=model,
             ),
             "load_config": _build_load_config(topology, engine.load),
+            "pooling_config": _build_pooling_config(stage_deploy, engine.pooling),
             "cache_config": _build_cache_config(
                 deploy,
                 engine.cache,
@@ -1756,6 +1794,20 @@ def _build_model_config(
     )
 
 
+def _build_pooling_config(
+    stage_deploy: StageDeployConfig | None,
+    engine: _PoolingEngineOverrides,
+) -> OmniStagePoolingConfig:
+    default_pooling_params = None
+    if stage_deploy is not None and stage_deploy.default_pooling_params:
+        default_pooling_params = PoolingParams(**dict(stage_deploy.default_pooling_params))
+    return cast(Any, OmniStagePoolingConfig)(
+        runner=_copy_value(engine.get("runner")),
+        pooling_output_decoder=_copy_value(engine.get("pooling_output_decoder")),
+        default_pooling_params=default_pooling_params,
+    )
+
+
 def _build_load_config(
     topology: StagePipelineConfig,
     engine: _LoadEngineOverrides,
@@ -1934,6 +1986,10 @@ class VllmOmniConfig:
     pipeline_config: PipelineConfig
     stage_configs: tuple[StageConfigType, ...]
     orchestrator_config: VllmOmniOrchestratorConfig = field(default_factory=VllmOmniOrchestratorConfig)
+    # Keep strategy provenance separate from the effective orchestrator value:
+    # an explicit CLI policy may be present even when the strategy declares no
+    # stage-replica axis.
+    strategy_omni_lb_policy: str | None = None
 
     def stage_by_id(self, stage_id: int) -> StageConfigType:
         for stage in self.stage_configs:
@@ -1949,6 +2005,7 @@ class VllmOmniConfig:
         user_deploy_config: DeployConfig | None = None,
         deploy_config_path: str | None = None,
         cli_overrides: dict[str, Any] | None = None,
+        strategy_specs: Mapping[Any, Any] | None = None,
     ) -> VllmOmniConfig:
         """Create a structured config from a resolved pipeline and deploy YAML."""
         if cli_overrides is None:
@@ -1973,6 +2030,62 @@ class VllmOmniConfig:
             deploy.async_chunk = False
         _validate_async_chunk_support(pipeline_cfg, deploy)
         validate_stage_async_chunk_edges(pipeline_cfg, deploy)
+
+        strategy_result = None
+        if strategy_specs:
+            from vllm_omni.config.composable_parallel import apply_strategy_specs
+
+            strategy_stages = merge_pipeline_deploy(pipeline_cfg, copy.deepcopy(deploy), {})
+            strategy_result = apply_strategy_specs(strategy_stages, strategy_specs)
+            strategy_overrides: dict[str, Any] = {}
+            axis_fields = {
+                "tp": "tensor_parallel_size",
+                "dp": "data_parallel_size",
+                "pp": "pipeline_parallel_size",
+            }
+            for stage_id, parallel in strategy_result.per_stage_config.items():
+                declared = set(parallel.l1_owners)
+                explicit = build_stage_runtime_overrides(stage_id, cli_overrides)
+                for axis, field_name in axis_fields.items():
+                    if axis not in declared:
+                        continue
+                    derived = getattr(parallel, field_name)
+                    value = explicit.get(field_name)
+                    if value is None:
+                        value = derived
+                    elif value != derived:
+                        logger.warning(
+                            "[composable_parallel] stage %s: CLI %s=%s overrides the "
+                            "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.",
+                            stage_id,
+                            field_name,
+                            value,
+                            field_name,
+                            derived,
+                        )
+                    strategy_overrides[f"stage_{stage_id}_{field_name}"] = value
+                if "ep" in declared and parallel.enable_expert_parallel:
+                    strategy_overrides[f"stage_{stage_id}_enable_expert_parallel"] = bool(
+                        explicit.get("enable_expert_parallel", True)
+                    )
+                if "stage_replica" in declared:
+                    derived = parallel.stage_replica_size
+                    value = explicit.get("num_replicas")
+                    if value is None:
+                        value = derived
+                    elif value != derived:
+                        logger.warning(
+                            "[composable_parallel] stage %s: CLI num_replicas=%s overrides the "
+                            "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.",
+                            stage_id,
+                            value,
+                            derived,
+                        )
+                    strategy_overrides[f"stage_{stage_id}_num_replicas"] = value
+            cli_overrides = {**strategy_overrides, **cli_overrides}
+            if strategy_result.omni_lb_policy is not None and cli_overrides.get("omni_lb_policy") is None:
+                cli_overrides["omni_lb_policy"] = strategy_result.omni_lb_policy
+
         deploy_by_id = {stage.stage_id: stage for stage in deploy.stages}
         model = cli_overrides.get("model")
 
@@ -1995,6 +2108,20 @@ class VllmOmniConfig:
             )
             for topology in pipeline_cfg.stages
         )
+        if strategy_result is not None:
+            from vllm_omni.config.composable_parallel import check_device_layout
+
+            by_id = {stage.stage_id: stage for stage in stage_configs}
+            for stage_id in strategy_result.per_stage_config:
+                stage = by_id[stage_id]
+                check_device_layout(
+                    stage.runtime_config.devices,
+                    tensor_parallel_size=stage.parallel_config.tensor_parallel_size,
+                    data_parallel_size=stage.parallel_config.data_parallel_size,
+                    pipeline_parallel_size=stage.parallel_config.pipeline_parallel_size,
+                    num_replicas=stage.runtime_config.num_replicas,
+                    role=stage.model_stage,
+                )
 
         orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)(
             deploy_config_path=loaded_deploy_config_path,
@@ -2004,6 +2131,7 @@ class VllmOmniConfig:
             pipeline_config=pipeline_cfg,
             stage_configs=stage_configs,
             orchestrator_config=orchestrator_config,
+            strategy_omni_lb_policy=(strategy_result.omni_lb_policy if strategy_result is not None else None),
         )
 
 

vllm_omni/config/resolver.py

@@ -14,7 +14,6 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot
 from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.omni_config import VllmOmniConfig
 from vllm_omni.config.stage_config import PipelineConfig
-from vllm_omni.config.yaml_util import create_config
 from vllm_omni.diffusion.data import resolve_model_class_name
 from vllm_omni.diffusion.registry import DiffusionModelRegistry
 from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
@@ -24,19 +23,10 @@ logger = init_logger(__name__)
 
 @dataclass(frozen=True)
 class OmniConfigResolution:
-    """Migration envelope returned by the production config resolver.
-
-    ``stage_configs`` intentionally carries the current OmegaConf-compatible
-    runtime ABI only until stage startup consumes ``VllmOmniConfig`` directly.
-    ``pipeline_config`` is the effective topology after runtime stage injection,
-    so both fields always describe the same set of stages.
-    It is not a stable authoring or extension API; new production callers
-    should resolve through :func:`resolve_omni_config` and must not construct or
-    merge this compatibility shape themselves.
-    """
+    """Canonical typed result returned by the production config resolver."""
 
     config_path: str | None
-    stage_configs: tuple[Any, ...]  # Temporary StageConfig/OmegaConf bridge.
+    stage_configs: tuple[Any, ...]
     pipeline_config: PipelineConfig | None = None
     omni_lb_policy: str | None = None
 
@@ -176,27 +166,13 @@ def _load_strategy_specs(strategy_config_path: str | None) -> Mapping[Any, Any]
 
 def _build_registered_resolution(
     structured_config: VllmOmniConfig,
-    *,
-    cli_overrides: dict[str, Any],
-    strategy_config_path: str | None,
 ) -> OmniConfigResolution:
-    """Build the temporary runtime view for an already-resolved pipeline."""
-    # Runtime consumers have not yet moved to typed per-stage configs. Use the
-    # factory-owned compatibility bridge instead of reimplementing legacy YAML
-    # discovery and merging in this resolver.
-    effective_deploy_path = structured_config.orchestrator_config.deploy_config_path
-    legacy_resolution = StageConfigFactory._resolve_legacy_from_registry(
-        structured_config.pipeline_config,
-        cli_overrides,
-        effective_deploy_path,
-        strategy_specs=_load_strategy_specs(strategy_config_path),
-    )
-
+    """Return the canonical structured runtime view for a resolved pipeline."""
     return OmniConfigResolution(
-        config_path=effective_deploy_path,
-        stage_configs=tuple(stage.to_omegaconf() for stage in legacy_resolution.stage_configs),
-        pipeline_config=legacy_resolution.pipeline_config,
-        omni_lb_policy=legacy_resolution.omni_lb_policy,
+        config_path=structured_config.orchestrator_config.deploy_config_path,
+        stage_configs=tuple(structured_config.stage_configs),
+        pipeline_config=structured_config.pipeline_config,
+        omni_lb_policy=getattr(structured_config, "strategy_omni_lb_policy", None),
     )
 
 
@@ -231,18 +207,16 @@ def resolve_omni_config(
     registry_overrides = dict(normalized_overrides)
     _flatten_stage_overrides(registry_overrides, stage_overrides)
 
+    strategy_specs = _load_strategy_specs(strategy_config_path)
     structured_config = StageConfigFactory.create_from_model(
         model,
         trust_remote_code=trust_remote_code,
         cli_overrides=registry_overrides,
         deploy_config_path=deploy_config_path,
+        strategy_specs=strategy_specs,
     )
     if structured_config is not None:
-        return _build_registered_resolution(
-            structured_config,
-            cli_overrides=registry_overrides,
-            strategy_config_path=strategy_config_path,
-        )
+        return _build_registered_resolution(structured_config)
 
     _apply_generic_stage_overrides(normalized_overrides, stage_overrides)
     supported, model_class_name = _resolve_generic_diffusion_model_class(model, normalized_overrides)
@@ -252,11 +226,12 @@ def resolve_omni_config(
         )
     if model_class_name is not None:
         normalized_overrides.setdefault("model_class_name", model_class_name)
-    default_stages = StageConfigFactory.create_default_diffusion(normalized_overrides)
+    default_config = StageConfigFactory.create_typed_default_diffusion(model, normalized_overrides)
 
     return OmniConfigResolution(
         config_path=deploy_config_path,
-        stage_configs=tuple(create_config(_convert_dataclasses_to_dict(default_stages))),
+        stage_configs=tuple(default_config.stage_configs),
+        pipeline_config=default_config.pipeline_config,
     )
 
 

vllm_omni/engine/omni_engine_base.py

@@ -278,7 +278,14 @@ class OmniEngineBase:
 
         self.num_stages = len(self.stage_configs)
         self.async_chunk = any(
-            bool(getattr(getattr(stage, "engine_args", None), "async_chunk", False)) for stage in self.stage_configs
+            bool(
+                getattr(
+                    getattr(stage, "connector_config", None),
+                    "async_chunk",
+                    getattr(getattr(stage, "engine_args", None), "async_chunk", False),
+                )
+            )
+            for stage in self.stage_configs
         )
         self.stage_pools: list[StagePool] = []
         self.stage_clients: list[StageClient] = []  # logical-stage view for external readers

vllm_omni/engine/stage_admission.py

@@ -287,8 +287,13 @@ def check_admission(
 
 
 def _diffusion_utilization(replica: Any) -> float | None:
-    """Best-effort gpu_memory_utilization for a diffusion replica from raw args."""
+    """Resolve a diffusion replica's memory budget from typed or legacy config."""
     stage_cfg = getattr(replica, "stage_cfg", None)
+    cache_config = getattr(stage_cfg, "cache_config", None)
+    util = getattr(cache_config, "gpu_memory_utilization", None)
+    if util is not None:
+        return float(util)
+
     engine_args = getattr(stage_cfg, "engine_args", None)
     if isinstance(engine_args, dict):
         util = engine_args.get("gpu_memory_utilization")

vllm_omni/engine/stage_engine_startup.py

@@ -18,6 +18,7 @@ from types import SimpleNamespace
 from typing import Any
 
 import msgspec
+import torch
 import zmq
 from omegaconf import OmegaConf
 from vllm.config import VllmConfig
@@ -34,6 +35,7 @@ from vllm.v1.engine.utils import (
 )
 from vllm.v1.executor import Executor
 
+from vllm_omni.config.omni_config import BaseVllmOmniStageConfig
 from vllm_omni.distributed.omni_connectors.utils import initialization
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.stage_init_utils import (
@@ -68,11 +70,13 @@ def _serialize_stage_config(stage_config: Any) -> Any:
     """Convert a stage config to msgpack-friendly builtins."""
     if stage_config is None or isinstance(stage_config, (str, bytes, int, float, bool)):
         return stage_config
+    if isinstance(stage_config, torch.dtype):
+        return str(stage_config).removeprefix("torch.")
 
     if OmegaConf.is_config(stage_config):
         return _serialize_stage_config(OmegaConf.to_container(stage_config, resolve=True))
 
-    if dataclasses.is_dataclass(stage_config):
+    if dataclasses.is_dataclass(stage_config) and not isinstance(stage_config, type):
         return _serialize_stage_config(dataclasses.asdict(stage_config))
 
     if isinstance(stage_config, dict):
@@ -1420,7 +1424,7 @@ def get_headless_replica_devices(
     omni_dp_size_local: int,
 ) -> list[str | None]:
     """Return per-replica device slices for a headless stage."""
-    runtime_cfg = getattr(stage_cfg, "runtime", None)
+    runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None))
     devices_str: str | None = None
     if runtime_cfg is not None:
         devices_str = (
@@ -1526,10 +1530,11 @@ def launch_headless_diffusion_replicas(
         stage_id,
     )
 
-    # Headless diffusion startup and its downstream helpers still consume the
-    # legacy StageConfig shape; switch this with the coordinated RFC #4021
-    # stage-init cutover.
-    metadata = stage_init_utils.extract_legacy_stage_metadata(stage_cfg)
+    metadata = (
+        stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage_cfg)
+        if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+        else stage_init_utils.extract_legacy_stage_metadata(stage_cfg)
+    )
     if omni_conn_cfg:
         inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
     # Headless single-stage launch must still infer cross-stage TP topology

vllm_omni/engine/stage_init_utils.py

@@ -402,6 +402,11 @@ def _tp_size_for_stage(stage_configs: Sequence[Any], stage_id: Any) -> int | Non
     for stage_cfg in stage_configs:
         if str(getattr(stage_cfg, "stage_id", None)) not in id_strs:
             continue
+        if isinstance(stage_cfg, BaseVllmOmniStageConfig):
+            try:
+                return max(1, int(stage_cfg.parallel_config.tensor_parallel_size))
+            except (TypeError, ValueError):
+                return 1
         engine_args = getattr(stage_cfg, "engine_args", None)
         if engine_args is None:
             return 1
@@ -489,11 +494,15 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[
     rank mappings automatically.
     """
     try:
-        engine_args = stage_cfg.engine_args
-        if hasattr(engine_args, "get"):
-            omni_kv = engine_args.get("omni_kv_config", None)
+        connector_config = getattr(stage_cfg, "connector_config", None)
+        if connector_config is not None:
+            omni_kv = connector_config.omni_kv_config
         else:
-            omni_kv = getattr(engine_args, "omni_kv_config", None)
+            engine_args = stage_cfg.engine_args
+            if hasattr(engine_args, "get"):
+                omni_kv = engine_args.get("omni_kv_config", None)
+            else:
+                omni_kv = getattr(engine_args, "omni_kv_config", None)
 
         if omni_kv is None:
             return
@@ -504,7 +513,11 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[
             if "stage_id" not in omni_kv:
                 omni_kv["stage_id"] = stage_id
 
-        engine_input_source = getattr(stage_cfg, "engine_input_source", None)
+        engine_input_source = (
+            getattr(stage_cfg, "input_sources", None)
+            if connector_config is not None
+            else getattr(stage_cfg, "engine_input_source", None)
+        )
         if engine_input_source is not None:
             if hasattr(omni_kv, "setdefault"):
                 omni_kv.setdefault("engine_input_source", list(engine_input_source))
@@ -703,17 +716,16 @@ def _resolve_omni_metadata_hook(path: str | None) -> Callable | None:
 def extract_stage_metadata_from_omni_stage_config(
     stage_config: BaseVllmOmniStageConfig,
 ) -> StageMetadata:
-    """Project one typed stage config into metadata for a future cutover.
-
-    This projection is not used by production startup yet. Current replica
-    layout, engine-argument, remote-diffusion, and platform setup paths still
-    require the legacy StageConfig/OmegaConf shape.
-    """
+    """Project one typed stage config into production runtime metadata."""
     stage_type: Literal["llm", "diffusion"] = "diffusion" if stage_config.stage_type == StageType.DIFFUSION else "llm"
-    sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
-    sampling_params: OmniSamplingParams = sampling_params_cls(
-        **(stage_config.model_config.default_sampling_params or {})
-    )
+    pooling_config = stage_config.pooling_config
+    if stage_type == "llm" and (pooling_config.runner or "").lower() == "pooling":
+        sampling_params: OmniSamplingParams | PoolingParams = (
+            copy.deepcopy(pooling_config.default_pooling_params) or PoolingParams()
+        )
+    else:
+        sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
+        sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {}))
     custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func)
 
     if stage_type == "diffusion":
@@ -878,7 +890,9 @@ def _get_local_llm_parallel_sizes(
     unset.
     """
     if engine_args is None:
-        engine_args = getattr(stage_cfg, "engine_args", {})
+        engine_args = getattr(stage_cfg, "parallel_config", None)
+        if engine_args is None:
+            engine_args = getattr(stage_cfg, "engine_args", {})
     tp_size = int(_get_attr_or_item(engine_args, "tensor_parallel_size", 1) or 1)
     pp_size = int(_get_attr_or_item(engine_args, "pipeline_parallel_size", 1) or 1)
     local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size_local", None)
@@ -889,9 +903,12 @@ def _get_local_llm_parallel_sizes(
 
 def get_stage_devices_per_replica(stage_cfg: Any, engine_args: Any | None = None) -> int:
     """Return the number of devices consumed by one replica of *stage_cfg*."""
-    if engine_args is None:
-        engine_args = getattr(stage_cfg, "engine_args", {})
     if getattr(stage_cfg, "stage_type", "llm") == "diffusion":
+        typed_parallel_config = getattr(stage_cfg, "parallel_config", None)
+        if typed_parallel_config is not None:
+            return max(1, int(typed_parallel_config.world_size))
+        if engine_args is None:
+            engine_args = getattr(stage_cfg, "engine_args", {})
         parallel_config = _get_attr_or_item(engine_args, "parallel_config")
         if parallel_config is None:
             return 1
@@ -933,7 +950,7 @@ def compute_replica_layout(
     """
     replicas_per_stage: list[int] = []
     for stage_cfg in stage_configs:
-        runtime_cfg = getattr(stage_cfg, "runtime", {})
+        runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {}))
         num_replicas = int(
             runtime_cfg.get("num_replicas", 1)
             if hasattr(runtime_cfg, "get")
@@ -948,7 +965,7 @@ def compute_replica_layout(
         num_replicas = replicas_per_stage[stage_id]
         if num_replicas <= 1:
             continue
-        runtime_cfg = getattr(stage_cfg, "runtime", {})
+        runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {}))
         devices_str = (
             runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
         )
@@ -1086,15 +1103,31 @@ def _project_omni_stage_engine_args(
         diffusion_stage = cast(VllmOmniDiffusionStageConfig, stage_config)
         engine_args.update(_project_omni_config_fields(diffusion_stage.diffusion_config))
 
+    model_excluded_fields = {
+        "default_sampling_params",
+        "has_sampling_extra_args",
+    }
+    runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus"}
+    if not is_diffusion:
+        # These values configure OmniDiffusionConfig or its worker process;
+        # OmniEngineArgs has no matching fields for LLM stages.
+        model_excluded_fields.update(
+            {
+                "disable_autocast",
+                "enable_multithread_weight_load",
+                "num_weight_load_threads",
+            }
+        )
+        runtime_excluded_fields.add("log_level")
+
     for config, excluded_fields in (
         (
             stage_config.model_config,
-            frozenset({"default_sampling_params", "has_sampling_extra_args"}),
+            frozenset(model_excluded_fields),
         ),
         (
             stage_config.runtime_config,
-            frozenset({"devices", "num_replicas", "env", "num_gpus"})
-            | (frozenset({"additional_config"}) if is_diffusion else frozenset()),
+            frozenset(runtime_excluded_fields) | (frozenset({"additional_config"}) if is_diffusion else frozenset()),
         ),
     ):
         engine_args.update(
@@ -1103,6 +1136,11 @@ def _project_omni_stage_engine_args(
                 exclude=excluded_fields,
             )
         )
+    pooling_config = stage_config.pooling_config
+    if pooling_config.runner is not None:
+        engine_args["runner"] = pooling_config.runner
+    if pooling_config.pooling_output_decoder is not None:
+        engine_args["pooling_output_decoder"] = pooling_config.pooling_output_decoder
 
     if is_diffusion:
         for config, field_map in (
@@ -1330,11 +1368,9 @@ def build_engine_args_dict(
     stage_connector_spec: dict[str, Any] | None = None,
     cli_tokenizer: str | None = None,
 ) -> dict[str, Any]:
-    """Build engine arguments through the stable production entry point.
+    """Preserve the legacy engine-argument API for compatibility callers.
 
-    Production inputs still use the legacy stage representation. Keep that
-    compatibility choice behind this function so callers do not bind directly
-    to a representation-specific implementation.
+    Typed runtime stages use ``build_engine_args_dict_from_omni_stage_config``.
     """
     return build_legacy_engine_args_dict(
         stage_config,
@@ -1350,13 +1386,7 @@ def build_engine_args_dict_from_omni_stage_config(
     stage_connector_spec: dict[str, Any] | None = None,
     cli_tokenizer: str | None = None,
 ) -> dict[str, Any]:
-    """Project one typed stage config into backend engine arguments.
-
-    This projection is prepared for the RFC #4021 stage-init cutover. Current
-    production startup reaches the legacy implementation through
-    ``build_engine_args_dict`` while strategy and startup-plan inputs still
-    use the legacy representation.
-    """
+    """Project one typed production stage config into backend engine arguments."""
     engine_args_dict = _project_omni_stage_engine_args(stage_config)
     _apply_rocm_attention_backend(engine_args_dict, stage_config.stage_type)
     return _finalize_engine_args_dict(
@@ -1391,7 +1421,7 @@ def _check_stage_device_layout(stage_config: Any, engine_args_dict: dict[str, An
     """
     from vllm_omni.config.composable_parallel import StrategyApplyError, check_device_layout
 
-    runtime = getattr(stage_config, "runtime", None)
+    runtime = getattr(stage_config, "runtime_config", getattr(stage_config, "runtime", None))
     devices = _get_attr_or_item(runtime, "devices", None) if runtime is not None else None
     if devices is None:
         # No explicit placement -> vLLM assigns devices itself; nothing to check.
@@ -1454,10 +1484,12 @@ def build_vllm_config(
         (vllm_config, executor_class)
     """
     if engine_args_dict is None:
-        engine_args_dict = build_engine_args_dict(
-            stage_config,
-            model,
-            stage_connector_spec=stage_connector_spec,
+        engine_args_dict = (
+            build_engine_args_dict_from_omni_stage_config(
+                stage_config, model, stage_connector_spec=stage_connector_spec
+            )
+            if isinstance(stage_config, BaseVllmOmniStageConfig)
+            else build_engine_args_dict(stage_config, model, stage_connector_spec=stage_connector_spec)
         )
 
     filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict)
@@ -1934,7 +1966,11 @@ def build_diffusion_config(
 ) -> Any:
     """Build diffusion config for a stage."""
 
-    engine_args_dict = build_engine_args_dict(stage_cfg, model)
+    engine_args_dict = (
+        build_engine_args_dict_from_omni_stage_config(stage_cfg, model)
+        if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+        else build_engine_args_dict(stage_cfg, model)
+    )
     od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict)
 
     num_devices_per_stage = od_config.parallel_config.world_size

vllm_omni/engine/stage_runtime.py

@@ -17,9 +17,9 @@ from multiprocessing.process import BaseProcess
 from typing import Any, TypedDict, cast
 
 import janus
-from omegaconf import OmegaConf
 from vllm.logger import init_logger
 
+from vllm_omni.config.omni_config import BaseVllmOmniStageConfig
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
     resolve_omni_kv_config_for_stage,
 )
@@ -50,11 +50,13 @@ from vllm_omni.engine.stage_init_utils import (
     _inject_inferred_kv_tp_topology,
     acquire_device_locks,
     build_engine_args_dict,
+    build_engine_args_dict_from_omni_stage_config,
     build_llm_stage_output_processor,
     build_vllm_config,
     compute_replica_layout,
     device_overlap_group_keys,
     extract_legacy_stage_metadata,
+    extract_stage_metadata_from_omni_stage_config,
     get_stage_connector_spec,
     inject_kv_stage_info,
     inject_omni_kv_connector_config,
@@ -592,12 +594,12 @@ class StageRuntime:
         """Build startup plans for every logical stage and replica."""
         stage_plans: list[LogicalStageInitPlan] = []
 
-        # RFC #4021 transition boundary: stage planning still relies on legacy
-        # StageConfig.runtime and StageConfig.engine_args for replica and engine
-        # setup. Keep metadata extraction on the legacy path until the
-        # coordinated stage-init cutover.
         for stage_idx, stage_cfg in enumerate(self._stage_configs):
-            base_metadata = extract_legacy_stage_metadata(stage_cfg)
+            base_metadata = (
+                extract_stage_metadata_from_omni_stage_config(stage_cfg)
+                if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+                else extract_legacy_stage_metadata(stage_cfg)
+            )
             stage_id = int(base_metadata.stage_id)
             if stage_id != stage_idx:
                 raise ValueError(
@@ -621,14 +623,20 @@ class StageRuntime:
             executor_class = None
             engine_args_dict = None
             if base_metadata.stage_type != "diffusion":
-                # The stable adapter entry point still receives the same
-                # legacy stage object as replica planning. Its implementation
-                # switches only at the coordinated RFC #4021 cutover.
-                engine_args_dict = build_engine_args_dict(
-                    stage_cfg,
-                    self._model,
-                    stage_connector_spec=stage_connector_spec,
-                    cli_tokenizer=self._tokenizer,
+                engine_args_dict = (
+                    build_engine_args_dict_from_omni_stage_config(
+                        stage_cfg,
+                        self._model,
+                        stage_connector_spec=stage_connector_spec,
+                        cli_tokenizer=self._tokenizer,
+                    )
+                    if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+                    else build_engine_args_dict(
+                        stage_cfg,
+                        self._model,
+                        stage_connector_spec=stage_connector_spec,
+                        cli_tokenizer=self._tokenizer,
+                    )
                 )
                 inject_omni_kv_connector_config(
                     engine_args_dict,
@@ -650,11 +658,20 @@ class StageRuntime:
                 )
 
             for replica_id in range(num_replicas):
-                replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg
+                # Keep the logical stage's device pool intact; each replica owns
+                # the same config throughout planning, metadata and launch.
+                replica_cfg = copy.deepcopy(stage_cfg) if num_replicas > 1 else stage_cfg
                 if stage_idx in replica_devices_map:
-                    replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id]
+                    devices = replica_devices_map[stage_idx][replica_id]
+                    runtime_cfg = getattr(replica_cfg, "runtime_config", getattr(replica_cfg, "runtime", None))
+                    if runtime_cfg is not None:
+                        runtime_cfg.devices = devices
 
-                replica_metadata = extract_legacy_stage_metadata(replica_cfg)
+                replica_metadata = (
+                    extract_stage_metadata_from_omni_stage_config(replica_cfg)
+                    if isinstance(replica_cfg, BaseVllmOmniStageConfig)
+                    else extract_legacy_stage_metadata(replica_cfg)
+                )
                 replica_metadata.replica_id = replica_id
                 if launch_mode == "remote" and replica_metadata.stage_type != "diffusion":
                     replica_metadata.runtime_cfg = None
@@ -1056,17 +1073,25 @@ class StageRuntime:
                         raise RuntimeError("Omni KV connector requires source and destination stages")
                     inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to)
                 inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs)
-                engine_args = getattr(plan.stage_cfg, "engine_args", {})
-                inline_diffusion = (
-                    engine_args.get("inline_diffusion", False)
-                    if hasattr(engine_args, "get")
-                    else getattr(engine_args, "inline_diffusion", False)
-                )
-                custom_pipeline_args = (
-                    engine_args.get("custom_pipeline_args")
-                    if hasattr(engine_args, "get")
-                    else getattr(engine_args, "custom_pipeline_args", None)
-                )
+                if isinstance(plan.stage_cfg, BaseVllmOmniStageConfig):
+                    inline_diffusion = plan.stage_cfg.stage_pipeline_config.inline_diffusion
+                    custom_pipeline_args = getattr(
+                        getattr(plan.stage_cfg, "diffusion_config", None),
+                        "custom_pipeline_args",
+                        None,
+                    )
+                else:
+                    engine_args = getattr(plan.stage_cfg, "engine_args", {})
+                    inline_diffusion = (
+                        engine_args.get("inline_diffusion", False)
+                        if hasattr(engine_args, "get")
+                        else getattr(engine_args, "inline_diffusion", False)
+                    )
+                    custom_pipeline_args = (
+                        engine_args.get("custom_pipeline_args")
+                        if hasattr(engine_args, "get")
+                        else getattr(engine_args, "custom_pipeline_args", None)
+                    )
                 client, resources = launch_diffusion_stage_replica(
                     model=self._model,
                     stage_config=plan.stage_cfg,
@@ -1225,7 +1250,7 @@ class DistStageRuntime(StageRuntime):
 
         for idx, stage_cfg in enumerate(self._stage_configs):
             stage_id = int(getattr(stage_cfg, "stage_id", idx))
-            runtime_cfg = getattr(stage_cfg, "runtime", None)
+            runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None))
             if runtime_cfg is None:
                 continue
             if stage_id == target_stage_id:
@@ -1297,13 +1322,9 @@ class DistStageRuntime(StageRuntime):
         if registered_stage_cfg is None:
             raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config")
 
-        # Remote diffusion registration still transports the legacy mapping
-        # shape. Reconstruct and project that shape until its RFC #4021 cutover.
-        metadata = (
-            extract_legacy_stage_metadata(OmegaConf.create(registered_stage_cfg))
-            if plan.metadata.stage_type == "diffusion"
-            else copy.deepcopy(plan.metadata)
-        )
+        # Registration is transport-only: the head-side typed plan remains
+        # authoritative for metadata and topology.
+        metadata = extract_stage_metadata_from_omni_stage_config(plan.stage_cfg)
         metadata.replica_id = plan.replica_id
         ctx = StageRemoteFactoryContext(
             stage_id=plan.metadata.stage_id,

vllm_omni/entrypoints/cli/serve.py

@@ -1108,7 +1108,13 @@ def _build_multi_api_stage_runtime(args: TrackingNamespace, num_api_servers: int
     sleep_stages = [
         int(getattr(stage_config, "stage_id", stage_index))
         for stage_index, stage_config in enumerate(stage_configs)
-        if bool(getattr(getattr(stage_config, "engine_args", None), "enable_sleep_mode", False))
+        if bool(
+            getattr(
+                getattr(stage_config, "model_config", getattr(stage_config, "engine_args", None)),
+                "enable_sleep_mode",
+                False,
+            )
+        )
     ]
     if sleep_stages:
         raise ValueError(
@@ -1116,8 +1122,16 @@ def _build_multi_api_stage_runtime(args: TrackingNamespace, num_api_servers: int
             f"disable enable_sleep_mode for stage(s) {sleep_stages}"
         )
 
-    stage0_args = getattr(stage_configs[0], "engine_args", None) if stage_configs else None
-    async_chunk = bool(getattr(stage0_args, "async_chunk", False))
+    async_chunk = any(
+        bool(
+            getattr(
+                getattr(stage, "connector_config", None),
+                "async_chunk",
+                getattr(getattr(stage, "engine_args", None), "async_chunk", False),
+            )
+        )
+        for stage in stage_configs
+    )
     return StageRuntime(
         stage_configs=stage_configs,
         model=model,
@@ -1282,6 +1296,7 @@ def run_headless(args: TrackingNamespace) -> None:
     )
     from vllm_omni.engine.stage_init_utils import (
         build_engine_args_dict,
+        build_engine_args_dict_from_omni_stage_config,
         build_vllm_config,
         get_stage_connector_spec,
         inject_omni_kv_connector_config,
@@ -1372,21 +1387,21 @@ def run_headless(args: TrackingNamespace) -> None:
     stage_connector_spec = get_stage_connector_spec(
         omni_transfer_config=omni_transfer_config,
         stage_id=stage_id,
-        async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)),
+        async_chunk=bool(
+            getattr(getattr(stage_cfg, "connector_config", None), "async_chunk", None)
+            if hasattr(stage_cfg, "connector_config")
+            else stage_cfg.engine_args.get("async_chunk", False)
+        ),
     )
 
-    # ``runtime_cfg`` is mostly inherited from the parent's
-    # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally
-    # bracket each replica's spawn below with setup_stage_devices so they
-    # don't all stack on cuda:0 (see ``per_replica_devices`` above).
-    # Headless startup still supplies the legacy OmegaConf stage shape through
-    # the stable adapter entry point. The implementation switches only when
-    # RFC #4021 threads structured stage configs through the launch plan.
-    engine_args_dict = build_engine_args_dict(
-        stage_cfg,
-        model,
-        stage_connector_spec=stage_connector_spec,
-        cli_tokenizer=getattr(args, "tokenizer", None),
+    engine_args_dict = (
+        build_engine_args_dict_from_omni_stage_config(
+            stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None)
+        )
+        if hasattr(stage_cfg, "connector_config")
+        else build_engine_args_dict(
+            stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None)
+        )
     )
 
     inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id)

vllm_omni/entrypoints/openai/serving_chat.py

@@ -248,10 +248,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             return obj.get(key, default)
         return getattr(obj, key, default)
 
+    @classmethod
+    def _stage_model_metadata(cls, stage: Any) -> tuple[str | None, str | None]:
+        """Return ``(model_arch, model_stage)`` for typed or legacy stages."""
+        model_config = cls._stage_get(stage, "model_config")
+        model_arch = cls._stage_get(model_config, "model_arch")
+        model_stage = cls._stage_get(stage, "model_stage")
+
+        topology = cls._stage_get(stage, "stage_pipeline_config")
+        if model_stage is None:
+            model_stage = cls._stage_get(topology, "model_stage")
+
+        engine_args = cls._stage_get(stage, "engine_args")
+        if model_arch is None:
+            model_arch = cls._stage_get(engine_args, "model_arch")
+        if model_stage is None:
+            model_stage = cls._stage_get(engine_args, "model_stage")
+        return model_arch, model_stage
+
     def _has_minicpmo45_stage(self) -> bool:
         for stage in getattr(self.engine_client, "stage_configs", []) or []:
-            engine_args = self._stage_get(stage, "engine_args")
-            model_arch = self._stage_get(engine_args, "model_arch")
+            model_arch, _ = self._stage_model_metadata(stage)
             if model_arch == "MiniCPMO45OmniForConditionalGeneration":
                 return True
         return False
@@ -324,11 +341,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             if not hasattr(sp, "output_kind"):
                 continue
 
-            engine_args = self._stage_get(stage, "engine_args")
-            if self._stage_get(engine_args, "model_arch") != "MiniCPMO45OmniForConditionalGeneration":
+            model_arch, model_stage = self._stage_model_metadata(stage)
+            if model_arch != "MiniCPMO45OmniForConditionalGeneration":
                 continue
 
-            model_stage = self._stage_get(engine_args, "model_stage")
             if model_stage == "llm":
                 sp.output_kind = RequestOutputKind.FINAL_ONLY
             elif model_stage == "tts":

vllm_omni/entrypoints/openai/serving_speech.py

@@ -36,6 +36,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.utils import random_uuid
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
+from vllm_omni.config.stage_config import StagePipelineConfig
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler
 from vllm_omni.entrypoints.openai.protocol.audio import (
     AudioResponse,
@@ -76,6 +77,24 @@ _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS"
 _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS"
 _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS"
 
+
+def _stage_speech_metadata(stage: Any) -> tuple[str | None, str | None, str | None]:
+    """Return speech routing metadata from typed or legacy stage configs."""
+    topology = getattr(stage, "stage_pipeline_config", None)
+    if isinstance(topology, StagePipelineConfig):
+        return (
+            getattr(topology, "model_stage", None),
+            getattr(getattr(stage, "model_config", None), "model_arch", None),
+            getattr(stage, "worker_type", None),
+        )
+    engine_args = getattr(stage, "engine_args", None)
+    return (
+        getattr(engine_args, "model_stage", None),
+        getattr(engine_args, "model_arch", None),
+        getattr(engine_args, "worker_type", None),
+    )
+
+
 # TTS Configuration
 #
 # The stage-key -> model-type mapping is NOT declared here: it is derived from
@@ -472,12 +491,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         """Find and return the TTS stage config, or None if not found."""
         tts_stage_keys = all_tts_stage_keys()
         entry_stage_archs = tts_entry_stage_archs()
-        all_stages = frozenset(stage.engine_args.model_stage for stage in self.engine_client.stage_configs)
+        all_stages = frozenset(
+            model_stage
+            for stage in self.engine_client.stage_configs
+            if (model_stage := _stage_speech_metadata(stage)[0]) is not None
+        )
         for stage in self.engine_client.stage_configs:
-            engine_args = stage.engine_args
-            model_stage = engine_args.model_stage
-            model_arch = getattr(engine_args, "model_arch", None)
-            worker_type = getattr(engine_args, "worker_type", None)
+            model_stage, model_arch, worker_type = _stage_speech_metadata(stage)
             if model_stage in tts_stage_keys:
                 # Owning the stage key is not always enough: a model may be
                 # speech-capable only in some deployment topologies. Ask the
@@ -501,10 +521,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         """
         if self._tts_stage is None:
             return None
-        return detect_tts_model_type(
-            getattr(self._tts_stage.engine_args, "model_stage", None),
-            getattr(self._tts_stage.engine_args, "model_arch", None),
-        )
+        return detect_tts_model_type(*_stage_speech_metadata(self._tts_stage)[:2])
 
     def _compute_max_instructions_length(self) -> int:
         """Compute max instructions length with precedence: CLI > stage config > default.
@@ -518,7 +535,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         # 2. Try to get from TTS stage config
         if self._tts_stage is not None:
-            tts_args = getattr(self._tts_stage, "tts_args", {})
+            topology = getattr(self._tts_stage, "stage_pipeline_config", None)
+            tts_args = (
+                getattr(topology, "extras", {}).get("tts_args", {})
+                if isinstance(topology, StagePipelineConfig)
+                else getattr(self._tts_stage, "tts_args", {})
+            )
             if "max_instructions_length" in tts_args:
                 return tts_args["max_instructions_length"]
 
@@ -1896,9 +1918,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             # chat-template markers (im_start_token_id 151644).  A raw-text
             # prompt produces a 1-token placeholder that crashes the talker's
             # prefill/decode handoff.  Reject early with an actionable message.
-            stage_names = {
-                getattr(getattr(s, "engine_args", None), "model_stage", None) for s in self.engine_client.stage_configs
-            }
+            stage_names = {_stage_speech_metadata(s)[0] for s in self.engine_client.stage_configs}
             if "talker" in stage_names:
                 raise ValueError(
                     "The /v1/audio/speech endpoint is only supported for "

vllm_omni/entrypoints/openai/serving_video.py

@@ -53,6 +53,30 @@ logger = init_logger(__name__)
 
 _VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS = 8
 
+
+def _config_value(config: Any, key: str, default: Any = None) -> Any:
+    if isinstance(config, Mapping):
+        return config.get(key, default)
+    return getattr(config, key, default)
+
+
+def _stage_diffusion_model_class_name(stage_config: Any) -> str | None:
+    """Resolve a stage's diffusion class across typed and legacy shapes."""
+    diffusion_config = _config_value(stage_config, "diffusion_config")
+    model_class_name = _config_value(diffusion_config, "model_class_name")
+    if model_class_name:
+        return str(model_class_name)
+
+    model_config = _config_value(stage_config, "model_config")
+    model_arch = _config_value(model_config, "model_arch") or _config_value(stage_config, "model_arch")
+    if model_arch:
+        return str(model_arch)
+
+    engine_args = _config_value(stage_config, "engine_args", {})
+    model_class_name = _config_value(engine_args, "model_class_name")
+    return str(model_class_name) if model_class_name else None
+
+
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -195,18 +219,7 @@ class OmniOpenAIServingVideo:
         model_class_name = getattr(od_config, "model_class_name", None)
         model_archs = [model_class_name]
         for stage_config in self.stage_configs or ():
-            stage_get = (
-                stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None)
-            )
-            engine_args = stage_get("engine_args") or {}
-            model_archs.extend(
-                (
-                    stage_get("model_arch"),
-                    engine_args.get("model_class_name")
-                    if isinstance(engine_args, Mapping)
-                    else getattr(engine_args, "model_class_name", None),
-                )
-            )
+            model_archs.append(_stage_diffusion_model_class_name(stage_config))
         metadata_capability = any(
             get_diffusion_model_metadata(model_arch).supports_mixed_reference_inputs for model_arch in model_archs
         )
@@ -223,18 +236,7 @@ class OmniOpenAIServingVideo:
         od_config = self._resolve_diffusion_od_config()
         model_archs = [None if od_config is None else getattr(od_config, "model_class_name", None)]
         for stage_config in self.stage_configs or ():
-            stage_get = (
-                stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None)
-            )
-            engine_args = stage_get("engine_args") or {}
-            model_archs.extend(
-                (
-                    stage_get("model_arch"),
-                    engine_args.get("model_class_name")
-                    if isinstance(engine_args, Mapping)
-                    else getattr(engine_args, "model_class_name", None),
-                )
-            )
+            model_archs.append(_stage_diffusion_model_class_name(stage_config))
 
         supported: set[str] = set()
         for model_arch in model_archs:

vllm_omni/entrypoints/openai/tts_adapters/audex.py

@@ -42,7 +42,11 @@ class AudexAdapter(ARTTSAdapter):
 
             from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
 
-            stage = getattr(getattr(self.ctx.server._tts_stage, "engine_args", None), "model_stage", None)
+            stage_config = self.ctx.server._tts_stage
+            topology = getattr(stage_config, "stage_pipeline_config", None)
+            stage = getattr(topology, "model_stage", None)
+            if stage is None:
+                stage = getattr(getattr(stage_config, "engine_args", None), "model_stage", None)
             profile, folder = (
                 ("full", "checkpoint_folder_full") if stage == "audex_omni" else ("tts", "checkpoint_folder_audiogen")
             )

vllm_omni/entrypoints/openai/tts_adapters/base.py

@@ -39,6 +39,21 @@ DEFAULT_TTS_LANGUAGES = frozenset(
 )
 
 
+def resolve_stage_model_path(engine_client: Any) -> str | None:
+    """Prefer typed stage model overrides, then legacy and served-model values."""
+    stages = getattr(engine_client, "stage_configs", ()) or ()
+    for stage in stages:
+        model_path = getattr(getattr(stage, "model_config", None), "model", None)
+        if model_path:
+            return str(model_path)
+    for stage in stages:
+        model_path = getattr(getattr(stage, "engine_args", None), "model", None)
+        if model_path:
+            return str(model_path)
+    model_path = getattr(engine_client, "model", None)
+    return str(model_path) if model_path else None
+
+
 def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict | None = None) -> str:
     """Stable hash of the real Stage 0 conditioning for the prefix cache.
 

vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py

@@ -9,7 +9,12 @@ import numpy as np
 from vllm.inputs import tokens_input
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    ARTTSAdapter,
+    PreparedRequest,
+    apply_max_new_tokens,
+    resolve_stage_model_path,
+)
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -81,17 +86,7 @@ class HiggsAudioV2Adapter(ARTTSAdapter):
 
         from transformers import AutoProcessor
 
-        model_path = None
-        for stage in self.engine_client.stage_configs:
-            model_path = getattr(getattr(stage, "engine_args", None), "model", None)
-            if model_path:
-                break
-        if model_path is None:
-            # Fallback: the orchestrator stores the served model id on the engine
-            # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args
-            # may not surface ``model`` when the deploy yaml doesn't set it per
-            # stage (the CLI-passed model id is the single source of truth).
-            model_path = getattr(self.engine_client, "model", None)
+        model_path = resolve_stage_model_path(self.engine_client)
         if model_path is None:
             raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs")
         processor = AutoProcessor.from_pretrained(model_path)

vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py

@@ -16,6 +16,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import (
     PreparedRequest,
     apply_max_new_tokens,
     conditioning_cache_salt,
+    resolve_stage_model_path,
 )
 
 _REF_CODE_CACHE_MAX_ENTRIES = 256
@@ -170,13 +171,7 @@ class HiggsAudioV3Adapter(ARTTSAdapter):
             HiggsAudioV3TokenizerAdapter,
         )
 
-        model_path = None
-        for stage in self.engine_client.stage_configs:
-            model_path = getattr(getattr(stage, "engine_args", None), "model", None)
-            if model_path:
-                break
-        if model_path is None:
-            model_path = getattr(self.engine_client, "model", None)
+        model_path = resolve_stage_model_path(self.engine_client)
         if model_path is None:
             raise RuntimeError("higgs_audio_v3 serving could not resolve model path")
         tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

vllm_omni/entrypoints/openai/video/generation/helpers.py

@@ -58,6 +58,7 @@ from vllm_omni.entrypoints.openai.serving_video import (
     ReferenceAudio,
     ReferenceImage,
     ReferenceVideo,
+    _stage_diffusion_model_class_name,
 )
 from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
 from vllm_omni.entrypoints.openai.stores import VIDEO_STORE
@@ -136,10 +137,6 @@ def _config_get(config: Any, key: str, default: Any = None) -> Any:
     return getattr(config, key, default)
 
 
-def _stage_engine_args(stage_cfg: Any) -> Any:
-    return _config_get(stage_cfg, "engine_args", {}) or {}
-
-
 def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]:
     if not stage_configs:
         return []
@@ -150,7 +147,7 @@ def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]:
     for stage_cfg in stage_configs:
         if get_stage_type(stage_cfg) != "diffusion":
             continue
-        model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name")
+        model_class_name = _stage_diffusion_model_class_name(stage_cfg)
         if not model_class_name:
             continue
         model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)

vllm_omni/entrypoints/openpi/serving.py

@@ -103,6 +103,14 @@ class ServingRealtimeRobotOpenPI:
             for stage_config in getattr(engine_client, "stage_configs", []) or []:
                 if getattr(stage_config, "stage_type", None) != "diffusion":
                     continue
+                # Typed diffusion stages keep model-owned OpenPI handshake
+                # metadata in diffusion_config.model_config. The out-of-process
+                # head only has this stage view because full od_config lives in
+                # the worker.
+                diffusion_config = getattr(stage_config, "diffusion_config", None)
+                model_config = getattr(diffusion_config, "model_config", None)
+                if model_config is not None:
+                    break
                 engine_args = getattr(stage_config, "engine_args", None)
                 model_config = getattr(engine_args, "model_config", None)
                 if model_config is not None:

vllm_omni/entrypoints/serve/profile/utils.py

@@ -4,18 +4,20 @@
 
 
 def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool:
-    """Check if any stage has profiler_config set in its engine_args."""
+    """Check whether any typed or legacy stage enables a profiler."""
     if not stage_configs:
         return False
     for stage in stage_configs:
-        engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None)
-        if engine_args is None:
-            continue
         profiler_config = (
-            engine_args.get("profiler_config")
-            if isinstance(engine_args, dict)
-            else getattr(engine_args, "profiler_config", None)
+            stage.get("profiler_config") if isinstance(stage, dict) else getattr(stage, "profiler_config", None)
         )
+        if profiler_config is None:
+            engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None)
+            profiler_config = (
+                engine_args.get("profiler_config")
+                if isinstance(engine_args, dict)
+                else getattr(engine_args, "profiler_config", None)
+            )
         if profiler_config is not None:
             profiler = (
                 profiler_config.get("profiler")

vllm_omni/entrypoints/stage_utils.py

@@ -1,7 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import logging
 import os
+from dataclasses import asdict, is_dataclass
 from multiprocessing import shared_memory as _shm
 from typing import Any
 
@@ -274,6 +278,8 @@ def _to_dict(x: Any) -> dict[str, Any]:
     try:
         if isinstance(x, dict):
             return dict(x)
+        if is_dataclass(x) and not isinstance(x, type):
+            return asdict(x)
         return _omega_to_dict(x)
     except Exception:
         try:

vllm_omni/entrypoints/utils.py

@@ -20,6 +20,13 @@ logger = init_logger(__name__)
 
 def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None:
     """Inject connector configuration into stage engine arguments."""
+    typed_connector_config = getattr(stage, "connector_config", None)
+    if typed_connector_config is not None:
+        omni_conf_dict = dict(typed_connector_config.omni_kv_config or {})
+        omni_conf_dict.update(connector_config=omni_conn_cfg, omni_from_stage=omni_from, omni_to_stage=omni_to)
+        typed_connector_config.omni_kv_config = omni_conf_dict
+        return
+
     # Prepare omni_kv_config dict
     omni_conf_dict = {}
     try: