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

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

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