[Test] L4 complete diffusion feature test for Qwen-Image-Edit models (#1682)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/test-nightly.yml

@@ -54,7 +54,6 @@ steps:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
           always-pull: true
-          shm-size: "8gb"
           propagate-environment: true
           shm-size: "8gb"
           environment:
@@ -62,6 +61,42 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: ":full_moon: Diffusion Model Test with H100"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    # if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
   - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
     timeout_in_minutes: 30
     depends_on: upload-nightly-pipeline

docs/contributing/ci/CI_5levels.md

@@ -508,7 +508,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu
     **2.4.5 Request Execution**
 
     ```python
-    openai_client.send_request(request_config, request_num=1)
+    openai_client.send_omni_request(request_config, request_num=1)  # for omni-understanding models
+    # or
+    openai_client.send_diffusion_request(request_config, request_num=1)  # for diffusion models
     ```
 
     **Explanation**:
@@ -517,7 +519,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
 
     **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
 
-    **Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
+    **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
 
 -   ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
 

docs/contributing/ci/tests_markers.md

@@ -32,7 +32,7 @@ Defined in `pyproject.toml`:
 | `slow`             | Slow tests (may skip in quick CI)                         |
 | `benchmark`        | Benchmark tests                                           |
 
-\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
+\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
 
 ### Example usage for markers
 
@@ -49,6 +49,7 @@ from tests.utils import hardware_test
 def test_video_to_audio()
     ...
 ```
+
 ### Decorator: `@hardware_test`
 
 This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
@@ -98,6 +99,24 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
     - ROCm only: `pytest -m "distributed_rocm and MI325"`
     - NPU only: `pytest -m "distributed_npu"`
 
+### Function: `hardware_marks`
+
+`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
+
+```python
+from tests.utils import hardware_marks
+
+MULTI_CARD_MARKS = hardware_marks(
+    res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
+)
+
+@pytest.mark.parametrize("omni_server", [
+    pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
+], indirect=True)
+def test_feature(omni_server):
+    ...
+```
+
 ## Add Support for a New Platform
 
 If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
@@ -132,12 +151,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
            # Optionally: add skipif_tpu when implemented
            return [test_platform, test_resource, test_distributed]
    ```
-3. **Update `hardware_test` to recognize your new platform**:
-    In the relevant place (see the `hardware_test` implementation), add:
+3. **Update `hardware_marks` to recognize your new platform**:
+    In the relevant place (see the `hardware_marks` implementation), add:
     ```python
     if platform == "tpu":
         marks = tpu_marks(res=resource, num_cards=cards)
     ```
+    (`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
 4. **(Recommended) Add a test using your new markers**:
    ```python
    @hardware_test(
@@ -151,7 +171,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
 **Summary**:  
 - Add pytest markers for your new platform/resources  
 - Implement a marker function (`xxx_marks`)  
-- Plug into `hardware_test`  
-- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
+- Plug into `hardware_marks`  
+- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
 
 See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).

docs/contributing/ci/tests_style.md

@@ -158,6 +158,7 @@ vllm_omni/                                    tests/
                                                    ├── test_diffusion_layerwise_offload.py
                                                    ├── test_diffusion_lora.py
                                                    ├── test_sequence_parallel.py
+                                                   ├── test_qwen_image_edit_expansion.py
                                                    └── stage_configs/
                                                        ├── qwen2_5_omni_ci.yaml
                                                        ├── qwen3_omni_ci.yaml

docs/contributing/model/adding_diffusion_model.md

@@ -651,8 +651,27 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
 
 #### 5.2 Functionality Check in CI
 
-To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
-For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
+To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
+
+For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
+The following tests are required to add:
+
+- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
+  - Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
+    - 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
+    - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
+    - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
+    - 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
+    - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
+    - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
+  - Validation: test that the multimodal output files of your model have the correct shapes.
+  - Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
+  - To maximize code reuse, you may refer to
+    - `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
+    - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
+    - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
+  - Doc: add a concise dostring for each test function.
+  - Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
 
 ---
 

tests/benchmarks/test_serve_cli.py

@@ -3,13 +3,16 @@ from pathlib import Path
 
 import pytest
 
+from tests.conftest import OmniServerParams
 from tests.utils import hardware_test
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
 stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
 
 # Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
+]
 
 
 @pytest.mark.core_model

tests/conftest.py

@@ -18,16 +18,21 @@ import subprocess
 import sys
 import threading
 import time
+from collections.abc import Generator
 from dataclasses import dataclass
+from io import BytesIO
 from pathlib import Path
-from typing import Any
+from typing import Any, NamedTuple
 
+import imageio.v3 as iio
 import numpy as np
 import psutil
 import pytest
+import soundfile as sf
 import torch
 import yaml
-from openai import OpenAI
+from openai import OpenAI, omit
+from PIL import Image
 from vllm import TextPrompt
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 from vllm.logger import init_logger
@@ -44,6 +49,65 @@ PromptImageInput = list[Any] | Any | None
 PromptVideoInput = list[Any] | Any | None
 
 
+class OmniServerParams(NamedTuple):
+    model: str
+    port: int | None = None
+    stage_config_path: str | None = None
+    server_args: list[str] | None = None
+
+
+def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None):
+    """Assert the file is a loadable image with optional exact dimensions."""
+    if isinstance(image, Path):
+        assert image.exists(), f"Image not found: {image}"
+        image = Image.open(image)
+        image.load()
+    assert image.width > 0 and image.height > 0
+    if width is not None:
+        assert image.width == width, f"Expected width={width}, got {image.width} in {image.name}"
+    if height is not None:
+        assert image.height == height, f"Expected height={height}, got {image.height} in {image.name}"
+    return image
+
+
+def assert_video_valid(frames: Path | np.ndarray, *, width: int, height: int, num_frames: int) -> None:
+    """Assert the MP4 has the expected resolution and exact frame count."""
+    if isinstance(frames, Path):
+        assert frames.exists(), f"Video not found: {frames}"
+        frames = iio.imread(str(frames), plugin="pyav", index=None)
+    assert frames.shape[0] == num_frames, f"Expected {num_frames} frames, got {frames.shape[0]}"
+    assert frames.shape[1] == height, f"Expected height={height}, got {frames.shape[1]}"
+    assert frames.shape[2] == width, f"Expected width={width}, got {frames.shape[2]}"
+
+
+def assert_audio_valid(path: Path, *, sample_rate: int, channels: int, duration_s: float) -> None:
+    """Assert the WAV has the expected sample rate, channel count, and duration."""
+
+    assert path.exists(), f"Audio not found: {path}"
+    info = sf.info(str(path))
+    assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}"
+    assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}"
+    expected_frames = int(duration_s * sample_rate)
+    assert info.frames == expected_frames, (
+        f"Expected {expected_frames} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}"
+    )
+
+
+def decode_b64_image(b64: str):
+    img = Image.open(BytesIO(base64.b64decode(b64)))
+    img.load()
+    return img
+
+
+@pytest.fixture(scope="session")
+def model_prefix() -> str:
+    """Optional model-path prefix from MODEL_PREFIX env var.
+    Useful if models are downloaded to non-default local directories.
+    """
+    prefix = os.environ.get("MODEL_PREFIX", "")
+    return f"{prefix.rstrip('/')}/" if prefix else ""
+
+
 @pytest.fixture(autouse=True)
 def default_vllm_config():
     """Set a default VllmConfig for all tests.
@@ -1151,7 +1215,9 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope="session")
-def run_level(request):
+def run_level(request) -> str:
+    """A command-line argument that specifies the level of tests to run in this session.
+    See https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/CI_5levels/"""
     return request.config.getoption("--run-level")
 
 
@@ -1159,15 +1225,17 @@ _omni_server_lock = threading.Lock()
 
 
 @pytest.fixture(scope="module")
-def omni_server(request, run_level):
+def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]:
     """Start vLLM-Omni server as a subprocess with actual model weights.
     Uses session scope so the server starts only once for the entire test session.
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        *port, model, stage_config_path = request.param
-        port = port[0] if port else None
-        if run_level == "advanced_model":
+        params: OmniServerParams = request.param
+        model = model_prefix + params.model
+        port = params.port
+        stage_config_path = params.stage_config_path
+        if run_level == "advanced_model" and stage_config_path is not None:
             stage_config_path = modify_stage_config(
                 stage_config_path,
                 deletes={
@@ -1179,7 +1247,11 @@ def omni_server(request, run_level):
                 },
             )
 
-        server_args = ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]
+        server_args = params.server_args or []
+        server_args = ["--stage-init-timeout", "120", *server_args]
+        if stage_config_path is not None:
+            server_args += ["--stage-configs-path", stage_config_path]
+
         with OmniServer(model, server_args, port=port) if port else OmniServer(model, server_args) as server:
             print("OmniServer started successfully")
             yield server
@@ -1199,6 +1271,17 @@ class OmniResponse:
     error_message: str | None = None
 
 
+@dataclass
+class DiffusionResponse:
+    text_content: str | None = None
+    images: list[Image.Image] | None = None
+    audios: list[Any] | None = None
+    videos: list[Any] | None = None
+    e2e_latency: float | None = None
+    success: bool = False
+    error_message: str | None = None
+
+
 def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level):
     """
     Validate response results.
@@ -1247,6 +1330,51 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
             print(f"similarity is: {response.similarity}")
 
 
+def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None):
+    """
+    Validate diffusion response results.
+
+    Args:
+        response: DiffusionResponse object. Any not-None content will be validated based on the request_config.
+        request_config: Request configuration dictionary containing parameters like model, messages, extra_body.
+            When validating a certain modality, the corresponding params in request_config['extra_body'] must present.
+            It will be used to check against the multimedia file in the response.
+        run_level: Test run level (e.g., "core_model", "advanced_model")
+
+    Raises:
+        AssertionError: When the response does not meet validation criteria
+        KeyError: When the request_config does not contain necessary parameters for validation
+    """
+    assert response.success, "The request failed."
+
+    e2e_latency = response.e2e_latency
+    if e2e_latency is not None:
+        print(f"the avg e2e is: {e2e_latency}")
+
+    extra_body = request_config.get("extra_body", {})
+
+    num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
+
+    if response.images is not None:
+        assert len(response.images) > 0, "No images in response"
+        assert len(response.images) == num_outputs_per_prompt, (
+            f"Expected {num_outputs_per_prompt} images, got {len(response.images)}"
+        )
+        if run_level == "advanced_model":
+            expected_width = extra_body["width"]  # intend to raise KeyError
+            expected_height = extra_body["height"]  # intend to raise KeyError
+            for img in response.images:
+                assert_image_valid(img, width=expected_width, height=expected_height)
+    if response.videos is not None:
+        raise NotImplementedError(
+            "Video validation is not implemented yet"
+        )  # consider using assert_video_valid defined above
+    if response.audios is not None:
+        raise NotImplementedError(
+            "Audio validation is not implemented yet"
+        )  # consider using assert_audio_valid defined above
+
+
 class OpenAIClientHandler:
     """
     OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic.
@@ -1269,7 +1397,7 @@ class OpenAIClientHandler:
         self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key)
         self.run_level = run_level
 
-    def _process_stream_response(self, chat_completion) -> OmniResponse:
+    def _process_stream_omni_response(self, chat_completion) -> OmniResponse:
         """
         Process streaming responses.
 
@@ -1330,7 +1458,7 @@ class OpenAIClientHandler:
 
         return result
 
-    def _process_non_stream_response(self, chat_completion) -> OmniResponse:
+    def _process_non_stream_omni_response(self, chat_completion) -> OmniResponse:
         """
         Process non-streaming responses.
 
@@ -1384,7 +1512,46 @@ class OpenAIClientHandler:
 
         return result
 
-    def send_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+    def _process_diffusion_response(self, chat_completion) -> DiffusionResponse:
+        """
+        Process diffusion responses (image generation/editing).
+
+        Args:
+            chat_completion: OpenAI response object
+
+        Returns:
+            DiffusionResponse: Processed response object
+        """
+        result = DiffusionResponse()
+        start_time = time.perf_counter()
+
+        try:
+            images = []
+            # [TODO] reading video and audio output from API response for later validation
+
+            for choice in chat_completion.choices:
+                if hasattr(choice.message, "content") and choice.message.content is not None:
+                    content = choice.message.content
+                    if isinstance(content, list):
+                        for item in content:
+                            if hasattr(item, "image_url") and item.image_url is not None:
+                                image_url = item.image_url.url
+                                if image_url.startswith("data:image"):
+                                    b64_data = image_url.split(",", 1)[1]
+                                    img = decode_b64_image(b64_data)
+                                    images.append(img)
+
+            result.e2e_latency = time.perf_counter() - start_time
+            result.images = images if images else None
+            result.success = True
+
+        except Exception as e:
+            result.error_message = f"Diffusion response processing error: {str(e)}"
+            print(f"Error: {result.error_message}")
+
+        return result
+
+    def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
         """
         Send OpenAI requests.
 
@@ -1410,9 +1577,9 @@ class OpenAIClientHandler:
             )
 
             if stream:
-                response = self._process_stream_response(chat_completion)
+                response = self._process_stream_omni_response(chat_completion)
             else:
-                response = self._process_non_stream_response(chat_completion)
+                response = self._process_non_stream_omni_response(chat_completion)
 
             assert_omni_response(response, request_config, run_level=self.run_level)
             responses.append(response)
@@ -1438,25 +1605,82 @@ class OpenAIClientHandler:
                     chat_completion = future.result()
 
                     if stream:
-                        response = self._process_stream_response(chat_completion)
+                        response = self._process_stream_omni_response(chat_completion)
                     else:
-                        response = self._process_non_stream_response(chat_completion)
+                        response = self._process_non_stream_omni_response(chat_completion)
 
                     assert_omni_response(response, request_config, run_level=self.run_level)
                     responses.append(response)
 
         return responses
 
+    def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+        """
+        Send OpenAI requests for diffusion models.
+
+        Args:
+            request_config: Request configuration dictionary containing parameters like model, messages
+            request_num: Number of requests to send concurrently, defaults to 1 (single request)
+        Returns:
+            List[OmniResponse]: List of response objects
+        """
+        responses = []
+        stream = request_config.get("stream", False)
+        modalities = request_config.get("modalities", omit)  # Most diffusion models don't require modalities param
+        extra_body = request_config.get("extra_body", None)
+
+        if stream:
+            raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test")
+
+        if request_num == 1:
+            # Send single request
+            chat_completion = self.client.chat.completions.create(
+                model=request_config.get("model"),
+                messages=request_config.get("messages"),
+                extra_body=extra_body,
+                modalities=modalities,
+            )
+
+            response = self._process_diffusion_response(chat_completion)
+            assert_diffusion_response(response, request_config, run_level=self.run_level)
+            responses.append(response)
+
+        else:
+            # Send concurrent requests
+            with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
+                futures = []
+
+                # Submit all request tasks
+                for _ in range(request_num):
+                    future = executor.submit(
+                        self.client.chat.completions.create,
+                        model=request_config.get("model"),
+                        messages=request_config.get("messages"),
+                        modalities=modalities,
+                        extra_body=extra_body,
+                    )
+                    futures.append(future)
+
+                # Process completed tasks
+                for future in concurrent.futures.as_completed(futures):
+                    chat_completion = future.result()
+                    response = self._process_diffusion_response(chat_completion)
+                    assert_diffusion_response(response, request_config, run_level=self.run_level)
+                    responses.append(response)
+
+        return responses
+
 
 @pytest.fixture
-def openai_client(omni_server, run_level):
-    """Create OpenAIClientHandler fixture"""
+def openai_client(omni_server: OmniServer, run_level: str):
+    """Create OpenAIClientHandler fixture to facilitate communication with OmniServer
+    with encapsulated request sending, concurrent requests, response handling, and validation."""
     return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level)
 
 
 class OmniRunner:
     """
-    Test runner for Omni models.
+    Offline test runner for Omni models.
     """
 
     def __init__(
@@ -1727,9 +1951,10 @@ class OmniRunner:
 
 
 @pytest.fixture(scope="module")
-def omni_runner(request):
+def omni_runner(request, model_prefix):
     with _omni_server_lock:
         model, stage_config_path = request.param
+        model = model_prefix + model
         with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
             print("OmniRunner started successfully")
             yield runner

tests/e2e/online_serving/test_qwen3_omni.py

@@ -8,6 +8,7 @@ from pathlib import Path
 import pytest
 
 from tests.conftest import (
+    OmniServerParams,
     dummy_messages_from_mix_data,
     generate_synthetic_audio,
     generate_synthetic_image,
@@ -51,7 +52,9 @@ else:
     stage_configs = [get_chunk_config()]
 
 # Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
+]
 
 
 def get_system_prompt():
@@ -120,7 +123,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
     }
 
     # Test single completion
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -146,4 +149,4 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -12,6 +12,7 @@ from pathlib import Path
 import pytest
 
 from tests.conftest import (
+    OmniServerParams,
     dummy_messages_from_mix_data,
     generate_synthetic_audio,
     generate_synthetic_image,
@@ -50,7 +51,9 @@ def get_chunk_config(default_path):
 default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
 stage_configs = [default_path, get_chunk_config(default_path)]
 # Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
+]
 
 
 def get_system_prompt():
@@ -106,7 +109,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -128,7 +131,7 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
 @pytest.mark.advanced_model
@@ -153,7 +156,7 @@ def test_image_to_text_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -177,7 +180,7 @@ def test_image_to_audio_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -201,7 +204,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
 @pytest.mark.advanced_model
@@ -225,7 +228,7 @@ def test_video_to_text_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -249,7 +252,7 @@ def test_video_to_audio_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -273,7 +276,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
 @pytest.mark.advanced_model
@@ -298,7 +301,7 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"audio": AUDIO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -324,7 +327,7 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"image": IMAGE_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.advanced_model
@@ -351,7 +354,7 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_omni_request(request_config)
 
 
 @pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
@@ -383,4 +386,4 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
         "stream": True,
         "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
     }
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -0,0 +1,164 @@
+"""
+Comprehensive tests of diffusion features that are available in online serving mode
+and are supported by the following models:
+- Qwen-Image-Edit: single image input
+- Qwen-Image-Edit-2509: two image inputs
+"""
+
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+    generate_synthetic_image,
+)
+from tests.utils import hardware_marks
+
+EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
+MULTI_EDIT_PROMPT = (
+    "Transform the first image into a Dadaism collage art. "
+    "Transform the second image into a Vincent van Gogh style painting. "
+    "Then juxtapose the two transformed images into a single artwork for visual contrast."
+)
+NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+# This test file targets two models, so I write a helper function.
+# If a similar test only involves one model, one can just define a global list variable.
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "tea_cache",  # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together.
+                    "--enable-layerwise-offload",
+                ],
+            ),
+            id="single_card_001",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--ulysses-degree",
+                    "2",
+                ],
+            ),
+            id="parallel_001",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="parallel_002",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "tea_cache",
+                    "--cfg-parallel-size",
+                    "2",
+                ],
+            ),
+            id="parallel_003",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--tensor-parallel-size",
+                    "2",
+                    "--vae-patch-parallel-size",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_004",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
+    indirect=True,
+)
+def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test all diffusion features with Qwen-Image-Edit in regular end-user scenarios."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
+
+    # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    indirect=True,
+)
+def test_qwen_image_edit_2509(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test all diffusion features with Qwen-Image-Edit-2509 in regular end-user scenarios."""
+    image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+
+    messages = dummy_messages_from_mix_data(
+        image_data_url=[image_data_url_1, image_data_url_2], content_text=MULTI_EDIT_PROMPT
+    )
+
+    # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/examples/online_serving/test_qwen2_5_omni.py

@@ -12,7 +12,7 @@ from pathlib import Path
 
 import pytest
 
-from tests.conftest import convert_audio_file_to_text, cosine_similarity_text
+from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text
 from tests.utils import hardware_test
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
@@ -22,7 +22,11 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config
 
 example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen2_5_omni")
 # Create parameter combinations for model and stage config
-test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, port=8091, stage_config_path=stage_config)
+    for model in models
+    for stage_config in stage_configs
+]
 
 
 def run_cmd(command):

tests/examples/online_serving/test_qwen3_omni.py

@@ -12,7 +12,7 @@ from pathlib import Path
 
 import pytest
 
-from tests.conftest import convert_audio_file_to_text, cosine_similarity_text
+from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text
 from tests.utils import hardware_test
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
@@ -22,7 +22,11 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config
 
 example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni")
 # Create parameter combinations for model and stage config
-test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs]
+test_params = [
+    OmniServerParams(model=model, port=8091, stage_config_path=stage_config)
+    for model in models
+    for stage_config in stage_configs
+]
 
 
 def run_cmd(command):

tests/utils.py

@@ -443,27 +443,11 @@ def npu_marks(*, res: str, num_cards: int):
         return [mark for mark in [test_platform, test_resource, test_distributed] if mark is not None]
 
 
-def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
+def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
     """
-    Decorate a test for multiple hardware platforms with a single call.
-    Automatically wraps the test with @create_new_process_for_each_test() for distributed tests.
-
-    Args:
-        res: Mapping from platform to resource type. Supported platforms/resources:
-            - cuda: L4, H100
-            - rocm: MI325
-            - npu: A2, A3
-        num_cards: Number of cards required. Can be:
-            - int: same card count for all platforms (default: 1)
-            - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2}
-
-    Example:
-        @hardware_test(
-            res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
-            num_cards={"cuda": 2, "rocm": 2, "npu": 2},
-        )
-        def test_multi_platform():
-            ...
+    Get a collection of pytest marks to apply for `@hardware_test`,
+    including CUDA, ROCm, and NPU,
+    based on the specified platforms and resources.
     """
     # Validate platforms
     # Don't validate platform details in this decorator
@@ -486,7 +470,7 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
                 num_cards_dict[platform] = 1
 
     # Collect marks from all platforms
-    all_marks: list[Callable[[Callable[_P, None]], Callable[_P, None]]] = []
+    all_marks: list[pytest.MarkDecorator] = []
     for platform, resource in res.items():
         cards = num_cards_dict[platform]
         if platform == "cuda" or platform == "rocm":
@@ -496,6 +480,32 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
         else:
             raise ValueError(f"Unsupported platform: {platform}")
         all_marks.extend(marks)
+    return all_marks
+
+
+def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
+    """
+    Decorate a test for multiple hardware platforms with a single call.
+    Automatically wraps the test with @create_new_process_for_each_test() for distributed tests.
+
+    Args:
+        res: Mapping from platform to resource type. Supported platforms/resources:
+            - cuda: L4, H100
+            - rocm: MI325
+            - npu: A2, A3
+        num_cards: Number of cards required. Can be:
+            - int: same card count for all platforms (default: 1)
+            - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2}
+
+    Example:
+        @hardware_test(
+            res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+            num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+        )
+        def test_multi_platform():
+            ...
+    """
+    all_marks = hardware_marks(res=res, num_cards=num_cards)
 
     def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
         func = f