Merge branch 'main' into feat/moss-tts-nano

Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/e2e/accuracy/conftest.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
 import os
-import shutil
 import subprocess
 from contextlib import contextmanager
 from dataclasses import dataclass
+from io import BytesIO
 from pathlib import Path
 
 import pytest
+import requests
 import torch
+from PIL import Image
 
-from tests.conftest import OmniServer, OmniServerParams
+from tests.helpers.runtime import OmniServer, OmniServerParams
 
 
 def pytest_addoption(parser):
@@ -183,16 +185,26 @@ def accuracy_artifact_root() -> Path:
     return root
 
 
-def reset_artifact_dir(path: Path) -> Path:
-    if path.exists():
-        shutil.rmtree(path)
-    path.mkdir(parents=True, exist_ok=True)
-    return path
+@pytest.fixture(scope="session")
+def qwen_bear_image(accuracy_artifact_root: Path) -> Image.Image:
+    """Download the Qwen bear image from the URL and save it to the accuracy artifact root."""
+    QWEN_BEAR_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
+    response = requests.get(QWEN_BEAR_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content)).convert("RGB")
+    image.save(accuracy_artifact_root / "qwen_bear.png")
+    return image
 
 
-def infer_model_label(model: str) -> str:
-    label = Path(model.rstrip("/\\")).name or "model"
-    return "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in label)
+@pytest.fixture(scope="session")
+def rabbit_image(accuracy_artifact_root: Path) -> Image.Image:
+    """Download the rabbit image from the URL and save it to the accuracy artifact root."""
+    RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
+    response = requests.get(RABBIT_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content)).convert("RGB")
+    image.save(accuracy_artifact_root / "rabbit.png")
+    return image
 
 
 def _build_accuracy_server_config(

tests/e2e/accuracy/helpers.py

@@ -0,0 +1,115 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+
+
+def reset_artifact_dir(path: Path) -> Path:
+    import shutil
+
+    if path.exists():
+        shutil.rmtree(path)
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def infer_model_label(model: str) -> str:
+    label = Path(model.rstrip("/\\")).name or "model"
+    return "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in label)
+
+
+def model_output_dir(parent_dir: Path, model: str) -> Path:
+    safe_model_name = model.split("/")[-1].replace(".", "_")
+    path = parent_dir / safe_model_name
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def assert_similarity(
+    *,
+    model_name: str,
+    vllm_image: Image.Image,
+    diffusers_image: Image.Image,
+    ssim_threshold: float,
+    psnr_threshold: float,
+    width: int | None = None,
+    height: int | None = None,
+    compare_mode: str = "RGB",
+) -> None:
+    requested_size = (width, height) if width is not None and height is not None else None
+    if requested_size is not None and diffusers_image.size != requested_size:
+        pytest.skip(
+            "Skipping as diffusers baseline output is corrupt and not comparable: "
+            f"dimensions do not match requested size; requested={requested_size}, got={diffusers_image.size}."
+        )
+
+    assert vllm_image.size == diffusers_image.size, (
+        f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}"
+    )
+
+    ssim_score, psnr_score = compute_image_ssim_psnr(
+        prediction=vllm_image,
+        reference=diffusers_image,
+        compare_mode=compare_mode,
+    )
+    print(f"{model_name} similarity metrics:")
+    print(f"  SSIM: value={ssim_score:.6f}, threshold>={ssim_threshold:.6f}, range=[-1, 1], higher_is_better=True")
+    print(
+        f"  PSNR: value={psnr_score:.6f} dB, threshold>={psnr_threshold:.6f} dB, range=[0, +inf), higher_is_better=True"
+    )
+
+    assert ssim_score >= ssim_threshold, (
+        f"SSIM below threshold for {model_name}: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}."
+    )
+    assert psnr_score >= psnr_threshold, (
+        f"PSNR below threshold for {model_name}: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}."
+    )
+
+
+def assert_image_sequence_similarity(
+    *,
+    model_name: str,
+    vllm_images: list[Image.Image],
+    diffusers_images: list[Image.Image],
+    ssim_threshold: float,
+    psnr_threshold: float,
+    compare_mode: str = "RGB",
+) -> None:
+    assert len(vllm_images) == len(diffusers_images), (
+        f"Output image count mismatch for {model_name}: online={len(vllm_images)}, diffusers={len(diffusers_images)}"
+    )
+    for index, (vllm_image, diffusers_image) in enumerate(zip(vllm_images, diffusers_images, strict=True), start=1):
+        assert_similarity(
+            model_name=f"{model_name}[layer={index}]",
+            vllm_image=vllm_image,
+            diffusers_image=diffusers_image,
+            ssim_threshold=ssim_threshold,
+            psnr_threshold=psnr_threshold,
+            compare_mode=compare_mode,
+        )
+
+
+def compute_image_ssim_psnr(
+    *,
+    prediction: Image.Image,
+    reference: Image.Image,
+    compare_mode: str = "RGB",
+) -> tuple[float, float]:
+    pred_tensor = _pil_to_batched_tensor(prediction, compare_mode=compare_mode)
+    ref_tensor = _pil_to_batched_tensor(reference, compare_mode=compare_mode)
+
+    ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0)
+    psnr_metric = PeakSignalNoiseRatio(data_range=1.0)
+
+    ssim_value = float(ssim_metric(pred_tensor, ref_tensor).item())
+    psnr_value = float(psnr_metric(pred_tensor, ref_tensor).item())
+    return ssim_value, psnr_value
+
+
+def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Tensor:
+    array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor

tests/e2e/accuracy/test_gebench_h100_smoke.py

@@ -6,8 +6,8 @@ from pathlib import Path
 import pytest
 
 from benchmarks.accuracy.text_to_image.gbench import main as gbench_main
-from tests.e2e.accuracy.conftest import infer_model_label, reset_artifact_dir
-from tests.utils import hardware_test
+from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
+from tests.helpers.mark import hardware_test
 
 
 @pytest.mark.advanced_model

tests/e2e/accuracy/test_gedit_bench_h100_smoke.py

@@ -7,8 +7,8 @@ import pytest
 
 from benchmarks.accuracy.image_to_image.gedit_bench import GROUPS
 from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main
-from tests.e2e.accuracy.conftest import infer_model_label, reset_artifact_dir
-from tests.utils import hardware_test
+from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
+from tests.helpers.mark import hardware_test
 
 
 @pytest.mark.advanced_model

tests/e2e/accuracy/test_qwen_image.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import base64
+import gc
+import io
+import os
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+MODEL_ID = "Qwen/Qwen-Image"
+MODEL_ENV_VAR = "QWEN_IMAGE_MODEL"
+PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style."
+NEGATIVE_PROMPT = "blurry, low quality"
+WIDTH = 512
+HEIGHT = 512
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 42
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
+
+
+def _model_name() -> str:
+    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
+def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
+    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
+            json={
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "size": f"{WIDTH}x{HEIGHT}",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "true_cfg_scale": TRUE_CFG_SCALE,
+                "seed": SEED,
+            },
+            timeout=600,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == 1
+        image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image.load()
+        image.save(output_path)
+        return image
+
+
+def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
+    run_pre_test_cleanup(enable_force=True)
+    pipe: DiffusionPipeline | None = None
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            prompt=PROMPT,
+            negative_prompt=NEGATIVE_PROMPT,
+            width=WIDTH,
+            height=HEIGHT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        run_post_test_cleanup(enable_force=True)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+
+    vllm_output = _run_vllm_omni_qwen_image(model=model, output_path=output_dir / "vllm_omni.png")
+    diffusers_output = _run_diffusers_qwen_image(model=model, output_path=output_dir / "diffusers.png")
+
+    assert_similarity(
+        model_name=MODEL_ID,
+        vllm_image=vllm_output,
+        diffusers_image=diffusers_output,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import gc
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline
+from PIL import Image
+
+from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
+MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509"
+WIDTH = 512
+HEIGHT = 512
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 42
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."
+PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting."
+NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion"
+SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+
+
+def _run_vllm_omni_image_edit(
+    *,
+    omni_server: OmniServer,
+    prompt: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    response = requests.post(
+        f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+        data={
+            "model": omni_server.model,
+            "prompt": prompt,
+            "size": f"{WIDTH}x{HEIGHT}",
+            "n": 1,
+            "response_format": "b64_json",
+            "negative_prompt": NEGATIVE_PROMPT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "true_cfg_scale": TRUE_CFG_SCALE,
+            "seed": SEED,
+        },
+        files=[
+            ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png"))
+            for index, image in enumerate(input_images)
+        ],
+        timeout=600,
+    )
+    response.raise_for_status()
+    payload = response.json()
+    assert len(payload["data"]) == 1
+    image = decode_base64_image(payload["data"][0]["b64_json"])
+    image.load()
+    image.save(output_path)
+    return image
+
+
+def _run_diffusers_image_edit(
+    *,
+    model: str,
+    pipeline_class: type[QwenImageEditPipeline] | type[QwenImageEditPlusPipeline],
+    prompt: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    run_pre_test_cleanup(enable_force=True)
+    pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
+    device = torch.device("cuda:0")
+    torch.cuda.set_device(device)
+    try:
+        images = input_images[0] if len(input_images) == 1 else input_images
+        pipe = pipeline_class.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+        ).to(device)
+        pipe.set_progress_bar_config(disable=False)
+        generator = torch.Generator(device=device).manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            prompt=prompt,
+            image=images,
+            negative_prompt=NEGATIVE_PROMPT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            width=WIDTH,
+            height=HEIGHT,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")  # pyright: ignore[reportAttributeAccessIssue]
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        run_post_test_cleanup(enable_force=True)
+
+
+def _vllm_omni_output_single_image(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
+    output_path = output_dir / "vllm_omni_single.png"
+    with OmniServer(model=SINGLE_MODEL, serve_args=SERVER_ARGS) as server:
+        output = _run_vllm_omni_image_edit(
+            omni_server=server,
+            prompt=PROMPT_SINGLE_IMAGE,
+            input_images=[qwen_bear_image],
+            output_path=output_path,
+        )
+    return output
+
+
+def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
+    output_path = output_dir / "diffusers_single.png"
+    return _run_diffusers_image_edit(
+        model=SINGLE_MODEL,
+        pipeline_class=QwenImageEditPipeline,
+        prompt=PROMPT_SINGLE_IMAGE,
+        input_images=[qwen_bear_image],
+        output_path=output_path,
+    )
+
+
+def _vllm_omni_output_multiple_image(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "vllm_omni_multiple.png"
+    with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
+        output = _run_vllm_omni_image_edit(
+            omni_server=server,
+            prompt=PROMPT_MULTIPLE_IMAGE,
+            input_images=[qwen_bear_image, rabbit_image],
+            output_path=output_path,
+        )
+    return output
+
+
+def _diffusers_output_multiple_image(
+    accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "diffusers_multiple.png"
+    return _run_diffusers_image_edit(
+        model=MULTIPLE_MODEL,
+        pipeline_class=QwenImageEditPlusPipeline,
+        prompt=PROMPT_MULTIPLE_IMAGE,
+        input_images=[qwen_bear_image, rabbit_image],
+        output_path=output_path,
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_edit_single_matches_diffusers(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+) -> None:
+    vllm_image = _vllm_omni_output_single_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+    )
+    diffusers_image = _diffusers_output_single_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+    )
+    assert_similarity(
+        model_name=SINGLE_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skip(
+    reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
+)
+def test_qwen_image_edit_multiple_matches_diffusers(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> None:
+    vllm_image = _vllm_omni_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    diffusers_image = _diffusers_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    assert_similarity(
+        model_name=MULTIPLE_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/test_qwen_image_layered.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import base64
+import gc
+import io
+import os
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import assert_image_sequence_similarity, model_output_dir
+from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+MODEL_ID = "Qwen/Qwen-Image-Layered"
+MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL"
+PROMPT = "decompose into layers"
+NEGATIVE_PROMPT = " "
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 777
+LAYERS = 3
+RESOLUTION = 640
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
+
+
+def _model_name() -> str:
+    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
+def _normalize_layered_images(images: object) -> list[Image.Image]:
+    if not isinstance(images, list) or not images:
+        raise AssertionError(f"Unexpected layered output container: {type(images).__name__}")
+
+    first_item = images[0]
+    if isinstance(first_item, Image.Image):
+        return [image.convert("RGBA") for image in images if isinstance(image, Image.Image)]
+    if isinstance(first_item, (list, tuple)):
+        return [image.convert("RGBA") for image in first_item if isinstance(image, Image.Image)]
+    raise AssertionError(f"Unexpected layered image element type: {type(first_item).__name__}")
+
+
+def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]:
+    input_image.save(output_dir / "input.png")
+    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        buffer = io.BytesIO()
+        input_image.save(buffer, format="PNG")
+        buffer.seek(0)
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+            data={
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "size": "auto",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "true_cfg_scale": TRUE_CFG_SCALE,
+                "seed": SEED,
+                "layers": LAYERS,
+                "resolution": RESOLUTION,
+            },
+            files=[("image", ("input.png", buffer, "image/png"))],
+            timeout=600,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == LAYERS
+        output_images = []
+        for item in payload["data"]:
+            image_bytes = base64.b64decode(item["b64_json"])
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGBA")
+            image.load()
+            output_images.append(image)
+        for index, image in enumerate(output_images, start=1):
+            image.save(output_dir / f"vllm_omni_layer_{index}.png")
+        return output_images
+
+
+def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]:
+    run_pre_test_cleanup(enable_force=True)
+    pipe: DiffusionPipeline | None = None
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            image=input_image,
+            prompt=PROMPT,
+            negative_prompt=NEGATIVE_PROMPT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            generator=generator,
+            num_images_per_prompt=1,
+            layers=LAYERS,
+            resolution=RESOLUTION,
+        )
+        output_images = _normalize_layered_images(result.images)
+        assert len(output_images) == LAYERS, f"Expected {LAYERS} diffusers layers, got {len(output_images)}"
+        for index, image in enumerate(output_images, start=1):
+            image.save(output_dir / f"diffusers_layer_{index}.png")
+        return output_images
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        run_post_test_cleanup(enable_force=True)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+    input_image = qwen_bear_image.convert("RGBA")
+
+    vllm_outputs = _run_vllm_omni_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir)
+    diffusers_outputs = _run_diffusers_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir)
+
+    assert_image_sequence_similarity(
+        model_name=MODEL_ID,
+        vllm_images=vllm_outputs,
+        diffusers_images=diffusers_outputs,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+        compare_mode="RGBA",
+    )

tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py

@@ -22,7 +22,6 @@ import torch
 from diffusers import UniPCMultistepScheduler
 from PIL import Image
 
-from tests.conftest import OmniServerParams
 from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
     _configure_scheduler,
     _ensure_wan_ftfy_fallback,
@@ -48,7 +47,8 @@ from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import (
     SSIM_THRESHOLD,
     WIDTH,
 )
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
 
 
 def test_parse_video_metadata_extracts_dimensions_and_fps() -> None:

tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py

@@ -26,7 +26,7 @@ from contextlib import ExitStack
 
 import pytest
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -19,7 +19,7 @@ import numpy as np
 import pytest
 from transformers import AutoTokenizer
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py

@@ -10,7 +10,7 @@ import pytest
 from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
     vLLMOmniColocateWorkerExtensionForTest,
 )
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -64,9 +64,6 @@ stage_args:
 # Top-level runtime config with Mooncake connector
 runtime:
   enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
   connectors:
     mooncake_connector:
       name: MooncakeConnector
@@ -80,4 +77,3 @@ runtime:
   edges:
     - from: 0
       to: 1
-      window_size: -1

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -62,10 +62,6 @@ stage_args:
 # Runtime edges
 runtime:
   enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
-
   # Distributed connectors configuration (optional)
   # More connectors will be supported in the future.
   connectors:
@@ -78,4 +74,3 @@ runtime:
   edges:
     - from: 0
       to: 1
-      window_size: -1

tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml

@@ -1,103 +0,0 @@
-# stage config for running qwen2.5-omni for multi-stage omni runtime.
-
-# This config is optimized for CI e2e tests.
-stage_args:
-  - stage_id: 0
-    runtime:
-      process: true            # Run this stage in a separate process
-      devices: "0"
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 896
-      max_num_batched_tokens: 896
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      enforce_eager: true  # Now we only support eager mode
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-      mm_processor_cache_gb: 0
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-  - stage_id: 1
-    runtime:
-      process: true
-      devices: "1"
-    engine_args:
-      model_stage: talker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 896
-      max_num_batched_tokens: 896
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
-    default_sampling_params:
-      temperature: 0.9
-      top_p: 0.8
-      top_k: 40
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-      stop_token_ids: [8294]
-  - stage_id: 2
-    runtime:
-      process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-    engine_args:
-      model_stage: code2wav
-      max_num_seqs: 1
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.15
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-    engine_input_source: [1]
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1             # Simplified: trigger downstream only after full upstream completion
-    max_inflight: 1             # Simplified: process serially within each stage
-  edges:
-    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
-      to: 1
-      window_size: -1
-    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
-      to: 2
-      window_size: -1

tests/e2e/offline_inference/test_bagel_img2img.py

@@ -22,9 +22,10 @@ import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
-from vllm_omni import Omni
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.platforms import current_omni_platform
 
 # Reference pixel data extracted from the known-good output image
@@ -32,30 +33,30 @@ from vllm_omni.platforms import current_omni_platform
 # prompt='Change the grass color to red',
 # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
 REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (157, 172, 217)},
-    {"position": (400, 50), "rgb": (105, 144, 218)},
-    {"position": (700, 100), "rgb": (118, 159, 233)},
-    {"position": (150, 400), "rgb": (195, 34, 60)},
-    {"position": (512, 336), "rgb": (222, 214, 193)},
-    {"position": (700, 400), "rgb": (197, 15, 43)},
-    {"position": (100, 600), "rgb": (105, 13, 18)},
-    {"position": (400, 600), "rgb": (169, 33, 44)},
-    {"position": (700, 600), "rgb": (101, 86, 93)},
-    {"position": (256, 256), "rgb": (181, 202, 222)},
+    {"position": (100, 100), "rgb": (156, 172, 217)},
+    {"position": (400, 50), "rgb": (105, 144, 217)},
+    {"position": (700, 100), "rgb": (118, 159, 232)},
+    {"position": (150, 400), "rgb": (180, 22, 52)},
+    {"position": (512, 336), "rgb": (221, 211, 194)},
+    {"position": (700, 400), "rgb": (192, 10, 46)},
+    {"position": (100, 600), "rgb": (102, 12, 22)},
+    {"position": (400, 600), "rgb": (161, 28, 47)},
+    {"position": (700, 600), "rgb": (100, 87, 94)},
+    {"position": (256, 256), "rgb": (181, 201, 221)},
 ]
 
 if current_omni_platform.is_rocm():
     REFERENCE_PIXELS = [
-        {"position": (100, 100), "rgb": (156, 172, 215)},
-        {"position": (400, 50), "rgb": (106, 144, 216)},
-        {"position": (700, 100), "rgb": (118, 158, 231)},
-        {"position": (150, 400), "rgb": (183, 23, 48)},
-        {"position": (512, 336), "rgb": (218, 215, 191)},
-        {"position": (700, 400), "rgb": (194, 14, 42)},
-        {"position": (100, 600), "rgb": (105, 10, 16)},
-        {"position": (400, 600), "rgb": (167, 33, 46)},
-        {"position": (700, 600), "rgb": (102, 86, 92)},
-        {"position": (256, 256), "rgb": (181, 201, 220)},
+        {"position": (100, 100), "rgb": (156, 172, 217)},
+        {"position": (400, 50), "rgb": (105, 144, 217)},
+        {"position": (700, 100), "rgb": (118, 159, 232)},
+        {"position": (150, 400), "rgb": (180, 22, 52)},
+        {"position": (512, 336), "rgb": (221, 211, 194)},
+        {"position": (700, 400), "rgb": (192, 10, 46)},
+        {"position": (100, 600), "rgb": (102, 12, 22)},
+        {"position": (400, 600), "rgb": (161, 28, 47)},
+        {"position": (700, 600), "rgb": (100, 87, 94)},
+        {"position": (256, 256), "rgb": (181, 201, 221)},
     ]
 
 PIXEL_TOLERANCE = 10

tests/e2e/offline_inference/test_bagel_lora.py

@@ -31,9 +31,10 @@ import torch
 from PIL import Image
 from safetensors.torch import save_file
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
-from vllm_omni import Omni
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -27,9 +27,10 @@ from typing import Any
 import pytest
 from PIL import Image
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
-from vllm_omni import Omni
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.platforms import current_omni_platform
 
 # Reference pixel data extracted from the known-good output image
@@ -37,30 +38,30 @@ from vllm_omni.platforms import current_omni_platform
 # "Generated with seed=52, num_inference_steps=15,
 # prompt='A futuristic city skyline at twilight, cyberpunk style'"
 REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (121, 118, 100)},
-    {"position": (400, 50), "rgb": (163, 162, 143)},
-    {"position": (700, 100), "rgb": (170, 156, 127)},
-    {"position": (150, 400), "rgb": (129, 127, 112)},
-    {"position": (512, 512), "rgb": (135, 61, 59)},
-    {"position": (700, 400), "rgb": (205, 107, 43)},
-    {"position": (100, 700), "rgb": (197, 177, 157)},
-    {"position": (400, 700), "rgb": (139, 107, 86)},
-    {"position": (700, 700), "rgb": (247, 205, 146)},
-    {"position": (256, 256), "rgb": (171, 160, 153)},
+    {"position": (100, 100), "rgb": (115, 113, 94)},
+    {"position": (400, 50), "rgb": (159, 160, 144)},
+    {"position": (700, 100), "rgb": (164, 151, 123)},
+    {"position": (150, 400), "rgb": (120, 121, 107)},
+    {"position": (512, 512), "rgb": (165, 133, 127)},
+    {"position": (700, 400), "rgb": (217, 130, 66)},
+    {"position": (100, 700), "rgb": (191, 168, 152)},
+    {"position": (400, 700), "rgb": (130, 96, 77)},
+    {"position": (700, 700), "rgb": (247, 203, 140)},
+    {"position": (256, 256), "rgb": (167, 156, 150)},
 ]
 
 if current_omni_platform.is_rocm():
     REFERENCE_PIXELS = [
-        {"position": (100, 100), "rgb": (123, 119, 100)},
-        {"position": (400, 50), "rgb": (162, 161, 142)},
-        {"position": (700, 100), "rgb": (171, 156, 127)},
-        {"position": (150, 400), "rgb": (131, 128, 112)},
-        {"position": (512, 512), "rgb": (134, 61, 59)},
-        {"position": (700, 400), "rgb": (204, 107, 43)},
-        {"position": (100, 700), "rgb": (201, 180, 165)},
-        {"position": (400, 700), "rgb": (140, 108, 87)},
-        {"position": (700, 700), "rgb": (247, 205, 145)},
-        {"position": (256, 256), "rgb": (171, 160, 153)},
+        {"position": (100, 100), "rgb": (115, 113, 94)},
+        {"position": (400, 50), "rgb": (159, 160, 144)},
+        {"position": (700, 100), "rgb": (164, 151, 123)},
+        {"position": (150, 400), "rgb": (120, 121, 107)},
+        {"position": (512, 512), "rgb": (165, 133, 127)},
+        {"position": (700, 400), "rgb": (217, 130, 66)},
+        {"position": (100, 700), "rgb": (191, 168, 152)},
+        {"position": (400, 700), "rgb": (130, 96, 77)},
+        {"position": (700, 700), "rgb": (247, 203, 140)},
+        {"position": (256, 256), "rgb": (167, 156, 150)},
     ]
 
 # Maximum allowed difference per color channel

tests/e2e/offline_inference/test_bagel_understanding.py

@@ -26,8 +26,9 @@ from pathlib import Path
 import pytest
 from vllm.assets.image import ImageAsset
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
 
 MODEL_NAME = "ByteDance-Seed/BAGEL-7B-MoT"
 STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")

tests/e2e/offline_inference/test_cache_dit.py

@@ -11,8 +11,8 @@ It uses minimal settings to keep test time short for CI.
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_cosyvoice3.py

@@ -15,19 +15,18 @@ from __future__ import annotations
 import functools
 import io
 import os
-import tempfile
 from pathlib import Path
 from urllib.request import urlopen
 
 import numpy as np
 import pytest
 import soundfile as sf
-import yaml
 from huggingface_hub import snapshot_download
 from vllm.sampling_params import SamplingParams
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 
@@ -44,13 +43,9 @@ REFERENCE_STAGE0_TEMPERATURE = 1.0
 REFERENCE_STAGE0_REPETITION_PENALTY = 2.0
 
 
-def _stage_config(name: str) -> str:
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
-STAGE_CONFIGS = [
-    _stage_config("cosyvoice3.yaml"),
-    _stage_config("cosyvoice3_async_chunk.yaml"),
+ASYNC_CHUNK_MODES = [
+    pytest.param(False, id="sync"),
+    pytest.param(True, id="async_chunk"),
 ]
 
 
@@ -138,19 +133,6 @@ def _get_stage_engine_outputs(omni_runner: OmniRunner, stage_id: int):
     raise AttributeError("Unable to locate stage outputs on Omni runner")
 
 
-def _patched_stage_config(base_stage_config: str, model_dir: Path, tmp_dir: Path) -> str:
-    cfg = yaml.safe_load(Path(base_stage_config).read_text(encoding="utf-8"))
-    tokenizer_path = str(model_dir / "CosyVoice-BlankEN")
-    for stage in cfg.get("stage_args", []):
-        engine_args = stage.setdefault("engine_args", {})
-        engine_args["tokenizer"] = tokenizer_path
-        engine_args["enforce_eager"] = True
-        engine_args["hf_overrides"] = {"architectures": ["CosyVoice3Model"]}
-    out_path = tmp_dir / Path(base_stage_config).name
-    out_path.write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
-    return str(out_path)
-
-
 def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[str, object]]:
     return [
         {
@@ -165,54 +147,54 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
 @pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("base_stage_config", STAGE_CONFIGS)
-def test_cosyvoice3_offline_reference_zero_shot(base_stage_config: str) -> None:
+@pytest.mark.parametrize("async_chunk", ASYNC_CHUNK_MODES)
+def test_cosyvoice3_offline_reference_zero_shot(async_chunk: bool) -> None:
     """CosyVoice3 zero-shot reference inference should stop cleanly and produce sane audio."""
     prompt_audio, prompt_sr = _load_reference_prompt_wav()
     model_dir = _resolve_model_dir()
     expected_stop_token = int(CosyVoice3Config().llm["eos_token_id"])
 
-    with tempfile.TemporaryDirectory(prefix="cv3-e2e-") as tmp:
-        stage_config = _patched_stage_config(base_stage_config, model_dir, Path(tmp))
-        with OmniRunner(
-            str(model_dir), seed=42, stage_configs_path=stage_config, stage_init_timeout=300
-        ) as omni_runner:
-            sampling_params_list = omni_runner.get_default_sampling_params_list()
-            sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT)
+    with OmniRunner(
+        str(model_dir),
+        seed=42,
+        stage_configs_path=get_deploy_config_path("cosyvoice3.yaml"),
+        async_chunk=async_chunk,
+        stage_init_timeout=300,
+    ) as omni_runner:
+        sampling_params_list = omni_runner.get_default_sampling_params_list()
+        sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT)
 
-            outputs = omni_runner.omni.generate(
-                _build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list
+        outputs = omni_runner.omni.generate(_build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list)
+
+        assert outputs, "No outputs returned"
+        audio_mm = outputs[0].multimodal_output
+        assert "audio" in audio_mm, "No audio output found"
+
+        audio = _concat_audio(audio_mm["audio"])
+        assert audio.size > 0, "Generated audio is empty"
+
+        sr_val = audio_mm.get("sr", 24000)
+        if isinstance(sr_val, list) and sr_val:
+            sr_val = sr_val[-1]
+        if hasattr(sr_val, "item"):
+            sr_val = sr_val.item()
+        sr = int(sr_val)
+        assert sr == 24000, f"Unexpected sample_rate={sr}"
+
+        duration_s = audio.size / sr
+        assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
+
+        stage0_outputs = _get_stage_engine_outputs(omni_runner, 0)
+        if stage0_outputs:
+            completion = stage0_outputs[0].outputs[0]
+            finish_reason = getattr(completion, "finish_reason", None)
+            stop_reason = getattr(completion, "stop_reason", None)
+            num_tokens = len(getattr(completion, "token_ids", []) or [])
+
+            assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'"
+            assert int(stop_reason) == expected_stop_token, (
+                f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}"
             )
-
-            assert outputs, "No outputs returned"
-            audio_mm = outputs[0].multimodal_output
-            assert "audio" in audio_mm, "No audio output found"
-
-            audio = _concat_audio(audio_mm["audio"])
-            assert audio.size > 0, "Generated audio is empty"
-
-            sr_val = audio_mm.get("sr", 24000)
-            if isinstance(sr_val, list) and sr_val:
-                sr_val = sr_val[-1]
-            if hasattr(sr_val, "item"):
-                sr_val = sr_val.item()
-            sr = int(sr_val)
-            assert sr == 24000, f"Unexpected sample_rate={sr}"
-
-            duration_s = audio.size / sr
-            assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
-
-            stage0_outputs = _get_stage_engine_outputs(omni_runner, 0)
-            if stage0_outputs:
-                completion = stage0_outputs[0].outputs[0]
-                finish_reason = getattr(completion, "finish_reason", None)
-                stop_reason = getattr(completion, "stop_reason", None)
-                num_tokens = len(getattr(completion, "token_ids", []) or [])
-
-                assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'"
-                assert int(stop_reason) == expected_stop_token, (
-                    f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}"
-                )
-                assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range"
-            else:
-                assert "async_chunk" in Path(base_stage_config).name, "Stage-0 produced no engine outputs"
+            assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range"
+        else:
+            assert async_chunk, "Stage-0 produced no engine outputs"

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -4,8 +4,9 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import OmniRunner
-from tests.utils import DeviceMemoryMonitor, hardware_test
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -2,8 +2,8 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import OmniRunner
-from tests.utils import DeviceMemoryMonitor
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -7,7 +7,7 @@ import pytest
 import torch
 from safetensors.torch import save_file
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_dynin_omni.py

@@ -18,7 +18,7 @@ import pytest
 import torch
 from transformers import AutoTokenizer
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"

tests/e2e/offline_inference/test_expert_parallel.py

@@ -18,8 +18,8 @@ import torch
 import torch.distributed as dist
 from PIL import Image
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_flux.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for Flux1 Schnell."""
+
+import pytest
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "black-forest-labs/FLUX.1-schnell"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+def test_flux_schnell_text_to_image():
+    """Test FLUX.1-schnell text-to-image generation."""
+    omni = Omni(model=MODEL)
+
+    omni_outputs = list(
+        omni.generate(
+            prompts=["A photo of a cat sitting on a laptop"],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=512,
+                width=512,
+                num_inference_steps=2,
+                seed=42,
+            ),
+        )
+    )
+
+    assert len(omni_outputs) > 0
+    images = omni_outputs[0].images
+    assert len(images) == 1
+    assert isinstance(images[0], Image.Image)
+    assert images[0].size == (512, 512)

tests/e2e/offline_inference/test_flux2_klein_inpaint.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for Flux2 Klein inpainting.
+
+"""
+
+# ruff: noqa: E402
+
+import os
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from PIL import Image, ImageDraw
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+MODEL = "black-forest-labs/FLUX.2-klein-4B"
+
+_HEIGHT = 512
+_WIDTH = 512
+_NUM_INFERENCE_STEPS = 4
+
+
+def _create_test_image(width: int = _WIDTH, height: int = _HEIGHT, color: tuple = (128, 128, 128)) -> Image.Image:
+    return Image.new("RGB", (width, height), color)
+
+
+def _create_test_mask(width: int = _WIDTH, height: int = _HEIGHT) -> Image.Image:
+    mask = Image.new("L", (width, height), 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
+    return mask
+
+
+def _create_test_inputs(color: tuple = (100, 150, 200)):
+    return _create_test_image(_WIDTH, _HEIGHT, color), _create_test_mask(_WIDTH, _HEIGHT)
+
+
+def _extract_images_from_output(outputs: list) -> list[Image.Image]:
+    images = []
+    for req_output in outputs:
+        if hasattr(req_output, "images") and req_output.images:
+            images.extend(req_output.images)
+        elif hasattr(req_output, "request_output") and req_output.request_output:
+            stage_out = req_output.request_output
+            if isinstance(stage_out, OmniRequestOutput) and hasattr(stage_out, "images"):
+                images.extend(stage_out.images)
+            elif isinstance(stage_out, list):
+                for s in stage_out:
+                    if hasattr(s, "images") and s.images:
+                        images.extend(s.images)
+    return images
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+def test_flux2_klein_inpaint_basic():
+    m = None
+    try:
+        m = Omni(model=MODEL)
+        input_image, mask_image = _create_test_inputs()
+
+        outputs = m.generate(
+            prompts=[
+                {
+                    "prompt": "Fill in the masked area with a beautiful garden",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        images = _extract_images_from_output(list(outputs))
+        assert len(images) == 1
+        assert images[0].size == (_WIDTH, _HEIGHT)
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()
+
+
+@pytest.mark.diffusion
+def test_flux2_klein_inpaint_deterministic():
+    m = None
+    try:
+        m = Omni(model=MODEL)
+        input_image, mask_image = _create_test_inputs()
+        seed = 12345
+
+        gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
+        gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
+
+        outputs1 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A red flower in a field",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen1,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        outputs2 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A red flower in a field",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen2,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        images1 = _extract_images_from_output(list(outputs1))
+        images2 = _extract_images_from_output(list(outputs2))
+
+        assert len(images1) == 1
+        assert len(images2) == 1
+
+        assert list(images1[0].getdata()) == list(images2[0].getdata()), (
+            "Same input with same seed should produce identical output. "
+            "This is critical for offline/online consistency."
+        )
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()
+
+
+@pytest.mark.diffusion
+def test_flux2_klein_inpaint_different_seeds_different_output():
+    m = None
+    try:
+        m = Omni(model=MODEL)
+        input_image, mask_image = _create_test_inputs()
+
+        gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(42)
+        gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(99999)
+
+        outputs1 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A beautiful landscape",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen1,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        outputs2 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A beautiful landscape",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen2,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        images1 = _extract_images_from_output(list(outputs1))
+        images2 = _extract_images_from_output(list(outputs2))
+
+        assert len(images1) == 1
+        assert len(images2) == 1
+
+        different_pixel_count = sum(1 for p1, p2 in zip(images1[0].getdata(), images2[0].getdata()) if p1 != p2)
+        assert different_pixel_count > 0, "Different seeds should produce different outputs"
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()

tests/e2e/offline_inference/test_flux_autoround_w4a16.py

@@ -14,8 +14,9 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import OmniRunner
-from tests.utils import DeviceMemoryMonitor, hardware_test
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_flux_kontext.py

@@ -13,7 +13,7 @@ import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 

tests/e2e/offline_inference/test_hunyuanimage3_text2img.py

@@ -8,7 +8,7 @@ import torch.nn.functional as F
 from PIL import Image
 from transformers import CLIPModel, CLIPProcessor
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -17,7 +17,7 @@ PROMPT = "A brown and white dog is running on the grass"
 MODEL_NAME = "tencent/HunyuanImage-3.0"
 LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32"
 REPO_ROOT = Path(__file__).resolve().parents[3]
-STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image_3_moe.yaml"
+STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_t2i.yaml"
 
 pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion]
 

tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py

@@ -11,7 +11,7 @@ import numpy as np
 import pytest
 from PIL import Image
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 
 REPO_ROOT = Path(__file__).resolve().parents[3]
 T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"

tests/e2e/offline_inference/test_magi_human.py

@@ -8,8 +8,8 @@ import av
 import numpy as np
 import pytest
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 

tests/e2e/offline_inference/test_mammoth_moda2.py

@@ -23,8 +23,10 @@ import pytest
 import torch
 from vllm.sampling_params import SamplingParams
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
 # ---------------------------------------------------------------------------
 # Constants

tests/e2e/offline_inference/test_ming_flash_omni.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import (
+    generate_synthetic_audio,
+    generate_synthetic_image,
+    generate_synthetic_video,
+)
+from tests.helpers.stage_config import modify_stage_config
+
+models = ["Jonathan1909/Ming-flash-omni-2.0"]
+
+# Ming-specific
+SYSTEM_PROMPT = "你是一个友好的AI助手。\n\ndetailed thinking off"
+EOS_TOKEN = "<|role_end|>"
+IMAGE_TOKEN = "<IMAGE>"
+VIDEO_TOKEN = "<VIDEO>"
+AUDIO_TOKEN = "<AUDIO>"
+
+
+def build_prompt(user_text: str) -> str:
+    """Build a Ming chat prompt."""
+    return (
+        f"<role>SYSTEM</role>{SYSTEM_PROMPT}{EOS_TOKEN}<role>HUMAN</role>{user_text}{EOS_TOKEN}<role>ASSISTANT</role>"
+    )
+
+
+def get_eager_config():
+    path = modify_stage_config(
+        str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_ci.yaml"),
+        updates={
+            "stage_args": {
+                0: {
+                    "engine_args.enforce_eager": "true",
+                },
+            },
+        },
+    )
+    return path
+
+
+stage_configs = [get_eager_config()]
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_text_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test text-only input processing and text output generation.
+    Input Modal: text
+    Output Modal: text
+    """
+    prompt = build_prompt("请详细介绍鹦鹉的生活习性。")
+    request_config = {"prompts": prompt, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_image_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test image understanding with text output.
+    Input Modal: image + text
+    Output Modal: text
+    """
+    image = generate_synthetic_image(224, 224)["np_array"]
+    prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.")
+    request_config = {"prompts": prompt, "images": image, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test audio understanding with text output.
+    Input Modal: audio + text
+    Output Modal: text
+    """
+    audio = generate_synthetic_audio(2, 1, 16000)["np_array"]
+    if len(audio.shape) == 2:
+        audio = audio.squeeze()
+    prompt = build_prompt(f"{AUDIO_TOKEN}Please recognize the language of this speech and transcribe it. Format: oral.")
+    request_config = {"prompts": prompt, "audios": audio, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_video_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test video understanding with text output.
+    Input Modal: video + text
+    Output Modal: text
+    """
+    video = generate_synthetic_video(224, 224, 30)["np_array"]
+    prompt = build_prompt(f"{VIDEO_TOKEN}Describe what is happening in this video.")
+    request_config = {"prompts": prompt, "videos": video, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test mixed modality input (image + audio) with text output.
+    Input Modal: image + audio + text
+    Output Modal: text
+    """
+    image = generate_synthetic_image(224, 224)["np_array"]
+    audio = generate_synthetic_audio(2, 1, 16000)["np_array"]
+    if len(audio.shape) == 2:
+        audio = audio.squeeze()
+    prompt = build_prompt(f"{IMAGE_TOKEN}{AUDIO_TOKEN}Describe the image and transcribe the audio.")
+    request_config = {"prompts": prompt, "images": image, "audios": audio, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)

tests/e2e/offline_inference/test_omni_sleep_mode.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Omni Sleep Mode across various model architectures.
+"""
+
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
+MODEL_DIFF = "riverclouds/qwen_image_random"
+
+
+def get_ack_info(ack, key, default=None):
+    if hasattr(ack, key):
+        return getattr(ack, key)
+    if isinstance(ack, dict):
+        return ack.get(key, default)
+    return default
+
+
+def get_dynamic_devices(stage_idx, num_stages, tp_size):
+    total_gpus = torch.cuda.device_count()
+    gpus_per_stage = tp_size
+    start_idx = stage_idx * gpus_per_stage
+    if start_idx + gpus_per_stage > total_gpus:
+        start_idx = start_idx % total_gpus
+    device_ids = [str(start_idx + i) for i in range(gpus_per_stage)]
+    return ",".join(device_ids)
+
+
+# Test 1: Diffusion Model (2-Stage BAGEL)
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_diffusion_model_sleep_tp(tp_size):
+    num_gpus = torch.cuda.device_count()
+    if num_gpus < tp_size:
+        pytest.skip(f"Skipping TP={tp_size}")
+
+    engine_args = {
+        "model": MODEL,
+        "enable_sleep_mode": True,
+        "tensor_parallel_size": tp_size,
+        "enforce_eager": True,
+        "trust_remote_code": True,
+        "dtype": "bfloat16",
+        "gpu_memory_utilization": 0.5,
+    }
+
+    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
+    try:
+        # BAGEL requires 2 params
+        diff_sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256)
+        llm_sp = SamplingParams()
+
+        # Warmup
+        async for _ in engine.generate("test", sampling_params=[llm_sp, diff_sp]):
+            pass
+
+        # Sleep all
+        acks = await engine.sleep(level=2)
+        statuses = [get_ack_info(ack, "status") for ack in acks]
+        assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
+
+        # Wakeup & Verify
+        await engine.wake_up()
+        async for _ in engine.generate("verify", sampling_params=[llm_sp, diff_sp]):
+            pass
+
+        print(f"Diffusion TP={tp_size} Lifecycle OK")
+    finally:
+        engine.shutdown()
+
+
+# Test 2: Multi-stage Manual Config
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_multistage_sleep_h100(tp_size):
+    num_gpus = torch.cuda.device_count()
+    if num_gpus < tp_size * 2:
+        pytest.skip("Not enough GPUs")
+
+    stages = []
+    for i in range(2):
+        devs = get_dynamic_devices(i, 2, tp_size)
+        stages.append(
+            {
+                "stage_id": i,
+                "stage_type": "llm" if i == 0 else "diffusion",
+                "runtime": {"process": True, "devices": devs},
+                "engine_args": {
+                    "model": MODEL,
+                    "model_stage": "thinker" if i == 0 else "base",
+                    "tensor_parallel_size": tp_size,
+                    "gpu_memory_utilization": 0.4,
+                    "dtype": "bfloat16",
+                    "enable_sleep_mode": True,
+                    "trust_remote_code": True,
+                },
+            }
+        )
+
+    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
+
+    engine = AsyncOmni(
+        model=MODEL, stages=stages, connectors=connectors, enable_sleep_mode=True, stage_init_timeout=1200
+    )
+    try:
+        sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+        async for _ in engine.generate("warmup", sampling_params=[SamplingParams(), sp]):
+            pass
+
+        acks = await engine.sleep(stage_ids=[0, 1], level=2)
+        assert len(acks) == 2 * tp_size
+
+        await engine.wake_up(stage_ids=[0, 1])
+        async for _ in engine.generate("verify", sampling_params=[SamplingParams(), sp]):
+            pass
+    finally:
+        engine.shutdown()
+
+
+# Test 3: Pure Diffusion Single-Stage
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_pure_diffusion_scenario(tp_size):
+    engine_args = {
+        "model": MODEL_DIFF,
+        "enable_sleep_mode": True,
+        "tensor_parallel_size": tp_size,
+        "enforce_eager": True,
+        "dtype": "bfloat16",
+        "gpu_memory_utilization": 0.5,
+    }
+
+    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
+    try:
+        await engine.sleep(level=1)
+        await engine.wake_up()
+        async for _ in engine.generate("test", sampling_params=[SamplingParams()]):
+            pass
+        print("Pure Diffusion OK")
+    finally:
+        engine.shutdown()

tests/e2e/offline_inference/test_omnivoice.py

@@ -16,8 +16,8 @@ from pathlib import Path
 import numpy as np
 import pytest
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 
 MODEL = "k2-fsa/OmniVoice"
 

tests/e2e/offline_inference/test_ovis_image.py

@@ -16,7 +16,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
 
 # Mock the OvisImageTransformer2DModel to avoid complex init if needed,

tests/e2e/offline_inference/test_quantization_fp8.py

@@ -36,8 +36,8 @@ from typing import Any
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -2,46 +2,39 @@
 E2E tests for Qwen2.5-Omni model with mixed modality inputs, audio and text output.
 """
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import (
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import (
     generate_synthetic_audio,
     generate_synthetic_image,
     generate_synthetic_video,
-    modify_stage_config,
 )
-from tests.utils import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
 
+# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
+# platforms: section. NPU still uses the legacy per-platform YAML until it
+# also migrates to the new schema.
+_CI_DEPLOY = get_deploy_config_path("ci/qwen2_5_omni.yaml")
+
 
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml"),
+    return modify_stage_config(
+        _CI_DEPLOY,
         updates={
-            "stage_args": {
-                0: {
-                    "engine_args.enforce_eager": "true",
-                },
-                1: {"engine_args.enforce_eager": "true"},
+            "stages": {
+                0: {"enforce_eager": True},
+                1: {"enforce_eager": True},
             },
         },
     )
-    return path
 
 
-# CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
-if current_omni_platform.is_npu():
-    stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
-elif current_omni_platform.is_rocm():
-    # ROCm stage config optimized for MI325 GPU
-    stage_config = str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
-elif current_omni_platform.is_xpu():
-    # Intel XPU stage config optimized for B60 GPU
-    stage_config = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")
+if current_omni_platform.is_rocm() or current_omni_platform.is_xpu() or current_omni_platform.is_npu():
+    stage_config = _CI_DEPLOY
 else:
     stage_config = get_cuda_graph_config()
 

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -7,41 +7,35 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import (
-    generate_synthetic_video,
-    modify_stage_config,
-)
-from tests.utils import hardware_test
+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.platforms import current_omni_platform
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 
+# 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")
+
+
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml"),
+    return modify_stage_config(
+        _CI_DEPLOY,
         updates={
-            "stage_args": {
-                0: {
-                    "engine_args.enforce_eager": "true",
-                },
-                1: {"engine_args.enforce_eager": "true"},
+            "stages": {
+                0: {"enforce_eager": True},
+                1: {"enforce_eager": True},
             },
         },
     )
-    return path
 
 
-# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if current_omni_platform.is_rocm():
-    # ROCm stage config optimized for MI325 GPU
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
-elif current_omni_platform.is_xpu():
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
+if current_omni_platform.is_rocm() or current_omni_platform.is_xpu():
+    stage_configs = [_CI_DEPLOY]
 else:
     stage_configs = [get_cuda_graph_config()]
 

tests/e2e/offline_inference/test_qwen3_tts_base.py

@@ -13,12 +13,10 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import modify_stage_config
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
@@ -26,23 +24,31 @@ REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know wh
 
 
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        get_stage_config(),
+    """Build a temp deploy yaml mirroring the deleted qwen3_tts_no_async_chunk.yaml.
+
+    Composes the synchronous (no-async-chunk) variant on top of the bundled
+    qwen3_tts.yaml prod default, with cudagraphs disabled. Replaces the deleted
+    standalone variant yaml; same effective config, no checked-in file needed.
+    """
+    return modify_stage_config(
+        get_deploy_config_path("qwen3_tts.yaml"),
         updates={
-            "stage_args": {
+            "async_chunk": False,
+            "stages": {
                 0: {
-                    "engine_args.enforce_eager": "true",
+                    "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
+                },
+                1: {
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
                 },
-                1: {"engine_args.enforce_eager": "true"},
             },
         },
     )
-    return path
-
-
-def get_stage_config(name: str = "qwen3_tts_no_async_chunk.yaml"):
-    """Get the no_async_chunk stage config path (async_chunk disable, cuda_graph disabled)."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
 
 
 # Same structure as test_qwen3_omni: models, stage_configs, test_params

tests/e2e/offline_inference/test_qwen3_tts_customvoice.py

@@ -13,34 +13,40 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import modify_stage_config
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 
 
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        get_stage_config(),
+    """Build a temp deploy yaml mirroring the deleted qwen3_tts_no_async_chunk.yaml.
+
+    Composes the synchronous (no-async-chunk) variant on top of the bundled
+    qwen3_tts.yaml prod default, with cudagraphs disabled. Replaces the deleted
+    standalone variant yaml; same effective config, no checked-in file needed.
+    """
+    return modify_stage_config(
+        get_deploy_config_path("qwen3_tts.yaml"),
         updates={
-            "stage_args": {
+            "async_chunk": False,
+            "stages": {
                 0: {
-                    "engine_args.enforce_eager": "true",
+                    "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
+                },
+                1: {
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
                 },
-                1: {"engine_args.enforce_eager": "true"},
             },
         },
     )
-    return path
-
-
-def get_stage_config(name: str = "qwen3_tts_no_async_chunk.yaml"):
-    """Get the no_async_chunk stage config path (async_chunk disable, cuda_graph disabled)."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
 
 
 # Same structure as test_qwen3_omni: models, stage_configs, test_params

tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py

@@ -36,8 +36,8 @@ from pathlib import Path
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -20,8 +20,8 @@ import torch
 import torch.distributed as dist
 from PIL import Image
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_stable_audio_expansion.py

@@ -15,8 +15,8 @@ import numpy as np
 import pytest
 import torch
 
-from tests.conftest import assert_audio_valid
-from tests.utils import hardware_test
+from tests.helpers.assertions import assert_audio_valid
+from tests.helpers.mark import hardware_test
 from vllm_omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/e2e/offline_inference/test_t2i_model.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_t2v_model.py

@@ -3,7 +3,7 @@ import os
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_teacache.py

@@ -11,8 +11,8 @@ It uses minimal settings to keep test time short for CI.
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_vae_decode_parallelism.py

@@ -18,7 +18,7 @@ if str(REPO_ROOT) not in sys.path:
 
 import time
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.platforms import current_omni_platform
 
 # os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"

tests/e2e/offline_inference/test_voxcpm.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test for VoxCPM offline inference."""
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+
+import tests.helpers.runtime as omni_runtime
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.model_executor.models.voxcpm.voxcpm_runtime_utils import (
+    prepare_voxcpm_hf_config_dir,
+    resolve_voxcpm_model_dir,
+)
+
+VOXCPM_MODEL = os.environ.get("VOXCPM_MODEL", "OpenBMB/VoxCPM1.5")
+STAGE_CONFIG = str(
+    Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
+)
+SAMPLE_RATE = 24000
+
+
+@pytest.fixture(autouse=True)
+def _patch_npu_cleanup_for_voxcpm(monkeypatch: pytest.MonkeyPatch):
+    """Limit the NPU cleanup workaround to this VoxCPM test module only."""
+    original_cleanup = omni_runtime.cleanup_dist_env_and_memory
+
+    def _safe_cleanup() -> None:
+        try:
+            original_cleanup()
+        except RuntimeError as exc:
+            if "Allocator for npu is not a DeviceAllocator" in str(exc):
+                return
+            raise
+
+    monkeypatch.setattr(omni_runtime, "cleanup_dist_env_and_memory", _safe_cleanup)
+
+
+def _build_prompt(text: str) -> dict[str, Any]:
+    return {
+        "prompt_token_ids": [1],
+        "additional_information": {
+            "text": [text],
+            "cfg_value": [2.0],
+            "inference_timesteps": [10],
+            "min_len": [2],
+            "max_new_tokens": [1024],
+        },
+    }
+
+
+def _extract_audio_tensor(multimodal_output: dict[str, Any]) -> torch.Tensor:
+    audio = multimodal_output.get("audio", multimodal_output.get("model_outputs"))
+    assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}"
+
+    if isinstance(audio, list):
+        parts: list[torch.Tensor] = []
+        for item in audio:
+            if item is None:
+                continue
+            tensor = torch.as_tensor(item)
+            if tensor.numel() == 0:
+                continue
+            parts.append(tensor.float().cpu().reshape(-1))
+        return torch.cat(parts, dim=-1) if parts else torch.zeros((0,), dtype=torch.float32)
+
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def _extract_final_multimodal_output(outputs) -> dict[str, Any]:
+    for item in reversed(outputs):
+        request_output = getattr(item, "request_output", None)
+        if request_output is not None:
+            multimodal_output = getattr(request_output, "multimodal_output", None)
+            if isinstance(multimodal_output, dict):
+                return multimodal_output
+            completions = getattr(request_output, "outputs", None) or []
+            for completion in completions:
+                multimodal_output = getattr(completion, "multimodal_output", None)
+                if isinstance(multimodal_output, dict):
+                    return multimodal_output
+
+        multimodal_output = getattr(item, "multimodal_output", None)
+        if isinstance(multimodal_output, dict):
+            return multimodal_output
+
+    raise AssertionError("No multimodal audio output found in VoxCPM generate results")
+
+
+@pytest.fixture
+def voxcpm_model_path(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> str:
+    model_dir = resolve_voxcpm_model_dir(VOXCPM_MODEL)
+
+    hf_config_env = os.environ.get("VLLM_OMNI_VOXCPM_HF_CONFIG_PATH")
+    if hf_config_env:
+        hf_config_dir = Path(hf_config_env).expanduser()
+    else:
+        hf_config_dir = tmp_path / "voxcpm_hf_config"
+
+    if not (hf_config_dir / "config.json").exists():
+        prepare_voxcpm_hf_config_dir(model_dir, hf_config_dir)
+
+    monkeypatch.setenv("VLLM_OMNI_VOXCPM_HF_CONFIG_PATH", str(hf_config_dir))
+    return str(model_dir)
+
+
+def test_prepare_voxcpm_hf_config_dir(tmp_path: Path):
+    model_dir = tmp_path / "model"
+    model_dir.mkdir()
+    (model_dir / "config.json").write_text(json.dumps({"hidden_size": 1024}), encoding="utf-8")
+    (model_dir / "generation_config.json").write_text(json.dumps({"do_sample": False}), encoding="utf-8")
+
+    hf_config_dir = prepare_voxcpm_hf_config_dir(model_dir, tmp_path / "voxcpm_hf_config")
+
+    prepared_config = json.loads((hf_config_dir / "config.json").read_text(encoding="utf-8"))
+    assert prepared_config["model_type"] == "voxcpm"
+    assert prepared_config["architectures"] == ["VoxCPMForConditionalGeneration"]
+    assert (hf_config_dir / "generation_config.json").exists()
+
+
+def test_resolve_voxcpm_model_dir_local_path(tmp_path: Path):
+    model_dir = tmp_path / "OpenBMB" / "VoxCPM1.5"
+    model_dir.mkdir(parents=True)
+
+    assert resolve_voxcpm_model_dir(str(model_dir)) == model_dir
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_voxcpm_zero_shot_001(voxcpm_model_path: str):
+    with OmniRunner(voxcpm_model_path, stage_configs_path=STAGE_CONFIG) as runner:
+        outputs = list(runner.omni.generate(_build_prompt("Hello, this is a VoxCPM offline inference test.")))
+
+    assert outputs, "No outputs returned"
+
+    multimodal_output = _extract_final_multimodal_output(outputs)
+    audio = _extract_audio_tensor(multimodal_output)
+    assert audio.numel() > SAMPLE_RATE // 2, f"Audio too short: {audio.numel()} samples"
+
+    duration_s = audio.shape[0] / SAMPLE_RATE
+    assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
+
+    peak = float(torch.max(torch.abs(audio)).item()) if audio.numel() > 0 else 0.0
+    assert peak > 0.01, "Generated audio appears to be silence"
+
+    audio_np = audio.numpy()
+    rms = float(np.sqrt(np.mean(np.square(audio_np)))) if audio_np.size else 0.0
+    assert rms > 1e-4, "Generated audio RMS too low"

tests/e2e/offline_inference/test_voxcpm2.py

@@ -5,27 +5,19 @@ import os
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
 
 VOXCPM2_MODEL = "openbmb/VoxCPM2"
-STAGE_CONFIG = os.path.join(
-    os.path.dirname(__file__),
-    "..",
-    "..",
-    "..",
-    "vllm_omni",
-    "model_executor",
-    "stage_configs",
-    "voxcpm2.yaml",
-)
+DEPLOY_CONFIG = get_deploy_config_path("voxcpm2.yaml")
 SAMPLE_RATE = 48000
 
 
 @pytest.fixture(scope="module")
 def voxcpm2_engine():
     """Create VoxCPM2 engine for testing."""
-    with OmniRunner(VOXCPM2_MODEL, stage_configs_path=STAGE_CONFIG) as runner:
+    with OmniRunner(VOXCPM2_MODEL, stage_configs_path=DEPLOY_CONFIG) as runner:
         yield runner.omni
 
 
@@ -33,14 +25,16 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor:
     """Extract the final complete audio tensor from multimodal output."""
     assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
 
-    # Output processor accumulates per-step full audio under "audio".
-    audio = multimodal_output.get("audio") or multimodal_output.get("model_outputs")
+    # Output processor accumulates per-step audio chunks under "audio".
+    audio = multimodal_output.get("audio")
+    if audio is None:
+        audio = multimodal_output.get("model_outputs")
     assert audio is not None, f"No audio key, got {list(multimodal_output.keys())}"
 
     if isinstance(audio, list):
-        valid = [x for x in audio if isinstance(x, torch.Tensor) and x.numel() > 100]
+        valid = [torch.as_tensor(x).float().cpu().reshape(-1) for x in audio if x is not None]
         assert valid, "No valid audio tensors in output list"
-        audio = valid[-1]
+        audio = torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
 
     assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
     return audio
@@ -98,3 +92,31 @@ def test_voxcpm2_voice_clone_002(voxcpm2_engine):
     audio = _extract_audio(outputs[0].outputs[0].multimodal_output)
     duration_s = audio.shape[0] / SAMPLE_RATE
     assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_voxcpm2_prefill_decode_mixed_batch_003(voxcpm2_engine):
+    """Regression: prefill+decode mixed batch must not crash (PR #2903)."""
+    long_prompt = (
+        "This is a deliberately long prompt that will stay in the decode "
+        "phase for many steps so that subsequent shorter prompts keep "
+        "entering prefill alongside it, reproducing the prefill plus "
+        "decode mixed batch scheduling pattern."
+    )
+    short_prompts = [
+        "Hello one.",
+        "Hello two.",
+        "Hello three.",
+        "Hello four.",
+    ]
+    requests = [{"prompt": long_prompt}] + [{"prompt": p} for p in short_prompts]
+
+    outputs = voxcpm2_engine.generate(requests)
+    assert len(outputs) == len(requests)
+
+    for i, out in enumerate(outputs):
+        audio = _extract_audio(out.outputs[0].multimodal_output)
+        duration_s = audio.shape[0] / SAMPLE_RATE
+        assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"

tests/e2e/offline_inference/test_voxtral_tts.py

@@ -29,14 +29,13 @@ from mistral_common.protocol.speech.request import SpeechRequest
 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
 from vllm import SamplingParams
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 MODEL = "mistralai/Voxtral-4B-TTS-2603"
-STAGE_CONFIG = str(
-    Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxtral_tts.yaml"
-)
+STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml")
 SAMPLE_RATE = 24000
 # Minimum expected audio samples for a short sentence (~0.04s of 24kHz audio)
 MIN_AUDIO_SAMPLES = 1000
@@ -65,9 +64,9 @@ def _resolve_stage_config(run_level: str) -> str:
         return modify_stage_config(
             STAGE_CONFIG,
             deletes={
-                "stage_args": {
-                    0: ["engine_args.load_format"],
-                    1: ["engine_args.load_format"],
+                "stages": {
+                    0: ["load_format"],
+                    1: ["load_format"],
                 }
             },
         )

tests/e2e/offline_inference/test_zimage_parallelism.py

@@ -20,8 +20,9 @@ import pytest
 import torch
 from PIL import Image
 
-from tests.conftest import OmniRunner
-from tests.utils import DeviceMemoryMonitor, hardware_test
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -213,7 +214,10 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     )
 
 
-@pytest.mark.integration
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
 def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
     if current_omni_platform.is_npu():
         pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA and ROCm for now.")

tests/e2e/offline_inference/utils.py

@@ -1,210 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import contextlib
-import functools
-import os
-import signal
-import subprocess
-import sys
-import tempfile
-from collections.abc import Callable
-from contextlib import ExitStack, suppress
-from pathlib import Path
-from typing import Any, Literal
-
-import cloudpickle
-from typing_extensions import ParamSpec
-from vllm.platforms import current_platform
-
-VLLM_PATH = Path(__file__).parent.parent.parent
-"""Path to root of the vLLM repository."""
-
-
-_P = ParamSpec("_P")
-
-
-def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
-    """Decorator to fork a new process for each test function.
-    See https://github.com/vllm-project/vllm/issues/7053 for more details.
-    """
-
-    @functools.wraps(func)
-    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
-        # Make the process the leader of its own process group
-        # to avoid sending SIGTERM to the parent process
-        os.setpgrp()
-        from _pytest.outcomes import Skipped
-
-        # Create a unique temporary file to store exception info from child
-        # process. Use test function name and process ID to avoid collisions.
-        with (
-            tempfile.NamedTemporaryFile(
-                delete=False,
-                mode="w+b",
-                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
-                suffix=".exc",
-            ) as exc_file,
-            ExitStack() as delete_after,
-        ):
-            exc_file_path = exc_file.name
-            delete_after.callback(os.remove, exc_file_path)
-
-            pid = os.fork()
-            print(f"Fork a new process to run a test {pid}")
-            if pid == 0:
-                # Parent process responsible for deleting, don't delete
-                # in child.
-                delete_after.pop_all()
-                try:
-                    func(*args, **kwargs)
-                except Skipped as e:
-                    # convert Skipped to exit code 0
-                    print(str(e))
-                    os._exit(0)
-                except Exception as e:
-                    import traceback
-
-                    tb_string = traceback.format_exc()
-
-                    # Try to serialize the exception object first
-                    exc_to_serialize: dict[str, Any]
-                    try:
-                        # First, try to pickle the actual exception with
-                        # its traceback.
-                        exc_to_serialize = {"pickled_exception": e}
-                        # Test if it can be pickled
-                        cloudpickle.dumps(exc_to_serialize)
-                    except (Exception, KeyboardInterrupt):
-                        # Fall back to string-based approach.
-                        exc_to_serialize = {
-                            "exception_type": type(e).__name__,
-                            "exception_msg": str(e),
-                            "traceback": tb_string,
-                        }
-                    try:
-                        with open(exc_file_path, "wb") as f:
-                            cloudpickle.dump(exc_to_serialize, f)
-                    except Exception:
-                        # Fallback: just print the traceback.
-                        print(tb_string)
-                    os._exit(1)
-                else:
-                    os._exit(0)
-            else:
-                pgid = os.getpgid(pid)
-                _pid, _exitcode = os.waitpid(pid, 0)
-                # ignore SIGTERM signal itself
-                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
-                # kill all child processes
-                os.killpg(pgid, signal.SIGTERM)
-                # restore the signal handler
-                signal.signal(signal.SIGTERM, old_signal_handler)
-                if _exitcode != 0:
-                    # Try to read the exception from the child process
-                    exc_info = {}
-                    if os.path.exists(exc_file_path):
-                        with (
-                            contextlib.suppress(Exception),
-                            open(exc_file_path, "rb") as f,
-                        ):
-                            exc_info = cloudpickle.load(f)
-
-                    original_exception = exc_info.get("pickled_exception")
-                    if original_exception is not None and isinstance(original_exception, Exception):
-                        # Re-raise the actual exception object if it was
-                        # successfully pickled.
-                        raise original_exception
-
-                    if (original_tb := exc_info.get("traceback")) is not None:
-                        # Use string-based traceback for fallback case
-                        raise AssertionError(
-                            f"Test {func.__name__} failed when called with"
-                            f" args {args} and kwargs {kwargs}"
-                            f" (exit code: {_exitcode}):\n{original_tb}"
-                        ) from None
-
-                    # Fallback to the original generic error
-                    raise AssertionError(
-                        f"function {func.__name__} failed when called with"
-                        f" args {args} and kwargs {kwargs}"
-                        f" (exit code: {_exitcode})"
-                    ) from None
-
-    return wrapper
-
-
-def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
-    """Decorator to spawn a new process for each test function."""
-
-    @functools.wraps(f)
-    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
-        # Check if we're already in a subprocess
-        if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
-            # If we are, just run the function directly
-            return f(*args, **kwargs)
-
-        import torch.multiprocessing as mp
-
-        with suppress(RuntimeError):
-            mp.set_start_method("spawn")
-
-        # Get the module
-        module_name = f.__module__
-
-        # Create a process with environment variable set
-        env = os.environ.copy()
-        env["RUNNING_IN_SUBPROCESS"] = "1"
-
-        with tempfile.TemporaryDirectory() as tempdir:
-            output_filepath = os.path.join(tempdir, "new_process.tmp")
-
-            # `cloudpickle` allows pickling complex functions directly
-            input_bytes = cloudpickle.dumps((f, output_filepath))
-
-            repo_root = str(VLLM_PATH.resolve())
-
-            env = dict(env or os.environ)
-            env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")
-
-            cmd = [sys.executable, "-m", f"{module_name}"]
-
-            returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env)
-
-            # check if the subprocess is successful
-            try:
-                returned.check_returncode()
-            except Exception as e:
-                # wrap raised exception to provide more information
-                raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e
-
-    return wrapper
-
-
-def create_new_process_for_each_test(
-    method: Literal["spawn", "fork"] | None = None,
-) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
-    """Creates a decorator that runs each test function in a new process.
-
-    Args:
-        method: The process creation method. Can be either "spawn" or "fork".
-               If not specified, it defaults to "spawn" on ROCm and XPU
-               platforms and "fork" otherwise.
-
-    Returns:
-        A decorator to run test functions in separate processes.
-    """
-    if method is None:
-        # TODO: Find out why spawn is not working correctly on ROCm
-        # The test content will not run and tests passed immediately.
-        # For now, using `fork` for ROCm as it can run with `fork`
-        # and tests are running correctly.
-        use_spawn = current_platform.is_xpu()
-        method = "spawn" if use_spawn else "fork"
-
-    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
-
-    if method == "fork":
-        return fork_new_process_for_each_test
-
-    return spawn_new_process_for_each_test

tests/e2e/online_serving/test_bagel_expansion.py

@@ -16,13 +16,8 @@ assert_diffusion_response validates successful generation and the expected
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
 NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
@@ -88,7 +83,7 @@ def _get_diffusion_feature_cases(model: str):
                 ],
             ),
             id="parallel_tp_2",
-            marks=PARALLEL_FEATURE_MARKS,
+            marks=[*PARALLEL_FEATURE_MARKS, pytest.mark.skip(reason="issue: #2862")],
         ),
         # Ulysses-SP degree=2 (2 GPUs)
         pytest.param(
@@ -138,7 +133,7 @@ def test_bagel(
     - Ulysses-SP (degree=2)
     - Ring-Attention (degree=2)
 
-    Validation is delegated to assert_diffusion_response in tests.conftest,
+    Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py,
     which checks output dimensions and basic correctness.
     """
 

tests/e2e/online_serving/test_bagel_online.py

@@ -28,8 +28,8 @@ from pathlib import Path
 import pytest
 from vllm.assets.image import ImageAsset
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"

tests/e2e/online_serving/test_cosyvoice3_tts.py

@@ -12,12 +12,11 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
 
@@ -27,8 +26,8 @@ REF_TEXT = "希望你以后能够做的比我还好呦。"
 
 
 def get_stage_config(name: str = "cosyvoice3.yaml"):
-    """Get the stage config path from vllm_omni model_executor stage_configs."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+    """Get the deploy config path for CosyVoice3."""
+    return get_deploy_config_path(name)
 
 
 def get_prompt(prompt_type="zh"):
@@ -44,7 +43,7 @@ tts_server_params = [
         OmniServerParams(
             model=MODEL,
             stage_config_path=get_stage_config(),
-            server_args=["--trust-remote-code", "--disable-log-stats"],
+            server_args=["--trust-remote-code", "--disable-log-stats", "--no-async-chunk"],
         ),
         id="cosyvoice3",
     )
@@ -54,7 +53,7 @@ tts_async_chunk_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_stage_config("cosyvoice3_async_chunk.yaml"),
+            stage_config_path=get_stage_config(),
             server_args=["--trust-remote-code", "--disable-log-stats"],
         ),
         id="cosyvoice3_async_chunk",
@@ -93,7 +92,7 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None:
 def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     """
     Test voice cloning TTS with Chinese text via async_chunk streaming.
-    Deploy Setting: cosyvoice3_async_chunk.yaml
+    Deploy Setting: cosyvoice3.yaml with default ``async_chunk: true``
     Input Modal: text + ref_audio + ref_text
     Output Modal: audio (streamed)
     Input Setting: stream=True

tests/e2e/online_serving/test_dynin_omni_expansion.py

@@ -15,9 +15,9 @@ import pytest
 import soundfile as sf
 from vllm.assets.image import ImageAsset
 
-from tests import conftest as tests_conftest
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import convert_audio_bytes_to_text
+from tests.helpers.runtime import OmniServerParams
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
@@ -87,7 +87,7 @@ def _convert_audio_bytes_to_text_without_ffmpeg(raw_bytes: bytes) -> str:
 @pytest.fixture
 def dynin_t2s_openai_client(openai_client, monkeypatch):
     monkeypatch.setattr(
-        tests_conftest,
+        convert_audio_bytes_to_text,
         "convert_audio_bytes_to_text",
         _convert_audio_bytes_to_text_without_ffmpeg,
     )

tests/e2e/online_serving/test_flux2_expansion.py

@@ -11,12 +11,8 @@ Coverage:
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
 FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
 POSITIVE_PROMPT = "A cat sitting on a windowsill"

tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Flux2 Klein inpainting in online serving mode.
+
+Uses /v1/images/edits endpoint which is the correct API for image inpainting.
+"""
+
+import base64
+from io import BytesIO
+
+import httpx
+import pytest
+from PIL import Image, ImageDraw
+
+from tests.helpers.runtime import OmniServer, OmniServerParams
+
+MODEL = "black-forest-labs/FLUX.2-klein-4B"
+
+_HEIGHT = 512
+_WIDTH = 512
+_NUM_INFERENCE_STEPS = 4
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--tensor-parallel-size", "2"],
+            ),
+            id="tp2_basic",
+        ),
+    ]
+
+
+def _image_to_base64_jpeg(image: Image.Image) -> str:
+    buffer = BytesIO()
+    image.save(buffer, format="JPEG")
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
+    mask = Image.new("L", (width, height), 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
+    return _image_to_base64_jpeg(mask)
+
+
+def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
+    return list(img1.getdata()) == list(img2.getdata())
+
+
+def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
+    url = f"http://{host}:{port}/v1/images/edits"
+    files = {
+        "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
+        "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
+    }
+    data = {"prompt": prompt, "model": model, **kwargs}
+    with httpx.Client(timeout=60.0) as client:
+        response = client.post(url, files=files, data=data)
+        response.raise_for_status()
+        return response.json()
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+
+    result = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt="Fill in the masked area with a beautiful garden",
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=42,
+    )
+
+    assert "data" in result and len(result["data"]) == 1
+    img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
+    img = Image.open(BytesIO(base64.b64decode(img_data)))
+    assert img.size == (_WIDTH, _HEIGHT)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+    prompt = "A red flower in a field"
+
+    result1 = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt=prompt,
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=12345,
+    )
+
+    result2 = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt=prompt,
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=12345,
+    )
+
+    img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
+    img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
+
+    img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
+    img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
+
+    assert _compare_images(img1, img2), (
+        "Same input with same seed should produce identical output. This is critical for offline/online consistency."
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+
+    result = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt="A beautiful landscape",
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=2,
+        seed=42,
+    )
+
+    assert "data" in result and len(result["data"]) == 2

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -14,13 +14,8 @@ image resolution.
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 MODEL = "black-forest-labs/FLUX.2-dev"
 PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography."

tests/e2e/online_serving/test_flux_kontext_expansion.py

@@ -5,13 +5,8 @@ and are supported by the FluxKontext model.
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-    generate_synthetic_image,
-)
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
 NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"

tests/e2e/online_serving/test_hunyuan_video_15_expansion.py

@@ -11,12 +11,8 @@ Coverage (H100, since model cannot fit L4):
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
 PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion."
 NEGATIVE_PROMPT = "low quality, blurry, distorted"

tests/e2e/online_serving/test_image_gen_edit.py

@@ -22,7 +22,7 @@ from PIL import Image
 from vllm.assets.image import ImageAsset
 from vllm.utils.network_utils import get_open_port
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 # Increase timeout for downloading assets from S3 (default 5s is too short for CI)

tests/e2e/online_serving/test_images_generations_lora.py

@@ -22,8 +22,8 @@ import torch
 from PIL import Image
 from safetensors.torch import save_file
 
-from tests.conftest import OmniServer
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 

tests/e2e/online_serving/test_longcat_image_edit_expansion.py

@@ -13,14 +13,9 @@ This validates:
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-    generate_synthetic_image,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 EDIT_PROMPT = "Transform this modern image into a cinematic animation style with vibrant colors and soft lighting."
 NEGATIVE_PROMPT = "blurry, low quality, distorted, oversaturated"

tests/e2e/online_serving/test_longcat_image_expansion.py

@@ -13,13 +13,8 @@ This validates:
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 TEXT_TO_IMAGE_PROMPT = (
     "A cinematic illustration of a cat typing on a silver laptop, soft window light, highly detailed."

tests/e2e/online_serving/test_mimo_audio.py

@@ -9,13 +9,10 @@ from pathlib import Path
 
 import pytest
 
-from tests.conftest import (
-    OmniServerParams,
-    dummy_messages_from_mix_data,
-    generate_synthetic_audio,
-    modify_stage_config,
-)
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -29,26 +26,6 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 models = ["XiaomiMiMo/MiMo-Audio-7B-Instruct"]
 
 
-def get_chunk_config():
-    path = modify_stage_config(
-        str(Path(__file__).parent.parent / "stage_configs" / "mimo_audio_ci.yaml"),
-        updates={
-            "async_chunk": True,
-            "stage_args": {
-                0: {
-                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav_async_chunk"
-                },
-                1: {
-                    "engine_args.max_model_len": 8192,
-                    "engine_args.max_num_batched_tokens": 8192,
-                },
-            },
-        },
-        deletes={"stage_args": {1: ["custom_process_input_func"]}},
-    )
-    return path
-
-
 def download_tokenizer():
     tokenizer_path = os.environ.get("MIMO_AUDIO_TOKENIZER_PATH", MIMO_AUDIO_TOKENIZER_REPO)
     if os.path.exists(tokenizer_path):
@@ -62,19 +39,21 @@ def download_tokenizer():
     return local_path
 
 
-# CI stage config for H100 / MI325
 # Guard module-level setup so test collection doesn't fail in environments
 # where the model cache is read-only or models aren't available.
 try:
-    stage_configs = [get_chunk_config()]
+    stage_configs = [get_deploy_config_path("mimo_audio.yaml")]
     tokenizer_path = download_tokenizer()
     os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
 
+    # --load-format dummy applies to every stage pipeline-wide, avoiding a
+    # per-stage yaml rewrite (the old approach wrote a tempfile + atexit-unlink
+    # which raced with CI's process lifecycle).
     test_params = [
         OmniServerParams(
             model=model,
             stage_config_path=stage_config,
-            server_args=["--chat-template", CHAT_TEMPLATE_PATH],
+            server_args=["--chat-template", CHAT_TEMPLATE_PATH, "--load-format", "dummy"],
         )
         for model in models
         for stage_config in stage_configs

tests/e2e/online_serving/test_ming_flash_omni.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online serving tests for Ming-flash-omni-2.0 model (Thinker stage).
+Tests multimodal understanding via OpenAI-compatible API.
+"""
+
+import os
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import (
+    generate_synthetic_audio,
+    generate_synthetic_image,
+    generate_synthetic_video,
+)
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import modify_stage_config
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+models = ["Jonathan1909/Ming-flash-omni-2.0"]
+
+
+def get_eager_config():
+    path = modify_stage_config(
+        str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_ci.yaml"),
+        updates={
+            "stage_args": {
+                0: {
+                    "engine_args.enforce_eager": "true",
+                },
+            },
+        },
+    )
+    return path
+
+
+stage_configs = [get_eager_config()]
+
+# Create parameter combinations for model and stage config
+test_params = [
+    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
+]
+
+
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": "你是一个友好的AI助手。\n\ndetailed thinking off",
+            }
+        ],
+    }
+
+
+def get_prompt(prompt_type="text_only"):
+    prompts = {
+        "text_only": "What is the capital of China? Answer in 20 words.",
+        "text_image": "What is in this image?",
+        "text_audio": "What is in this audio?",
+        "text_video": "What is in this video?",
+        "mix": "What is recited in the audio? What is in this image? What is in this video?",
+    }
+    return prompts.get(prompt_type, prompts["text_only"])
+
+
+def get_max_batch_size(size_type="few"):
+    batch_sizes = {"few": 5, "medium": 100, "large": 256}
+    return batch_sizes.get(size_type, 5)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_text_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text
+    Output Modal: text
+    Input Setting: stream=False
+    Datasets: single request
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt("text_only"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+
+    openai_client.send_omni_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_text_stream_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text
+    Output Modal: text
+    Input Setting: stream=True
+    Datasets: few requests
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt("text_only"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_image_to_text_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: image + text
+    Output Modal: text
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        image_data_url=image_data_url,
+        content_text=get_prompt("text_image"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "modalities": ["text"],
+    }
+
+    openai_client.send_omni_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_to_text_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: audio + text
+    Output Modal: text
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(2, 1)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        audio_data_url=audio_data_url,
+        content_text=get_prompt("text_audio"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "modalities": ["text"],
+    }
+
+    openai_client.send_omni_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_video_to_text_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: video + text
+    Output Modal: text
+    Input Setting: stream=False
+    Datasets: single request
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        content_text=get_prompt("text_video"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+    }
+
+    openai_client.send_omni_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_mix_to_text_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text + audio + image + video
+    Output Modal: text
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(2, 1)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        image_data_url=image_data_url,
+        audio_data_url=audio_data_url,
+        content_text=get_prompt("mix"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "modalities": ["text"],
+    }
+
+    openai_client.send_omni_request(request_config)

tests/e2e/online_serving/test_nextstep_expansion.py

@@ -0,0 +1,71 @@
+"""
+Online serving E2E for NextStep-1.1 text-to-image (tensor parallel).
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+)
+
+# L4: 4 GPUs + TP=4; XPU B60: 2 cards (use num_cards={"cuda": 4, "xpu": 4} if needed)
+FOUR_CARD_MARKS = hardware_marks(
+    res={"cuda": "L4", "xpu": "B60"},
+    num_cards={"cuda": 2, "xpu": 2},
+)
+
+POSITIVE_PROMPT = "A small red barn in a snowy field, simple illustration."
+NEGATIVE_PROMPT = "blurry, low quality"
+
+_DEFAULT_MODEL = "stepfun-ai/NextStep-1.1"
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Single online config: TP=4, explicit pipeline class."""
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--tensor-parallel-size",
+                    "2",
+                    "--model-class-name",
+                    "NextStep11Pipeline",
+                ],
+            ),
+            id="nextstep_tp4_pipeline",
+            marks=FOUR_CARD_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(model=os.environ.get("VLLM_TEST_NEXTSTEP_MODEL", _DEFAULT_MODEL)),
+    indirect=True,
+)
+def test_nextstep_11(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "guidance_scale": 5.0,
+            "guidance_scale_2": 1.0,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_omnivoice.py

@@ -17,8 +17,9 @@ from pathlib import Path
 import httpx
 import pytest
 
-from tests.conftest import OmniServerParams, generate_synthetic_audio
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio
+from tests.helpers.runtime import OmniServerParams
 
 try:
     from transformers import HiggsAudioV2TokenizerModel  # noqa: F401

tests/e2e/online_serving/test_qwen2_5_omni.py

@@ -3,20 +3,13 @@ E2E Online tests for Qwen2.5-Omni model with video input and audio output.
 """
 
 import os
-from pathlib import Path
 
 import pytest
 
-from tests.conftest import (
-    OmniServerParams,
-    dummy_messages_from_mix_data,
-    generate_synthetic_audio,
-    generate_synthetic_image,
-    generate_synthetic_video,
-    modify_stage_config,
-)
-from tests.utils import hardware_test
-from vllm_omni.platforms import current_omni_platform
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
@@ -24,20 +17,9 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
 
-
-def get_config():
-    path = modify_stage_config(
-        str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml"),
-    )
-    return path
-
-
-# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if current_omni_platform.is_rocm():
-    # ROCm stage config optimized for MI325 GPU
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")]
-else:
-    stage_configs = [get_config()]
+# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
+# platforms: section in vllm_omni/deploy/ci/qwen2_5_omni.yaml.
+stage_configs = [modify_stage_config(get_deploy_config_path("ci/qwen2_5_omni.yaml"))]
 
 # Create parameter combinations for model and stage config
 test_params = [

tests/e2e/online_serving/test_qwen3_omni.py

@@ -3,19 +3,13 @@ E2E Online tests for Qwen3-Omni model with video input and audio output.
 """
 
 import os
-from pathlib import Path
 
 import pytest
 
-from tests.conftest import (
-    OmniServerParams,
-    dummy_messages_from_mix_data,
-    generate_synthetic_audio,
-    generate_synthetic_image,
-    generate_synthetic_video,
-    modify_stage_config,
-)
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -24,35 +18,64 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
+# Set VLLM_TEST_PD_MODE=1 to test PD disaggregation (follow-up — deploy overlay not yet migrated).
+_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
 
-def get_chunk_config():
+_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
+
+
+def get_chunk_config(config_path: str | None = None):
+    """Load the qwen3_omni CI deploy yaml with async_chunk modifications for streaming mode."""
+    if config_path is None:
+        config_path = _CI_DEPLOY
+    # TODO: remove this workaround once legacy `stage_args` path is deleted.
+    # The pipeline (qwen3_omni/pipeline.py) already wires
+    # thinker2talker_async_chunk / talker2code2wav_async_chunk on stage 0/1,
+    # so only async_chunk needs flipping. Writing nested `engine_args:` into
+    # the new-schema overlay trips _parse_stage_deploy's legacy branch and
+    # drops flat fields (load_format, max_num_seqs, ...).
+    return modify_stage_config(config_path, updates={"async_chunk": True})
+
+
+def get_prefix_caching_config(config_path: str):
+    """Create a stage config with prefix caching enabled on the thinker (stage 0)."""
     path = modify_stage_config(
-        str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml"),
+        config_path,
         updates={
-            "async_chunk": True,
             "stage_args": {
-                0: {
-                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
-                },
-                1: {
-                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk"
-                },
+                0: {"engine_args.enable_prefix_caching": True},
             },
         },
-        deletes={"stage_args": {2: ["custom_process_input_func"]}},
     )
     return path
 
 
+# Platform-specific overrides live inside the new deploy yaml's ``platforms:``
+# section, so a single ``_CI_DEPLOY`` path serves CUDA, ROCm, and XPU.
+# TODO: re-add VLLM_TEST_PD_MODE branch once the PD-disaggregation deploy
+# overlay has been migrated to the new schema (previously used the deleted
+# ``qwen3_omni_moe_pd_ci.yaml`` stage-configs file).
 if current_omni_platform.is_xpu():
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
-else:  # MI325 GPU should share the same config as H100
+    stage_configs = [_CI_DEPLOY]
+else:  # CUDA + ROCm MI325 share the same deploy config
     stage_configs = [get_chunk_config()]
+prefix_caching_stage_configs = [get_prefix_caching_config(_CI_DEPLOY)]
 
 # Create parameter combinations for model and stage config
 test_params = [
     OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
 ]
+# For prefix caching, we need to enable prompt token details so that we
+# can determine if any tokens were cached.
+prefix_test_params = [
+    OmniServerParams(
+        model=model,
+        stage_config_path=stage_config,
+        server_args=["--enable-prompt-tokens-details"],  # Enable prompt tokens details to get cached_tokens
+    )
+    for model in models
+    for stage_config in prefix_caching_stage_configs
+]
 
 
 def get_system_prompt():
@@ -75,6 +98,7 @@ def get_prompt(prompt_type="text_only"):
     prompts = {
         "text_only": "What is the capital of China? Answer in 20 words.",
         "mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
+        "text_image": "What color are the squares in this image?",
     }
     return prompts.get(prompt_type, prompts["text_only"])
 
@@ -87,7 +111,8 @@ def get_max_batch_size(size_type="few"):
 @pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
     """
@@ -126,7 +151,8 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_001(omni_server, openai_client) -> None:
     """
@@ -147,3 +173,42 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
     }
 
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True)
+@pytest.mark.skip(reason="issue: #2833")
+def test_thinker_prefix_caching(omni_server, openai_client) -> None:
+    """
+    Test thinker prefix caching by sending identical requests with an image (i.e.,
+    a large shared prefix) and verifying that the second request uses cached tokens
+    & produces the same output.
+    """
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        image_data_url=image_data_url,
+        content_text=get_prompt("text_image"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+    }
+
+    response_1 = openai_client.send_omni_request(request_config, request_num=1)[0]
+    response_2 = openai_client.send_omni_request(request_config, request_num=1)[0]
+
+    assert response_1.success
+    assert response_2.success
+    assert response_2.cached_tokens is not None
+    # We should cache the vast majority of the prompt (image + up to last full block),
+    # and set seed in the CI config, so the second request should give an identical
+    # response for the generated input image, even if we use dummy weights
+    assert response_2.cached_tokens > 0
+    assert response_1.text_content == response_2.text_content

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -6,22 +6,14 @@ E2E Online tests for Qwen3-Omni model.
 
 import os
 
-from vllm_omni.platforms import current_omni_platform
-
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-from pathlib import Path
 
 import pytest
 
-from tests.conftest import (
-    OmniServerParams,
-    dummy_messages_from_mix_data,
-    generate_synthetic_audio,
-    generate_synthetic_image,
-    generate_synthetic_video,
-    modify_stage_config,
-)
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
 
@@ -40,47 +32,57 @@ LARGE_IMAGE_HEIGHT = 1080
 LONG_AUDIO_DURATION_SEC = 120
 
 
-def get_chunk_config(default_path):
-    path = modify_stage_config(
-        default_path,
-        updates={
-            "async_chunk": True,
-            "stage_args": {
-                0: {
-                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk",
-                    "default_sampling_params.max_tokens": 2048,
-                },
-                1: {
-                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk"
-                },
-            },
-        },
-        deletes={"stage_args": {2: ["custom_process_input_func"]}},
-    )
-    return path
-
-
 def get_batch_token_config(default_path):
-    path = modify_stage_config(
+    """Override stage 1's max_num_batched_tokens to exercise small-batch paths.
+
+    Uses the new flat-stage schema (``stages.<id>.<field>``); the legacy
+    ``stage_args.<id>.engine_args.<field>`` path no longer applies because
+    the deploy YAML doesn't nest engine fields under ``engine_args:``.
+    """
+    return modify_stage_config(
         default_path,
         updates={
-            "stage_args": {1: {"engine_args.max_num_batched_tokens": 64}},
+            "stages": {1: {"max_num_batched_tokens": 64}},
         },
     )
-    return path
 
 
-# CI stage config for 2*H100-80G GPUs
-default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
+def get_async_chunk_config(default_path):
+    """Flip async_chunk on and bump stage 0 thinker output to 2048 tokens.
 
-if current_omni_platform.is_xpu():
-    default_path = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")
+    Pipeline registry (qwen3_omni/pipeline.py) already wires
+    thinker2talker_async_chunk / talker2code2wav_async_chunk on stages 0/1,
+    so no per-stage processor override is needed. Using only flat-schema
+    writes so _parse_stage_deploy stays in its flat branch (nested
+    ``engine_args:`` would drop other overlay fields).
+    """
+    return modify_stage_config(
+        default_path,
+        updates={
+            "stages": {0: {"default_sampling_params.max_tokens": 2048}},
+        },
+    )
+
+
+# CI deploy YAML (single file; xpu deltas applied via ``platforms:`` section).
+# The overlay explicitly sets ``async_chunk: False``, so ``default`` tests the
+# sync path and ``async_chunk`` tests the streaming path with a longer thinker
+# output — two distinct scenarios, kept as separate parametrizations.
+default_path = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
 
-# Create parameter combinations for model and stage config
 test_params = [
-    pytest.param(OmniServerParams(model=model, stage_config_path=default_path, use_stage_cli=True), id="default"),
     pytest.param(
-        OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path), use_stage_cli=True),
+        OmniServerParams(
+            model=model, stage_config_path=default_path, use_stage_cli=True, server_args=["--no-async-chunk"]
+        ),
+        id="default",
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=model,
+            stage_config_path=get_async_chunk_config(default_path),
+            use_stage_cli=True,
+        ),
         id="async_chunk",
     ),
 ]
@@ -371,6 +373,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.skip(reason="issue: #2827")
 def test_audio_in_video_001(omni_server, openai_client) -> None:
     """
     Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
@@ -447,7 +450,7 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["london"]},
     }
 
-    # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/conftest.py).
+    # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/helpers/assertions.py).
     _similarity_assert_msg = "The audio content is not same as the text"
     _max_retries = 3
     for attempt in range(_max_retries):
@@ -491,6 +494,7 @@ def test_speaker_001(omni_server, openai_client) -> None:
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.skip(reason="Known issue: occasional inaccuracy in voice recognition.")
 def test_speaker_002(omni_server, openai_client) -> None:
     """
     Input Modal: text only (one-word answer constraint).
@@ -511,7 +515,7 @@ def test_speaker_002(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    # Retry only when assert_omni_response fails on preset voice gender (see tests/conftest.py).
+    # Retry only when assert_omni_response fails on preset voice gender (see tests/helpers/assertions.py).
     _gender_assert_substr = "estimated gender"
     _max_retries = 3
     for attempt in range(_max_retries):

tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import io
+import json
+import os
+import wave
+
+import pytest
+import websockets
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import (
+    convert_audio_bytes_to_text,
+    cosine_similarity_text,
+    generate_synthetic_audio,
+)
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+
+# The new-schema CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU
+# via its ``platforms:`` section, so one path serves all three.
+default_stage_config = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
+
+
+def _realtime_stage_config_path() -> str:
+    """CI omni layout without async_chunk; stage 0 thinker max_tokens=10."""
+    return modify_stage_config(
+        default_stage_config,
+        updates={"stages": {0: {"default_sampling_params.max_tokens": 10}}},
+    )
+
+
+realtime_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=_realtime_stage_config_path(),
+            use_stage_cli=True,
+        ),
+        id="thinker_max_tokens_10",
+    ),
+]
+
+
+def _pcm16_mono_16k_from_wav_bytes(wav_bytes: bytes) -> bytes:
+    with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
+        if wf.getnchannels() != 1:
+            raise ValueError(f"Expected mono WAV, got {wf.getnchannels()} channels")
+        if wf.getsampwidth() != 2:
+            raise ValueError(f"Expected 16-bit PCM, sampwidth={wf.getsampwidth()}")
+        if wf.getframerate() != 16000:
+            raise ValueError(f"Expected 16 kHz input for /v1/realtime, got {wf.getframerate()} Hz")
+        if wf.getcomptype() != "NONE":
+            raise ValueError(f"Expected uncompressed PCM, comptype={wf.getcomptype()!r}")
+        return wf.readframes(wf.getnframes())
+
+
+def _wav_bytes_from_pcm16(pcm: bytes, sample_rate_hz: int) -> bytes:
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate_hz)
+        wf.writeframes(pcm)
+    return buf.getvalue()
+
+
+async def _run_realtime_audio_roundtrip(
+    host: str,
+    port: int,
+    model: str,
+    pcm16: bytes,
+    *,
+    chunk_ms: int = 100,
+) -> dict:
+    uri = f"ws://{host}:{port}/v1/realtime"
+    incremental: list[bytes] = []
+    output_sr = 24000
+    text_chunks: list[str] = []
+    final_text = ""
+    delta_events = 0
+
+    bytes_per_ms = 16000 * 2 // 1000
+    chunk_bytes = max(bytes_per_ms * chunk_ms, 2)
+
+    async with websockets.connect(uri, max_size=64 * 1024 * 1024) as ws:
+        await ws.send(json.dumps({"type": "session.update", "model": model}))
+        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))
+
+        for i in range(0, len(pcm16), chunk_bytes):
+            chunk = pcm16[i : i + chunk_bytes]
+            await ws.send(
+                json.dumps(
+                    {
+                        "type": "input_audio_buffer.append",
+                        "audio": base64.b64encode(chunk).decode("utf-8"),
+                    }
+                )
+            )
+
+        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
+
+        while True:
+            message = await asyncio.wait_for(ws.recv(), timeout=600)
+            if isinstance(message, bytes):
+                continue
+
+            event = json.loads(message)
+            event_type = event.get("type")
+
+            if event_type == "session.created":
+                continue
+
+            if event_type == "response.audio.delta":
+                delta_events += 1
+                sr = event.get("sample_rate_hz")
+                if isinstance(sr, int) and sr > 0:
+                    output_sr = sr
+                audio_b64 = event.get("audio", "")
+                if audio_b64:
+                    incremental.append(base64.b64decode(audio_b64))
+                continue
+
+            if event_type == "transcription.delta":
+                d = event.get("delta", "")
+                if d:
+                    text_chunks.append(d)
+                continue
+
+            if event_type == "transcription.done":
+                final_text = event.get("text", "") or "".join(text_chunks)
+                continue
+
+            if event_type == "response.audio.done":
+                break
+
+            if event_type == "error":
+                raise AssertionError(f"WebSocket error: {event}")
+
+            raise AssertionError(f"Unexpected WebSocket event: {event}")
+
+    out_pcm = b"".join(incremental)
+    return {
+        "output_pcm": out_pcm,
+        "output_sample_rate": output_sr,
+        "transcription_text": final_text if final_text else "".join(text_chunks),
+        "delta_events": delta_events,
+    }
+
+
+class TestQwen3OmniRealtimeWebSocket:
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    @pytest.mark.parametrize("omni_server", realtime_server_params, indirect=True)
+    def test_streaming_audio_input_pcm_output(self, omni_server) -> None:
+        """
+        Short streamed 16 kHz mono PCM16 input; expect streamed PCM16 audio deltas and
+        transcription. Verify Whisper(output audio) aligns with model text (same idea
+        as multimodal omni e2e).
+        """
+        syn = generate_synthetic_audio(10, 1, sample_rate=16000)
+        wav_bytes = base64.b64decode(syn["base64"])
+        pcm16 = _pcm16_mono_16k_from_wav_bytes(wav_bytes)
+
+        result = asyncio.run(
+            _run_realtime_audio_roundtrip(
+                omni_server.host,
+                omni_server.port,
+                omni_server.model,
+                pcm16,
+                chunk_ms=100,
+            )
+        )
+
+        out_pcm = result["output_pcm"]
+        assert result["delta_events"] >= 1
+        assert out_pcm, "No output PCM from response.audio.delta"
+        assert len(out_pcm) % 2 == 0
+        assert len(out_pcm) >= 4096, "Output audio unexpectedly small"
+        assert result["output_sample_rate"] > 0
+
+        final_text = (result["transcription_text"] or "").strip()
+        assert final_text, "Expected non-empty transcription (model text stream)"
+
+        wav_out = _wav_bytes_from_pcm16(out_pcm, result["output_sample_rate"])
+        whisper_text = convert_audio_bytes_to_text(wav_out).strip()
+        assert whisper_text, "Whisper returned empty string for synthesized output audio"
+
+        sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
+        assert sim > 0.9, (
+            f"Output audio transcript should match model text (sim={sim:.3f}): "
+            f"whisper={whisper_text!r}, model_text={final_text!r}"
+        )

tests/e2e/online_serving/test_qwen3_tts_base.py

@@ -12,12 +12,11 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 
@@ -25,11 +24,6 @@ REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/
 REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
 
 
-def get_stage_config(name: str = "qwen3_tts.yaml"):
-    """Get the stage config path from vllm_omni model_executor stage_configs."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
 def get_prompt(prompt_type="text"):
     """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
     prompts = {
@@ -48,7 +42,7 @@ tts_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
             server_args=["--trust-remote-code", "--disable-log-stats"],
         ),
         id="async_chunk",

tests/e2e/online_serving/test_qwen3_tts_base_expansion.py

@@ -12,12 +12,11 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 
@@ -25,11 +24,6 @@ REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/
 REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
 
 
-def get_stage_config(name: str = "qwen3_tts.yaml"):
-    """Get the stage config path from vllm_omni model_executor stage_configs."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
 def get_prompt(prompt_type="text"):
     """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
     prompts = {
@@ -48,16 +42,19 @@ tts_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
             server_args=["--trust-remote-code", "--disable-log-stats"],
         ),
         id="async_chunk",
     ),
+    # Synchronous (no async-chunk) variant — ``--no-async-chunk`` alone
+    # flips the deploy yaml's bool and the pipeline dispatches to the
+    # end-to-end codec processor. No variant yaml / pipeline needed.
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_stage_config("qwen3_tts_no_async_chunk.yaml"),
-            server_args=["--trust-remote-code", "--disable-log-stats"],
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats", "--no-async-chunk"],
         ),
         id="no_async_chunk",
     ),

tests/e2e/online_serving/test_qwen3_tts_batch.py

@@ -22,19 +22,18 @@ import httpx
 import pytest
 import yaml
 
-from tests.conftest import (
-    OmniServer,
-    convert_audio_file_to_text,
-    cosine_similarity_text,
-)
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import convert_audio_file_to_text, cosine_similarity_text
+from tests.helpers.runtime import OmniServer
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
 STAGE_INIT_TIMEOUT_S = 120
 
 
-def get_stage_config(name: str = "qwen3_tts.yaml"):
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+def get_stage_config(name: str = "qwen3_tts.yaml") -> str:
+    """Resolve a deploy config path under vllm_omni/deploy/."""
+    return get_deploy_config_path(name)
 
 
 @pytest.fixture(scope="module")

tests/e2e/online_serving/test_qwen3_tts_customvoice.py

@@ -12,21 +12,15 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 
 
-def get_stage_config(name: str = "qwen3_tts.yaml"):
-    """Get the stage config path from vllm_omni model_executor stage_configs."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
 def get_prompt(prompt_type="text"):
     """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
     prompts = {
@@ -45,7 +39,7 @@ tts_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
             server_args=["--trust-remote-code", "--disable-log-stats"],
         ),
         id="async_chunk",

tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py

@@ -12,21 +12,15 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 
 
-def get_stage_config(name: str = "qwen3_tts.yaml"):
-    """Get the stage config path from vllm_omni model_executor stage_configs."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
 def get_prompt(prompt_type="english"):
     """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
     prompts = {
@@ -46,16 +40,19 @@ tts_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
             server_args=["--trust-remote-code", "--disable-log-stats"],
         ),
         id="async_chunk",
     ),
+    # Synchronous (no async-chunk) variant — ``--no-async-chunk`` alone
+    # flips the deploy yaml's bool and the pipeline dispatches to the
+    # end-to-end codec processor. No variant yaml / pipeline needed.
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_stage_config("qwen3_tts_no_async_chunk.yaml"),
-            server_args=["--trust-remote-code", "--disable-log-stats"],
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats", "--no-async-chunk"],
         ),
         id="no_async_chunk",
     ),

tests/e2e/online_serving/test_qwen3_tts_speaker_embedding.py

@@ -13,13 +13,13 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 import struct
-from pathlib import Path
 
 import httpx
 import pytest
 
-from tests.conftest import OmniServer
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL_BASE = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 MODEL_BASE_1_7B = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
@@ -37,10 +37,8 @@ MIN_AUDIO_BYTES = 2000
 MAX_NEW_TOKENS = 256
 
 
-def get_stage_config():
-    return str(
-        Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml"
-    )
+def get_stage_config() -> str:
+    return get_deploy_config_path("qwen3_tts.yaml")
 
 
 def _server_args():

tests/e2e/online_serving/test_qwen3_tts_websocket.py

@@ -7,13 +7,13 @@ E2E online tests for Qwen3-TTS WebSocket streaming speech.
 import asyncio
 import json
 import os
-from pathlib import Path
 
 import pytest
 import websockets
 
-from tests.conftest import OmniServer
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+from tests.helpers.stage_config import get_deploy_config_path
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
@@ -23,9 +23,7 @@ STAGE_INIT_TIMEOUT_S = 120
 
 
 def get_stage_config() -> str:
-    return str(
-        Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml"
-    )
+    return get_deploy_config_path("qwen3_tts.yaml")
 
 
 @pytest.fixture(scope="module")

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -7,14 +7,9 @@ and are supported by the following models:
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-    generate_synthetic_image,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (

tests/e2e/online_serving/test_qwen_image_expansion.py

@@ -12,13 +12,8 @@ See docs/user_guide/diffusion_acceleration.md.
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 T2I_PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style."
 NEGATIVE_PROMPT = "blurry, low quality"

tests/e2e/online_serving/test_qwen_image_layered_expansion.py

@@ -14,15 +14,9 @@ Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    decode_b64_image,
-    dummy_messages_from_mix_data,
-    generate_synthetic_image,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import decode_b64_image, generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
 MODEL = "Qwen/Qwen-Image-Layered"
 EDIT_PROMPT = "Decompose this image into layers."

tests/e2e/online_serving/test_sd3_expansion.py

@@ -4,12 +4,8 @@ Tests for Stable Diffusion 3.5 medium model.
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
 FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
 POSITIVE_PROMPT = "A serene mountain landscape at sunset"

tests/e2e/online_serving/test_video_generation_api.py

@@ -16,8 +16,8 @@ from typing import Any
 import pytest
 import requests
 
-from tests.conftest import OmniServer
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"

tests/e2e/online_serving/test_voxtral_tts.py

@@ -12,19 +12,16 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import httpx
 import pytest
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "mistralai/Voxtral-4B-TTS-2603"
 
-STAGE_CONFIG = str(
-    Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxtral_tts.yaml"
-)
+STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml")
 EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"]
 TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, server_args=EXTRA_ARGS)]
 

tests/e2e/online_serving/test_wan22_expansion.py

@@ -19,13 +19,9 @@ assert_diffusion_response validates successful generation
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    generate_synthetic_image,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
 PROMPT = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
 NEGATIVE_PROMPT = "low quality, blurry, distorted face, extra limbs, bad anatomy, watermark, logo, text, ugly, deformed, mutated, jpeg artifacts"

tests/e2e/online_serving/test_wan_2_1_vace_expansion.py

@@ -23,12 +23,8 @@ Coverage:
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
 MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
 PROMPT = "A cat walking slowly across a sunlit garden path"

tests/e2e/online_serving/test_zimage_expansion.py

@@ -13,12 +13,8 @@ exercise Z-Image's supported feature combinations:
 
 import pytest
 
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-)
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
 MODEL = "Tongyi-MAI/Z-Image-Turbo"
 PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard."

tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml

@@ -0,0 +1,35 @@
+# Thinker stage only
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0,1,2,3"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: MingFlashOmniForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: false
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      max_model_len: 32768
+      tensor_parallel_size: 4
+      hf_config_name: llm_config
+      load_format: dummy
+      mm_processor_cache_gb: 0
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      max_tokens: 100
+      repetition_penalty: 1.05
+      seed: 42
+      detokenize: true
+      ignore_eos: false

tests/e2e/stage_configs/dynin_omni_ci.yaml

@@ -72,13 +72,8 @@ stage_args:
 
 runtime:
   enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
   edges:
     - from: 0
       to: 1
-      window_size: -1
     - from: 1
       to: 2
-      window_size: -1

tests/e2e/stage_configs/mimo_audio_ci.yaml

@@ -1,71 +0,0 @@
-# CI stage config for running MiMo-Audio in multi-stage omni runtime.
-# Based on mimo_audio.yaml with load_format: dummy for CI testing.
-
-async_chunk: false
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-    engine_args:
-      dtype: bfloat16
-      max_num_seqs: 1
-      model_stage: fused_thinker_talker
-      model_arch: MiMoAudioForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      tensor_parallel_size: 1
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-      max_model_len: 8192
-      max_num_batched_tokens: 8192
-      load_format: dummy
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.6
-      top_p: 0.95
-      top_k: 50
-      max_tokens: 18192
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-    engine_args:
-      model_stage: code2wav
-      max_num_seqs: 1
-      model_arch: MiMoAudioForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      tensor_parallel_size: 1
-      gpu_memory_utilization: 0.2
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-      max_model_len: 18192
-      max_num_batched_tokens: 18192
-      async_scheduling: false
-      load_format: dummy
-    engine_input_source: [ 0 ]
-    is_comprehension: false
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 18192
-      seed: 42
-      detokenize: false

tests/e2e/stage_configs/qwen2_5_omni_ci.yaml

@@ -1,109 +0,0 @@
-# stage config for running qwen2.5-omni for multi-stage omni runtime.
-
-# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
-# This config is optimized for CI e2e tests.
-stage_args:
-  - stage_id: 0
-    runtime:
-      process: true            # Run this stage in a separate process
-      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 16384
-      max_num_batched_tokens: 16384
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.9
-      skip_mm_profiling: true
-      enforce_eager: true  # Now we only support eager mode
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-      mm_processor_cache_gb: 0
-      load_format: dummy
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-  - stage_id: 1
-    runtime:
-      process: true
-      devices: "1"
-    engine_args:
-      model_stage: talker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 16384
-      max_num_batched_tokens: 16384
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.4
-      skip_mm_profiling: true
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-      load_format: dummy
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
-    default_sampling_params:
-      temperature: 0.9
-      top_p: 0.8
-      top_k: 40
-      max_tokens: 4096
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-      stop_token_ids: [8294]
-  - stage_id: 2
-    runtime:
-      process: true
-      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-    engine_args:
-      model_stage: code2wav
-      max_num_seqs: 1
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.5 #increase the gpu memory utilization to enable the test on H800
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-      max_num_batched_tokens: 8192
-      max_model_len: 8192
-      load_format: dummy
-    engine_input_source: [1]
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 8192
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1             # Simplified: trigger downstream only after full upstream completion
-    max_inflight: 1             # Simplified: process serially within each stage
-  edges:
-    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
-      to: 1
-      window_size: -1
-    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
-      to: 2
-      window_size: -1

tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml

@@ -1,31 +0,0 @@
-stage_args:
-  - stage_id: 0
-    runtime:
-      process: true            # Run this stage in a separate process
-      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 16384
-      max_num_batched_tokens: 16384
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.9
-      skip_mm_profiling: true
-      enforce_eager: true  # Now we only support eager mode
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-      mm_processor_cache_gb: 0
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -1,102 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
-# Stage 0: Thinker (multimodal understanding + text generation)
-# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
-# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
-# The following config has been verified on 2x H100-80G GPUs.
-stage_args:
-- stage_id: 0
-  runtime:
-    devices: "0"
-  engine_args:
-    model_stage: thinker
-    max_num_seqs: 5
-    model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.9
-    enforce_eager: false
-    trust_remote_code: true
-    engine_output_type: latent  # Output hidden states for talker
-    distributed_executor_backend: "mp"
-    max_num_batched_tokens: 32768
-    max_model_len: 32768
-    enable_prefix_caching: false
-    mm_processor_cache_gb: 0
-    hf_config_name: thinker_config
-    tensor_parallel_size: 1
-    load_format: dummy
-  final_output: true
-  final_output_type: text
-  is_comprehension: true
-  default_sampling_params:
-    temperature: 0.4
-    top_p: 0.9
-    top_k: 1
-    max_tokens: 150
-    seed: 42
-    ignore_eos: False
-    detokenize: True
-    repetition_penalty: 1.05
-
-- stage_id: 1
-  runtime:
-    devices: "1"
-  engine_args:
-    model_stage: talker
-    max_num_seqs: 5
-    model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.5
-    enforce_eager: false
-    trust_remote_code: true
-    engine_output_type: latent  # Output codec codes for code2wav
-    enable_prefix_caching: false
-    max_num_batched_tokens: 32768
-    max_model_len: 32768
-    distributed_executor_backend: "mp"
-    hf_config_name: talker_config
-    load_format: dummy
-  engine_input_source: [0]
-  custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-  default_sampling_params:
-    temperature: 0.9
-    top_k: 50
-    max_tokens: 1000
-    seed: 42
-    detokenize: False
-    repetition_penalty: 1.05
-    stop_token_ids: [2150]
-
-- stage_id: 2
-  runtime:
-    devices: "1"
-  engine_args:
-    model_stage: code2wav
-    max_num_seqs: 5
-    model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_type: generation
-    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    engine_output_type: audio  # Final output: audio waveform
-    gpu_memory_utilization: 0.1
-    distributed_executor_backend: "mp"
-    max_num_batched_tokens: 100000
-    hf_config_name: thinker_config
-    async_scheduling: false
-    load_format: dummy
-  engine_input_source: [1]
-  custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-  final_output: true
-  final_output_type: audio
-  default_sampling_params:
-    temperature: 0.0
-    top_p: 1.0
-    top_k: -1
-    max_tokens: 2000
-    seed: 42
-    detokenize: True
-    repetition_penalty: 1.1

tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml

@@ -1,106 +0,0 @@
-# stage config for running qwen2.5-omni for multi-stage omni runtime.
-
-# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
-# This config is optimized for CI e2e tests.
-stage_args:
-  - stage_id: 0
-    runtime:
-      process: true            # Run this stage in a separate process
-      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 16384
-      max_num_batched_tokens: 16384
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      enforce_eager: true  # Now we only support eager mode
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-      mm_processor_cache_gb: 0
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-  - stage_id: 1
-    runtime:
-      process: true
-      devices: "1"
-    engine_args:
-      model_stage: talker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 16384
-      max_num_batched_tokens: 16384
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
-    default_sampling_params:
-      temperature: 0.9
-      top_p: 0.8
-      top_k: 40
-      max_tokens: 4096
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-      stop_token_ids: [8294]
-  - stage_id: 2
-    runtime:
-      process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-    engine_args:
-      model_stage: code2wav
-      max_num_seqs: 1
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.15
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-      max_num_batched_tokens: 4096
-      max_model_len: 4096
-    engine_input_source: [1]
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 4096
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1             # Simplified: trigger downstream only after full upstream completion
-    max_inflight: 1             # Simplified: process serially within each stage
-  edges:
-    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
-      to: 1
-      window_size: -1
-    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
-      to: 2
-      window_size: -1

tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -1,100 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
-# Stage 0: Thinker (multimodal understanding + text generation)
-# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
-# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
-# The following config has been verified on 2x H100-80G GPUs.
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: thinker
-      max_num_seqs: 1
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: false
-      trust_remote_code: true
-      engine_output_type: latent  # Output hidden states for talker
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      mm_processor_cache_gb: 0
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-      load_format: dummy
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 100
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    runtime:
-       devices: "1"
-    engine_args:
-       model_stage: talker
-       max_num_seqs: 1
-       model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_type: ar
-       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.6
-       enforce_eager: true
-       trust_remote_code: true
-       engine_output_type: latent  # Output codec codes for code2wav
-      #  tensor_parallel_size: 2
-       enable_prefix_caching: false
-       distributed_executor_backend: "mp"
-       hf_config_name: talker_config
-       load_format: dummy
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-    # final_output: true
-    # final_output_type: text
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 100
-      seed: 42
-      detokenize: False
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 2
-    runtime:
-      devices: "1"
-    engine_args:
-      model_stage: code2wav
-      max_num_seqs: 1
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-      hf_config_name: thinker_config
-      load_format: dummy
-      async_scheduling: false
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 200
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1

tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml

@@ -1,108 +0,0 @@
-# stage config for running qwen2.5-omni for multi-stage omni runtime.
-
-# The following config is verified with 2 * Intel Arc Pro B60 XPU.
-stage_args:
-  - stage_id: 0
-    stage_type: llm # Use llm stage type for AR stages
-    runtime:
-      process: true # Run this stage in a separate process
-      devices: "0" # Visible devices for this stage
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 16384
-      max_num_batched_tokens: 16384
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.9 # thinker weight is around 16.74GB for Qwen2.5-Omni-7B
-      skip_mm_profiling: true
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-      mm_processor_cache_gb: 0
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-  - stage_id: 1
-    stage_type: llm # Use llm stage type for AR stages
-    runtime:
-      process: true
-      devices: "1"
-    engine_args:
-      model_stage: talker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 16384
-      max_num_batched_tokens: 16384
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.5  # talker weight is 6.03GB for Qwen2.5-Omni-7B
-      skip_mm_profiling: true
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
-    default_sampling_params:
-      temperature: 0.9
-      top_p: 0.8
-      top_k: 40
-      max_tokens: 4096
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-      stop_token_ids: [8294]
-
-  - stage_id: 2
-    stage_type: llm # Use llm stage type for AR stages
-    runtime:
-      process: true
-      devices: "2"
-    engine_args:
-      max_num_seqs: 1
-      model_stage: code2wav
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.3 # code2wav weight is around 1.46GB for Qwen2.5-Omni-7B
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-    engine_input_source: [1]
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1 # Simplified: trigger downstream only after full upstream completion
-    max_inflight: 1 # Simplified: process serially within each stage
-
-  edges:
-    - from: 0 # thinker → talker: trigger only after receiving full input (-1)
-      to: 1
-      window_size: -1
-    - from: 1 # talker → code2wav: trigger only after receiving full input (-1)
-      to: 2
-      window_size: -1