[Core] Retain structured configs through runtime startup (#6849)

Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com>

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

@@ -12,6 +12,7 @@ import types
 import pytest
 from omegaconf import OmegaConf
 
+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
@@ -21,12 +22,23 @@ 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
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+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)
@@ -223,6 +235,50 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
     assert captured["log_stats"] is True
 
 
+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,
@@ -1715,6 +1771,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
@@ -1852,3 +1953,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

@@ -156,6 +156,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

@@ -64,6 +64,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 # 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

@@ -1134,6 +1134,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
@@ -1191,6 +1217,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

@@ -7,10 +7,12 @@ from __future__ import annotations
 
 import argparse
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
 from vllm_omni.config.resolver import OmniConfigResolution
 from vllm_omni.entrypoints.cli.serve import (
     OmniServeCommand,
@@ -334,6 +336,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
@@ -706,6 +728,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