Merge branch 'main' into lance-model-support

tests/benchmarks/metrics/test_metrics.py

@@ -63,6 +63,38 @@ def test_total_input_aggregated_from_output_prompt_len():
     )
 
 
+def test_audio_continuity_aggregation():
+    """Continuity rate and underrun percentile must aggregate from per-output fields."""
+    bad = _make_output(100)
+    bad.audio_underrun_s = 0.5
+    bad.audio_continuity_ok = False
+    good_a = _make_output(100)
+    good_a.audio_underrun_s = 0.02
+    good_a.audio_continuity_ok = True
+    good_b = _make_output(100)
+    good_b.audio_underrun_s = 0.0
+    good_b.audio_continuity_ok = True
+
+    metrics, _ = calculate_metrics(
+        input_requests=[],
+        outputs=[bad, good_a, good_b],
+        dur_s=10.0,
+        tokenizer=None,
+        selected_percentiles=[50.0, 99.0],
+        goodput_config_dict={},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=["audio_underrun"],
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=10.0,
+    )
+
+    assert metrics.audio_continuity_ok_rate == pytest.approx(2 / 3, abs=1e-6)
+    # p99 of [0.5, 0.02, 0.0] is dominated by the 0.5 outlier.
+    p99 = dict(metrics.percentiles_audio_underrun_s).get(99.0)
+    assert p99 is not None and p99 > 0.4
+
+
 # ============================================================================
 # TTFT suppression for pure-audio (TTS) benchmarks
 # ============================================================================

tests/benchmarks/test_audio_continuity.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the audio-streaming continuity helper.
+
+The helper is intentionally vllm-free: it operates on a flat timeline of
+``(arrival_time_s, num_bytes)`` chunks plus PCM format params.  These tests
+exercise the underrun arithmetic against hand-crafted timelines so future
+refactors of the streaming backend cannot silently regress the metric.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.benchmarks.audio_continuity import (
+    ContinuityStats,
+    compute_continuity_stats,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
+
+
+# PCM s16le mono at 24 kHz -> 48 000 bytes/sec, the Qwen3-TTS / VoxCPM2 default.
+_SR = 24_000
+_BPS = _SR * 2  # 48 000
+
+
+def test_empty_timeline_is_continuous() -> None:
+    stats = compute_continuity_stats([], [], sample_rate=_SR)
+    assert stats == ContinuityStats(
+        max_underrun_s=0.0,
+        underrun_event_count=0,
+        is_continuous=True,
+    )
+
+
+def test_single_chunk_has_no_underrun() -> None:
+    # One large chunk, no later arrivals to check against.
+    stats = compute_continuity_stats(
+        chunk_arrival_times_s=[0.05],
+        chunk_bytes=[_BPS],  # one second of audio in a single chunk
+        sample_rate=_SR,
+    )
+    assert stats.max_underrun_s == 0.0
+    assert stats.underrun_event_count == 0
+    assert stats.is_continuous is True
+
+
+def test_chunks_faster_than_realtime_have_no_underrun() -> None:
+    # 100ms of audio per chunk, but arrive every 50ms — buffer always builds.
+    stats = compute_continuity_stats(
+        chunk_arrival_times_s=[0.0, 0.05, 0.10, 0.15],
+        chunk_bytes=[_BPS // 10] * 4,
+        sample_rate=_SR,
+    )
+    assert stats.max_underrun_s == 0.0
+    assert stats.underrun_event_count == 0
+    assert stats.is_continuous is True
+
+
+def test_single_gap_above_threshold_breaks_continuity() -> None:
+    # First chunk carries 100ms of audio at t=0.  Next chunk arrives at t=0.5s,
+    # so the player ran dry from t=0.1 to t=0.5 -> 400ms underrun.
+    stats = compute_continuity_stats(
+        chunk_arrival_times_s=[0.0, 0.5],
+        chunk_bytes=[_BPS // 10, _BPS // 10],  # 100ms each
+        sample_rate=_SR,
+        threshold_s=0.1,
+    )
+    assert stats.max_underrun_s == pytest.approx(0.4, abs=1e-6)
+    assert stats.underrun_event_count == 1
+    assert stats.is_continuous is False
+
+
+def test_underrun_is_max_across_gaps_not_sum() -> None:
+    # Two gaps: a 200ms shortfall and a 50ms shortfall.  max_underrun_s
+    # must surface the 200ms one, not the sum.
+    stats = compute_continuity_stats(
+        chunk_arrival_times_s=[0.0, 0.3, 0.45],
+        chunk_bytes=[_BPS // 10, _BPS // 10, _BPS // 10],  # 100ms chunks
+        sample_rate=_SR,
+        threshold_s=0.1,
+    )
+    # At t=0.3 the player consumed 300ms but only 100ms received -> 200ms deficit
+    # At t=0.45 player consumed 450ms, received 200ms -> 250ms deficit.
+    assert stats.max_underrun_s == pytest.approx(0.25, abs=1e-6)
+    assert stats.underrun_event_count == 2
+    assert stats.is_continuous is False
+
+
+def test_continuity_threshold_boundary_inclusive() -> None:
+    # First chunk = 100ms audio at t=0.  Next arrives at t=0.2 -> 100ms deficit
+    # exactly. With threshold=0.1, continuity_ok stays True (boundary).
+    stats = compute_continuity_stats(
+        chunk_arrival_times_s=[0.0, 0.2],
+        chunk_bytes=[_BPS // 10, _BPS // 10],
+        sample_rate=_SR,
+        threshold_s=0.1,
+    )
+    assert stats.max_underrun_s == pytest.approx(0.1, abs=1e-6)
+    assert stats.is_continuous is True
+
+
+def test_mismatched_lengths_returns_zero_stats() -> None:
+    # Defensive: never crash on malformed timeline.
+    stats = compute_continuity_stats(
+        chunk_arrival_times_s=[0.0, 0.1],
+        chunk_bytes=[100],
+        sample_rate=_SR,
+    )
+    assert stats == ContinuityStats(0.0, 0, True)
+
+
+def test_threshold_is_configurable() -> None:
+    # 200ms underrun: violates 100ms threshold, OK against 500ms threshold.
+    arrivals = [0.0, 0.3]
+    bytes_list = [_BPS // 10, _BPS // 10]
+    strict = compute_continuity_stats(arrivals, bytes_list, sample_rate=_SR, threshold_s=0.1)
+    lenient = compute_continuity_stats(arrivals, bytes_list, sample_rate=_SR, threshold_s=0.5)
+    assert strict.is_continuous is False
+    assert lenient.is_continuous is True
+    assert strict.max_underrun_s == lenient.max_underrun_s

tests/benchmarks/test_diffusion_backends_metrics.py

@@ -1,6 +1,11 @@
 import pytest
 
-from benchmarks.diffusion.backends import RequestFuncInput, async_request_chat_completions
+from benchmarks.diffusion.backends import (
+    RequestFuncInput,
+    async_request_chat_completions,
+    endpoint_filename_token,
+    normalize_endpoint,
+)
 
 
 class _MockResponse:
@@ -29,6 +34,32 @@ class _MockSession:
         return _MockResponse(self._payload)
 
 
+@pytest.mark.core_model
+@pytest.mark.benchmark
+@pytest.mark.cpu
+def test_endpoint_normalization_accepts_optional_leading_slash():
+    assert normalize_endpoint("v1/videos") == "/v1/videos"
+    assert normalize_endpoint("/v1/videos") == "/v1/videos"
+    assert normalize_endpoint("v1/chat/completions") == "/v1/chat/completions"
+    assert normalize_endpoint("v1/images/generations") == "/v1/images/generations"
+
+
+@pytest.mark.core_model
+@pytest.mark.benchmark
+@pytest.mark.cpu
+def test_endpoint_normalization_accepts_legacy_backend_aliases():
+    assert normalize_endpoint("vllm-omni") == "/v1/chat/completions"
+    assert normalize_endpoint("openai") == "/v1/images/generations"
+
+
+@pytest.mark.core_model
+@pytest.mark.benchmark
+@pytest.mark.cpu
+def test_endpoint_filename_token_drops_leading_slash():
+    assert endpoint_filename_token("/v1/videos") == "v1_videos"
+    assert endpoint_filename_token("v1/chat/completions") == "v1_chat_completions"
+
+
 @pytest.mark.core_model
 @pytest.mark.benchmark
 @pytest.mark.cpu

tests/conftest.py

@@ -36,6 +36,7 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config):
 _ASSERTION_EXPORT_NAMES = (
     "assert_audio_speech_response",
     "assert_diffusion_response",
+    "assert_http_error",
     "assert_image_diffusion_response",
     "assert_image_valid",
     "assert_omni_response",
@@ -54,6 +55,7 @@ _MEDIA_EXPORT_NAMES = (
 _STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",)
 _RUNTIME_EXPORT_NAMES = (
     "DiffusionResponse",
+    "HttpResponse",
     "OmniResponse",
     "OmniRunner",
     "OmniRunnerHandler",

tests/dfx/conftest.py

@@ -279,6 +279,10 @@ def create_benchmark_indices(
     return indices
 
 
+# Omni perf/stability: ``vllm bench serve --omni`` result template (used if bench exits without writing JSON).
+OMNI_RESULT_TEMPLATE_PATH = Path(__file__).resolve().parent / "perf" / "scripts" / "result_omni_template.json"
+
+
 def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     """Make a single path segment safe for result filenames on common filesystems."""
     if value is None:
@@ -289,7 +293,7 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     return s if s else default
 
 
-def _resolve_baseline_value(
+def resolve_baseline_value(
     baseline_raw: Any,
     *,
     sweep_index: int | None,
@@ -330,7 +334,7 @@ def _baseline_thresholds_for_step(
 ) -> dict[str, Any]:
     """Resolve baseline config to one threshold per metric for this iteration."""
     return {
-        metric_name: _resolve_baseline_value(
+        metric_name: resolve_baseline_value(
             baseline_raw,
             sweep_index=sweep_index,
             max_concurrency=max_concurrency,
@@ -354,7 +358,12 @@ def run_benchmark(
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
 ) -> dict[str, Any]:
-    """Run one ``vllm bench serve --omni`` iteration and return parsed metrics."""
+    """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
+
+    After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the resolved per-metric thresholds
+    (when ``baseline_config`` is provided). If the benchmark exits without writing a result file,
+    ``result_omni_template.json`` is used as a fallback.
+    """
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
     ri = _safe_filename_token(random_input_len)
     ro = _safe_filename_token(random_output_len)
@@ -403,8 +412,17 @@ def run_benchmark(
         result_dir = "./"
 
     result_path = os.path.join(result_dir, result_filename)
-    with open(result_path, encoding="utf-8") as f:
-        result = json.load(f)
+    if not os.path.exists(result_path):
+        with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f:
+            template_result: dict[str, Any] = json.load(f)
+        Path(result_path).parent.mkdir(parents=True, exist_ok=True)
+        with open(result_path, "w", encoding="utf-8") as f:
+            json.dump(template_result, f, ensure_ascii=False, indent=2)
+        print(f"Benchmark result file not generated, fallback to template: {result_path}")
+        result = template_result
+    else:
+        with open(result_path, encoding="utf-8") as f:
+            result = json.load(f)
 
     if baseline_config:
         result["baseline"] = _baseline_thresholds_for_step(
@@ -432,3 +450,12 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         default=None,
         help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
     )
+    parser.addoption(
+        "--assert-baseline",
+        action="store_true",
+        default=False,
+        help=(
+            "When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config "
+            "(default: off)."
+        ),
+    )

tests/dfx/perf/scripts/diffusion_result_template.json

@@ -1,7 +1,7 @@
 [
   {
     "test_name": null,
-    "backend": null,
+    "endpoint": null,
     "timestamp": null,
     "server_params": {
       "model": null,
@@ -43,7 +43,7 @@
       "stage_durations_mean": {},
       "stage_durations_p50": {},
       "stage_durations_p99": {},
-      "backend": null,
+      "endpoint": null,
       "model": null,
       "dataset": null,
       "task": null
@@ -51,6 +51,7 @@
     "log_file": null,
     "Model": null,
     "Framework": null,
+    "API Endpoint": null,
     "Hardware": null,
     "Deployment": null,
     "Task": null,

tests/dfx/perf/scripts/run_benchmark.py

@@ -12,12 +12,17 @@ from tests.dfx.conftest import (
     create_unique_server_params,
     get_benchmark_params_for_server,
     load_configs,
+    resolve_baseline_value,
     run_benchmark,
 )
 from tests.helpers.runtime import OmniServer
 
 pytestmark = [pytest.mark.full_model]
 
+# Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline``
+# (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are
+# defined in the same module.
+
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -46,7 +51,6 @@ if CONFIG_FILE_PATH is None:
     CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE
 
 BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-OMNI_RESULT_TEMPLATE_PATH = Path(__file__).parent / "result_omni_template.json"
 
 
 DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
@@ -87,16 +91,6 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
-def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
-    """Make a single path segment safe for result filenames on common filesystems."""
-    if value is None:
-        return default
-    s = str(value).strip()
-    for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
-        s = s.replace(bad, "_")
-    return s if s else default
-
-
 benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
 
 
@@ -126,76 +120,23 @@ def benchmark_params(request, omni_server):
     }
 
 
-def _resolve_baseline_value(
-    baseline_raw: Any,
-    *,
-    sweep_index: int | None,
-    max_concurrency: Any = None,
-    request_rate: Any = None,
-) -> Any:
-    """Pick the baseline threshold for this sweep step.
-
-    Supported shapes per metric:
-    - **Scalar** — same threshold for every concurrency / QPS.
-    - **List** — aligned with ``max_concurrency`` / ``request_rate`` sweep order; use ``sweep_index``.
-    - **Dict** — keyed by concurrency or rate, e.g. ``{"1": 500, "4": 800}`` (keys are strings in JSON).
-
-    For dict lookup, ``max_concurrency`` is preferred when both are set (concurrency sweep).
-    """
-    if baseline_raw is None:
-        # If no baseline is set, the maximum value will be used.
-        return 100000
-    if isinstance(baseline_raw, dict):
-        if max_concurrency is not None:
-            for key in (max_concurrency, str(max_concurrency)):
-                if key in baseline_raw:
-                    return baseline_raw[key]
-        if request_rate is not None:
-            for key in (request_rate, str(request_rate)):
-                if key in baseline_raw:
-                    return baseline_raw[key]
-        raise KeyError(
-            f"baseline dict has no key for max_concurrency={max_concurrency!r} "
-            f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
-        )
-    if isinstance(baseline_raw, (list, tuple)):
-        return baseline_raw[sweep_index]
-    return baseline_raw
-
-
-def _baseline_thresholds_for_step(
-    baseline_data: dict[str, Any],
-    *,
-    sweep_index: int | None = None,
-    max_concurrency: Any = None,
-    request_rate: Any = None,
-) -> dict[str, Any]:
-    """Resolve ``test.json`` ``baseline`` block to one threshold per metric (same as ``assert_result``)."""
-    return {
-        metric_name: _resolve_baseline_value(
-            baseline_raw,
-            sweep_index=sweep_index,
-            max_concurrency=max_concurrency,
-            request_rate=request_rate,
-        )
-        for metric_name, baseline_raw in baseline_data.items()
-    }
-
-
 def assert_result(
     result,
     params,
     num_prompt,
     *,
+    assert_baseline: bool,
     sweep_index: int | None = None,
-    max_concurrency: Any = None,
-    request_rate: Any = None,
+    max_concurrency: Any | None = None,
+    request_rate: Any | None = None,
 ) -> None:
     assert result["completed"] == num_prompt, "Request failures exist"
+    if not assert_baseline:
+        return
     baseline_data = params.get("baseline", {})
     for metric_name, baseline_raw in baseline_data.items():
         current_value = result[metric_name]
-        baseline_value = _resolve_baseline_value(
+        baseline_value = resolve_baseline_value(
             baseline_raw,
             sweep_index=sweep_index,
             max_concurrency=max_concurrency,
@@ -214,7 +155,7 @@ def assert_result(
 @pytest.mark.benchmark
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 @pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
-def test_performance_benchmark(omni_server, benchmark_params):
+def test_performance_benchmark(omni_server, benchmark_params, request):
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
     dataset_name = params.get("dataset_name", "")
@@ -226,6 +167,8 @@ def test_performance_benchmark(omni_server, benchmark_params):
     print(f"Running benchmark for model: {model}")
     print(f"Benchmark parameters: {benchmark_params}")
 
+    assert_baseline = request.config.getoption("--assert-baseline", default=False)
+
     def to_list(value, default=None):
         if value is None:
             return [] if default is None else [default]
@@ -264,7 +207,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
         elif not isinstance(value, bool):
             args.extend([arg_name, str(value)])
 
-    # QPS test (sweep_index aligns with qps_list / num_prompt_list for this loop)
+    # QPS / request-rate sweep
     for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
         result = run_benchmark(
@@ -283,12 +226,13 @@ def test_performance_benchmark(omni_server, benchmark_params):
         assert_result(
             result,
             params,
-            num_prompt=num_prompt,
+            num_prompt,
+            assert_baseline=assert_baseline,
             sweep_index=i,
             request_rate=qps,
         )
 
-    # concurrency test (sweep_index aligns with max_concurrency_list for separate thresholds per concurrency)
+    # concurrency test
     for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
         args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
         result = run_benchmark(
@@ -307,7 +251,8 @@ def test_performance_benchmark(omni_server, benchmark_params):
         assert_result(
             result,
             params,
-            num_prompt=num_prompt,
+            num_prompt,
+            assert_baseline=assert_baseline,
             sweep_index=i,
             max_concurrency=concurrency,
         )

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -5,16 +5,17 @@ This runner separates two concepts:
 
 1. ``server_type``: how the serving process is started.
    Currently only ``vllm-omni`` is supported here.
-2. ``benchmark_backend``: which serving API the benchmark client calls.
-   Examples: ``vllm-omni`` for ``/v1/chat/completions`` and ``v1/videos``
-   for async video jobs.
+2. ``benchmark_endpoint``: which serving API the benchmark client calls.
+   Examples: ``/v1/chat/completions`` and ``/v1/videos``.
 
 A config JSON file is REQUIRED via --test-config-file:
   pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
 
+Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off).
+
 All benchmark results for a session are consolidated into a single JSON file under
 BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable).
-Each entry in the file contains the test metadata (test_name, backend, benchmark_params,
+Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params,
 timestamp) together with the raw metrics returned by the benchmark script.
 """
 
@@ -33,6 +34,8 @@ from typing import Any, cast
 import psutil
 import pytest
 
+from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
+
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -400,7 +403,7 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]
                 "model": cfg["server_params"]["model"],
                 "serve_args_dict": serve_args_dict,
                 "serve_args": _build_serve_args(serve_args_dict),
-                "benchmark_backend": cfg.get("benchmark_backend"),
+                "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")),
                 "server_params": cfg["server_params"],
             }
         )
@@ -482,7 +485,7 @@ def run_benchmark(
     model: str,
     params: dict[str, Any],
     test_name: str,
-    backend: str = "vllm-omni",
+    endpoint: str = "/v1/chat/completions",
     server_cfg: dict[str, Any] | None = None,
     source_file: str = "",
 ) -> dict[str, Any]:
@@ -490,18 +493,19 @@ def run_benchmark(
 
     The raw metrics are written to a temporary file by the subprocess.  After
     the run completes the metrics are merged with full metadata (test_name,
-    backend, benchmark_params, timestamp, flat reporting fields) and appended
+    endpoint, benchmark_params, timestamp, flat reporting fields) and appended
     to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE).  The
     temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
     to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in
     the record.
     """
     timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+    endpoint = normalize_endpoint(endpoint)
 
     log_dir = BENCHMARK_RESULT_DIR / "logs"
     log_dir.mkdir(parents=True, exist_ok=True)
-    backend_label = backend.replace("/", "_")
-    log_file = log_dir / f"{test_name}_{backend_label}_{timestamp}.log"
+    endpoint_label = endpoint_filename_token(endpoint)
+    log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
 
     with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
         tmp_result_file = Path(tmp.name)
@@ -517,8 +521,8 @@ def run_benchmark(
         str(port),
         "--model",
         model,
-        "--backend",
-        backend,
+        "--endpoint",
+        endpoint,
         "--dataset",
         params.get("dataset", "random"),
         "--task",
@@ -545,7 +549,7 @@ def run_benchmark(
     # cause truncated or out-of-order log output when stdout is piped).
     cmd = [cmd[0], "-u"] + cmd[1:]
 
-    print(f"\nRunning benchmark (backend={backend}): {' '.join(cmd)}")
+    print(f"\nRunning benchmark (endpoint={endpoint}): {' '.join(cmd)}")
     print(f"  Log file: {log_file}")
 
     # Redirect stdout + stderr directly to the log file at the OS level
@@ -597,7 +601,7 @@ def run_benchmark(
 
     record: dict[str, Any] = {
         "test_name": test_name,
-        "backend": backend,
+        "endpoint": endpoint,
         "timestamp": timestamp,
         "server_params": server_cfg.get("server_params"),
         "benchmark_params": params,
@@ -605,7 +609,7 @@ def run_benchmark(
         "log_file": str(log_file),
         "Model": model,
         "Framework": server_type,
-        "API Backend": backend,
+        "API Endpoint": endpoint,
         "Hardware": "",
         "Deployment": "",
         "Task": params.get("task", "t2i"),
@@ -650,12 +654,14 @@ def run_benchmark(
 # ---------------------------------------------------------------------------
 
 
-def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None:
-    """Assert that benchmark metrics satisfy the configured baselines."""
+def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_baseline: bool) -> None:
+    """Assert request completion; optionally assert metrics against ``baseline`` in *params*."""
     num_prompts = params.get("num-prompts", 10)
     completed = result.get("completed_requests", result.get("completed", 0))
     assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
 
+    if not assert_baseline:
+        return
     if params.get("skip-performance-assertion", False):
         print("Skipping performance assertions.")
         return
@@ -669,21 +675,21 @@ def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None:
             assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}"
 
 
-def _default_benchmark_backend_for_task(task: str) -> str:
-    """Return the default client-side benchmark backend for a diffusion task."""
+def _default_benchmark_endpoint_for_task(task: str) -> str:
+    """Return the default client-side benchmark endpoint for a diffusion task."""
     if task in {"t2v", "i2v", "ti2v"}:
-        return "v1/videos"
+        return "/v1/videos"
     if task in {"t2i", "i2i", "ti2i"}:
-        return "vllm-omni"
-    raise ValueError(f"Unsupported task for benchmark backend resolution: {task}")
+        return "/v1/chat/completions"
+    raise ValueError(f"Unsupported task for benchmark endpoint resolution: {task}")
 
 
-def _resolve_benchmark_backend(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
+def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
     """Resolve which serving API the benchmark client should call."""
-    configured = server_cfg.get("benchmark_backend")
+    configured = server_cfg.get("benchmark_endpoint")
     if configured:
-        return cast(str, configured)
-    return _default_benchmark_backend_for_task(cast(str, params.get("task", "t2i")))
+        return normalize_endpoint(cast(str, configured))
+    return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i")))
 
 
 # ---------------------------------------------------------------------------
@@ -697,21 +703,19 @@ def _resolve_benchmark_backend(server_cfg: dict[str, Any], params: dict[str, Any
     indirect=True,
 )
 @pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
-def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
-    """Run the diffusion performance benchmark and assert against baselines.
+def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
+    """Run the diffusion performance benchmark and verify request completion.
 
     One server is started per unique parallel configuration (module scope).
     For each server, all benchmark parameter sets defined in the config JSON
-    are executed sequentially; results are asserted against the baselines.
+    are executed sequentially; metrics are recorded to the aggregated result file.
 
-    Tracked metrics:
-        - throughput_qps          (higher is better)
-        - latency_p50, latency_p99 (lower is better)
+    Pass ``--assert-baseline`` to compare ``throughput_qps``, ``latency_mean``, etc. to the JSON ``baseline`` block.
     """
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
     server_cfg = getattr(diffusion_server, "server_cfg", {})
-    backend = _resolve_benchmark_backend(server_cfg, params)
+    endpoint = _resolve_benchmark_endpoint(server_cfg, params)
 
     result = run_benchmark(
         host=diffusion_server.host,
@@ -719,13 +723,13 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
         model=diffusion_server.model,
         params=params,
         test_name=test_name,
-        backend=backend,
+        endpoint=endpoint,
         server_cfg=server_cfg,
         source_file=cast(str, CONFIG_FILE_PATH),
     )
 
     print(f"\n{'=' * 60}")
-    print(f"Results for {test_name} (server={diffusion_server.server_type}, backend={backend}):")
+    print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
     for key in (
         "throughput_qps",
         "latency_mean",
@@ -742,4 +746,5 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
     print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
     print("=" * 60)
 
-    assert_result(result, params)
+    assert_baseline = request.config.getoption("--assert-baseline", default=False)
+    assert_result(result, params, assert_baseline=assert_baseline)

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -3,31 +3,67 @@
         "test_name": "test_qwen3_omni_chunk",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": ["--async-chunk"]
+            "extra_cli_args": [
+                "--async-chunk"
+            ]
         },
         "benchmark_params": [
             {
                 "dataset_name": "random",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16, 32],
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
                 "random_input_len": 2500,
                 "random_output_len": 900,
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
-                    "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000, 9000],
-                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85, 0.9]
+                    "mean_ttft_ms": [
+                        400,
+                        300,
+                        400,
+                        600,
+                        900
+                    ],
+                    "mean_e2el_ms": [
+                        26000,
+                        43000,
+                        52000,
+                        72000,
+                        109000
+                    ],
+                    "mean_audio_rtf": [
+                        0.15,
+                        0.30,
+                        0.35,
+                        0.50,
+                        0.75
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -42,17 +78,27 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [2000],
-                    "mean_audio_ttfp_ms": [2000],
-                    "mean_audio_rtf": [0.25]
+                    "mean_ttft_ms": [
+                        600
+                    ],
+                    "mean_e2el_ms": [
+                        6000
+                    ],
+                    "mean_audio_rtf": [
+                        0.20
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -69,17 +115,27 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [6000],
-                    "mean_audio_ttfp_ms": [6000],
-                    "mean_audio_rtf": [0.7]
+                    "mean_ttft_ms": [
+                        300
+                    ],
+                    "mean_e2el_ms": [
+                        7000
+                    ],
+                    "mean_audio_rtf": [
+                        0.25
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -98,40 +154,79 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [12000],
-                    "mean_audio_ttfp_ms": [12000],
-                    "mean_audio_rtf": [1.0]
+                    "mean_ttft_ms": [
+                        700
+                    ],
+                    "mean_e2el_ms": [
+                        11000
+                    ],
+                    "mean_audio_rtf": [
+                        0.35
+                    ]
                 }
             },
             {
                 "dataset_name": "random",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16, 32],
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
                 "random_input_len": 2500,
                 "random_output_len": 900,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
+                    "mean_ttft_ms": [
+                        100,
+                        300,
+                        500,
+                        800,
+                        1500
+                    ],
+                    "mean_e2el_ms": [
+                        5000,
+                        8000,
+                        10000,
+                        14000,
+                        20000
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "random_mm_base_items_per_request": 1,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -143,21 +238,32 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [2000]
+                    "mean_ttft_ms": [
+                        500
+                    ],
+                    "mean_e2el_ms": [
+                        1000
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "random_mm_base_items_per_request": 2,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -171,21 +277,32 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [6000]
+                    "mean_ttft_ms": [
+                        200
+                    ],
+                    "mean_e2el_ms": [
+                        800
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "random_mm_base_items_per_request": 3,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -201,7 +318,12 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [6000]
+                    "mean_ttft_ms": [
+                        600
+                    ],
+                    "mean_e2el_ms": [
+                        1300
+                    ]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -3,31 +3,67 @@
         "test_name": "test_qwen3_omni",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": ["--no-async-chunk"]
+            "extra_cli_args": [
+                "--no-async-chunk"
+            ]
         },
         "benchmark_params": [
             {
                 "dataset_name": "random",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16, 32],
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
                 "random_input_len": 2500,
                 "random_output_len": 900,
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
-                    "mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000],
-                    "mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75]
+                    "mean_ttft_ms": [
+                        400,
+                        300,
+                        400,
+                        600,
+                        900
+                    ],
+                    "mean_e2el_ms": [
+                        32000,
+                        34000,
+                        45000,
+                        61000,
+                        90000
+                    ],
+                    "mean_audio_rtf": [
+                        0.20,
+                        0.25,
+                        0.30,
+                        0.45,
+                        0.60
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -42,17 +78,27 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [2000],
-                    "mean_audio_ttfp_ms": [10000],
-                    "mean_audio_rtf": [0.25]
+                    "mean_ttft_ms": [
+                        500
+                    ],
+                    "mean_e2el_ms": [
+                        6000
+                    ],
+                    "mean_audio_rtf": [
+                        0.20
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -69,17 +115,27 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [6000],
-                    "mean_audio_ttfp_ms": [15000],
-                    "mean_audio_rtf": [0.45]
+                    "mean_ttft_ms": [
+                        300
+                    ],
+                    "mean_e2el_ms": [
+                        7000
+                    ],
+                    "mean_audio_rtf": [
+                        0.20
+                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -98,9 +154,15 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [12000],
-                    "mean_audio_ttfp_ms": [18000],
-                    "mean_audio_rtf": [0.9]
+                    "mean_ttft_ms": [
+                        700
+                    ],
+                    "mean_e2el_ms": [
+                        9000
+                    ],
+                    "mean_audio_rtf": [
+                        0.30
+                    ]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -4,7 +4,9 @@
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
             "use_omni": false,
-            "extra_cli_args": ["--no-enable-prefix-caching"]
+            "extra_cli_args": [
+                "--no-enable-prefix-caching"
+            ]
         },
         "benchmark_params": [
             {
@@ -12,8 +14,20 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16, 32],
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
                 "random_input_len": 2500,
                 "random_output_len": 900,
                 "temperature": 0.4,
@@ -23,11 +37,26 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
+                    "mean_ttft_ms": [
+                        300,
+                        300,
+                        1000,
+                        600,
+                        800
+                    ],
+                    "mean_e2el_ms": [
+                        6000,
+                        8000,
+                        11000,
+                        13000,
+                        17000
+                    ]
                 }
             },
             {
@@ -35,8 +64,12 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -47,7 +80,9 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "random_mm_base_items_per_request": 1,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -59,7 +94,12 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [2000]
+                    "mean_ttft_ms": [
+                        30
+                    ],
+                    "mean_e2el_ms": [
+                        540
+                    ]
                 }
             },
             {
@@ -67,8 +107,12 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -79,7 +123,9 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "random_mm_base_items_per_request": 2,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -93,7 +139,12 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [6000]
+                    "mean_ttft_ms": [
+                        30
+                    ],
+                    "mean_e2el_ms": [
+                        600
+                    ]
                 }
             },
             {
@@ -101,8 +152,12 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -113,7 +168,9 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": ["text"]
+                    "modalities": [
+                        "text"
+                    ]
                 },
                 "random_mm_base_items_per_request": 3,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -129,7 +186,12 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [6000]
+                    "mean_ttft_ms": [
+                        40
+                    ],
+                    "mean_e2el_ms": [
+                        700
+                    ]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json

@@ -22,10 +22,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.05,
-                    "latency_mean": 18,
-                    "peak_memory_mb_max": 78500,
-                    "peak_memory_mb_mean": 78500
+                    "throughput_qps": 0.06,
+                    "latency_mean": 20,
+                    "peak_memory_mb_max": 57000,
+                    "peak_memory_mb_mean": 57000
                 }
             },
             {
@@ -41,9 +41,9 @@
                 "enable-negative-prompt": true,
                 "baseline": {
                     "throughput_qps": 0.01,
-                    "latency_mean": 70,
-                    "peak_memory_mb_max": 81000,
-                    "peak_memory_mb_mean": 81000
+                    "latency_mean": 60,
+                    "peak_memory_mb_max": 67000,
+                    "peak_memory_mb_mean": 67000
                 }
             }
         ]
@@ -75,10 +75,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.1,
-                    "latency_mean": 12,
-                    "peak_memory_mb_max": 69000,
-                    "peak_memory_mb_mean": 69000
+                    "throughput_qps": 0.15,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 55000,
+                    "peak_memory_mb_mean": 55000
                 }
             },
             {
@@ -93,10 +93,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.03,
-                    "latency_mean": 28,
-                    "peak_memory_mb_max": 69000,
-                    "peak_memory_mb_mean": 69000
+                    "throughput_qps": 0.05,
+                    "latency_mean": 20,
+                    "peak_memory_mb_max": 55000,
+                    "peak_memory_mb_mean": 55000
                 }
             }
         ]
@@ -138,10 +138,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.10,
-                    "latency_mean": 12,
-                    "peak_memory_mb_max": 73000,
-                    "peak_memory_mb_mean": 73000
+                    "throughput_qps": 0.21,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 57000,
+                    "peak_memory_mb_mean": 57000
                 }
             },
             {
@@ -156,10 +156,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.05,
+                    "throughput_qps": 0.08,
                     "latency_mean": 20,
-                    "peak_memory_mb_max": 81000,
-                    "peak_memory_mb_mean": 81000
+                    "peak_memory_mb_max": 66000,
+                    "peak_memory_mb_mean": 66000
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -21,10 +21,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.05,
-                    "latency_mean": 15.0,
-                    "peak_memory_mb_max": 72500,
-                    "peak_memory_mb_mean": 72500
+                    "throughput_qps": 0.1,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 57000,
+                    "peak_memory_mb_mean": 57000
                 }
             },
             {
@@ -38,10 +38,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.01,
-                    "latency_mean": 65.6,
-                    "peak_memory_mb_max": 80777,
-                    "peak_memory_mb_mean": 80777
+                    "throughput_qps": 0.02,
+                    "latency_mean": 50,
+                    "peak_memory_mb_max": 65000,
+                    "peak_memory_mb_mean": 65000
                 }
             }
         ]
@@ -72,10 +72,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.10,
-                    "latency_mean": 7.2,
-                    "peak_memory_mb_max": 68100,
-                    "peak_memory_mb_mean": 68100
+                    "throughput_qps": 0.21,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 56000,
+                    "peak_memory_mb_mean": 56000
                 }
             },
             {
@@ -89,10 +89,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.03,
-                    "latency_mean": 24.0,
-                    "peak_memory_mb_max": 68100,
-                    "peak_memory_mb_mean": 68100
+                    "throughput_qps": 0.06,
+                    "latency_mean": 20,
+                    "peak_memory_mb_max": 56000,
+                    "peak_memory_mb_mean": 56000
                 }
             }
         ]
@@ -133,10 +133,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.1,
-                    "latency_mean": 6.5,
-                    "peak_memory_mb_max": 72600,
-                    "peak_memory_mb_mean": 72600
+                    "throughput_qps": 0.26,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 58000,
+                    "peak_memory_mb_mean": 58000
                 }
             },
             {
@@ -150,10 +150,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.05,
-                    "latency_mean": 16.0,
-                    "peak_memory_mb_max": 81000,
-                    "peak_memory_mb_mean": 81000
+                    "throughput_qps": 0.1,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 67000,
+                    "peak_memory_mb_mean": 67000
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json

@@ -21,10 +21,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.02,
-                    "latency_mean": 40.0,
-                    "peak_memory_mb_max": 70000,
-                    "peak_memory_mb_mean": 70000
+                    "throughput_qps": 0.05,
+                    "latency_mean": 20,
+                    "peak_memory_mb_max": 63000,
+                    "peak_memory_mb_mean": 63000
                 }
             },
             {
@@ -38,10 +38,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.005,
-                    "latency_mean": 80.0,
-                    "peak_memory_mb_max": 70000,
-                    "peak_memory_mb_mean": 70000
+                    "throughput_qps": 0.03,
+                    "latency_mean": 30,
+                    "peak_memory_mb_max": 63000,
+                    "peak_memory_mb_mean": 63000
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -21,9 +21,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.30,
-                    "latency_mean": 3.50,
-                    "peak_memory_mb_mean": 67000
+                    "throughput_qps": 0.35,
+                    "latency_mean": 3,
+                    "peak_memory_mb_mean": 56000
                 }
             },
             {
@@ -37,9 +37,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.037,
-                    "latency_mean": 27.0,
-                    "peak_memory_mb_mean": 74000
+                    "throughput_qps": 0.03,
+                    "latency_mean": 26,
+                    "peak_memory_mb_mean": 65000
                 }
             }
         ]
@@ -67,9 +67,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.30,
-                    "latency_mean": 3.50,
-                    "peak_memory_mb_mean": 67000
+                    "throughput_qps": 0.35,
+                    "latency_mean": 3,
+                    "peak_memory_mb_mean": 56000
                 }
             },
             {
@@ -83,9 +83,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.037,
-                    "latency_mean": 27.0,
-                    "peak_memory_mb_mean": 74000
+                    "throughput_qps": 0.03,
+                    "latency_mean": 26,
+                    "peak_memory_mb_mean": 65000
                 }
             }
         ]
@@ -116,9 +116,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.1,
-                    "latency_mean": 9.1,
-                    "peak_memory_mb_mean": 61000
+                    "throughput_qps": 0.11,
+                    "latency_mean": 9,
+                    "peak_memory_mb_mean": 55000
                 }
             }
         ]
@@ -159,9 +159,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.1,
-                    "latency_mean": 3.50,
-                    "peak_memory_mb_mean": 67000
+                    "throughput_qps": 0.46,
+                    "latency_mean": 3,
+                    "peak_memory_mb_mean": 56000
                 }
             },
             {
@@ -175,9 +175,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.1,
-                    "latency_mean": 6.15,
-                    "peak_memory_mb_mean": 74000
+                    "throughput_qps": 0.18,
+                    "latency_mean": 6,
+                    "peak_memory_mb_mean": 67000
                 }
             }
         ]

tests/dfx/perf/tests/test_tts.json

@@ -4,7 +4,9 @@
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
       "stage_overrides": {
-        "1": {"devices": "1"}
+        "1": {
+          "devices": "1"
+        }
       }
     },
     "benchmark_params": [
@@ -15,13 +17,21 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [20],
-        "max_concurrency": [1],
+        "num_prompts": [
+          20
+        ],
+        "max_concurrency": [
+          1
+        ],
         "seed_tts_locale": "en",
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [350],
-          "median_audio_rtf": [0.25]
+          "median_audio_ttfp_ms": [
+            3000
+          ],
+          "median_audio_rtf": [
+            0.6
+          ]
         }
       },
       {
@@ -31,14 +41,34 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [80, 128, 128],
-        "max_concurrency": [8, 16, 64],
+        "num_prompts": [
+          80,
+          128,
+          128
+        ],
+        "max_concurrency": [
+          8,
+          16,
+          64
+        ],
         "seed_tts_locale": "en",
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [3500, 30000, 7000],
-          "median_audio_rtf": [0.75, 10.0, 1.5],
-          "audio_throughput": [10.0, 1.0, 40.0]
+          "median_audio_ttfp_ms": [
+            1000,
+            3000,
+            14000
+          ],
+          "median_audio_rtf": [
+            0.70,
+            1.15,
+            3.90
+          ],
+          "audio_throughput": [
+            12,
+            14,
+            14
+          ]
         }
       },
       {
@@ -49,13 +79,19 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [200],
-        "max_concurrency": [4],
+        "num_prompts": [
+          200
+        ],
+        "max_concurrency": [
+          4
+        ],
         "seed_tts_locale": "en",
         "seed_tts_wer_eval": true,
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "mean_audio_rtf": [0.45]
+          "mean_audio_rtf": [
+            0.45
+          ]
         }
       }
     ]
@@ -65,7 +101,9 @@
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
       "stage_overrides": {
-        "1": {"devices": "1"}
+        "1": {
+          "devices": "1"
+        }
       }
     },
     "benchmark_params": [
@@ -76,14 +114,26 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [20],
-        "max_concurrency": [1],
+        "num_prompts": [
+          20
+        ],
+        "max_concurrency": [
+          1
+        ],
         "seed_tts_locale": "en",
-        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
+        "extra_body": {
+          "voice": "Vivian",
+          "language": "English",
+          "task_type": "CustomVoice"
+        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [150],
-          "median_audio_rtf": [0.15]
+          "median_audio_ttfp_ms": [
+            2000
+          ],
+          "median_audio_rtf": [
+            0.40
+          ]
         }
       },
       {
@@ -93,15 +143,39 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [80, 128, 128],
-        "max_concurrency": [8, 16, 64],
+        "num_prompts": [
+          80,
+          128,
+          128
+        ],
+        "max_concurrency": [
+          8,
+          16,
+          64
+        ],
         "seed_tts_locale": "en",
-        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
+        "extra_body": {
+          "voice": "Vivian",
+          "language": "English",
+          "task_type": "CustomVoice"
+        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [1500, 30000, 3000],
-          "median_audio_rtf": [0.30, 10.0, 0.60],
-          "audio_throughput": [30.0, 1.0, 100.0]
+          "median_audio_ttfp_ms": [
+            300,
+            1000,
+            5000
+          ],
+          "median_audio_rtf": [
+            0.25,
+            0.35,
+            1.10
+          ],
+          "audio_throughput": [
+            30,
+            40,
+            40
+          ]
         }
       },
       {
@@ -112,14 +186,24 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [200],
-        "max_concurrency": [4],
+        "num_prompts": [
+          200
+        ],
+        "max_concurrency": [
+          4
+        ],
         "seed_tts_locale": "en",
-        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
+        "extra_body": {
+          "voice": "Vivian",
+          "language": "English",
+          "task_type": "CustomVoice"
+        },
         "seed_tts_wer_eval": true,
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "mean_audio_rtf": [0.35]
+          "mean_audio_rtf": [
+            0.35
+          ]
         }
       },
       {
@@ -129,15 +213,29 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [100],
-        "request_rate": [2.0],
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          2.0
+        ],
         "seed_tts_locale": "en",
-        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
+        "extra_body": {
+          "voice": "Vivian",
+          "language": "English",
+          "task_type": "CustomVoice"
+        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [3000],
-          "median_audio_rtf": [0.25],
-          "audio_throughput": [4.0]
+          "median_audio_ttfp_ms": [
+            100
+          ],
+          "median_audio_rtf": [
+            0.20
+          ],
+          "audio_throughput": [
+            10
+          ]
         }
       },
       {
@@ -147,14 +245,25 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_design",
-        "num_prompts": [20],
-        "max_concurrency": [1],
+        "num_prompts": [
+          20
+        ],
+        "max_concurrency": [
+          1
+        ],
         "seed_tts_locale": "en",
-        "extra_body": {"task_type": "VoiceDesign", "language": "English"},
+        "extra_body": {
+          "task_type": "VoiceDesign",
+          "language": "English"
+        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [150],
-          "median_audio_rtf": [0.15]
+          "median_audio_ttfp_ms": [
+            100
+          ],
+          "median_audio_rtf": [
+            0.15
+          ]
         }
       },
       {
@@ -164,15 +273,38 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_design",
-        "num_prompts": [80, 128, 128],
-        "max_concurrency": [8, 16, 64],
+        "num_prompts": [
+          80,
+          128,
+          128
+        ],
+        "max_concurrency": [
+          8,
+          16,
+          64
+        ],
         "seed_tts_locale": "en",
-        "extra_body": {"task_type": "VoiceDesign", "language": "English"},
+        "extra_body": {
+          "task_type": "VoiceDesign",
+          "language": "English"
+        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [1500, 30000, 3500],
-          "median_audio_rtf": [0.35, 10.0, 0.70],
-          "audio_throughput": [25.0, 1.0, 80.0]
+          "median_audio_ttfp_ms": [
+            100,
+            1000,
+            5000
+          ],
+          "median_audio_rtf": [
+            0.25,
+            0.40,
+            1.20
+          ],
+          "audio_throughput": [
+            30,
+            40,
+            40
+          ]
         }
       },
       {
@@ -182,15 +314,28 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_design",
-        "num_prompts": [100],
-        "request_rate": [2.0],
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          2.0
+        ],
         "seed_tts_locale": "en",
-        "extra_body": {"task_type": "VoiceDesign", "language": "English"},
+        "extra_body": {
+          "task_type": "VoiceDesign",
+          "language": "English"
+        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [3500],
-          "median_audio_rtf": [0.30],
-          "audio_throughput": [4.0]
+          "median_audio_ttfp_ms": [
+            100
+          ],
+          "median_audio_rtf": [
+            0.20
+          ],
+          "audio_throughput": [
+            9
+          ]
         }
       }
     ]

tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json

@@ -30,9 +30,9 @@
                     }
                 ],
                 "baseline": {
-                    "throughput_qps": 0.034,
-                    "latency_mean": 26.0,
-                    "peak_memory_mb_mean": 80000
+                    "throughput_qps": 0.04,
+                    "latency_mean": 30,
+                    "peak_memory_mb_mean": 75000
                 }
             }
         ]
@@ -72,9 +72,9 @@
                     }
                 ],
                 "baseline": {
-                    "throughput_qps": 0.042,
-                    "latency_mean": 21.6,
-                    "peak_memory_mb_mean": 55300
+                    "throughput_qps": 0.06,
+                    "latency_mean": 20,
+                    "peak_memory_mb_mean": 48000
                 }
             },
             {
@@ -97,9 +97,9 @@
                     }
                 ],
                 "baseline": {
-                    "throughput_qps": 0.0085,
-                    "latency_mean": 101.6,
-                    "peak_memory_mb_mean": 65200
+                    "throughput_qps": 0.01,
+                    "latency_mean": 90,
+                    "peak_memory_mb_mean": 56000
                 }
             }
         ]

tests/dfx/reliability/invalid_param_test/conftest.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared env for ``http_invalid`` (real :func:`omni_server`)."""
+
+from __future__ import annotations
+
+import io
+import os
+
+import pytest
+from PIL import Image
+
+# Match ``tests/e2e/online_serving/*`` module-level env for subprocess serve.
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+os.environ.setdefault("VLLM_TEST_CLEAN_GPU_MEMORY", "0")
+
+
+@pytest.fixture
+def tiny_png_bytes() -> bytes:
+    img = Image.new("RGB", (32, 32), color="gray")
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()

tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``POST /v1/audio/generate`` validation (live Stable Audio Open server)."""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+# https://github.com/vllm-project/vllm-omni/issues/3649 — invalid inputs still HTTP 200 / timeout / missing validation.
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+_PARAMS = [
+    pytest.param(
+        OmniServerParams(
+            model="stabilityai/stable-audio-open-1.0",
+            server_args=[
+                "--trust-remote-code",
+                "--enforce-eager",
+                "--model-class-name",
+                "StableAudioPipeline",
+            ],
+        ),
+        id="stable_audio_open",
+        marks=hardware_marks(res={"cuda": "L4"}),
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"input": ""}, ("cannot be empty", "input"), id="input_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"input": "   "}, ("cannot be empty", "input"), id="input_whitespace_only", marks=_SKIP_ISSUE_3649
+        ),
+        pytest.param({"input": 123}, ("input", "a valid string"), id="input_wrong_type"),
+        pytest.param(
+            {"response_format": "mpeg"},
+            ("response_format", "literal_error", "Input should be"),
+            id="response_format_invalid",
+        ),
+        pytest.param({"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="speed_too_low"),
+        pytest.param({"speed": 99.0}, ("speed", "less_than_equal", "4"), id="speed_too_high"),
+        pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"),
+        pytest.param(
+            {"stream_format": "sse"},
+            ("sse", "stream_format", "value_error", "not a supported stream_format"),
+            id="stream_format_sse_blocked",
+        ),
+        pytest.param(
+            {"audio_length": -1.0}, ("audio_end_in_s", "higher than", "audio_start_in_s"), id="audio_length_negative"
+        ),
+        pytest.param({"audio_length": 0}, "audio_length", id="audio_length_zero", marks=_SKIP_ISSUE_3649),
+        # ``OmniOpenAIServingAudioGenerate`` only forwards ``audio_start`` when ``audio_length`` is set.
+        pytest.param(
+            {"audio_start": -0.5, "audio_length": 5.0},
+            ("audio_start_in_s", "between", "0", "512"),
+            id="audio_start_negative",
+        ),
+        pytest.param(
+            {"negative_prompt": ["noise"]}, ("negative_prompt", "a valid string"), id="negative_prompt_wrong_type"
+        ),
+        pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649),
+        pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649
+        ),
+        pytest.param(
+            {"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative"
+        ),
+        pytest.param(
+            {"num_inference_steps": 6000},
+            "num_inference_steps",
+            id="num_inference_steps_above_max",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            {"guidance_scale": 1001.0}, "guidance_scale", id="guidance_scale_above_max", marks=_SKIP_ISSUE_3649
+        ),
+        pytest.param(
+            {"audio_length": 86401.0}, ("Requested audio length", "exceeds maximum"), id="audio_length_above_max"
+        ),
+        pytest.param(
+            {"audio_start": 86401.0, "audio_length": 1.0},
+            ("audio_start_in_s", "between", "0", "512"),
+            id="audio_start_above_max",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_audio_generate_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body = {"model": omni_server.model, "input": "ambient electronic pad"}
+    body.update(overrides)
+    openai_client.send_audio_generate_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_audio_generate_missing_input(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_generate_http_request(
+        {
+            "json": {"model": omni_server.model},
+            "timeout": 120,
+            "err_code": 400,
+            "err_message": ("missing", "input", "field required"),
+        }
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -0,0 +1,862 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HTTP + WebSocket validation for Qwen3-TTS Base: ``/v1/audio/speech``, ``/v1/audio/speech/stream``, batch, voices."""
+
+from __future__ import annotations
+
+import io
+import json
+import uuid
+import wave
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.stage_config import get_deploy_config_path
+
+# Speech / voice-upload numeric caps for these tests only (not exported from vllm_omni).
+# Align where possible with serving_speech defaults: _TTS_MAX_INSTRUCTIONS_LENGTH (500),
+# _TTS_MAX_NEW_TOKENS_MAX (4096). Voice form max lengths are not public constants in prod;
+# values below keep tests self-contained and deterministic.
+_SPEECH_API_MAX_INSTRUCTIONS_CHARS = 500
+_SPEECH_API_MAX_NEW_TOKENS = 4096
+_VOICE_UPLOAD_MAX_CONSENT_LEN = 1024
+_VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192
+_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS = 2048
+REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+# Qwen3-TTS 0.6B expects 1024-dim speaker embeddings (see serving_speech / talker hidden_size).
+_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM = 1024
+# Pop-only marker for ``test_speech_invalid_field_values``: omit ``ref_audio`` / ``ref_text`` so Base
+# ``speaker_embedding`` cases mirror standalone embedding-only JSON (no ref_audio clash).
+_SPEECH_INVALID_EMBEDDING_ONLY_BODY = "__speech_invalid_embedding_only_body__"
+# Batch-only: merged items inherit batch ``ref_audio`` / ``ref_text`` because ``_pick`` treats item ``None``
+# as "unset". Strip clone fields from the batch envelope before asserting CustomVoice preset errors.
+_SPEECH_BATCH_NO_CLONE_FIELDS = "__speech_batch_no_clone_fields__"
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+_L4_SPEECH_HW = hardware_marks(res={"cuda": "L4"})
+_SPEECH_SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
+
+_QWEN3_TTS_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="qwen3_tts",
+        marks=_L4_SPEECH_HW,
+    ),
+]
+
+
+def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None:
+    if loc == "batch":
+        body.update(overrides)
+        return
+    items = body["items"]
+    assert isinstance(items, list) and len(items) >= 1
+    first = items[0]
+    assert isinstance(first, dict)
+    body["items"] = [{**first, **overrides}]
+
+
+def _pcm_wav_mono_bytes(*, duration_s: float = 1.1, sample_rate: int = 44100) -> bytes:
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        nframes = max(1, int(sample_rate * duration_s))
+        wf.writeframes(b"\x00\x00" * nframes)
+    return buf.getvalue()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/audio/speech (JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_http_request(
+        {
+            "raw_body": "{",
+            "timeout": 120,
+            "err_code": 400,
+            "err_message": ("json", "decode", "invalid"),
+        }
+    )
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_http_request(
+        {
+            "json": {"model": omni_server.model},
+            "timeout": 120,
+            "err_code": 400,
+            "err_message": ("input", "field required", "missing"),
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"input": ""}, ("input", "empty"), id="input_empty"),
+        pytest.param({"input": "   "}, ("input", "empty"), id="input_whitespace_only"),
+        pytest.param({"voice": ""}, "Invalid voice", id="voice_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"instructions": 123}, ("instructions", "string_type", "valid string"), id="instructions_wrong_type"
+        ),
+        pytest.param(
+            {"response_format": "mpeg"}, ("response_format", "literal_error", "wav"), id="response_format_invalid"
+        ),
+        pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"),
+        pytest.param(
+            {"stream_format": "sse"}, ("stream_format", "value_error", "audio"), id="stream_format_sse_blocked"
+        ),
+        pytest.param({"stream": "wrong_type"}, ("stream", "bool_parsing", "validation error"), id="stream_wrong_type"),
+        pytest.param(
+            {
+                "voice": "vivian",
+                "task_type": "CustomVoice",
+                "ref_audio": None,
+                "ref_text": None,
+            },
+            ("CustomVoice", "does not support"),
+            id="customvoice_task_not_supported",
+        ),
+        pytest.param(
+            {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
+        ),
+        pytest.param({"language": ""}, ("language", "invalid language", "chinese"), id="language_empty"),
+        pytest.param({"ref_audio": "ftp://example.com/a.wav"}, ("ref_audio", "URL"), id="ref_audio_bad_scheme"),
+        pytest.param(
+            {"ref_audio": "not_a_valid_uri"},
+            ("ref_audio", "url"),
+            id="ref_audio_invalid_uri",
+        ),
+        pytest.param({"ref_text": ["x"]}, ("ref_text", "string_type", "valid string"), id="ref_text_wrong_type"),
+        pytest.param(
+            {"x_vector_only_mode": "wrong"},
+            ("x_vector_only_mode", "valid boolean"),
+            id="x_vector_only_mode_string_not_boolean",
+        ),
+        pytest.param(
+            {"speaker_embedding": [0.01] * _QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM},
+            ("mutually exclusive", "speaker_embedding", "ref_audio"),
+            id="speaker_embedding_with_ref_audio_mutually_exclusive",
+        ),
+        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
+        pytest.param(
+            {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
+            ("max_new_tokens", "exceed 4096"),
+            id="max_new_tokens_above_max",
+        ),
+        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
+        pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
+        pytest.param(
+            {"initial_codec_chunk_frames": -1},
+            ("initial_codec_chunk_frames", "greater_than_equal", "0"),
+            id="initial_codec_chunk_frames_negative",
+        ),
+        pytest.param(
+            {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
+            ("instructions", "too long"),
+            id="instructions_exceed_api_limit",
+        ),
+        # Base clone: whitespace-only ref_text fails serving validation (not pydantic).
+        pytest.param(
+            {
+                "task_type": "Base",
+                "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
+                "ref_text": "   ",
+            },
+            ("ref_text", "base task", "non-empty"),
+            id="ref_text_whitespace_only_base_clone",
+        ),
+        pytest.param(
+            {
+                _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
+                "task_type": "Base",
+                "speaker_embedding": [],
+            },
+            ("speaker_embedding", "non-empty"),
+            id="base_speaker_embedding_empty",
+        ),
+        pytest.param(
+            {
+                _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
+                "task_type": "Base",
+                "x_vector_only_mode": True,
+                "speaker_embedding": [0.01] * 512,
+            },
+            ("speaker_embedding", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)),
+            id="base_speaker_embedding_wrong_dim_short",
+        ),
+        pytest.param(
+            {
+                _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
+                "task_type": "Base",
+                "x_vector_only_mode": True,
+                "speaker_embedding": [0.01] * 2048,
+            },
+            ("speaker_embedding", "2048", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)),
+            id="base_speaker_embedding_wrong_dim_long",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    ov = dict(overrides)
+    embedding_only = bool(ov.pop(_SPEECH_INVALID_EMBEDDING_ONLY_BODY, False))
+    body = {
+        "model": omni_server.model,
+        "input": "Hello.",
+        "voice": "clone",
+        "language": "English",
+        "response_format": "wav",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    if embedding_only:
+        body.pop("ref_audio", None)
+        body.pop("ref_text", None)
+    body.update(ov)
+    openai_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/audio/speech/batch (JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_batch_empty_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_batch_http_request(
+        {"json": {"items": []}, "timeout": 120, "err_code": (400, 422), "err_message": ("items", "least 1 item")}
+    )
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_batch_http_request(
+        {
+            "json": {"model": omni_server.model},
+            "timeout": 120,
+            "err_code": (400, 422),
+            "err_message": ("items", "field required", "missing"),
+        }
+    )
+
+
+# Align field-level checks with ``POST /v1/audio/speech`` where ``BatchSpeechRequest`` /
+# ``SpeechBatchItem`` expose the same knobs. Not on batch schema (single-speech only):
+# ``stream``, ``stream_format``, ``seed``, ``speaker_embedding`` — see dedicated speech tests.
+@pytest.mark.parametrize(
+    "loc, overrides, err_message",
+    [
+        pytest.param("item", {"input": ""}, ("input", "empty"), id="item_input_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            "item",
+            {"input": "   "},
+            ("input", "empty"),
+            id="item_input_whitespace_only",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param("batch", {"voice": ""}, "Invalid voice", id="batch_voice_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param("item", {"voice": ""}, "Invalid voice", id="item_voice_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            "batch",
+            {
+                "voice": "nonexistent_voice_xyz",
+                "task_type": "CustomVoice",
+                "ref_audio": None,
+                "ref_text": None,
+            },
+            "Invalid voice",
+            id="batch_voice_unknown_preset",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {
+                _SPEECH_BATCH_NO_CLONE_FIELDS: True,
+                "voice": "nonexistent_voice_xyz",
+                "task_type": "CustomVoice",
+            },
+            "Invalid voice",
+            id="item_voice_unknown_preset",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"instructions": 123},
+            ("instructions", "string_type", "valid string"),
+            id="batch_instructions_wrong_type",
+        ),
+        pytest.param(
+            "item",
+            {"instructions": 123},
+            ("instructions", "string_type", "valid string"),
+            id="item_instructions_wrong_type",
+        ),
+        pytest.param(
+            "batch",
+            {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
+            ("instructions", "too long"),
+            id="batch_instructions_exceed_api_limit",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
+            ("instructions", "too long"),
+            id="item_instructions_exceed_api_limit",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"response_format": "mpeg"},
+            ("response_format", "literal_error", "wav"),
+            id="batch_response_format_invalid",
+        ),
+        pytest.param(
+            "item",
+            {"response_format": "mpeg"},
+            ("response_format", "literal_error", "wav"),
+            id="item_response_format_invalid",
+        ),
+        pytest.param("batch", {"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="batch_speed_too_low"),
+        pytest.param("item", {"speed": 10.0}, ("speed", "less_than_equal", "4"), id="item_speed_too_high"),
+        pytest.param(
+            "batch",
+            {"task_type": "InvalidEnum"},
+            ("task_type", "literal_error", "CustomVoice"),
+            id="batch_task_type_invalid",
+        ),
+        pytest.param(
+            "item",
+            {"task_type": "InvalidEnum"},
+            ("task_type", "literal_error", "CustomVoice"),
+            id="item_task_type_invalid",
+        ),
+        pytest.param(
+            "batch",
+            {"language": ""},
+            ("language", "invalid language", "chinese"),
+            id="batch_language_empty",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"language": ""},
+            ("language", "invalid language", "chinese"),
+            id="item_language_empty",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"ref_audio": "ftp://example.com/a.wav"},
+            ("ref_audio", "URL"),
+            id="batch_ref_audio_bad_scheme",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"ref_audio": "ftp://example.com/a.wav"},
+            ("ref_audio", "URL"),
+            id="item_ref_audio_bad_scheme",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"ref_audio": "not_a_valid_uri"},
+            ("ref_audio", "url"),
+            id="batch_ref_audio_invalid_uri",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"ref_audio": "not_a_valid_uri"},
+            ("ref_audio", "url"),
+            id="item_ref_audio_invalid_uri",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"ref_text": ["x"]},
+            ("ref_text", "string_type", "valid string"),
+            id="batch_ref_text_wrong_type",
+        ),
+        pytest.param(
+            "item",
+            {"ref_text": ["x"]},
+            ("ref_text", "string_type", "valid string"),
+            id="item_ref_text_wrong_type",
+        ),
+        pytest.param(
+            "batch",
+            {"x_vector_only_mode": "wrong_type"},
+            ("x_vector_only_mode", "bool_parsing", "validation error"),
+            id="batch_x_vector_only_mode_wrong_type",
+        ),
+        pytest.param(
+            "item",
+            {"x_vector_only_mode": "wrong_type"},
+            ("x_vector_only_mode", "bool_parsing", "validation error"),
+            id="item_x_vector_only_mode_wrong_type",
+        ),
+        pytest.param(
+            "batch",
+            {"max_new_tokens": 0},
+            ("max_new_tokens", "least 1"),
+            id="batch_max_new_tokens_below_min",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"max_new_tokens": 0},
+            ("max_new_tokens", "least 1"),
+            id="item_max_new_tokens_below_min",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
+            ("max_new_tokens", "exceed 4096"),
+            id="batch_max_new_tokens_above_max",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
+            ("max_new_tokens", "exceed 4096"),
+            id="item_max_new_tokens_above_max",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"initial_codec_chunk_frames": -1},
+            ("initial_codec_chunk_frames", "greater_than_equal", "0"),
+            id="batch_initial_codec_chunk_frames_negative",
+        ),
+        pytest.param(
+            "item",
+            {"initial_codec_chunk_frames": -1},
+            ("initial_codec_chunk_frames", "greater_than_equal", "0"),
+            id="item_initial_codec_chunk_frames_negative",
+        ),
+        # Base clone: whitespace-only ref_text (same serving check as single speech).
+        pytest.param(
+            "batch",
+            {
+                "task_type": "Base",
+                "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
+                "ref_text": "   ",
+            },
+            ("ref_text", "base task", "non-empty"),
+            id="batch_ref_text_whitespace_only_base_clone",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {
+                "task_type": "Base",
+                "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
+                "ref_text": "   ",
+            },
+            ("ref_text", "base task", "non-empty"),
+            id="item_ref_text_whitespace_only_base_clone",
+            marks=_SKIP_ISSUE_3649,
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_batch_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    loc: str,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    ov = dict(overrides)
+    strip_clone = bool(ov.pop(_SPEECH_BATCH_NO_CLONE_FIELDS, False))
+    body = {
+        "model": omni_server.model,
+        "voice": "clone",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+        "language": "English",
+        "response_format": "wav",
+        "items": [{"input": "Hello."}],
+    }
+    if strip_clone:
+        body.pop("ref_audio", None)
+        body.pop("ref_text", None)
+    _apply_batch_overrides(body, loc=loc, overrides=ov)
+    openai_client.send_audio_speech_batch_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WebSocket /v1/audio/speech/stream
+# ─────────────────────────────────────────────────────────────────────────────
+
+# ``test_speech_stream_invalid_requests``: build ``session.config`` JSON with ``omni_server.model``.
+_SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES = object()
+
+
+@pytest.mark.parametrize(
+    "send_frames_spec, err_message",
+    [
+        pytest.param("{not-json", "Invalid JSON", id="invalid_json"),
+        pytest.param(
+            json.dumps({"type": "input.text", "text": "hello"}),
+            "Expected session.config",
+            id="wrong_first_message_type",
+        ),
+        pytest.param(
+            _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES,
+            ("invalid session config", "response_format", "pcm"),
+            id="session_config_stream_audio_requires_pcm",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_stream_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    send_frames_spec: Any,
+    err_message: str | tuple[str, ...],
+) -> None:
+    """Bad first WebSocket text frames: non-JSON, wrong ``type``, or invalid ``session.config``."""
+    if send_frames_spec is _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES:
+        send_frames = json.dumps(
+            {
+                "type": "session.config",
+                "model": omni_server.model,
+                "response_format": "wav",
+                "stream_audio": True,
+            }
+        )
+    else:
+        send_frames = send_frames_spec
+    assert isinstance(send_frames, str)
+    openai_client.send_audio_speech_stream_ws_request(
+        {
+            "send_frames": send_frames,
+            "timeout": 120,
+            "ws_max_size": None,
+            "ws_json_type": "error",
+            "err_message": err_message,
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/audio/voices (multipart + embedding-only) and DELETE /v1/audio/voices/{name}
+# ─────────────────────────────────────────────────────────────────────────────
+
+# Templates for ``test_voices_create_invalid_requests``: resolved in ``_finalize_voices_form_data``.
+_VF_LONG_CONSENT = object()
+_VF_LONG_REF_TEXT = object()
+_VF_LONG_SPEAKER_DESC = object()
+_VF_EMBEDDING_NAN_JSON = object()
+
+
+def _finalize_voices_form_data(template: dict[str, Any], uuid_hex: str) -> dict[str, Any]:
+    out: dict[str, Any] = {}
+    for key, val in template.items():
+        if val is _VF_LONG_CONSENT:
+            out[key] = "x" * (_VOICE_UPLOAD_MAX_CONSENT_LEN + 1)
+        elif val is _VF_LONG_REF_TEXT:
+            out[key] = "y" * (_VOICE_UPLOAD_MAX_REF_TEXT_CHARS + 1)
+        elif val is _VF_LONG_SPEAKER_DESC:
+            out[key] = "z" * (_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS + 1)
+        elif val is _VF_EMBEDDING_NAN_JSON:
+            out[key] = json.dumps([0.1] * 1023 + [float("nan")])
+        elif isinstance(val, str):
+            out[key] = val.replace("{uuid}", uuid_hex)
+        else:
+            out[key] = val
+    return out
+
+
+def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
+    """``None`` => no multipart file (embedding-only or neither)."""
+    if kind is None:
+        return None
+    if kind == "wav_ok":
+        return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "audio/wav")}
+    if kind == "wav_short":
+        return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(duration_s=0.3), "audio/wav")}
+    if kind == "wav_bad_body":
+        return {"audio_sample": ("clip.wav", b"not valid wav content", "audio/wav")}
+    if kind == "wav_pdf_type":
+        return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "application/pdf")}
+    raise AssertionError(f"unknown multipart kind {kind!r}")
+
+
+@pytest.mark.parametrize(
+    "multipart_kind, form_template, err_message",
+    [
+        pytest.param(
+            None,
+            {"consent": "c1", "name": "v1"},
+            ("audio_sample", "speaker_embedding", "provided"),
+            id="neither_audio_nor_embedding",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"name": "v_miss_consent_{uuid}"},
+            ("missing", "consent", "field required"),
+            id="missing_consent",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok"},
+            ("missing", "name", "field required"),
+            id="missing_name",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "", "name": "v_empty_consent_{uuid}"},
+            ("missing", "consent", "field required"),
+            id="empty_consent",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "   ", "name": "v_ws_consent_{uuid}"},
+            "consent",
+            id="whitespace_consent",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "bad/consent", "name": "v_bad_cons_{uuid}"},
+            "consent",
+            id="consent_path_sep",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": _VF_LONG_CONSENT, "name": "v_long_cons_{uuid}"},
+            ("consent", "too long", "Failed to save"),
+            id="consent_too_long",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok", "name": ""},
+            ("missing", "name", "field required"),
+            id="empty_name",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok", "name": "evil/name"},
+            ("voice name", "invalid voice name"),
+            id="invalid_name_path_sep",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok", "name": "v_long_ref_{uuid}", "ref_text": _VF_LONG_REF_TEXT},
+            "ref_text",
+            id="ref_text_too_long",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {
+                "consent": "consent_ok",
+                "name": "v_long_desc_{uuid}",
+                "speaker_description": _VF_LONG_SPEAKER_DESC,
+            },
+            "speaker_description",
+            id="speaker_description_too_long",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_pdf_type",
+            {"consent": "consent_ok", "name": "v_bad_mime_{uuid}"},
+            ("mime", "unsupported", "audio/mp4"),
+            id="audio_unsupported_mime",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_short",
+            {"consent": "consent_ok", "name": "v_short_{uuid}"},
+            ("too short", "reference audio"),
+            id="audio_too_short",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_bad_body",
+            {"consent": "consent_ok", "name": "v_bad_wav_{uuid}"},
+            ("decode", "audio", "Format not recognised"),
+            id="audio_decode_error",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": "not valid json [[[",
+                "consent": "consent_emb",
+                "name": "v_bad_json_{uuid}",
+            },
+            ("speaker_embedding", "valid JSON"),
+            id="speaker_embedding_invalid_json",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": json.dumps([]),
+                "consent": "consent_emb",
+                "name": "v_empty_emb_{uuid}",
+            },
+            ("speaker_embedding", "non-empty"),
+            id="speaker_embedding_empty",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": _VF_EMBEDDING_NAN_JSON,
+                "consent": "consent_emb",
+                "name": "v_nan_emb_{uuid}",
+            },
+            ("speaker_embedding", "finite"),
+            id="speaker_embedding_nan",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": json.dumps([0.01] * 512),
+                "consent": "consent_emb",
+                "name": "v_wrong_dim_{uuid}",
+            },
+            ("expected", "dimensions", "1024"),
+            id="speaker_embedding_wrong_dim",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {
+                "consent": "consent_emb",
+                "name": "v_both_{uuid}",
+                "speaker_embedding": json.dumps([0.01] * 1024),
+            },
+            ("mutually exclusive", "audio_sample", "speaker_embedding"),
+            id="audio_and_embedding_mutually_exclusive",
+            marks=_SKIP_ISSUE_3649,
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_voices_create_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    multipart_kind: str | None,
+    form_template: dict[str, Any],
+    err_message: str | tuple[str, ...],
+) -> None:
+    """Invalid ``POST /v1/audio/voices`` bodies (multipart or embedding-only)."""
+    uid = uuid.uuid4().hex[:8]
+    cfg: dict[str, Any] = {
+        "data": _finalize_voices_form_data(form_template, uid),
+        "timeout": 120,
+        "err_code": 400,
+        "err_message": err_message,
+    }
+    files = _voices_upload_multipart_files(multipart_kind)
+    if files is not None:
+        cfg["files"] = files
+    openai_client.send_audio_voices_create_http_request(cfg)
+
+
+@pytest.mark.parametrize(
+    "voice_name, err_code, err_message",
+    [
+        pytest.param(
+            "missing-voice-404",
+            404,
+            "not found",
+            id="not_found",
+        ),
+        pytest.param(
+            "   ",
+            404,
+            "not found",
+            id="whitespace_only",
+        ),
+        pytest.param(
+            "evil/name",
+            404,
+            "not found",
+            id="path_separator",
+        ),
+        pytest.param(
+            ".",
+            405,
+            ("Method Not Allowed", "detail"),
+            id="dot_single",
+        ),
+        pytest.param(
+            "..",
+            404,
+            "not found",
+            id="dot_dot",
+        ),
+        pytest.param(
+            "bad\x00voice",
+            404,
+            "not found",
+            id="nul_byte",
+        ),
+        pytest.param(
+            "",
+            405,
+            "Method Not Allowed",
+            id="empty_name_path",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_voices_delete_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    voice_name: str,
+    err_code: int | tuple[int, ...],
+    err_message: str | tuple[str, ...],
+) -> None:
+    """DELETE ``/v1/audio/voices/{name}``: missing voice (404), odd segments (often 404), ``.``paths (405)."""
+    openai_client.send_audio_voices_delete_http_request(
+        {
+            "name": voice_name,
+            "timeout": 120,
+            "err_code": err_code,
+            "err_message": err_message,
+        }
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``POST /v1/images/edits`` multipart validation (live ``Qwen/Qwen-Image-Edit``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+_PARAMS = [
+    pytest.param(
+        OmniServerParams(model="Qwen/Qwen-Image-Edit"),
+        id="qwen_image_edit",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+# Populated in ``test_images_edits_invalid_requests`` from ``omni_server.model``.
+_IMAGE_EDITS_SERVER_MODEL = object()
+
+
+def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> dict[str, Any]:
+    return {k: (server_model if v is _IMAGE_EDITS_SERVER_MODEL else v) for k, v in template.items()}
+
+
+@pytest.mark.parametrize(
+    "include_image, data_template, err_message, err_code",
+    [
+        pytest.param(
+            False,
+            {"prompt": "make it brighter", "model": _IMAGE_EDITS_SERVER_MODEL},
+            ("image", "required", "Unprocessable Entity"),
+            422,
+            id="missing_image",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "model": "wrong-model-id"},
+            ("model", "mismatch"),
+            400,
+            id="model_mismatch",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "response_format": "url"},
+            ("response_format", "b64_json", "supported"),
+            400,
+            id="unsupported_response_format",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "output_compression": "101"},
+            ("less_than_equal", "output_compression", "100"),
+            400,
+            id="output_compression_above_max",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"},
+            ("Invalid layers", "3", "1"),
+            400,
+            id="invalid_layers",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_images_edits_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    tiny_png_bytes: bytes,
+    include_image: bool,
+    data_template: dict[str, Any],
+    err_message: str | tuple[str, ...],
+    err_code: int | tuple[int, ...],
+) -> None:
+    cfg: dict[str, Any] = {
+        "data": _finalize_edits_form_data(data_template, omni_server.model),
+        "timeout": 300,
+        "err_code": err_code,
+        "err_message": err_message,
+    }
+    if include_image:
+        cfg["files"] = {"image": ("x.png", tiny_png_bytes, "image/png")}
+    openai_client.send_images_edits_http_request(cfg)

tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``POST /v1/images/generations`` validation (live ``Qwen/Qwen-Image``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+_PARAMS = [
+    pytest.param(
+        OmniServerParams(model="Qwen/Qwen-Image"),
+        id="qwen_image",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+# Full replacement body for ``test_images_generations_invalid_requests`` (no minimal prompt/model merge).
+_IMAGES_GEN_MISSING_PROMPT_JSON = object()
+
+
+def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
+    return {"prompt": "a simple red apple icon", "model": omni_server.model}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/images/generations (JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "body_spec, err_message",
+    [
+        pytest.param(_IMAGES_GEN_MISSING_PROMPT_JSON, ("prompt", "Field required", "Missing"), id="missing_prompt"),
+        pytest.param(
+            {"prompt": "cat", "model": "not-the-loaded-checkpoint"},
+            ("model", "mismatch"),
+            id="model_mismatch",
+        ),
+        pytest.param({"prompt": ""}, ("prompt", "empty"), id="prompt_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param({"prompt": "   "}, ("prompt", "empty"), id="prompt_whitespace_only", marks=_SKIP_ISSUE_3649),
+        pytest.param({"prompt": 123}, ("prompt", "string_type", "valid string"), id="prompt_wrong_type"),
+        pytest.param({"n": 0}, ("n", "greater_than_equal", "1"), id="n_below_min"),
+        pytest.param({"n": 11}, ("n", "less_than_equal", "10"), id="n_above_max"),
+        pytest.param({"size": "1024"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_height_token"),
+        pytest.param({"size": "not-a-size"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_separator"),
+        pytest.param({"size": "1024xabc"}, ("Invalid size format", "integers"), id="size_non_integer_height"),
+        pytest.param({"size": "1024x"}, ("Invalid size format", "integers"), id="size_incomplete_height"),
+        pytest.param({"size": "0x1024"}, ("Invalid size", "positive integers"), id="size_non_positive_width"),
+        pytest.param(
+            {"response_format": "url"},
+            ("response_format", "value_error", "b64_json"),
+            id="response_format_not_b64_json",
+        ),
+        pytest.param(
+            {"use_system_prompt": "bogus"},
+            ("use_system_prompt", "value_error", "dynamic"),
+            id="use_system_prompt_invalid",
+        ),
+        pytest.param(
+            {"num_inference_steps": 0},
+            ("num_inference_steps", "'greater_than_equal", "1"),
+            id="num_inference_steps_below_min",
+        ),
+        pytest.param(
+            {"num_inference_steps": 201},
+            ("num_inference_steps", "less_than_equal", "200"),
+            id="num_inference_steps_above_max",
+        ),
+        pytest.param(
+            {"guidance_scale": -1.0}, ("guidance_scale", "greater_than_equal", "0"), id="guidance_scale_negative"
+        ),
+        pytest.param(
+            {"guidance_scale": 21.0}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max"
+        ),
+        pytest.param(
+            {"true_cfg_scale": -1.0}, ("true_cfg_scale", "greater_than_equal", "0"), id="true_cfg_scale_negative"
+        ),
+        pytest.param(
+            {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max"
+        ),
+        pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"),
+        pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"),
+        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"seed": 2**32},
+            ("seed", "less_than_equal", "4294967295"),
+            id="seed_above_uint32",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            {"output_format": "gif"},
+            ("output_format", "value_error", "b64_json"),
+            id="output_format_invalid",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            {"vae_use_slicing": "wrong_type"},
+            ("vae_use_slicing", "bool_parsing", "validation error"),
+            id="vae_use_slicing_wrong_type",
+        ),
+        pytest.param({"user": 123}, ("user", "string_type", "valid string"), id="user_wrong_type"),
+        pytest.param(
+            {"negative_prompt": ["noise"]},
+            ("negative_prompt", "string_type", "valid string"),
+            id="negative_prompt_wrong_type",
+        ),
+        pytest.param(
+            {"generator_device": 1},
+            ("generator_device", "string_type", "valid string"),
+            id="generator_device_wrong_type",
+        ),
+        pytest.param(
+            {"lora": {"foo": "bar"}}, ("lora", "both name and path", "required"), id="lora_missing_required_fields"
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_images_generations_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    body_spec: Any,
+    err_message: str | tuple[str, ...],
+) -> None:
+    """Invalid JSON bodies for ``POST /v1/images/generations`` (missing prompt, model mismatch, bad fields)."""
+    if body_spec is _IMAGES_GEN_MISSING_PROMPT_JSON:
+        body: dict[str, object] = {"size": "1024x1024"}
+    else:
+        assert isinstance(body_spec, dict)
+        body = _minimal_images_gen_json(omni_server)
+        body.update(body_spec)
+    openai_client.send_images_generations_http_request(
+        {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message}
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.omni]
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+
+def _minimal_chat_json(omni_server: OmniServer) -> dict[str, object]:
+    """Minimal valid chat body; individual tests override one offending field."""
+    return {
+        "model": omni_server.model,
+        "messages": [{"role": "user", "content": "hello"}],
+        "stream": False,
+    }
+
+
+_QWEN3_OMNI_SERVER = [
+    pytest.param(
+        OmniServerParams(
+            model="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            stage_config_path=get_deploy_config_path("qwen3_omni_moe.yaml"),
+            server_args=["--async-chunk"],
+        ),
+        id="qwen3_omni",
+    ),
+]
+
+
+def _chat_completions_request_without_expectations(omni_server: OmniServer, case_id: str) -> dict[str, Any]:
+    """Build ``send_chat_completions_http_request`` kwargs excluding ``err_code`` / ``err_message``."""
+    if case_id == "malformed_json":
+        return {"raw_body": "{not-json", "timeout": 120}
+    if case_id == "missing_messages":
+        return {"json": {"model": omni_server.model, "stream": False}, "timeout": 120}
+    if case_id == "model_mismatch":
+        return {
+            "json": {
+                "model": "this-model-is-not-served-on-this-instance",
+                "messages": [{"role": "user", "content": "hello"}],
+                "stream": False,
+            },
+            "timeout": 120,
+        }
+    body = _minimal_chat_json(omni_server)
+    if case_id == "stream_wrong_type":
+        body["stream"] = "not-a-boolean"  # type: ignore[assignment]
+    elif case_id == "temperature_wrong_type":
+        body["temperature"] = "hot"  # type: ignore[assignment]
+    elif case_id == "max_tokens_wrong_type":
+        body["max_tokens"] = "not-an-int"  # type: ignore[assignment] # type: ignore[assignment]
+    elif case_id == "modalities_list_bad":
+        body["modalities"] = [123]  # type: ignore[assignment]
+    elif case_id == "response_format_json_schema_incomplete":
+        body["response_format"] = {"type": "json_schema"}
+    elif case_id == "logprobs_wrong_type":
+        body["logprobs"] = "yes"  # type: ignore[assignment]
+    elif case_id == "logprobs_top_without_enabled":
+        body["logprobs"] = False
+        body["top_logprobs"] = 5
+    elif case_id == "speaker_unknown":
+        body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
+    else:
+        raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
+    return {"json": body, "timeout": 120}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/chat/completions
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize(
+    "case_id, err_code, err_message",
+    [
+        pytest.param(
+            "malformed_json",
+            400,
+            ("json_invalid", "decode error"),
+            id="malformed_json",
+        ),
+        pytest.param("missing_messages", 400, ("messages", "missing", "Field required"), id="missing_messages"),
+        pytest.param(
+            "model_mismatch",
+            404,
+            ("model", "does not exist", "NotFoundError"),
+            id="model_mismatch",
+        ),
+        pytest.param("stream_wrong_type", 400, ("stream", "bool_parsing", "not-a-boolean"), id="invalid_stream_type"),
+        pytest.param(
+            "temperature_wrong_type",
+            400,
+            ("temperature", "float_parsing", "a valid number"),
+            id="invalid_temperature_type",
+        ),
+        pytest.param(
+            "max_tokens_wrong_type", 400, ("max_tokens", "int_parsing", "integer"), id="invalid_max_tokens_type"
+        ),
+        pytest.param(
+            "modalities_list_bad",
+            400,
+            ("modalities", "value_error", ""),
+            id="modalities_list_bad_element",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "response_format_json_schema_incomplete",
+            400,
+            ("response_format", "value_error", "json_schema"),
+            id="invalid_response_format_json_schema",
+        ),
+        pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            "logprobs_top_without_enabled",
+            400,
+            ("top_logprobs", "logprobs", "True"),
+            id="logprobs_top_without_logprobs_true",
+        ),
+        pytest.param(
+            "speaker_unknown",
+            400,
+            ("Invalid speaker", "Supported"),
+            id="speaker_unknown_preset",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
+def test_chat_completions_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    case_id: str,
+    err_code: int | tuple[int, ...],
+    err_message: str | tuple[str, ...],
+) -> None:
+    cfg = _chat_completions_request_without_expectations(omni_server, case_id)
+    cfg["err_code"] = err_code
+    cfg["err_message"] = err_message
+    openai_client.send_chat_completions_http_request(cfg)[0]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WS /v1/video/chat/stream
+# ─────────────────────────────────────────────────────────────────────────────
+
+# First-frame JSON with ``omni_server.model`` for ``test_video_chat_stream_invalid_requests``.
+_VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR = object()
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize(
+    "send_frames_spec, err_message",
+    [
+        pytest.param("{not-json", "Invalid JSON", id="invalid_json"),
+        pytest.param(
+            _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR,
+            "Invalid session config",
+            id="session_modalities_scalar",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
+def test_video_chat_stream_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    send_frames_spec: Any,
+    err_message: str,
+) -> None:
+    if send_frames_spec is _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR:
+        send_frames = json.dumps({"type": "session.config", "model": omni_server.model, "modalities": "text"})
+    else:
+        send_frames = send_frames_spec
+    assert isinstance(send_frames, str)
+    openai_client.send_video_chat_stream_ws_request(
+        {
+            "send_frames": send_frames,
+            "timeout": 120,
+            "ws_max_size": None,
+            "ws_json_type": "error",
+            "err_message": err_message,
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WS /v1/realtime
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
+def test_realtime_rejected_when_async_chunk_enabled(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+) -> None:
+    openai_client.send_realtime_ws_request(
+        {
+            "timeout": 120,
+            "ws_max_size": None,
+            "ws_json_type": "error",
+            "ws_error_code": "unsupported",
+            "err_message": "async_chunk",
+        }
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Omni sleep/wakeup JSON validation (live ``Qwen/Qwen-Image`` process)."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+_QWEN_IMAGE = [
+    pytest.param(
+        OmniServerParams(model="Qwen/Qwen-Image"),
+        id="qwen_image",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST omni sleep / omni wakeup (invalid JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "endpoint, json_body, err_message",
+    [
+        pytest.param("sleep", {"level": 2}, ("stage_ids", "Field required", "Missing"), id="sleep_missing_stage_ids"),
+        pytest.param(
+            "sleep",
+            {"stage_ids": [], "level": 2},
+            ("stage_ids", "Field required", "Missing"),
+            id="sleep_empty_stage_ids",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "sleep",
+            {"stage_ids": "not-a-list", "level": 2},
+            ("stage_ids", "list_type", "a valid list"),
+            id="sleep_invalid_stage_ids_type",
+        ),
+        pytest.param(
+            "sleep",
+            {"stage_ids": [0], "level": "high"},
+            ("level", "int_parsing", "a valid integer"),
+            id="sleep_invalid_level_type",
+        ),
+        pytest.param(
+            "sleep",
+            {"stage_ids": [0], "level": -1},
+            ("level", "greater"),
+            id="sleep_negative_level",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param("wakeup", {}, ("stage_ids", "missing", "Field required"), id="wakeup_missing_stage_ids"),
+        pytest.param(
+            "wakeup",
+            {"stage_ids": None},
+            ("stage_ids", "list_type", "a valid list"),
+            id="wakeup_invalid_stage_ids_type",
+        ),
+        pytest.param(
+            "wakeup",
+            {"stage_ids": []},
+            ("stage_ids", "length"),
+            id="wakeup_empty_stage_ids",
+            marks=_SKIP_ISSUE_3649,
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
+def test_omni_sleep_wakeup_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    endpoint: Literal["sleep", "wakeup"],
+    json_body: dict[str, Any],
+    err_message: str | tuple[str, ...],
+) -> None:
+    cfg: dict[str, Any] = {
+        "json": json_body,
+        "timeout": 120,
+        "err_code": 400,
+        "err_message": err_message,
+    }
+    if endpoint == "sleep":
+        openai_client.send_omni_sleep_http_request(cfg)
+    else:
+        openai_client.send_omni_wakeup_http_request(cfg)

tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HTTP video job routes (live ``Wan-AI/Wan2.2-T2V-A14B-Diffusers``).
+
+WebSocket ``/v1/video/chat/stream`` and ``/v1/realtime`` invalid-input checks live in
+``test_invalid_omni_chat.py`` with the same Qwen3-Omni server module fixture (Wan cannot serve those paths).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+_WAN_T2V = [
+    pytest.param(
+        OmniServerParams(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers"),
+        id="wan22_t2v",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+
+def _minimal_video_form_data() -> dict[str, str]:
+    return {"prompt": "a cat walking on grass"}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/videos (async)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "form_extra, err_message",
+    [
+        pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"),
+        pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"),
+        pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"),
+        pytest.param(
+            {"num_inference_steps": "0"},
+            ("num_inference_steps", "greater_than_equal", "1"),
+            id="num_inference_steps_zero",
+        ),
+        pytest.param(
+            {"guidance_scale": "25"}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max"
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_post_videos_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    form_extra: dict[str, str] | None,
+    err_message: str | tuple[str, ...],
+) -> None:
+    data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
+    openai_client.send_videos_create_http_request(
+        {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/videos/sync
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "form_extra, err_message",
+    [
+        pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"),
+        pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"),
+        pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"),
+        pytest.param(
+            {"num_inference_steps": "201"},
+            ("num_inference_steps", "less_than_equal", "200"),
+            id="num_inference_steps_above_max",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_post_videos_sync_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    form_extra: dict[str, str] | None,
+    err_message: str | tuple[str, ...],
+) -> None:
+    data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
+    openai_client.send_videos_sync_http_request(
+        {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GET /v1/videos (list)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "params, err_message",
+    [
+        pytest.param({"limit": -1}, ("limit", "greater_than_equal", "0"), id="invalid_limit_negative"),
+        pytest.param({"limit": 101}, ("limit", "less_than_equal", "100"), id="limit_above_max"),
+        pytest.param({"order": "newest"}, ("order", "literal_error", "asc", "desc"), id="invalid_order"),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_get_videos_list_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    params: dict[str, Any],
+    err_message: str | tuple[str, ...],
+) -> None:
+    openai_client.send_videos_list_http_request(
+        {"params": params, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GET /v1/videos/{video_id}
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_video_retrieve_http_request(
+        {
+            "video_id": "does-not-exist",
+            "timeout": 120,
+            "err_code": 404,
+            "err_message": ("not found", "video"),
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# DELETE /v1/videos/{video_id}
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_video_delete_http_request(
+        {
+            "video_id": "does-not-exist",
+            "timeout": 120,
+            "err_code": 404,
+            "err_message": ("not found", "video"),
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GET /v1/videos/{video_id}/content
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_get_video_content_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_video_content_http_request(
+        {
+            "video_id": "does-not-exist",
+            "timeout": 120,
+            "err_code": 404,
+            "err_message": ("not found", "video"),
+        }
+    )

tests/diffusion/attention/test_attention_config.py

@@ -420,3 +420,102 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         assert attn.use_ring is True
         assert attn.ring_runner is not None
         assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA"
+
+
+class TestDiffusionKvCacheQuantization:
+    @staticmethod
+    def _install_attention_init_stubs(monkeypatch):
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            def forward(self, query, key, value, attn_metadata=None):
+                return query
+
+        class _FakeBackend:
+            @staticmethod
+            def get_name() -> str:
+                return "FAKE_BACKEND"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+            @staticmethod
+            def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
+                return True
+
+        class _FakeRingParallelAttention:
+            def __init__(self, sp_group, attn_backend_pref=None):
+                self.sp_group = sp_group
+                self.attn_backend_pref = attn_backend_pref
+
+        monkeypatch.setattr(
+            layer_mod,
+            "get_attn_backend_for_role",
+            lambda role, head_size, attention_config=None, role_category=None: (_FakeBackend, None),
+        )
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "get_sp_group", lambda: SimpleNamespace(ring_group="ring-group"))
+        monkeypatch.setattr(layer_mod, "RingParallelAttention", _FakeRingParallelAttention)
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+    def test_diffusion_kv_cache_dtype_none_does_not_trigger_ring_quantization_error(self, monkeypatch):
+        self._install_attention_init_stubs(monkeypatch)
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attn = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+            )
+
+        assert attn._kv_cache_dtype is None
+
+    def test_diffusion_kv_cache_dtype_auto_does_not_trigger_ring_quantization_error(self, monkeypatch):
+        self._install_attention_init_stubs(monkeypatch)
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype="auto",
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attn = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+            )
+
+        assert attn._kv_cache_dtype is None
+
+    def test_diffusion_kv_cache_dtype_fp8_raises_with_ring_attention(self, monkeypatch):
+        self._install_attention_init_stubs(monkeypatch)
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype="fp8",
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            with pytest.raises(ValueError, match="KV quantization is not compatible with ring attention"):
+                Attention(
+                    num_heads=4,
+                    head_size=64,
+                    causal=False,
+                    softmax_scale=1.0,
+                )

tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (
+    DistributedAutoencoderKLHunyuan,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan):
+    def __init__(self):
+        torch.nn.Module.__init__(self)
+        self.tile_latent_min_size = 8
+        self.tile_sample_min_size = 8
+        self.tile_overlap_factor = 0.25
+        self.use_spatial_tiling = False
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return torch.float32
+
+    def decoder(self, x: torch.Tensor) -> torch.Tensor:
+        return x + 10
+
+    def encoder(self, x: torch.Tensor) -> torch.Tensor:
+        return x + 20
+
+    def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
+        return b
+
+    def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
+        return b
+
+
+def test_hunyuan_vae_use_tiling_aliases_spatial_tiling():
+    # Verify use_tiling property maps to use_spatial_tiling.
+    vae = _DummyDistributedAutoencoderKLHunyuan()
+
+    assert not vae.use_tiling
+
+    vae.use_tiling = True
+
+    assert vae.use_spatial_tiling
+
+
+def test_hunyuan_vae_decode_tiles_round_trip():
+    # Validate decode tile split/exec/merge returns expected reconstructed tensor.
+    vae = _DummyDistributedAutoencoderKLHunyuan()
+    z = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
+
+    tile_tasks, grid_spec = vae.tile_split(z)
+    decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks}
+    output = vae.tile_merge(decoded_tiles, grid_spec)
+
+    assert grid_spec.split_dims == (3, 4)
+    assert grid_spec.grid_shape == (2, 2)
+    assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
+    assert len(tile_tasks) == 4
+    assert torch.equal(output, z + 10)
+
+
+def test_hunyuan_vae_encode_tiles_round_trip():
+    # Validate encode tile split/exec/merge returns expected latent tensor.
+    vae = _DummyDistributedAutoencoderKLHunyuan()
+    x = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
+
+    tile_tasks, grid_spec = vae.encode_tile_split(x)
+    encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks}
+    output = vae.encode_tile_merge(encoded_tiles, grid_spec)
+
+    assert grid_spec.split_dims == (3, 4)
+    assert grid_spec.grid_shape == (2, 2)
+    assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
+    assert len(tile_tasks) == 4
+    assert torch.equal(output, x + 20)

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -0,0 +1,802 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import socket
+
+import pytest
+import torch
+from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers
+from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
+
+import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    get_classifier_free_guidance_rank,
+    get_pp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.parallel]
+
+
+def _find_free_port() -> str:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return str(s.getsockname()[1])
+
+
+def update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for k, v in envs_dict.items():
+        os.environ[k] = v
+
+
+# ---------------------------------------------------------------------------
+# Shared stubs used by both unit and distributed tests
+# ---------------------------------------------------------------------------
+
+
+class FakeWork:
+    """Drop-in for torch.distributed.Work that records whether wait() was called."""
+
+    def __init__(self):
+        self.waited = False
+
+    def wait(self):
+        self.waited = True
+
+
+class SimpleScheduler:
+    """Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred."""
+
+    def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False):
+        return (latents - 0.1 * noise_pred,)
+
+
+class FakeVAE:
+    def __init__(self, distributed_enabled: bool = False):
+        self.calls = 0
+        self.distributed_enabled = distributed_enabled
+
+    def decode(self, z: torch.Tensor):
+        """Original decode docstring."""
+        self.calls += 1
+        return (z + 1,)
+
+    def is_distributed_enabled(self) -> bool:
+        return self.distributed_enabled
+
+
+class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
+    """Minimal pipeline used to exercise PipelineParallelMixin.
+
+    Uses vLLM's ``make_layers`` for layer partitioning — the same utility used
+    by real DiT models — so the PP layer-split logic is exercised faithfully.
+
+    Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i``
+    is initialized identically on every rank regardless of which ranks are active,
+    allowing the distributed output to be compared against the single-GPU baseline.
+
+    Args:
+        num_layers: Total number of Linear layers.
+        dim:        Input / hidden dimension.
+        seed:       Base RNG seed; layer ``i`` uses ``seed + i``.
+        device:     Target device for layer weights (default: CPU).
+        dtype:      Target dtype for layer weights (default: float32).
+    """
+
+    def __init__(
+        self,
+        num_layers: int = 4,
+        dim: int = 64,
+        seed: int = 42,
+        device: torch.device | None = None,
+        dtype: torch.dtype = torch.float32,
+    ):
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            num_layers,
+            lambda prefix: torch.nn.Linear(dim, dim, bias=False),
+            prefix="layers",
+        )
+
+        for i, layer in enumerate(self.layers):
+            if not isinstance(layer, PPMissingLayer):
+                torch.manual_seed(seed + i)
+                torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02)
+
+        self.layers.to(device=device, dtype=dtype)
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim)
+        self.scheduler = SimpleScheduler()
+
+    def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors:
+        """Layer-split forward pass.
+
+        * First PP rank: uses ``x`` from caller kwargs.
+        * Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]``
+          (which transparently waits for the async receive).
+        * Non-last PP ranks return ``IntermediateTensors``; the last rank
+          returns the plain noise-prediction tensor.
+        """
+        if intermediate_tensors is not None:
+            x = intermediate_tensors["hidden_states"]
+
+        for i in range(self.start_layer, self.end_layer):
+            x = self.layers[i](x)
+
+        pp_group = get_pp_group()
+        if not pp_group.is_last_rank:
+            return IntermediateTensors({"hidden_states": x})
+        return x
+
+
+# ---------------------------------------------------------------------------
+# 1.  AsyncLatents – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestAsyncLatents:
+    """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
+        return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
+
+    def test_resolve_returns_wrapped_tensor(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al._resolve() is t
+
+    def test_attribute_access_resolves(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al.shape == t.shape
+        assert al.dtype == t.dtype
+
+    def test_torch_function_protocol(self):
+        """torch ops that receive an AsyncLatents should see the underlying tensor."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        mask = torch.ones_like(t)
+        result = mask * al  # triggers __torch_function__
+        torch.testing.assert_close(result, mask * t)
+
+    def test_torch_function_with_list_arg(self):
+        """__torch_function__ must unwrap AsyncLatents inside list/tuple args."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.cat([al, al], dim=0)
+        torch.testing.assert_close(result, torch.cat([t, t], dim=0))
+
+    def test_torch_tensor_conversion(self):
+        """torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy)."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.as_tensor(al)
+        assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage"
+
+    def test_handles_are_waited_before_resolve(self):
+        t = torch.randn(2, 4)
+        h1, h2 = FakeWork(), FakeWork()
+        al = self._make(t, handles=[h1, h2])
+        _ = al.shape  # trigger resolution
+        assert h1.waited and h2.waited, "Not all handles were waited on"
+
+    def test_postproc_callbacks_invoked_on_resolve(self):
+        t = torch.randn(2, 4)
+        log: list[int] = []
+        al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)])
+        _ = al.shape
+        assert log == [1, 2], f"postproc not called in order: {log}"
+
+    def test_idempotent_resolve(self):
+        """handle.wait() must not be called twice if _resolve() is called twice."""
+        t = torch.randn(2, 4)
+        h = FakeWork()
+        al = self._make(t, handles=[h])
+        _ = al.shape  # first resolve
+        h.waited = False  # reset sentinel
+        _ = al.dtype  # second resolve
+        assert not h.waited, "handle.wait() was called a second time"
+
+
+# ---------------------------------------------------------------------------
+# 2.  _sync_pp_send / diffuse wrapper – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestSyncPPSend:
+    """Verifies PipelineParallelMixin's internal PP-send flush."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline() -> PipelineParallelMixin:
+        # Instantiate a bare mixin — no layers, no distributed env needed.
+        # _sync_pp_send only touches _pp_send_work, so this is sufficient.
+        class _BarePP(PipelineParallelMixin, CFGParallelMixin):
+            pass
+
+        return _BarePP()
+
+    def test_noop_when_work_list_empty(self):
+        pipeline = self._make_pipeline()
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+    def test_waits_all_pending_handles(self):
+        pipeline = self._make_pipeline()
+        works = [FakeWork(), FakeWork(), FakeWork()]
+        pipeline._pp_send_work = works
+        pipeline._sync_pp_send()
+        assert all(w.waited for w in works), "Some handles were not waited on"
+
+    def test_clears_work_list_after_sync(self):
+        pipeline = self._make_pipeline()
+        pipeline._pp_send_work = [FakeWork()]
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+
+class TestDiffuseWrapper:
+    """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def test_diffuse_flushes_pending_sends_on_success(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                return "done"
+
+        pipeline = _DiffusePP()
+
+        assert pipeline.diffuse() == "done"
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_flushes_pending_sends_on_exception(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                raise RuntimeError("boom")
+
+        pipeline = _DiffusePP()
+
+        with pytest.raises(RuntimeError, match="boom"):
+            pipeline.diffuse()
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_wrapper_preserves_metadata(self):
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                """Original diffuse docstring."""
+                return "done"
+
+        assert _DiffusePP.diffuse.__name__ == "diffuse"
+        assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring."
+
+
+class TestVaeDecodeGuard:
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
+        class _DecodePP(PipelineParallelMixin, CFGParallelMixin):
+            def __init__(self):
+                self.vae = FakeVAE(distributed_enabled=distributed_enabled)
+
+        return _DecodePP()
+
+    @staticmethod
+    def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None:
+        monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size)
+        monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage)
+
+    def test_calls_original_decode_when_pp_disabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=1, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_calls_original_decode_on_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_skips_decode_on_non_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)
+
+        assert pipeline.vae.calls == 0
+        assert output == (None,)
+
+    def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline(distributed_enabled=True)
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_decode_wrapper_preserves_metadata(self):
+        pipeline = self._make_pipeline()
+
+        assert pipeline.vae.decode.__name__ == "decode"
+        assert pipeline.vae.decode.__doc__ == "Original decode docstring."
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_cfg_mixin():
+    with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
+
+        class _MissingCFG(PipelineParallelMixin):
+            pass
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_mro_before_cfg_mixin():
+    with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
+
+        class _WrongOrder(CFGParallelMixin, PipelineParallelMixin):
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Distributed test helpers
+# ---------------------------------------------------------------------------
+
+
+def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
+    """Initialise the distributed environment for a spawned worker."""
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": master_port,
+        }
+    )
+    init_distributed_environment()
+    return device
+
+
+def make_pipeline_and_inputs(
+    test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
+) -> tuple["MockPipelineParallel", dict, dict | None]:
+    """Create a MockPipelineParallel and seeded inputs from a test_config dict.
+
+    Must be called after ``initialize_model_parallel`` so that ``make_layers``
+    can read the PP group to determine this rank's layer slice.
+
+    Returns ``(pipeline, positive_kwargs, negative_kwargs)``.
+    ``negative_kwargs`` is ``None`` when ``do_true_cfg=False``.
+    """
+    pipeline = MockPipelineParallel(
+        num_layers=test_config["num_layers"],
+        dim=test_config["dim"],
+        seed=test_config["model_seed"],
+        device=device,
+        dtype=dtype,
+    )
+
+    torch.manual_seed(test_config["input_seed"])
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(test_config["input_seed"])
+    pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
+
+    negative_kwargs = None
+    if do_true_cfg:
+        torch.manual_seed(test_config["input_seed"] + 1)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
+        neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
+        negative_kwargs = {"x": neg_x}
+
+    return pipeline, pos_x, negative_kwargs
+
+
+# ---------------------------------------------------------------------------
+# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
+# ---------------------------------------------------------------------------
+
+
+def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=world_size)
+    pp_group = get_pp_group()
+
+    if pp_group.is_first_rank:
+        torch.manual_seed(77)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(77)
+        tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
+        handles = pp_group.isend_tensor_dict({"t": tensor})
+        for h in handles:
+            h.wait()
+        result_queue.put(("sent", tensor.cpu()))
+    else:
+        received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
+        result_queue.put(("received", received["t"].cpu()))
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
+@pytest.mark.parametrize("pp_size", [2])
+def test_isend_irecv_tensor_dict(pp_size: int):
+    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
+
+    results = {label: tensor for label, tensor in [q.get(), q.get()]}
+    torch.testing.assert_close(
+        results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 4.  predict_noise_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+_baseline_cache: dict[tuple, torch.Tensor] = {}
+
+
+def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
+    """Compute expected single-GPU output using the same MockPipelineParallel.
+
+    Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
+    Results are cached so identical configs are only computed once.
+    """
+    key = (
+        test_config["num_layers"],
+        test_config["dim"],
+        test_config["batch_size"],
+        test_config["model_seed"],
+        test_config["input_seed"],
+        test_config["cfg_scale"],
+        dtype,
+        do_true_cfg,
+    )
+    if key in _baseline_cache:
+        return _baseline_cache[key]
+
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    destroy_distributed_env()
+
+    _baseline_cache[key] = noise_pred.cpu()
+    return _baseline_cache[key]
+
+
+def predict_noise_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    dtype: torch.dtype,
+    test_config: dict,
+    result_queue,
+):
+    """Generic predict-noise worker parameterized by PP and CFG topology."""
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+    # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the non-last PP rank's async send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_last_rank:
+        assert noise_pred is not None
+        if cfg_rank == 0:
+            result_queue.put(noise_pred.cpu())
+    else:
+        assert noise_pred is None
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.float32,
+            4,
+            100,
+            1e-5,
+            1e-5,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-float32",
+        ),
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            1,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-seq_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-bfloat16",
+        ),
+        pytest.param(
+            3,
+            1,
+            False,
+            torch.bfloat16,
+            6,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
+            id="pp3-no_cfg-bfloat16",
+        ),
+    ],
+)
+def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
+    """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": num_layers,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    pp_q = manager.Queue()
+
+    world_size = pp_size * cfg_size
+    port = _find_free_port()
+    torch.multiprocessing.spawn(
+        predict_noise_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
+        nprocs=world_size,
+    )
+
+    pp_out = pp_q.get()
+
+    assert baseline_out.shape == pp_out.shape
+    torch.testing.assert_close(
+        pp_out,
+        baseline_out,
+        rtol=rtol,
+        atol=atol,
+        msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})",
+    )
+
+
+# ---------------------------------------------------------------------------
+# 5.  scheduler_step_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+
+def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
+    """Single-GPU reference: predict_noise + scheduler_step."""
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        result = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+
+    destroy_distributed_env()
+    return result.cpu()
+
+
+def scheduler_step_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    test_config: dict,
+    result_queue,
+):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        latents = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+    # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the last PP rank's async latent send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_first_rank and cfg_rank == 0:
+        resolved = latents.contiguous()
+        result_queue.put(resolved.cpu())
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            300,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            600,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-true_cfg",
+        ),
+    ],
+)
+def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": 4,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    world_size = pp_size * cfg_size
+    torch.multiprocessing.spawn(
+        scheduler_step_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
+        nprocs=world_size,
+    )
+
+    result = q.get()
+    torch.testing.assert_close(
+        result,
+        baseline,
+        rtol=0,
+        atol=0,
+        msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
+    )

tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import (
+    Flux2KleinPipeline,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _StopAfterCheckInputsError(Exception):
+    pass
+
+
+def _make_pipeline():
+    pipeline = object.__new__(Flux2KleinPipeline)
+    pipeline.vae_scale_factor = 8
+    pipeline.is_distilled = True
+    pipeline._guidance_scale = 0.0
+    return pipeline
+
+
+def _make_minimal_request(
+    prompt="valid prompt",
+    *,
+    num_inference_steps=None,
+):
+    """Build a minimal OmniDiffusionRequest-like object that forward() reads from."""
+    params = OmniDiffusionSamplingParams(
+        height=512,
+        width=512,
+        num_inference_steps=num_inference_steps,
+        seed=42,
+    )
+    req = MagicMock()
+    req.sampling_params = params
+    req.prompts = [prompt]
+    req.multi_modal_data = {}
+    return req
+
+
+def _capture_num_inference_steps(pipe, steps):
+    """Run pipe.forward() and capture the num_inference_steps passed to check_inputs."""
+    captured = {}
+
+    def fake_check_inputs(**kwargs):
+        captured["num_inference_steps"] = kwargs.get("num_inference_steps")
+        raise _StopAfterCheckInputsError
+
+    original = pipe.check_inputs
+    pipe.check_inputs = fake_check_inputs
+    try:
+        req = _make_minimal_request(num_inference_steps=steps)
+        pipe.forward(req)
+    except _StopAfterCheckInputsError:
+        pass
+    finally:
+        pipe.check_inputs = original
+    return captured.get("num_inference_steps")
+
+
+# --- Forward resolution tests (the actual regression) ---
+
+
+def test_forward_preserves_zero_not_default():
+    """#3703: num_inference_steps=0 must reach validation as 0, not as default 50."""
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, 0)
+    assert captured == 0, f"Expected 0, got {captured}"
+
+
+def test_forward_preserves_negative():
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, -1)
+    assert captured == -1
+
+
+def test_forward_none_uses_default():
+    """None means unset — forward() should substitute the pipeline default."""
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, None)
+    assert captured == 50
+
+
+def test_forward_preserves_positive():
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, 9)
+    assert captured == 9
+
+
+# --- check_inputs validation tests ---
+
+
+@pytest.mark.parametrize("steps", [0, -1])
+def test_check_inputs_rejects_non_positive(steps):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError):
+        pipe.check_inputs(
+            prompt="valid prompt",
+            height=512,
+            width=512,
+            num_inference_steps=steps,
+        )
+
+
+def test_check_inputs_accepts_none():
+    pipe = _make_pipeline()
+    pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=None)
+
+
+@pytest.mark.parametrize("steps", [1, 2])
+def test_check_inputs_accepts_positive(steps):
+    pipe = _make_pipeline()
+    pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=steps)

tests/diffusion/models/flux2/test_flux2_klein_validation.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import (
+    Flux2KleinPipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_pipeline():
+    pipeline = object.__new__(Flux2KleinPipeline)
+    pipeline.vae_scale_factor = 8
+    pipeline.is_distilled = True
+    pipeline._guidance_scale = 0.0
+    return pipeline
+
+
+def _check(pipeline, prompt, prompt_embeds=None):
+    pipeline.check_inputs(
+        prompt=prompt,
+        height=512,
+        width=512,
+        prompt_embeds=prompt_embeds,
+        num_inference_steps=4,
+        guidance_scale=0.0,
+    )
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        "",
+        "   ",
+    ],
+)
+def test_rejects_empty_or_whitespace_prompt(prompt):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError, match="empty or whitespace-only"):
+        _check(pipe, prompt)
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        ["valid prompt", ""],
+        ["valid prompt", "   "],
+        ["   "],
+    ],
+)
+def test_rejects_list_with_empty_or_whitespace_element(prompt):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError, match="empty or whitespace-only"):
+        _check(pipe, prompt)
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        "valid prompt",
+        ["valid prompt", "another valid prompt"],
+    ],
+)
+def test_accepts_valid_prompts(prompt):
+    pipe = _make_pipeline()
+    _check(pipe, prompt)
+
+
+def test_rejects_none_without_embeds():
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError):
+        _check(pipe, None)
+
+
+def test_accepts_none_with_prompt_embeds():
+    pipe = _make_pipeline()
+    _check(pipe, None, prompt_embeds=torch.randn(1, 256, 4096))

tests/diffusion/models/wan2_2/test_create_transformer_quant.py

@@ -0,0 +1,383 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for transformer quant-config auto-detection.
+
+The loader path at pipeline_wan2_2.py carries two quantization contracts:
+
+  create_transformer_from_config (~L137)
+    - Reads quantization_config from config.json (injected by the merge scripts)
+    - Auto-detects the quant method when no CLI quant_config is provided
+    - Rejects method mismatches (CLI vs disk)
+    - Upgrades online → offline when disk marks is_checkpoint_*_serialized=True
+    - Rebuilds when the active ignored_layers differs from the disk value
+
+  Wan22Pipeline._create_transformer (~L456)
+    - Passes od_config.quantization_config (set by CLI or externally) to
+      create_transformer_from_config for each transformer independently.
+    - When od_config.quantization_config is None, each transformer auto-detects
+      from its own config.json; od_config is NOT modified by this call.
+
+All tests are pure-CPU and do not load model weights.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+# Minimum config that create_transformer_from_config accepts without raising.
+_MIN_CFG: dict = {"patch_size": [1, 2, 2]}
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_fake_transformer():
+    """Return (FakeTransformer class, captured list).
+
+    Each _FakeTransformer.__init__ call appends its **kwargs to captured,
+    letting tests inspect what quant_config was passed per transformer.
+    """
+    captured: list[dict] = []
+
+    class _FakeTransformer:
+        def __init__(self, **kwargs):
+            captured.append(kwargs)
+
+    return _FakeTransformer, captured
+
+
+class _FakePipeline:
+    """Minimal stand-in exposing only what _create_transformer needs from self."""
+
+    def __init__(self, od_config: SimpleNamespace) -> None:
+        self.od_config = od_config
+
+    # Bind the real unbound method so the tests exercise production code.
+    _create_transformer = wan22_module.Wan22Pipeline._create_transformer
+
+
+# ---------------------------------------------------------------------------
+# create_transformer_from_config — auto-detection
+# ---------------------------------------------------------------------------
+
+
+def test_create_transformer_detects_mxfp8_serialized_from_config_json(monkeypatch):
+    """When config.json carries MXFP8 quant and no CLI quant_config is provided,
+    the transformer must receive a DiffusionMXFP8Config with serialized=True."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp8",
+            "is_checkpoint_mxfp8_serialized": True,
+        },
+    }
+    create_transformer_from_config(config)
+
+    qc = captured[0].get("quant_config")
+    assert isinstance(qc, DiffusionMXFP8Config)
+    assert qc.is_checkpoint_mxfp8_serialized is True
+
+
+def test_create_transformer_detects_mxfp4_dualscale_from_config_json(monkeypatch):
+    """config.json with mxfp4_dualscale + ignored_layers must produce
+    a DiffusionMXFP4DualScaleMixedConfig with the correct ignored_layers and
+    serialized flag."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp4_dualscale",
+            "ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"],
+            "is_checkpoint_serialized": True,
+        },
+    }
+    create_transformer_from_config(config)
+
+    qc = captured[0].get("quant_config")
+    assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig)
+    assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.0.attn1.to_k"}
+    assert qc.is_checkpoint_serialized is True
+
+
+def test_create_transformer_without_quantization_config_passes_no_quant(monkeypatch):
+    """A plain BF16 config.json (no quantization_config key) must result in no
+    quant_config being passed to WanTransformer3DModel."""
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    create_transformer_from_config(_MIN_CFG)
+
+    assert "quant_config" not in captured[0]
+
+
+# ---------------------------------------------------------------------------
+# create_transformer_from_config — method-mismatch guard
+# ---------------------------------------------------------------------------
+
+
+def test_create_transformer_rejects_method_mismatch(monkeypatch):
+    """Passing a CLI quant_config whose get_name() differs from the config.json
+    quant_method must raise ValueError immediately (prevents silent weight corruption).
+
+    fp8 (vLLM built-in, get_name()=='fp8') vs disk 'mxfp8' triggers the guard.
+    These are distinct methods; using the same method for both would not trigger it.
+    """
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    FakeTransformer, _ = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    fp8_cli = Fp8Config(is_checkpoint_fp8_serialized=True, activation_scheme="dynamic")
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp8",
+            "is_checkpoint_mxfp8_serialized": True,
+        },
+    }
+    with pytest.raises(ValueError, match="quant_method"):
+        create_transformer_from_config(config, quant_config=fp8_cli)
+
+
+# ---------------------------------------------------------------------------
+# create_transformer_from_config — online → offline upgrade
+# ---------------------------------------------------------------------------
+
+
+def test_create_transformer_upgrades_to_serialized_when_disk_marks_it(monkeypatch):
+    """CLI passes online (is_checkpoint_mxfp8_serialized=False) but config.json
+    marks is_checkpoint_mxfp8_serialized=True → must switch to offline (serialized)
+    so that pre-quantized FP8 tensors are loaded correctly."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    online_cli = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp8",
+            "is_checkpoint_mxfp8_serialized": True,
+        },
+    }
+    create_transformer_from_config(config, quant_config=online_cli)
+
+    qc = captured[0].get("quant_config")
+    assert isinstance(qc, DiffusionMXFP8Config)
+    assert qc.is_checkpoint_mxfp8_serialized is True
+
+
+# ---------------------------------------------------------------------------
+# create_transformer_from_config — ignored_layers rebuild
+# ---------------------------------------------------------------------------
+
+
+def test_create_transformer_rebuilds_when_ignored_layers_differ(monkeypatch):
+    """When the active quant_config has different ignored_layers than config.json,
+    the config must be rebuilt from disk so per-layer BF16 routing is authoritative
+    for this specific transformer."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    stale = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=True,
+        ignored_layers=["blocks.0.attn1.to_q"],
+    )
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp4_dualscale",
+            "ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"],
+            "is_checkpoint_serialized": True,
+        },
+    }
+    create_transformer_from_config(config, quant_config=stale)
+
+    qc = captured[0].get("quant_config")
+    assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig)
+    assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.1.attn1.to_q"}
+
+
+def test_create_transformer_does_not_rebuild_when_ignored_layers_match(monkeypatch):
+    """When the active quant_config already has the same ignored_layers,
+    the same instance must be passed through unchanged (no unnecessary rebuild)."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    matching = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=True,
+        ignored_layers=["blocks.0.attn1.to_q"],
+    )
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp4_dualscale",
+            "ignored_layers": ["blocks.0.attn1.to_q"],
+            "is_checkpoint_serialized": True,
+        },
+    }
+    create_transformer_from_config(config, quant_config=matching)
+
+    assert captured[0].get("quant_config") is matching
+
+
+# ---------------------------------------------------------------------------
+# Wan22Pipeline._create_transformer — od_config passthrough
+# ---------------------------------------------------------------------------
+
+
+def test_pipeline_create_transformer_auto_detects_from_config_json(monkeypatch):
+    """When od_config.quantization_config is None, _create_transformer must
+    auto-detect the quant method from config.json and pass it to the transformer.
+    od_config itself must remain unchanged (None)."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    od_config = SimpleNamespace(quantization_config=None)
+    pipeline = _FakePipeline(od_config)
+
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp8",
+            "is_checkpoint_mxfp8_serialized": True,
+        },
+    }
+    pipeline._create_transformer(config)
+
+    qc = captured[0].get("quant_config")
+    assert isinstance(qc, DiffusionMXFP8Config)
+    assert qc.is_checkpoint_mxfp8_serialized is True
+    # od_config must NOT be modified — each transformer auto-detects independently.
+    assert od_config.quantization_config is None
+
+
+def test_pipeline_create_transformer_does_not_overwrite_existing_od_config(monkeypatch):
+    """If od_config.quantization_config is already set (e.g. via CLI), _create_transformer
+    must pass it through to the transformer unchanged and leave od_config unmodified."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    FakeTransformer, _ = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    existing = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
+    od_config = SimpleNamespace(quantization_config=existing)
+    pipeline = _FakePipeline(od_config)
+
+    config = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp8",
+            "is_checkpoint_mxfp8_serialized": True,
+        },
+    }
+    pipeline._create_transformer(config)
+
+    assert od_config.quantization_config is existing
+
+
+# ---------------------------------------------------------------------------
+# Wan22Pipeline._create_transformer — cascade contracts
+# ---------------------------------------------------------------------------
+
+
+def test_pipeline_cascade_both_transformers_get_mxfp8_serialized_config(monkeypatch):
+    """Cascade model (transformer + transformer_2) with MXFP8 checkpoint:
+    - Each transformer auto-detects from its own config.json independently.
+    Both must receive is_checkpoint_mxfp8_serialized=True."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    od_config = SimpleNamespace(quantization_config=None)
+    pipeline = _FakePipeline(od_config)
+
+    mxfp8_qc = {"quant_method": "mxfp8", "is_checkpoint_mxfp8_serialized": True}
+    pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc})
+    pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc})
+
+    assert len(captured) == 2
+    for i, kwargs in enumerate(captured):
+        qc = kwargs.get("quant_config")
+        assert isinstance(qc, DiffusionMXFP8Config), f"transformer[{i}]: expected DiffusionMXFP8Config, got {type(qc)}"
+        assert qc.is_checkpoint_mxfp8_serialized is True, f"transformer[{i}]: expected serialized=True"
+
+
+def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_layers(monkeypatch):
+    """Cascade with mxfp4_dualscale where transformer and transformer_2 have
+    different ignored_layers in their config.json.
+
+    Expected outcome:
+      transformer   → ignored_layers=["blocks.0.attn1.to_q"]  (auto-detected from its config.json)
+      transformer_2 → ignored_layers=["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"]
+                      (auto-detected from its own config.json independently)
+      od_config     → quantization_config remains None (not modified by _create_transformer)
+    """
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    od_config = SimpleNamespace(quantization_config=None)
+    pipeline = _FakePipeline(od_config)
+
+    cfg1 = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp4_dualscale",
+            "ignored_layers": ["blocks.0.attn1.to_q"],
+            "is_checkpoint_serialized": True,
+        },
+    }
+    cfg2 = {
+        **_MIN_CFG,
+        "quantization_config": {
+            "quant_method": "mxfp4_dualscale",
+            "ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"],
+            "is_checkpoint_serialized": True,
+        },
+    }
+
+    pipeline._create_transformer(cfg1)
+    pipeline._create_transformer(cfg2)
+
+    assert len(captured) == 2
+    qc1 = captured[0].get("quant_config")
+    qc2 = captured[1].get("quant_config")
+
+    assert isinstance(qc1, DiffusionMXFP4DualScaleMixedConfig)
+    assert isinstance(qc2, DiffusionMXFP4DualScaleMixedConfig)
+    assert set(qc1.ignored_layers) == {"blocks.0.attn1.to_q"}, f"transformer expected 1 layer, got {qc1.ignored_layers}"
+    assert set(qc2.ignored_layers) == {
+        "blocks.0.attn1.to_q",
+        "blocks.1.attn1.to_q",
+    }, f"transformer_2 expected 2 layers, got {qc2.ignored_layers}"
+
+    # od_config must remain unchanged — _create_transformer does not modify it.
+    assert od_config.quantization_config is None

tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py

@@ -0,0 +1,301 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Wan2.2 quant_config propagation through transformer creation.
+
+Tests cover:
+- create_transformer_from_config passes quant_config and prefix
+- create_vace_transformer_from_config passes quant_config and prefix
+- set_tf_model_config propagates quant_config to OmniDiffusionConfig
+- patch_wan_rms_norm safely iterates sys.modules with concurrent modifications
+- I2V transformer_2 quant_config is built from config dict
+"""
+
+import sys
+from types import SimpleNamespace
+
+import pytest
+from pytest_mock import MockerFixture
+
+import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
+import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as wan22_vace_module
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    create_transformer_from_config,
+)
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
+    create_vace_transformer_from_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+# ---------------------------------------------------------------------------
+# create_transformer_from_config: quant_config / prefix forwarding
+# ---------------------------------------------------------------------------
+
+
+class TestCreateTransformerQuant:
+    """Verify quant_config and prefix are forwarded to WanTransformer3DModel."""
+
+    def test_quant_config_passed_through(self, mocker: MockerFixture):
+        captured = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+        fake_qc = mocker.MagicMock()
+        create_transformer_from_config(
+            {"patch_size": [1, 2, 2], "num_layers": 2},
+            quant_config=fake_qc,
+        )
+        assert captured.get("quant_config") is fake_qc
+
+    def test_prefix_passed_through(self, mocker: MockerFixture):
+        captured = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+        create_transformer_from_config(
+            {"patch_size": [1, 2, 2]},
+            prefix="model.transformer.",
+        )
+        assert captured.get("prefix") == "model.transformer."
+
+    def test_quant_config_none_by_default(self, mocker: MockerFixture):
+        captured = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+        create_transformer_from_config({"patch_size": [1, 2, 2]})
+        # When quant_config is None and prefix is "", they are not added
+        assert "quant_config" not in captured or captured["quant_config"] is None
+
+    def test_quant_config_and_prefix_together(self, mocker: MockerFixture):
+        captured = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+        fake_qc = mocker.MagicMock()
+        create_transformer_from_config(
+            {"patch_size": [1, 2, 2], "num_attention_heads": 4},
+            quant_config=fake_qc,
+            prefix="blocks.",
+        )
+        assert captured["quant_config"] is fake_qc
+        assert captured["prefix"] == "blocks."
+
+
+# ---------------------------------------------------------------------------
+# create_vace_transformer_from_config: quant_config / prefix forwarding
+# ---------------------------------------------------------------------------
+
+
+class TestCreateVaceTransformerQuant:
+    """Verify quant_config and prefix are forwarded to WanVACETransformer3DModel."""
+
+    def test_quant_config_passed_through(self, mocker: MockerFixture):
+        captured = {}
+
+        class FakeVACETransformer:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        mocker.patch.object(wan22_vace_module, "WanVACETransformer3DModel", FakeVACETransformer)
+
+        fake_qc = mocker.MagicMock()
+        create_vace_transformer_from_config(
+            {"patch_size": [1, 2, 2], "num_layers": 2},
+            quant_config=fake_qc,
+        )
+        assert captured.get("quant_config") is fake_qc
+
+    def test_prefix_passed_through(self, mocker: MockerFixture):
+        captured = {}
+
+        class FakeVACETransformer:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        mocker.patch.object(wan22_vace_module, "WanVACETransformer3DModel", FakeVACETransformer)
+
+        create_vace_transformer_from_config(
+            {"patch_size": [1, 2, 2]},
+            prefix="vace.",
+        )
+        assert captured.get("prefix") == "vace."
+
+
+# ---------------------------------------------------------------------------
+# set_tf_model_config: propagation of quant_config
+# ---------------------------------------------------------------------------
+
+
+class TestSetTfModelConfig:
+    """Test that set_tf_model_config propagates quant_config correctly."""
+
+    def _make_od_config(self):
+        """Create a minimal OmniDiffusionConfig-like object for testing."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        cfg = object.__new__(OmniDiffusionConfig)
+        cfg.quantization_config = None
+        cfg.tf_model_config = None
+        return cfg
+
+    def test_propagates_quant_config_when_none(self, mocker: MockerFixture):
+        cfg = self._make_od_config()
+        fake_qc = mocker.MagicMock()
+        tf_config = SimpleNamespace(quant_config=fake_qc, quant_method="auto-round")
+
+        cfg.set_tf_model_config(tf_config)
+
+        assert cfg.tf_model_config is tf_config
+        assert cfg.quantization_config is fake_qc
+
+    def test_does_not_overwrite_existing_quantization_config(self, mocker: MockerFixture):
+        cfg = self._make_od_config()
+        existing_qc = mocker.MagicMock()
+        cfg.quantization_config = existing_qc
+        tf_config = SimpleNamespace(quant_config=mocker.MagicMock())
+
+        cfg.set_tf_model_config(tf_config)
+
+        assert cfg.tf_model_config is tf_config
+        assert cfg.quantization_config is existing_qc  # not overwritten
+
+    def test_no_propagation_when_tf_quant_config_is_none(self, mocker: MockerFixture):
+        cfg = self._make_od_config()
+        tf_config = SimpleNamespace(quant_config=None)
+
+        cfg.set_tf_model_config(tf_config)
+
+        assert cfg.tf_model_config is tf_config
+        assert cfg.quantization_config is None
+
+
+# ---------------------------------------------------------------------------
+# patch_wan_rms_norm: sys.modules snapshot safety
+# ---------------------------------------------------------------------------
+
+
+class TestPatchWanRmsNorm:
+    """Test that patch_wan_rms_norm doesn't raise on concurrent module registration."""
+
+    def test_patches_modules_with_wan_rms_norm(self):
+        from vllm_omni.diffusion.layers.norm import RMSNormVAE
+        from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
+
+        # Create a fake module that has WanRMS_norm
+        fake_module = SimpleNamespace(WanRMS_norm=lambda x: x)
+        sys.modules["_test_fake_wan_module"] = fake_module
+
+        try:
+            patch_wan_rms_norm()
+            assert fake_module.WanRMS_norm is RMSNormVAE
+        finally:
+            del sys.modules["_test_fake_wan_module"]
+
+    def test_no_error_when_modules_change_during_iteration(self):
+        """Regression test: list() snapshot prevents RuntimeError."""
+        from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
+
+        # Simulate a module being added during iteration by a side effect
+        original_items = sys.modules.items
+
+        def items_with_side_effect():
+            # This would cause RuntimeError without list() snapshot
+            result = list(original_items())
+            # Add a new module to simulate concurrent modification
+            sys.modules["_test_dynamic_module"] = SimpleNamespace()
+            return result
+
+        try:
+            # The function uses list(sys.modules.items()) so it takes a snapshot
+            # Just verify it doesn't raise
+            patch_wan_rms_norm()
+        finally:
+            sys.modules.pop("_test_dynamic_module", None)
+
+
+# ---------------------------------------------------------------------------
+# I2V transformer_2 quant_config extraction
+# ---------------------------------------------------------------------------
+
+
+class TestI2VTransformer2QuantConfig:
+    """Test the transformer_2 quant_config build logic from pipeline_wan2_2_i2v."""
+
+    def test_transformer_2_quant_config_built_from_dict(self):
+        """When transformer_2 config has quantization_config dict, build_quant_config is called."""
+        from vllm_omni.quantization.factory import build_quant_config
+
+        t2_config = {
+            "patch_size": [1, 2, 2],
+            "num_layers": 2,
+            "quantization_config": {
+                "quant_method": "auto-round",
+                "bits": 4,
+                "group_size": 128,
+                "sym": True,
+                "packing_format": "auto_round:auto_gptq",
+            },
+        }
+
+        # Replicate the logic from pipeline_wan2_2_i2v.py
+        t2_quant = t2_config.get("quantization_config")
+        if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
+            method = t2_quant["quant_method"]
+            kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
+            t2_quant = build_quant_config(method, **kwargs)
+        else:
+            t2_quant = None
+
+        from vllm.model_executor.layers.quantization.inc import INCConfig
+
+        assert isinstance(t2_quant, INCConfig)
+        assert t2_quant.weight_bits == 4
+        assert t2_quant.group_size == 128
+
+    def test_transformer_2_quant_config_none_when_missing(self):
+        """When transformer_2 config has no quantization_config, result is None."""
+        t2_config = {
+            "patch_size": [1, 2, 2],
+            "num_layers": 2,
+        }
+
+        t2_quant = t2_config.get("quantization_config")
+        if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
+            pass  # won't enter
+        else:
+            t2_quant = None
+
+        assert t2_quant is None
+
+    def test_transformer_2_quant_config_none_when_dict_lacks_method(self):
+        """When quantization_config is a dict but missing quant_method, result is None."""
+        t2_config = {
+            "patch_size": [1, 2, 2],
+            "quantization_config": {"bits": 4},  # no quant_method key
+        }
+
+        t2_quant = t2_config.get("quantization_config")
+        if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
+            pass
+        else:
+            t2_quant = None
+
+        assert t2_quant is None

tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py

@@ -1,97 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from PIL import Image
-from torch import nn
-
-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
-    Wan22TI2VPipeline,
-    get_wan22_ti2v_pre_process_func,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def _make_ti2v_pipeline() -> Wan22TI2VPipeline:
-    pipeline = object.__new__(Wan22TI2VPipeline)
-    nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
-    pipeline.vae = StubVAE(z_dim=4)
-    pipeline.vae_scale_factor_temporal = 4
-    pipeline.vae_scale_factor_spatial = 8
-    pipeline.progress_bar = noop_progress_bar
-    return pipeline
-
-
-def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None:
-    preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace())
-    request = SimpleNamespace(
-        prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}],
-        sampling_params=SimpleNamespace(height=None, width=None),
-    )
-
-    result = preprocess(request)
-
-    assert result.sampling_params.height == 672
-    assert result.sampling_params.width == 1344
-    assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672)
-
-
-def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None:
-    pipeline = _make_ti2v_pipeline()
-    latents = torch.zeros(1, 4, 2, 4, 4)
-    calls = []
-
-    def fake_predict_noise_maybe_with_cfg(**kwargs):
-        calls.append(kwargs)
-        return torch.ones_like(latents)
-
-    pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg  # type: ignore[method-assign]
-    pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise  # type: ignore[method-assign]
-
-    result = pipeline.diffuse(
-        latents=latents,
-        timesteps=torch.tensor([7]),
-        prompt_embeds=torch.zeros(1, 2, 3),
-        negative_prompt_embeds=torch.zeros(1, 2, 3),
-        guidance_scale=3.0,
-        dtype=torch.float32,
-        attention_kwargs={"a": "b"},
-        num_latent_frames=2,
-        latent_height=4,
-        latent_width=4,
-    )
-
-    positive = calls[0]["positive_kwargs"]
-    assert calls[0]["do_true_cfg"] is True
-    assert positive["timestep"].shape == (1, 8)
-    torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype))
-    torch.testing.assert_close(positive["hidden_states"], latents)
-    torch.testing.assert_close(result, torch.ones_like(latents))
-
-
-def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None:
-    pipeline = _make_ti2v_pipeline()
-    latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents(
-        image=torch.zeros(1, 3, 16, 16),
-        batch_size=1,
-        num_channels_latents=4,
-        height=16,
-        width=16,
-        num_frames=5,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-        generator=None,
-        latents=torch.zeros(1, 4, 2, 2, 2),
-    )
-
-    torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2))
-    assert latent_condition.shape == (1, 4, 1, 2, 2)
-    assert first_frame_mask[:, :, 0].sum() == 0
-    assert first_frame_mask[:, :, 1].sum() == 4

tests/diffusion/quantization/test_mxfp4_config.py

@@ -0,0 +1,539 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _patch_tp_state(monkeypatch):
+    """Patch TP rank/world_size so ModelWeightParameter can be instantiated on CPU
+    without an initialized distributed group.  Returns TP=1 rank=0 for all tests."""
+    monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP4Config
+# ---------------------------------------------------------------------------
+
+
+def test_mxfp4_config_get_name():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    assert DiffusionMXFP4Config.get_name() == "mxfp4"
+
+
+def test_mxfp4_config_from_config_defaults():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({})
+    assert cfg.is_checkpoint_mxfp4_serialized is False
+    assert cfg.ignored_layers == []
+
+
+def test_mxfp4_config_from_config_serialized():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({"is_checkpoint_mxfp4_serialized": True})
+    assert cfg.is_checkpoint_mxfp4_serialized is True
+
+
+def test_mxfp4_config_from_config_ignored_layers():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({"ignored_layers": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
+    """modules_to_not_convert must be accepted as an alias for ignored_layers."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({"modules_to_not_convert": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+# ---------------------------------------------------------------------------
+# build_quant_config integration
+# ---------------------------------------------------------------------------
+
+
+def test_build_quant_config_mxfp4_string():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = build_quant_config("mxfp4")
+    assert isinstance(cfg, DiffusionMXFP4Config)
+    assert cfg.get_name() == "mxfp4"
+    assert cfg.is_checkpoint_mxfp4_serialized is False
+
+
+def test_build_quant_config_mxfp4_dict():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = build_quant_config({"method": "mxfp4", "is_checkpoint_mxfp4_serialized": True})
+    assert isinstance(cfg, DiffusionMXFP4Config)
+    assert cfg.is_checkpoint_mxfp4_serialized is True
+
+
+def test_build_quant_config_mxfp4_dualscale_string():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = build_quant_config("mxfp4_dualscale")
+    assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
+    assert cfg.is_checkpoint_serialized is False
+    assert cfg.num_bf16_fallback_layers == 5
+    assert cfg.ignored_layers == []
+
+
+def test_build_quant_config_mxfp4_dualscale_dict_offline():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = build_quant_config(
+        {
+            "method": "mxfp4_dualscale",
+            "is_checkpoint_serialized": True,
+            "ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"],
+        }
+    )
+    assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
+    assert cfg.is_checkpoint_serialized is True
+    assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"]
+
+
+def test_build_quant_config_mxfp4_dualscale_dict_online_custom_fallback():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = build_quant_config({"method": "mxfp4_dualscale", "num_bf16_fallback_layers": 10})
+    assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
+    assert cfg.num_bf16_fallback_layers == 10
+
+
+# ---------------------------------------------------------------------------
+# Block-index dispatch (_parse_block_idx)
+# ---------------------------------------------------------------------------
+
+
+def test_parse_block_idx_valid():
+    from vllm_omni.quantization.mxfp4_config import _parse_block_idx
+
+    assert _parse_block_idx("blocks.0.attn1.to_q") == 0
+    assert _parse_block_idx("blocks.5.ffn.net.0.proj") == 5
+    assert _parse_block_idx("blocks.40.norm1.weight") == 40
+
+
+def test_parse_block_idx_non_block_prefixes():
+    """Prefixes that do not start with 'blocks.N.' must return None."""
+    from vllm_omni.quantization.mxfp4_config import _parse_block_idx
+
+    assert _parse_block_idx("condition_embedder.time_embedder.linear_1") is None
+    assert _parse_block_idx("proj_out.weight") is None
+    assert _parse_block_idx("model.layers.0.self_attn.q_proj") is None
+    assert _parse_block_idx("scale_shift_table") is None
+
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_QUANTIZATION_METHODS
+# ---------------------------------------------------------------------------
+
+
+def test_supported_methods_include_mxfp4_variants():
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "mxfp4" in SUPPORTED_QUANTIZATION_METHODS
+    assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS
+    assert "mxfp4_dualscale" in SUPPORTED_QUANTIZATION_METHODS
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP4DualScaleMixedConfig — config roundtrips
+# ---------------------------------------------------------------------------
+
+
+def test_mixed_dualscale_config_get_name():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    assert DiffusionMXFP4DualScaleMixedConfig.get_name() == "mxfp4_dualscale"
+
+
+def test_mixed_dualscale_config_no_args_defaults():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig()
+    assert cfg.is_checkpoint_serialized is False
+    assert cfg.ignored_layers == []
+    assert cfg.num_bf16_fallback_layers == 5
+
+
+def test_mixed_dualscale_config_from_config_offline():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig.from_config(
+        {
+            "quant_method": "mxfp4_dualscale",
+            "is_checkpoint_serialized": True,
+            "ignored_layers": ["blocks.0.attn1.to_q", "proj_out"],
+        }
+    )
+    assert cfg.is_checkpoint_serialized is True
+    assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "proj_out"]
+    assert cfg.num_bf16_fallback_layers == 5  # default
+
+
+def test_mixed_dualscale_config_from_config_online_custom_fallback():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig.from_config({"num_bf16_fallback_layers": 10})
+    assert cfg.is_checkpoint_serialized is False
+    assert cfg.num_bf16_fallback_layers == 10
+
+
+def test_mixed_dualscale_config_from_config_modules_to_not_convert_fallback():
+    """modules_to_not_convert must be accepted as an alias for ignored_layers."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig.from_config(
+        {"is_checkpoint_serialized": True, "modules_to_not_convert": ["proj_out"]}
+    )
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP4DualScaleMixedConfig — get_quant_method dispatch
+# ---------------------------------------------------------------------------
+
+
+def test_mixed_dualscale_offline_ignored_layer_returns_unquantized(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Offline: a prefix in ignored_layers must return UnquantizedLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=True,
+        ignored_layers=["blocks.0.attn1.to_q"],
+    )
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = cfg.get_quant_method(layer, "blocks.0.attn1.to_q")
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+def test_mixed_dualscale_offline_non_ignored_returns_mxfp4(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Offline: a prefix NOT in ignored_layers must return NPUMxfp4DualScaleLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=True,
+        ignored_layers=["blocks.0.attn1.to_q"],
+    )
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = cfg.get_quant_method(layer, "blocks.1.attn1.to_q")
+    assert isinstance(method, NPUMxfp4DualScaleLinearMethod)
+
+
+def test_mixed_dualscale_online_fallback_block_returns_unquantized(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: blocks < num_bf16_fallback_layers must return UnquantizedLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    assert isinstance(cfg.get_quant_method(layer, "blocks.0.attn1.to_q"), UnquantizedLinearMethod)
+    assert isinstance(cfg.get_quant_method(layer, "blocks.4.ffn.net.0.proj"), UnquantizedLinearMethod)
+
+
+def test_mixed_dualscale_online_quantized_block_returns_mxfp4(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: blocks >= num_bf16_fallback_layers must return NPUMxfp4DualScaleOnlineLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import (
+        DiffusionMXFP4DualScaleMixedConfig,
+        NPUMxfp4DualScaleOnlineLinearMethod,
+    )
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    assert isinstance(cfg.get_quant_method(layer, "blocks.5.attn1.to_q"), NPUMxfp4DualScaleOnlineLinearMethod)
+    assert isinstance(cfg.get_quant_method(layer, "blocks.40.ffn.net.0.proj"), NPUMxfp4DualScaleOnlineLinearMethod)
+
+
+def test_mixed_dualscale_online_non_block_prefix_returns_mxfp4(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: layers outside 'blocks.N.*' (condition_embedder etc.) always use MXFP4 online."""
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import (
+        DiffusionMXFP4DualScaleMixedConfig,
+        NPUMxfp4DualScaleOnlineLinearMethod,
+    )
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = cfg.get_quant_method(layer, "condition_embedder.time_embedder.linear_1")
+    assert isinstance(method, NPUMxfp4DualScaleOnlineLinearMethod)
+
+
+def test_mixed_dualscale_online_ignored_layers_override(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: explicit ignored_layers must return UnquantizedLinearMethod regardless of block index.
+
+    A layer that is NOT in the leading-block range (block 10 >= num_bf16_fallback_layers=5)
+    but IS listed in ignored_layers must still fall back to BF16.  This lets power users
+    pin specific interleaved layers to BF16 during online quantization without needing an
+    offline checkpoint.
+    """
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=False,
+        num_bf16_fallback_layers=5,
+        ignored_layers=["blocks.10.attn1.to_q"],
+    )
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    # block 10 is above the leading-block threshold but is in ignored_layers → BF16
+    assert isinstance(cfg.get_quant_method(layer, "blocks.10.attn1.to_q"), UnquantizedLinearMethod)
+
+
+def test_mixed_dualscale_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch):
+    """Non-LinearBase layers (norms, embeddings) must return None → no quantization."""
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig()
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    norm_layer = torch.nn.LayerNorm(64)
+    assert cfg.get_quant_method(norm_layer, "blocks.0.norm1") is None
+
+
+# ---------------------------------------------------------------------------
+# TP=2 create_weights: parameter shapes and input_dim/output_dim
+#
+# Two scenarios mirror real Wan2.2 A14B linear layer types:
+#   Column-parallel (to_q, ffn.net_0): output is sharded (N/TP), input is full (K).
+#   Row-parallel   (to_out, ffn.net_2): input is sharded (K/TP), output is full (N).
+#
+# Tests verify:
+#   1. Registered parameter shapes are correct for each partition configuration.
+#   2. input_dim/output_dim attributes are set so RowParallelLinear.weight_loader
+#      can shard scale tensors correctly (the fix for the TP>1 shape-mismatch bug).
+#   3. Simulated loader slicing: slicing the full checkpoint tensor along the
+#      declared input_dim produces the exact shape stored in the parameter —
+#      proving the dim declaration is consistent with the allocation.
+# ---------------------------------------------------------------------------
+
+# K must be divisible by 32 (fine groups) and 512 (coarse groups).
+_TP2_K, _TP2_N, _TP2 = 1024, 512, 2
+
+
+class _FakeLayer(torch.nn.Module):
+    """Bare nn.Module that accepts register_parameter without a real weight_loader."""
+
+
+def _create_weights(method, *, input_size_per_partition, output_partition_sizes):
+    layer = _FakeLayer()
+    method.create_weights(
+        layer=layer,
+        input_size_per_partition=input_size_per_partition,
+        output_partition_sizes=output_partition_sizes,
+        input_size=_TP2_K,
+        output_size=_TP2_N,
+        params_dtype=torch.bfloat16,
+    )
+    return layer
+
+
+def _shard(tensor, param, rank, tp, dim_attr):
+    """Slice `tensor` along the dimension given by `param.<dim_attr>` for `rank`."""
+    dim = getattr(param, dim_attr)
+    if dim is None:
+        return tensor  # not sharded along this axis
+    shard_size = param.shape[dim]
+    slices = [slice(None)] * tensor.ndim
+    slices[dim] = slice(rank * shard_size, (rank + 1) * shard_size)
+    return tensor[tuple(slices)]
+
+
+# ── DualScale method ─────────────────────────────────────────────────────────
+
+
+def test_dualscale_column_parallel_tp2_shapes():
+    """Column-parallel TP=2: output halved, fine/coarse groups stay full, mul_scale full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
+
+    assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K)
+    assert layer.weight_scale.shape == (_TP2_N // _TP2, _TP2_K // 32)
+    assert layer.weight_dual_scale.shape == (_TP2_N // _TP2, _TP2_K // 512, 1)
+    assert layer.mul_scale.shape == (_TP2_K,)
+
+
+def test_dualscale_row_parallel_tp2_shapes():
+    """Row-parallel TP=2: input halved, fine/coarse groups halved, mul_scale halved."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
+    assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32)
+    assert layer.weight_dual_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 512, 1)
+    assert layer.mul_scale.shape == (_TP2_K // _TP2,)
+
+
+def test_dualscale_scale_parameter_input_dims():
+    """weight_scale/weight_dual_scale must have input_dim=1; mul_scale must have input_dim=0.
+
+    RowParallelLinear.weight_loader only shards a parameter when input_dim is set.
+    Without these, loading a full checkpoint tensor into a per-rank shape causes a
+    shape mismatch for TP>1 row-parallel layers (to_out, ffn.net_2).
+    """
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight_scale.input_dim == 1
+    assert layer.weight_scale.output_dim == 0
+    assert layer.weight_dual_scale.input_dim == 1
+    assert layer.weight_dual_scale.output_dim == 0
+    assert layer.mul_scale.input_dim == 0
+    assert layer.mul_scale.output_dim is None
+
+
+def test_dualscale_row_parallel_tp2_loader_simulation():
+    """Slicing full checkpoint tensors along input_dim must match row-parallel parameter shapes.
+
+    Simulates what RowParallelLinear.weight_loader does: for each scale parameter,
+    take the slice at rank*shard_size:(rank+1)*shard_size along input_dim.
+    The resulting shape must equal the per-rank parameter shape allocated by create_weights.
+    """
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    # Full checkpoint tensors (what the loader reads from disk).
+    ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
+    ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1)
+    ckpt_mul_scale = torch.zeros(_TP2_K)
+
+    for rank in range(_TP2):
+        assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape
+        assert (
+            _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "input_dim").shape
+            == layer.weight_dual_scale.shape
+        )
+        assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "input_dim").shape == layer.mul_scale.shape
+
+
+def test_dualscale_column_parallel_tp2_loader_simulation():
+    """Slicing full checkpoint tensors along output_dim must match column-parallel parameter shapes.
+
+    For column-parallel layers, the loader shards along output_dim (rows).
+    mul_scale has output_dim=None → not sharded (full tensor, same for all ranks).
+    """
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
+
+    ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
+    ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1)
+    ckpt_mul_scale = torch.zeros(_TP2_K)
+
+    for rank in range(_TP2):
+        assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "output_dim").shape == layer.weight_scale.shape
+        assert (
+            _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "output_dim").shape
+            == layer.weight_dual_scale.shape
+        )
+        # mul_scale: output_dim=None → no sharding → full tensor fits the column-parallel parameter
+        assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "output_dim").shape == layer.mul_scale.shape
+
+
+# ── Single-scale method ───────────────────────────────────────────────────────
+
+
+def test_single_scale_row_parallel_tp2_shapes():
+    """Row-parallel TP=2: input halved → weight_scale groups halved."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
+
+    method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
+    assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32)
+
+
+def test_single_scale_scale_parameter_input_dims():
+    """Single-scale weight_scale must have input_dim=1 for RowParallel TP sharding."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
+
+    method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight_scale.input_dim == 1
+    assert layer.weight_scale.output_dim == 0
+
+
+def test_single_scale_row_parallel_tp2_loader_simulation():
+    """Slicing full checkpoint weight_scale along input_dim matches row-parallel parameter shape."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
+
+    method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
+
+    for rank in range(_TP2):
+        assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape

tests/diffusion/quantization/test_mxfp4_key_remap.py

@@ -0,0 +1,399 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for merge_mxfp4_dualscale_checkpoint.py key-remapping helpers.
+
+These are pure-Python unit tests that exercise the transformation functions
+without loading any actual checkpoint files or requiring NPU hardware.
+"""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_MODEL_TYPES
+# ---------------------------------------------------------------------------
+
+
+def test_supported_model_types_includes_a14b():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES
+
+    assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES
+    assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES
+
+
+def test_supported_model_types_excludes_ti2v_5b():
+    """TI2V-5B is explicitly NOT supported under W4A4 quantization."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES
+
+    assert "Wan2.2-TI2V-5B" not in SUPPORTED_MODEL_TYPES
+
+
+# ---------------------------------------------------------------------------
+# _apply_rename_dict
+# ---------------------------------------------------------------------------
+
+
+def test_apply_rename_dict_self_attn_qkvo():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.0.self_attn.q.weight") == "blocks.0.attn1.to_q.weight"
+    assert _apply_rename_dict("blocks.0.self_attn.k.weight") == "blocks.0.attn1.to_k.weight"
+    assert _apply_rename_dict("blocks.0.self_attn.v.weight") == "blocks.0.attn1.to_v.weight"
+    assert _apply_rename_dict("blocks.0.self_attn.o.weight") == "blocks.0.attn1.to_out.0.weight"
+
+
+def test_apply_rename_dict_ffn():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.1.ffn.0.weight") == "blocks.1.ffn.net.0.proj.weight"
+    assert _apply_rename_dict("blocks.1.ffn.2.weight") == "blocks.1.ffn.net.2.weight"
+
+
+def test_apply_rename_dict_norm_order_swap():
+    """norm2↔norm3 swap: quant tool uses norm1/norm3/norm2 order,
+    Diffusers uses norm1/norm2/norm3."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.0.norm2.weight") == "blocks.0.norm3.weight"
+    assert _apply_rename_dict("blocks.0.norm3.weight") == "blocks.0.norm2.weight"
+
+
+def test_apply_rename_dict_cross_attn():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.0.cross_attn.q.weight") == "blocks.0.attn2.to_q.weight"
+    assert _apply_rename_dict("blocks.0.cross_attn.k.weight") == "blocks.0.attn2.to_k.weight"
+    assert _apply_rename_dict("blocks.0.cross_attn.v.weight") == "blocks.0.attn2.to_v.weight"
+    assert _apply_rename_dict("blocks.0.cross_attn.o.weight") == "blocks.0.attn2.to_out.0.weight"
+
+
+def test_apply_rename_dict_head_and_modulation():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("head.head.weight") == "proj_out.weight"
+
+
+# ---------------------------------------------------------------------------
+# _strip_mxfp4_wrapper
+# ---------------------------------------------------------------------------
+
+
+def test_strip_mxfp4_wrapper_linear_weight():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight") == "blocks.0.attn1.to_q.weight"
+
+
+def test_strip_mxfp4_wrapper_linear_weight_scale():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_scale") == "blocks.0.attn1.to_q.weight_scale"
+
+
+def test_strip_mxfp4_wrapper_linear_weight_dual_scale():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert (
+        _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_dual_scale") == "blocks.0.attn1.to_q.weight_dual_scale"
+    )
+
+
+def test_strip_mxfp4_wrapper_linear_bias():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.bias") == "blocks.0.attn1.to_q.bias"
+
+
+def test_strip_mxfp4_wrapper_div_mul_scale():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.div.mul_scale") == "blocks.0.attn1.to_q.mul_scale"
+
+
+def test_strip_mxfp4_wrapper_noop_for_plain_weight():
+    """MXFP8 / FLOAT tensors have no wrapper — must be returned unchanged."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.weight") == "blocks.0.attn1.to_q.weight"
+    assert _strip_mxfp4_wrapper("blocks.0.norm_q.weight") == "blocks.0.norm_q.weight"
+    assert _strip_mxfp4_wrapper("condition_embedder.time_embedder.linear_1.weight") == (
+        "condition_embedder.time_embedder.linear_1.weight"
+    )
+
+
+# ---------------------------------------------------------------------------
+# _classify_blocks
+# ---------------------------------------------------------------------------
+
+
+def test_classify_blocks_mixed_mxfp8_and_mxfp4():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
+
+    quant_meta = {
+        "blocks.0.attn1.to_q.weight": "W8A8_MXFP8",
+        "blocks.1.attn1.to_q.weight": "W8A8_MXFP8",
+        "blocks.2.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
+        "blocks.3.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
+        "condition_embedder.time_embedder.linear_1.weight": "FLOAT",
+    }
+    block_types = _classify_blocks(quant_meta)
+    assert block_types[0] == "mxfp8"
+    assert block_types[1] == "mxfp8"
+    assert block_types[2] == "mxfp4_dualscale"
+    assert block_types[3] == "mxfp4_dualscale"
+    # Non-block key must not produce an entry
+    assert None not in block_types
+
+
+def test_classify_blocks_float_entries_skipped():
+    """FLOAT-typed tensors must not contribute to block classification."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
+
+    quant_meta = {
+        "blocks.0.bias": "FLOAT",
+        "blocks.1.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
+    }
+    block_types = _classify_blocks(quant_meta)
+    assert 0 not in block_types
+    assert block_types[1] == "mxfp4_dualscale"
+
+
+def test_classify_blocks_empty():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
+
+    assert _classify_blocks({}) == {}
+
+
+# ---------------------------------------------------------------------------
+# _is_mxfp4_tensor
+# ---------------------------------------------------------------------------
+
+
+def test_is_mxfp4_tensor_quantized_weight():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight", "W4A4_MXFP4_DUALSCALE") is True
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_scale", "W4A4_MXFP4_DUALSCALE") is True
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_dual_scale", "W4A4_MXFP4_DUALSCALE") is True
+
+
+def test_is_mxfp4_tensor_companion_bias_and_mul_scale():
+    """Companion tensors (bias, mul_scale) are FLOAT but belong to MXFP4 layers
+    and must be included in the merged checkpoint."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    # bias inside .linear. wrapper → must be included
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.bias", "FLOAT") is True
+    # mul_scale inside .div. wrapper → must be included
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.div.mul_scale", "FLOAT") is True
+
+
+def test_is_mxfp4_tensor_bf16_fallback_weight():
+    """BF16 fallback linear layers have plain .weight keys (no wrapper) → must NOT be included."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.weight", "FLOAT") is False
+    assert _is_mxfp4_tensor("condition_embedder.time_embedder.linear_1.weight", "FLOAT") is False
+
+
+def test_is_mxfp4_tensor_norm_weight():
+    """Norm layers are always FLOAT with plain .weight → not MXFP4 related."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    assert _is_mxfp4_tensor("blocks.0.norm1.weight", "FLOAT") is False
+
+
+# ---------------------------------------------------------------------------
+# _collect_ignored_layers
+# ---------------------------------------------------------------------------
+
+
+def test_collect_ignored_layers_basic():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
+
+    # block 0: all three of Q/K/V are BF16 → fused as to_qkv in ignored_layers
+    # block 1: to_qkv is MXFP4 → not in ignored_layers
+    merged = {
+        "blocks.0.attn1.to_q.weight": torch.zeros(4, 4),
+        "blocks.0.attn1.to_k.weight": torch.zeros(4, 4),
+        "blocks.0.attn1.to_v.weight": torch.zeros(4, 4),
+        "blocks.1.attn1.to_q.weight": torch.zeros(4, 4),
+        "blocks.1.attn1.to_q.weight_scale": torch.zeros(4, 1),
+    }
+    mxfp4_prefixes = {"blocks.1.attn1.to_q"}
+
+    ignored = _collect_ignored_layers(merged, mxfp4_prefixes)
+    assert "blocks.0.attn1.to_qkv" in ignored  # fused vllm-omni name
+    assert "blocks.0.attn1.to_q" not in ignored  # diffusers name must not appear
+    assert "blocks.0.attn1.to_k" not in ignored
+    assert "blocks.1.attn1.to_qkv" not in ignored  # MXFP4, not ignored
+
+
+def test_collect_ignored_layers_empty_mxfp4():
+    """When no layers are MXFP4 (all BF16), all .weight prefixes become ignored_layers."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
+
+    # Use non-attn layers to avoid the Q/K/V completeness requirement.
+    merged = {
+        "blocks.0.ffn.net.0.proj.weight": torch.zeros(4, 4),
+        "proj_out.weight": torch.zeros(4, 4),
+    }
+    ignored = _collect_ignored_layers(merged, set())
+    assert sorted(ignored) == ["blocks.0.ffn.net_0.proj", "proj_out"]
+
+
+def test_collect_ignored_layers_returns_sorted():
+    """ignored_layers must be sorted for deterministic config.json output."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
+
+    # Use FFN layers (no Q/K/V fusion concern) to isolate the sort guarantee.
+    merged = {
+        "blocks.2.ffn.net.0.proj.weight": torch.zeros(1),
+        "blocks.0.ffn.net.0.proj.weight": torch.zeros(1),
+        "blocks.1.ffn.net.0.proj.weight": torch.zeros(1),
+    }
+    ignored = _collect_ignored_layers(merged, set())
+    assert ignored == sorted(ignored)
+
+
+# ---------------------------------------------------------------------------
+# _build_quant_config (new mxfp4_dualscale format)
+# ---------------------------------------------------------------------------
+
+
+def test_build_quant_config_new_format():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config
+
+    config = _build_quant_config(["blocks.0.attn1.to_q", "proj_out"])
+    assert config["quant_method"] == "mxfp4_dualscale"
+    assert config["is_checkpoint_serialized"] is True
+    assert config["ignored_layers"] == ["blocks.0.attn1.to_q", "proj_out"]
+    assert "num_mxfp8_blocks" not in config
+
+
+def test_build_quant_config_empty_ignored_layers():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config
+
+    config = _build_quant_config([])
+    assert config["ignored_layers"] == []
+
+
+# ---------------------------------------------------------------------------
+# _diffusers_to_vllm_ignored
+# ---------------------------------------------------------------------------
+
+
+def test_remap_self_attn_qkv_fusion():
+    """All three of attn1.to_q/k/v in BF16 → fused attn1.to_qkv."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_k", "blocks.0.attn1.to_q", "blocks.0.attn1.to_v"])
+    assert result == ["blocks.0.attn1.to_qkv"]
+
+
+def test_remap_self_attn_qkv_partial_raises():
+    """Partial Q/K/V in ignored_layers is invalid — must raise ValueError.
+
+    Self-attention Q/K/V are fused into a single to_qkv layer at runtime;
+    partial precision (some BF16, some MXFP4) cannot be expressed and must
+    be caught early in the merge script.
+    """
+    import pytest
+
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    with pytest.raises(ValueError, match="Partial BF16 fallback"):
+        _diffusers_to_vllm_ignored(["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"])
+
+
+def test_remap_cross_attn_kept_separate():
+    """Cross-attention (attn2) to_q/k/v must NOT be fused."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.attn2.to_k", "blocks.0.attn2.to_q", "blocks.0.attn2.to_v"])
+    assert "blocks.0.attn2.to_qkv" not in result
+    assert "blocks.0.attn2.to_q" in result
+    assert "blocks.0.attn2.to_k" in result
+    assert "blocks.0.attn2.to_v" in result
+
+
+def test_remap_ffn_dot_to_underscore():
+    """ffn.net.0.proj → ffn.net_0.proj; ffn.net.2 → ffn.net_2."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.ffn.net.0.proj", "blocks.0.ffn.net.2"])
+    assert "blocks.0.ffn.net_0.proj" in result
+    assert "blocks.0.ffn.net_2" in result
+    assert "blocks.0.ffn.net.0.proj" not in result
+    assert "blocks.0.ffn.net.2" not in result
+
+
+def test_remap_to_out_strips_index():
+    """attn1.to_out.0 and attn2.to_out.0 must have the .0 suffix removed."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_out.0", "blocks.0.attn2.to_out.0"])
+    assert "blocks.0.attn1.to_out" in result
+    assert "blocks.0.attn2.to_out" in result
+    assert "blocks.0.attn1.to_out.0" not in result
+    assert "blocks.0.attn2.to_out.0" not in result
+
+
+def test_remap_noop_for_non_block_layers():
+    """Layers outside 'blocks.N.*' (condition_embedder, proj_out) pass through unchanged."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    inputs = [
+        "condition_embedder.time_embedder.linear_1",
+        "condition_embedder.text_embedder.linear_1",
+        "proj_out",
+    ]
+    result = _diffusers_to_vllm_ignored(inputs)
+    assert result == sorted(inputs)
+
+
+def test_remap_mixed_block_layers():
+    """Typical mixed block: QKV fused, to_out stripped, FFN renamed, cross-attn separate."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    inputs = [
+        "blocks.0.attn1.to_k",
+        "blocks.0.attn1.to_out.0",
+        "blocks.0.attn1.to_q",
+        "blocks.0.attn1.to_v",
+        "blocks.0.attn2.to_k",
+        "blocks.0.attn2.to_out.0",
+        "blocks.0.attn2.to_q",
+        "blocks.0.attn2.to_v",
+        "blocks.0.ffn.net.0.proj",
+        "blocks.0.ffn.net.2",
+    ]
+    result = _diffusers_to_vllm_ignored(inputs)
+    assert "blocks.0.attn1.to_qkv" in result
+    assert "blocks.0.attn1.to_out" in result
+    assert "blocks.0.attn2.to_q" in result
+    assert "blocks.0.attn2.to_k" in result
+    assert "blocks.0.attn2.to_v" in result
+    assert "blocks.0.attn2.to_out" in result
+    assert "blocks.0.ffn.net_0.proj" in result
+    assert "blocks.0.ffn.net_2" in result
+    # Old names must be gone
+    assert "blocks.0.attn1.to_q" not in result
+    assert "blocks.0.ffn.net.0.proj" not in result
+
+
+def test_remap_returns_sorted():
+    """Output of _diffusers_to_vllm_ignored must always be sorted."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    inputs = [
+        "blocks.2.ffn.net.0.proj",
+        "blocks.0.attn1.to_q",
+        "blocks.0.attn1.to_k",
+        "blocks.0.attn1.to_v",
+        "blocks.1.attn2.to_out.0",
+    ]
+    result = _diffusers_to_vllm_ignored(inputs)
+    assert result == sorted(result)

tests/diffusion/quantization/test_mxfp8_config.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for MXFP8 quantization config and linear method dispatch.
+
+Coverage:
+- DiffusionMXFP8Config.from_config roundtrips (CPU, no NPU required)
+- get_quant_method dispatch (mocked platform)
+- MXFPLinearMethodBase.apply() reshape skeleton (CPU)
+- Weight / scale shape-transform arithmetic from process_weights_after_loading (CPU)
+- build_quant_config integration
+- MXFP8_QUANT_CONFIG structure as the auto-detection contract
+"""
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available")
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP8Config — from_config roundtrips
+# ---------------------------------------------------------------------------
+
+
+def test_mxfp8_config_get_name():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    assert DiffusionMXFP8Config.get_name() == "mxfp8"
+
+
+def test_mxfp8_config_from_config_defaults():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({})
+    assert cfg.is_checkpoint_mxfp8_serialized is False
+    assert cfg.ignored_layers == []
+
+
+def test_mxfp8_config_from_config_serialized():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({"is_checkpoint_mxfp8_serialized": True})
+    assert cfg.is_checkpoint_mxfp8_serialized is True
+
+
+def test_mxfp8_config_from_config_ignored_layers():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({"ignored_layers": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+def test_mxfp8_config_from_config_modules_to_not_convert_fallback():
+    """modules_to_not_convert must be accepted as an alias for ignored_layers."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({"modules_to_not_convert": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+# ---------------------------------------------------------------------------
+# build_quant_config integration
+# ---------------------------------------------------------------------------
+
+
+def test_build_quant_config_mxfp8_string():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = build_quant_config("mxfp8")
+    assert isinstance(cfg, DiffusionMXFP8Config)
+    assert cfg.get_name() == "mxfp8"
+    assert cfg.is_checkpoint_mxfp8_serialized is False
+
+
+def test_build_quant_config_mxfp8_dict():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = build_quant_config({"method": "mxfp8", "is_checkpoint_mxfp8_serialized": True})
+    assert isinstance(cfg, DiffusionMXFP8Config)
+    assert cfg.is_checkpoint_mxfp8_serialized is True
+
+
+def test_build_quant_config_mxfp8_config_json_format():
+    """Verify that the exact quantization_config injected by merge_mxfp8_checkpoint.py
+    is accepted by build_quant_config and selects the offline (serialized) path.
+
+    This is the critical auto-detection contract: TransformerConfig.from_dict()
+    reads quant_method + is_checkpoint_mxfp8_serialized to pick NPUMxfp8LinearMethod.
+    """
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
+
+    cfg = build_quant_config(MXFP8_QUANT_CONFIG)
+    assert isinstance(cfg, DiffusionMXFP8Config)
+    assert cfg.is_checkpoint_mxfp8_serialized is True
+
+
+def test_mxfp8_quant_config_structure():
+    """MXFP8_QUANT_CONFIG must contain exactly the keys that auto-detection reads."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
+
+    assert MXFP8_QUANT_CONFIG.get("quant_method") == "mxfp8"
+    assert MXFP8_QUANT_CONFIG.get("is_checkpoint_mxfp8_serialized") is True
+
+
+# ---------------------------------------------------------------------------
+# get_quant_method dispatch
+# ---------------------------------------------------------------------------
+
+
+def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Offline (serialized) path on NPU must return NPUMxfp8LinearMethod."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod
+
+    config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
+    assert isinstance(method, NPUMxfp8LinearMethod)
+
+
+def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod
+
+    config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
+    assert isinstance(method, NPUMxfp8OnlineLinearMethod)
+
+
+def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Non-NPU platform must raise NotImplementedError (CUDA not yet supported)."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config()
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+
+    with pytest.raises(NotImplementedError):
+        config.get_quant_method(layer, "blocks.0.attn1.to_q")
+
+
+def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """A prefix in ignored_layers must bypass quantization → UnquantizedLinearMethod."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config(ignored_layers=["proj_out"])
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = config.get_quant_method(layer, "proj_out")
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+def test_get_quant_method_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch):
+    """Non-LinearBase layers (norms, embeddings) must get None → no quantization."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config()
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    norm_layer = torch.nn.LayerNorm(64)
+    assert config.get_quant_method(norm_layer, "blocks.0.norm1") is None
+
+
+# ---------------------------------------------------------------------------
+# MXFPLinearMethodBase.apply() — reshape skeleton (CPU, no NPU)
+# ---------------------------------------------------------------------------
+
+
+def test_apply_reshape_skeleton():
+    """apply() must flatten batch dims → _apply_inner → restore original leading dims."""
+    from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase
+
+    OUT_FEATURES = 4
+
+    class _StubMethod(MXFPLinearMethodBase):
+        def create_weights(
+            self,
+            layer,
+            input_size_per_partition,
+            output_partition_sizes,
+            input_size,
+            output_size,
+            params_dtype,
+            **extra_weight_attrs,
+        ):
+            pass
+
+        def _quantize_activation(self, x):
+            return x, None
+
+        def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype):
+            return torch.zeros(x_q.shape[0], OUT_FEATURES, dtype=ori_dtype)
+
+    method = _StubMethod()
+    x = torch.randn(2, 3, 8)  # (batch=2, seq=3, K=8)
+    out = method.apply(None, x)
+    assert out.shape == (2, 3, OUT_FEATURES)
+
+
+def test_apply_reshape_with_bias():
+    """apply() must pass bias through to _apply_inner unchanged."""
+    from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase
+
+    received_bias = []
+
+    class _StubMethod(MXFPLinearMethodBase):
+        def create_weights(
+            self,
+            layer,
+            input_size_per_partition,
+            output_partition_sizes,
+            input_size,
+            output_size,
+            params_dtype,
+            **extra_weight_attrs,
+        ):
+            pass
+
+        def _quantize_activation(self, x):
+            return x, None
+
+        def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype):
+            received_bias.append(bias)
+            return torch.zeros(x_q.shape[0], 4, dtype=ori_dtype)
+
+    method = _StubMethod()
+    bias = torch.zeros(4)
+    method.apply(None, torch.randn(2, 8), bias=bias)
+    assert received_bias[0] is bias
+
+
+# ---------------------------------------------------------------------------
+# process_weights_after_loading shape arithmetic (pure torch, no NPU ops)
+#
+# These tests replicate the CPU-safe portions of process_weights_after_loading
+# to guard the key layout contract: (N,K) weight → (K,N) and (N,S) scale →
+# (S/2,N,2). They do NOT call NPU ops; they test only the math.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "n, k",
+    [(64, 128), (32, 64), (16, 96)],
+    ids=["64x128", "32x64", "16x96"],
+)
+def test_weight_transpose_contract(n: int, k: int):
+    """Weight must be transposed from (N, K) to (K, N) and be contiguous."""
+    w = torch.zeros(n, k, dtype=torch.uint8)
+    w = w.transpose(0, 1).contiguous()
+    assert w.shape == (k, n)
+    assert w.is_contiguous()
+
+
+@pytest.mark.parametrize(
+    "n, k_groups, expected_groups",
+    [
+        (64, 4, 2),  # even K_groups — no padding
+        (64, 3, 2),  # odd  K_groups — padded to 4
+        (32, 1, 1),  # odd  K_groups — padded to 2
+    ],
+    ids=["even", "odd-to-4", "odd-to-2"],
+)
+def test_weight_scale_reshape_contract(n: int, k_groups: int, expected_groups: int):
+    """Scale must be reshaped from (N, K_groups) to (K_groups_even//2, N, 2).
+
+    Odd K_groups must be padded to even before the reshape.
+    """
+    s = torch.zeros(n, k_groups, dtype=torch.uint8)
+    if k_groups % 2 == 1:
+        s = torch.cat([s, torch.zeros(n, 1, dtype=s.dtype)], dim=1)
+        k_groups += 1
+    s = s.reshape(n, k_groups // 2, 2).transpose(0, 1).contiguous()
+    assert s.shape == (expected_groups, n, 2)
+    assert s.is_contiguous()
+
+
+def test_num_groups_formula():
+    """K_groups formula: ceil(K / 32) — spot-check boundary values."""
+    assert (31 + 31) // 32 == 1  # K=31 → 1 group
+    assert (32 + 31) // 32 == 1  # K=32 → 1 group
+    assert (33 + 31) // 32 == 2  # K=33 → 2 groups
+    assert (128 + 31) // 32 == 4  # K=128 → 4 groups (even)
+    assert (96 + 31) // 32 == 3  # K=96  → 3 groups (odd → needs padding)
+
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_QUANTIZATION_METHODS
+# ---------------------------------------------------------------------------
+
+
+def test_supported_methods_include_mxfp8():
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS

tests/diffusion/quantization/test_mxfp8_key_remap.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for merge_mxfp8_checkpoint.py key-remapping helpers and model metadata.
+
+These are pure-Python unit tests that exercise the transformation functions
+without loading any actual checkpoint files or requiring NPU hardware.
+
+Key contracts verified:
+- SUPPORTED_MODEL_TYPES includes Wan2.2-TI2V-5B (MXFP8 supports it; MXFP4 does not)
+- MXFP8_QUANT_CONFIG structure matches what TransformerConfig.from_dict() reads
+- _remap_keys correctly translates msModelSlim naming → Diffusers naming
+- _get_transformer_dirs routes cascade vs single-transformer models
+- _get_quant_subdir maps high/low noise subdirs for cascade models
+"""
+
+import pathlib
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_MODEL_TYPES and CASCADE_MODEL_TYPES
+# ---------------------------------------------------------------------------
+
+
+def test_supported_model_types_includes_all_wan22():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES
+
+    assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES
+    assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES
+
+
+def test_supported_model_types_includes_ti2v_5b():
+    """MXFP8 supports TI2V-5B (contrast: MXFP4 explicitly excludes it)."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES
+
+    assert "Wan2.2-TI2V-5B" in SUPPORTED_MODEL_TYPES
+
+
+def test_cascade_model_types_excludes_ti2v():
+    """TI2V-5B is a single-transformer model — must not be in CASCADE_MODEL_TYPES."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import CASCADE_MODEL_TYPES
+
+    assert "Wan2.2-TI2V-5B" not in CASCADE_MODEL_TYPES
+    assert "Wan2.2-T2V-A14B" in CASCADE_MODEL_TYPES
+    assert "Wan2.2-I2V-A14B" in CASCADE_MODEL_TYPES
+
+
+# ---------------------------------------------------------------------------
+# MXFP8_QUANT_CONFIG — auto-detection contract
+# ---------------------------------------------------------------------------
+
+
+def test_mxfp8_quant_config_has_required_keys():
+    """MXFP8_QUANT_CONFIG must carry exactly the keys that auto-detection reads:
+    quant_method (selects DiffusionMXFP8Config) and is_checkpoint_mxfp8_serialized
+    (selects NPUMxfp8LinearMethod over the online path)."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
+
+    assert MXFP8_QUANT_CONFIG["quant_method"] == "mxfp8"
+    assert MXFP8_QUANT_CONFIG["is_checkpoint_mxfp8_serialized"] is True
+
+
+# ---------------------------------------------------------------------------
+# _get_transformer_dirs
+# ---------------------------------------------------------------------------
+
+
+def test_get_transformer_dirs_cascade():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs
+
+    assert _get_transformer_dirs("Wan2.2-T2V-A14B") == ["transformer", "transformer_2"]
+    assert _get_transformer_dirs("Wan2.2-I2V-A14B") == ["transformer", "transformer_2"]
+
+
+def test_get_transformer_dirs_single():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs
+
+    assert _get_transformer_dirs("Wan2.2-TI2V-5B") == ["transformer"]
+
+
+# ---------------------------------------------------------------------------
+# _get_quant_subdir
+# ---------------------------------------------------------------------------
+
+
+def test_get_quant_subdir_cascade_high_noise():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
+
+    base = pathlib.Path("/quant")
+    result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer")
+    assert result == base / "high_noise_model"
+
+
+def test_get_quant_subdir_cascade_low_noise():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
+
+    base = pathlib.Path("/quant")
+    result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer_2")
+    assert result == base / "low_noise_model"
+
+
+def test_get_quant_subdir_non_cascade():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
+
+    base = pathlib.Path("/quant")
+    result = _get_quant_subdir("Wan2.2-TI2V-5B", base, "transformer")
+    assert result == base
+
+
+# ---------------------------------------------------------------------------
+# _remap_keys — msModelSlim naming → Diffusers naming
+# ---------------------------------------------------------------------------
+
+
+def test_remap_keys_self_attn_q():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"blocks.0.self_attn.q.weight": torch.zeros(1)}
+    meta = {"blocks.0.self_attn.q.weight": "W8A8_MXFP8"}
+    new_state, new_meta = _remap_keys(state, meta)
+    assert "blocks.0.attn1.to_q.weight" in new_state
+    assert new_meta.get("blocks.0.attn1.to_q.weight") == "W8A8_MXFP8"
+
+
+def test_remap_keys_self_attn_all_heads():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    pairs = {
+        "self_attn.q": "attn1.to_q",
+        "self_attn.k": "attn1.to_k",
+        "self_attn.v": "attn1.to_v",
+        "self_attn.o": "attn1.to_out.0",
+    }
+    for src_part, dst_part in pairs.items():
+        src_key = f"blocks.0.{src_part}.weight"
+        state = {src_key: torch.zeros(1)}
+        new_state, _ = _remap_keys(state, {})
+        expected = f"blocks.0.{dst_part}.weight"
+        assert expected in new_state, f"{src_key} → expected {expected}, got {list(new_state)}"
+
+
+def test_remap_keys_ffn():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {
+        "blocks.1.ffn.0.weight": torch.zeros(1),
+        "blocks.1.ffn.2.weight": torch.zeros(1),
+    }
+    new_state, _ = _remap_keys(state, {})
+    assert "blocks.1.ffn.net.0.proj.weight" in new_state
+    assert "blocks.1.ffn.net.2.weight" in new_state
+
+
+def test_remap_keys_norm_order_swap():
+    """norm2↔norm3 swap: msModelSlim uses norm1/norm3/norm2 order,
+    Diffusers uses norm1/norm2/norm3."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {
+        "blocks.0.norm2.weight": torch.zeros(1),
+        "blocks.0.norm3.weight": torch.zeros(1),
+    }
+    new_state, _ = _remap_keys(state, {})
+    # norm2 → norm3 and norm3 → norm2
+    assert "blocks.0.norm3.weight" in new_state
+    assert "blocks.0.norm2.weight" in new_state
+    # Both must be present (swap, not collapse)
+    assert len([k for k in new_state if "norm" in k and "norm_q" not in k and "norm_k" not in k]) == 2
+
+
+def test_remap_keys_head():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"head.head.weight": torch.zeros(1)}
+    new_state, _ = _remap_keys(state, {})
+    assert "proj_out.weight" in new_state
+
+
+def test_remap_keys_cross_attn():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"blocks.0.cross_attn.q.weight": torch.zeros(1)}
+    new_state, _ = _remap_keys(state, {})
+    assert "blocks.0.attn2.to_q.weight" in new_state
+
+
+def test_remap_keys_meta_only_mapped_for_existing_state_keys():
+    """quant_meta entries are only emitted for keys present in state_dict."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"blocks.0.self_attn.q.weight": torch.zeros(1)}
+    # meta has an extra key not in state_dict
+    meta = {
+        "blocks.0.self_attn.q.weight": "W8A8_MXFP8",
+        "blocks.0.self_attn.q.weight_scale": "W8A8_MXFP8",
+    }
+    _, new_meta = _remap_keys(state, meta)
+    assert "blocks.0.attn1.to_q.weight" in new_meta
+    # weight_scale was in meta but not state_dict → must be absent
+    assert "blocks.0.attn1.to_q.weight_scale" not in new_meta
+
+
+def test_remap_keys_preserves_tensors():
+    """Tensor values must survive the key rename unchanged."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    t = torch.randn(4, 8)
+    state = {"blocks.0.self_attn.q.weight": t}
+    new_state, _ = _remap_keys(state, {})
+    assert torch.equal(new_state["blocks.0.attn1.to_q.weight"], t)

tests/diffusion/quantization/test_quantization_fp8.py

@@ -109,6 +109,85 @@ def _generate_single_stage_image(
         return images, peak_mem
 
 
+def _generate_single_stage_video(
+    model: str,
+    quantization: str | None = None,
+    height: int = 256,
+    width: int = 256,
+    num_frames: int = 25,
+    num_inference_steps: int = 8,
+    guidance_scale: float = 4.0,
+    seed: int = 42,
+    prompt: str = "A serene lakeside sunrise with mist over the water",
+    **extra_omni_kwargs: Any,
+) -> tuple[int, float]:
+    """Generate a t2v output with a single-stage diffusion model.
+
+    Returns (num_frames_produced, peak_memory_gib)
+    """
+    omni_kwargs: dict[str, Any] = dict(extra_omni_kwargs)
+    if quantization:
+        omni_kwargs["quantization"] = quantization
+
+    with OmniRunner(model, **omni_kwargs) as runner:
+        torch.accelerator.reset_peak_memory_stats()
+
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+        outputs = runner.omni.generate(
+            {"prompt": prompt, "negative_prompt": ""},
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+            ),
+        )
+
+        first = outputs[0]
+
+        # Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output).
+        frames: Any = None
+        if hasattr(first, "request_output") and isinstance(first.request_output, list):
+            inner = first.request_output[0]
+            if isinstance(inner, OmniRequestOutput) and inner.images:
+                frames = inner.images[0]
+        if frames is None and hasattr(first, "images") and first.images:
+            frames = first.images[0]
+        assert frames is not None, "No video frames returned from generate()"
+
+        # LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts.
+        if isinstance(frames, dict):
+            frames = frames.get("video") or frames.get("frames")
+        elif isinstance(frames, tuple) and len(frames) == 2:
+            frames = frames[0]
+        assert frames is not None, "Could not extract video frames from output"
+
+        if isinstance(frames, torch.Tensor):
+            video = frames.detach().cpu()
+            if video.dim() == 5:
+                video = video[0]
+            if video.dim() == 4 and video.shape[0] in (3, 4):
+                video = video.permute(1, 2, 3, 0)
+            num_frames_produced = int(video.shape[0])
+        else:
+            import numpy as np
+
+            arr = np.asarray(frames)
+            if arr.ndim == 5:
+                arr = arr[0]
+            num_frames_produced = int(arr.shape[0])
+
+        peak_mem_mb = getattr(first, "peak_memory_mb", None)
+        if peak_mem_mb:
+            peak_mem = float(peak_mem_mb) / 1024.0
+        else:
+            peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
+
+        return num_frames_produced, peak_mem
+
+
 def _generate_bagel_image(
     quantization_config: str | None = None,
     num_inference_steps: int = 15,
@@ -277,6 +356,25 @@ def test_single_stage_flux_fp8_uses_less_memory():
     assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
 
 
+@hardware_test(res={"cuda": "H100"})
+def test_single_stage_ltx2_fp8_uses_less_memory():
+    """FP8 should use less peak memory than BF16 for LTX-2."""
+    _, mem_bf16 = _generate_single_stage_video(
+        model="Lightricks/LTX-2",
+        quantization=None,
+    )
+    torch.accelerator.empty_cache()
+
+    _, mem_fp8 = _generate_single_stage_video(
+        model="Lightricks/LTX-2",
+        quantization="fp8",
+    )
+
+    print(f"LTX-2 BF16 peak memory: {mem_bf16:.2f} GiB")
+    print(f"LTX-2 FP8 peak memory:  {mem_fp8:.2f} GiB")
+    assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
+
+
 # ─── Multi-stage model tests (BAGEL) ─────────────────────────────────────────
 
 

tests/diffusion/quantization/test_quantization_quality.py

@@ -143,6 +143,18 @@ QUALITY_CONFIGS = [
         seed=142,
         num_inference_steps=20,
     ),
+    QualityTestConfig(
+        id="fp8_ltx2",
+        model="Lightricks/LTX-2",
+        quantization="fp8",
+        task="t2v",
+        prompt="A serene lakeside sunrise with mist over the water",
+        max_lpips=0.10,
+        height=256,
+        width=256,
+        num_frames=25,
+        num_inference_steps=8,
+    ),
 ]
 
 
@@ -261,10 +273,19 @@ def _generate_video(omni, config: QualityTestConfig):
         else:
             frames = inner
     elif hasattr(first, "images") and first.images:
-        frames = first.images
+        frames = first.images[0]
     else:
         raise ValueError("Could not extract video frames from output.")
 
+    # LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts
+    if isinstance(frames, dict):
+        frames = frames.get("video") or frames.get("frames")
+    elif isinstance(frames, tuple) and len(frames) == 2:
+        frames = frames[0]
+
+    if frames is None:
+        raise ValueError("Could not extract video frames from output.")
+
     if isinstance(frames, torch.Tensor):
         video = frames.detach().cpu()
         if video.dim() == 5:
@@ -275,7 +296,11 @@ def _generate_video(omni, config: QualityTestConfig):
             video = video.clamp(-1, 1) * 0.5 + 0.5
         return video.float().numpy(), peak_mem
 
-    return np.asarray(frames), peak_mem
+    frames_array = np.asarray(frames)
+    if frames_array.ndim == 5:
+        # strip the leading batch dim
+        frames_array = frames_array[0]
+    return frames_array, peak_mem
 
 
 def _compute_lpips(baseline, quantized, task: str) -> float:
@@ -318,8 +343,7 @@ def _compute_psnr_and_mae(baseline, quantized, task: str) -> tuple[float, float]
     return psnr, mae
 
 
-def _unload(omni):
-    del omni
+def _free_gpu_memory():
     gc.collect()
     if torch.cuda.is_available():
         torch.accelerator.empty_cache()
@@ -395,7 +419,9 @@ def test_quantization_quality(config: QualityTestConfig):
     # --- BF16 baseline ---
     omni_bl = Omni(model=config.baseline_ref())
     baseline_out, bl_mem = generate_fn(omni_bl, config)
-    _unload(omni_bl)
+    omni_bl.shutdown()
+    del omni_bl
+    _free_gpu_memory()
     _maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out)
 
     # --- Quantized ---
@@ -405,7 +431,9 @@ def test_quantization_quality(config: QualityTestConfig):
     else:
         omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization)
     quant_out, qt_mem = generate_fn(omni_qt, config)
-    _unload(omni_qt)
+    omni_qt.shutdown()
+    del omni_qt
+    _free_gpu_memory()
     _maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out)
 
     # --- Similarity metrics ---

tests/diffusion/quantization/test_trajectory_similarity_tool.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import (
+    VariantRun,
+    _build_variant_config,
+    _request_peak_memory_mb,
+    _run_summary,
+    compute_tensor_metrics,
+    compute_uint8_image_metrics,
+    metric_guidance,
+    parse_args,
+    summarize_output_image_metrics,
+)
+
+
+def test_compute_tensor_metrics_identical_tensors():
+    metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3))
+
+    assert metrics["cosine_similarity"] == pytest.approx(1.0)
+    assert metrics["mae"] == 0.0
+    assert metrics["mse"] == 0.0
+    assert metrics["rmse"] == 0.0
+    assert metrics["max_abs"] == 0.0
+    assert metrics["l2"] == 0.0
+    assert metrics["relative_l2"] == 0.0
+
+
+def test_compute_uint8_image_metrics_adds_psnr():
+    lhs = np.zeros((2, 2, 3), dtype=np.uint8)
+    rhs = np.zeros((2, 2, 3), dtype=np.uint8)
+
+    metrics = compute_uint8_image_metrics(lhs, rhs)
+
+    assert math.isinf(metrics["psnr_db"])
+
+
+def test_summarize_output_image_metrics_stacks_pil_images():
+    reference = [Image.fromarray(np.zeros((2, 2, 3), dtype=np.uint8))]
+    candidate = [Image.fromarray(np.ones((2, 2, 3), dtype=np.uint8))]
+
+    summary = summarize_output_image_metrics(reference, candidate)
+
+    assert summary["num_images"] == 1
+    assert summary["image0_metrics"]["mae"] == 1.0
+    assert summary["all_images_metrics"]["mse"] == 1.0
+
+
+def test_run_summary_reports_worker_peak_memory():
+    summary = _run_summary(
+        VariantRun(
+            label="candidate",
+            result=object(),
+            generation_times_s=[1.0, 3.0],
+            peak_memory_mb=[100.0, 150.0],
+        )
+    )
+
+    assert summary["peak_memory_mb"] == 150.0
+    assert summary["avg_peak_memory_mb"] == 125.0
+    assert summary["max_peak_memory_mb"] == 150.0
+    assert summary["per_run_peak_memory_mb"] == [100.0, 150.0]
+
+
+def test_request_peak_memory_prefers_inner_diffusion_output():
+    inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})()
+    outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})()
+
+    assert _request_peak_memory_mb(outer) == 321.0
+
+
+def test_metric_guidance_describes_thresholds():
+    guidance = metric_guidance()
+
+    assert "cosine_similarity" in guidance["descriptions"]
+    assert guidance["recommended_thresholds"]["output_images_or_frames_uint8"]["psnr_db"]["recommended_min"] == 20.0
+    assert (
+        guidance["recommended_thresholds"]["performance"]["max_peak_memory_ratio_candidate_over_reference"][
+            "recommended_max"
+        ]
+        == 1.00
+    )
+
+
+def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantization():
+    args = SimpleNamespace(
+        model="Qwen/Qwen-Image",
+        candidate_model="Qwen/Qwen-Image-FP8",
+        candidate_quantization=None,
+        candidate_quantization_config_json=None,
+        candidate_ignored_layers=None,
+        ignored_layers=None,
+        candidate_gguf_model=None,
+        candidate_load_format="default",
+    )
+
+    config = _build_variant_config(args, "candidate")
+
+    assert config.model == "Qwen/Qwen-Image-FP8"
+    assert config.quantization is None
+    assert config.quantization_config is None
+
+
+def test_step_execution_defaults_to_false(monkeypatch):
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "compare_diffusion_trajectory_similarity.py",
+            "--output-json",
+            "result.json",
+        ],
+    )
+
+    args = parse_args()
+
+    assert args.step_execution is False

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -136,6 +136,7 @@ def _make_engine_with_loop(
     calls and a real ``RequestScheduler``.
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._closed = False
     engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
 
     sched = RequestScheduler()
@@ -371,6 +372,7 @@ def test_collective_rpc_before_loop_starts_calls_executor_directly():
     on the caller's thread without enqueueing.
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._closed = False
     engine._loop_started = False
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
@@ -391,6 +393,7 @@ def test_collective_rpc_before_loop_starts_serializes_concurrent_callers():
     serialize them so they cannot race on the shared executor MQ pair.
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._closed = False
     engine._loop_started = False
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)

tests/diffusion/test_diffusion_scheduler.py

@@ -153,6 +153,54 @@ class _StubScheduler(SchedulerInterface):
         return None
 
 
+class TestGetSamplingParamsKey:
+    """Pure-function tests for the batch-compatibility key builder."""
+
+    @staticmethod
+    def _make(lora_int_id: int | None = None, lora_scale: float = 1.0) -> OmniDiffusionRequest:
+        from vllm_omni.lora.request import LoRARequest
+
+        sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+        if lora_int_id is not None:
+            sp.lora_request = LoRARequest(
+                lora_name=f"adapter-{lora_int_id}",
+                lora_int_id=lora_int_id,
+                lora_path=f"/tmp/lora-{lora_int_id}",
+            )
+        sp.lora_scale = lora_scale
+        return OmniDiffusionRequest(
+            prompts=["prompt"],
+            sampling_params=sp,
+            request_ids=[f"req-{lora_int_id}-{lora_scale}"],
+        )
+
+    def test_distinguishes_lora_id(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        assert get_sampling_params_key(self._make(lora_int_id=1)) != get_sampling_params_key(self._make(lora_int_id=2))
+
+    def test_distinguishes_lora_scale(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        assert get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) != get_sampling_params_key(
+            self._make(lora_int_id=1, lora_scale=1.0)
+        )
+
+    def test_treats_no_lora_as_distinct_bucket(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        assert get_sampling_params_key(self._make(lora_int_id=None)) != get_sampling_params_key(
+            self._make(lora_int_id=1)
+        )
+
+    def test_equal_for_same_lora_identity(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        a = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
+        b = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
+        assert a == b
+
+
 class TestRequestScheduler:
     def setup_method(self) -> None:
         self.scheduler: RequestScheduler = RequestScheduler()
@@ -728,6 +776,114 @@ class TestStepScheduler:
         assert second.num_running_reqs == 1
         assert second.num_waiting_reqs == 1
 
+    def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
+        """Multiple requests with the same LoRA (id + scale) co-batch."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=3))
+
+        lora = LoRARequest(lora_name="adapter", lora_int_id=42, lora_path="/tmp/lora")
+
+        def _with_lora(req_id: str) -> OmniDiffusionRequest:
+            sp = OmniDiffusionSamplingParams(num_inference_steps=4)
+            sp.lora_request = lora
+            sp.lora_scale = 0.5
+            return _make_step_request(req_id, sampling_params=sp)
+
+        req_a = scheduler.add_request(_with_lora("a"))
+        req_b = scheduler.add_request(_with_lora("b"))
+        req_c = scheduler.add_request(_with_lora("c"))
+
+        sched_output = scheduler.schedule()
+
+        assert _new_ids(sched_output) == [req_a, req_b, req_c]
+        assert sched_output.num_running_reqs == 3
+        assert sched_output.num_waiting_reqs == 0
+
+    def test_step_batch_separates_requests_with_different_lora_ids(self) -> None:
+        """Different LoRA adapters → distinct batches admitted in FIFO order."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=4))
+
+        lora_a = LoRARequest(lora_name="adapter-A", lora_int_id=1, lora_path="/tmp/lora-a")
+        lora_b = LoRARequest(lora_name="adapter-B", lora_int_id=2, lora_path="/tmp/lora-b")
+
+        def _build(req_id: str, lora: LoRARequest) -> OmniDiffusionRequest:
+            sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+            sp.lora_request = lora
+            return _make_step_request(req_id, sampling_params=sp)
+
+        req_a1 = scheduler.add_request(_build("a1", lora_a))
+        req_b1 = scheduler.add_request(_build("b1", lora_b))
+        req_a2 = scheduler.add_request(_build("a2", lora_a))
+
+        # Strict FIFO admission: a1 starts; b1 (different LoRA) blocks the
+        # queue head, so a2 (compatible with a1) is *not* skipped ahead.
+        first = scheduler.schedule()
+        assert _new_ids(first) == [req_a1]
+        assert first.num_waiting_reqs == 2
+
+        # Drain a1 → b1 becomes head-of-line and is admitted with its LoRA.
+        scheduler.update_from_output(first, _make_step_output(req_a1, step_index=2, finished=True))
+        second = scheduler.schedule()
+        assert _new_ids(second) == [req_b1]
+        assert second.num_waiting_reqs == 1
+
+        # Drain b1 → a2 is admitted next; LoRA-A is re-activated for it.
+        scheduler.update_from_output(second, _make_step_output(req_b1, step_index=2, finished=True))
+        third = scheduler.schedule()
+        assert _new_ids(third) == [req_a2]
+        assert third.num_waiting_reqs == 0
+
+    def test_step_batch_separates_requests_with_different_lora_scale(self) -> None:
+        """Same adapter id but different scales → still separate batches."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=4))
+
+        lora = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora")
+
+        def _build(req_id: str, scale: float) -> OmniDiffusionRequest:
+            sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+            sp.lora_request = lora
+            sp.lora_scale = scale
+            return _make_step_request(req_id, sampling_params=sp)
+
+        req_full = scheduler.add_request(_build("full", 1.0))
+        req_half = scheduler.add_request(_build("half", 0.5))
+
+        sched_output = scheduler.schedule()
+
+        admitted = _new_ids(sched_output)
+        assert admitted == [req_full]
+        assert req_half not in admitted
+        assert sched_output.num_waiting_reqs == 1
+
+    def test_step_batch_separates_lora_from_no_lora(self) -> None:
+        """A LoRA request and a no-LoRA request do not share a batch."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=4))
+
+        lora = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora")
+
+        sp_with = OmniDiffusionSamplingParams(num_inference_steps=2)
+        sp_with.lora_request = lora
+        req_with = scheduler.add_request(_make_step_request("with", sampling_params=sp_with))
+        req_without = scheduler.add_request(_make_step_request("without", num_inference_steps=2))
+
+        sched_output = scheduler.schedule()
+
+        admitted = _new_ids(sched_output)
+        assert admitted == [req_with]
+        assert req_without not in admitted
+        assert sched_output.num_waiting_reqs == 1
+
     def test_preempt_request_preserves_step_index(self) -> None:
         request = _make_step_request("preempt", num_inference_steps=3)
         req_id = self.scheduler.add_request(request)

tests/diffusion/test_diffusion_step_pipeline.py

@@ -285,6 +285,19 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id
     )
 
 
+def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
+    """Scheduler output for a homogeneous batch (one NewRequestData per req)."""
+    new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
+    return DiffusionSchedulerOutput(
+        step_id=step_id,
+        scheduled_new_reqs=new_reqs,
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
+        num_running_reqs=len(new_reqs),
+        num_waiting_reqs=0,
+    )
+
+
 def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
@@ -489,71 +502,128 @@ class TestRunner:
             DiffusionModelRunner.load_model(runner)
 
 
+class _RecordingLoRAManager:
+    def __init__(self) -> None:
+        self.calls: list[tuple[object | None, float]] = []
+
+    def set_active_adapter(self, adapter, scale: float = 1.0) -> None:
+        self.calls.append((adapter, scale))
+
+
+def _make_step_worker(lora_manager=None, *, expected_output=None):
+    """Build a bare DiffusionWorker primed for execute_stepwise tests."""
+    worker = object.__new__(DiffusionWorker)
+    worker.lora_manager = lora_manager
+    worker._step_lora_state = {}
+    output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
+    worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
+    return worker
+
+
 @pytest.mark.cpu
 class TestWorker:
     """DiffusionWorker.execute_stepwise"""
 
     def test_delegates_to_model_runner(self):
-        worker = object.__new__(DiffusionWorker)
         expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
-        scheduler_output = SimpleNamespace(
-            scheduled_new_reqs=[
-                SimpleNamespace(
-                    req=SimpleNamespace(
-                        sampling_params=SimpleNamespace(lora_request=None),
-                    )
-                )
-            ]
-        )
-        worker.lora_manager = None
-        worker.model_runner = SimpleNamespace(
-            execute_stepwise=lambda arg: expected if arg is scheduler_output else None
-        )
+        worker = _make_step_worker(expected_output=expected)
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
 
         output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
         assert output is expected
 
-    def test_clears_active_lora_before_stepwise_execution(self):
-        worker = object.__new__(DiffusionWorker)
-        scheduler_output = SimpleNamespace(
-            scheduled_new_reqs=[
-                SimpleNamespace(
-                    req=SimpleNamespace(
-                        sampling_params=SimpleNamespace(lora_request=None),
-                    )
-                )
-            ]
-        )
-        calls: list[object | None] = []
-
-        class _FakeLoRAManager:
-            def set_active_adapter(self, adapter):
-                calls.append(adapter)
-
-        worker.lora_manager = _FakeLoRAManager()
-        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
+    def test_deactivates_lora_when_request_has_no_adapter(self):
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
-        assert calls == [None]
+        assert manager.calls == [(None, 1.0)]
 
-    def test_rejects_lora_requests_in_step_mode(self):
-        worker = object.__new__(DiffusionWorker)
-        scheduler_output = SimpleNamespace(
-            scheduled_new_reqs=[
-                SimpleNamespace(
-                    req=SimpleNamespace(
-                        sampling_params=SimpleNamespace(lora_request=object()),
-                    )
-                )
-            ]
+    def test_activates_lora_for_step_requests(self):
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora")
+        request = _make_engine_request("req-1")
+        request.sampling_params.lora_request = lora_request
+        request.sampling_params.lora_scale = 0.75
+
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
+
+        DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert manager.calls == [(lora_request, 0.75)]
+
+    def test_recovers_lora_for_cached_step_requests(self):
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=11, lora_path="/tmp/lora")
+        request = _make_engine_request("req-1")
+        request.sampling_params.lora_request = lora_request
+        request.sampling_params.lora_scale = 0.5
+
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        first = _make_scheduler_output(request, sched_req_id="req-1")
+        second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
+
+        DiffusionWorker.execute_stepwise(worker, first)
+        DiffusionWorker.execute_stepwise(worker, second)
+
+        assert manager.calls == [(lora_request, 0.5), (lora_request, 0.5)]
+
+    def test_activates_single_lora_for_homogeneous_batch(self):
+        """Multiple requests sharing the same LoRA → exactly one activation,
+        and every request id is registered in ``_step_lora_state``."""
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=9, lora_path="/tmp/lora")
+        reqs = []
+        for rid in ("req-1", "req-2", "req-3"):
+            r = _make_engine_request(rid)
+            r.sampling_params.lora_request = lora_request
+            r.sampling_params.lora_scale = 0.6
+            reqs.append(r)
+
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        scheduler_output = _make_batch_scheduler_output(reqs)
+
+        DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert manager.calls == [(lora_request, 0.6)]
+        assert set(worker._step_lora_state) == {"req-1", "req-2", "req-3"}
+        for entry in worker._step_lora_state.values():
+            assert entry == (lora_request, 0.6)
+
+    def test_evicts_step_lora_state_for_finished_requests(self):
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora")
+        finishing = _make_engine_request("req-1")
+        finishing.sampling_params.lora_request = lora_request
+        next_request = _make_engine_request("req-2")
+        next_request.sampling_params.lora_request = lora_request
+
+        worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
+        first = _make_scheduler_output(finishing, sched_req_id="req-1")
+        next_batch = _make_scheduler_output(
+            next_request,
+            sched_req_id="req-2",
+            step_id=1,
+            finished_req_ids={"req-1"},
         )
-        worker.lora_manager = None
-        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
 
-        with pytest.raises(ValueError, match="does not support LoRA"):
-            DiffusionWorker.execute_stepwise(worker, scheduler_output)
+        DiffusionWorker.execute_stepwise(worker, first)
+        assert "req-1" in worker._step_lora_state
+
+        DiffusionWorker.execute_stepwise(worker, next_batch)
+        assert "req-1" not in worker._step_lora_state
+        assert worker._step_lora_state == {"req-2": (lora_request, 1.0)}
 
 
 @pytest.mark.cpu

tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py

@@ -15,8 +15,9 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
     BufferAllocator,
     ManagedBuffer,
 )
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import KVCacheTransferData
 
-# All tests in this file are pure-CPU unit tests for the memory allocator.
+# Most tests in this file are pure-CPU unit tests; CUDA smoke tests are skip-guarded.
 pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model]
 
 
@@ -225,3 +226,70 @@ class TestManagedBuffer:
         # Space should be reusable
         new_offset = self.allocator.alloc(128)
         assert new_offset == offset
+
+    def test_to_bytes_cpu_pool(self):
+        """Smoke-test Mooncake's ManagedBuffer byte extraction on CPU pools."""
+        expected = bytes(range(64))
+        offset = self.allocator.alloc(len(expected))
+        buf = ManagedBuffer(self.allocator, offset, len(expected), self.pool)
+        self.pool[offset : offset + len(expected)] = torch.tensor(list(expected), dtype=torch.uint8)
+
+        try:
+            assert buf.to_bytes() == expected
+        finally:
+            buf.release()
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for GPU pool smoke test")
+    def test_to_bytes_cuda_pool(self):
+        """Smoke-test Mooncake's existing CUDA pool D2H byte path."""
+        expected = bytes(range(64))
+        allocator = BufferAllocator(total_size=4096, alignment=64)
+        pool = torch.zeros(4096, dtype=torch.uint8, device="cuda:0")
+        offset = allocator.alloc(len(expected))
+        buf = ManagedBuffer(allocator, offset, len(expected), pool)
+        pool[offset : offset + len(expected)].copy_(torch.tensor(list(expected), dtype=torch.uint8, device="cuda:0"))
+
+        try:
+            assert buf.to_bytes() == expected
+        finally:
+            buf.release()
+
+
+def _make_kv_payload(device: str | torch.device = "cpu") -> tuple[KVCacheTransferData, torch.Tensor, torch.Tensor]:
+    key_tensor = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    value_tensor = key_tensor + 100
+    payload = KVCacheTransferData(
+        request_id="moon-smoke",
+        layer_blocks={"key_cache": [key_tensor], "value_cache": [value_tensor]},
+        block_ids=[1],
+        metadata={"seq_len": 3},
+    )
+    return payload, key_tensor, value_tensor
+
+
+class TestMooncakePackedPayloadSmoke:
+    """Smoke tests for Mooncake raw payload helpers without real Mooncake/RDMA."""
+
+    def test_load_header_from_cpu_tensor(self):
+        payload, _, _ = _make_kv_payload()
+        raw = payload.to_bytes()
+        packed = torch.tensor(list(raw), dtype=torch.uint8)
+        header, data_start = KVCacheTransferData._load_header_from_tensor(packed)
+
+        assert header["rid"] == "moon-smoke"
+        assert data_start == 4 + int.from_bytes(raw[:4], "big")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for packed CUDA smoke test")
+    def test_from_bytes_device_cuda_round_trip(self):
+        payload, key_tensor, value_tensor = _make_kv_payload(device="cuda:0")
+        packed = payload.to_gpu_tensor()
+
+        header, data_start = KVCacheTransferData._load_header_from_tensor(packed)
+        decoded = KVCacheTransferData.from_bytes_device(packed)
+
+        assert header["rid"] == "moon-smoke"
+        assert data_start == 4 + int.from_bytes(payload.to_bytes()[:4], "big")
+        assert decoded["layer_blocks"]["key_cache"][0].is_cuda
+        assert decoded["layer_blocks"]["value_cache"][0].is_cuda
+        assert torch.equal(decoded["layer_blocks"]["key_cache"][0].cpu(), key_tensor.cpu())
+        assert torch.equal(decoded["layer_blocks"]["value_cache"][0].cpu(), value_tensor.cpu())

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -6,18 +6,18 @@ from time import time
 import pytest
 
 from vllm_omni.distributed.omni_coordinator import (
-    InstanceInfo,
     LeastQueueLengthBalancer,
     RandomBalancer,
+    ReplicaInfo,
+    ReplicaStatus,
     RoundRobinBalancer,
-    StageStatus,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_load_balancer_select_returns_valid_index():
-    """Verify RandomBalancer.select() returns a valid index for instances."""
+    """Verify RandomBalancer.select() returns a valid index for replicas."""
     # Task structure mirrors async_omni; RandomBalancer ignores task contents.
     task: dict = {
         "request_id": "test",
@@ -26,30 +26,30 @@ def test_load_balancer_select_returns_valid_index():
     }
 
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=1,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
@@ -58,38 +58,38 @@ def test_load_balancer_select_returns_valid_index():
 
     balancer = RandomBalancer()
 
-    index = balancer.select(task, instances)
+    index = balancer.select(task, replicas)
 
     assert isinstance(index, int)
-    assert 0 <= index < len(instances)
+    assert 0 <= index < len(replicas)
 
 
-def test_round_robin_balancer_cycles_instances():
+def test_round_robin_balancer_cycles_replicas():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=1,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
@@ -97,35 +97,35 @@ def test_round_robin_balancer_cycles_instances():
     ]
 
     balancer = RoundRobinBalancer()
-    results = [balancer.select({}, instances) for _ in range(5)]
+    results = [balancer.select({}, replicas) for _ in range(5)]
 
     # Default start_index=0 => 0,1,2,0,1
     assert results == [0, 1, 2, 0, 1]
 
 
-def test_round_robin_balancer_empty_instances_raises():
-    with pytest.raises(ValueError, match="instances must not be empty"):
+def test_round_robin_balancer_empty_replicas_raises():
+    with pytest.raises(ValueError, match="replicas must not be empty"):
         RoundRobinBalancer().select({}, [])
 
 
 def test_round_robin_balancer_after_large_index_and_shorter_list():
-    """Large start_index % len(instances) then counter wraps with shorter list."""
+    """Large start_index % len(replicas) then counter wraps with shorter list."""
     now = time()
     two = [
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
@@ -138,30 +138,30 @@ def test_round_robin_balancer_after_large_index_and_shorter_list():
 
 def test_least_queue_length_balancer_picks_min_queue():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=5,
             last_heartbeat=now,
             registered_at=now,
@@ -169,41 +169,41 @@ def test_least_queue_length_balancer_picks_min_queue():
     ]
 
     balancer = LeastQueueLengthBalancer()
-    index = balancer.select({}, instances)
+    index = balancer.select({}, replicas)
     assert index == 1
 
 
-def test_least_queue_length_balancer_empty_instances_raises():
-    with pytest.raises(ValueError, match="instances must not be empty"):
+def test_least_queue_length_balancer_empty_replicas_raises():
+    with pytest.raises(ValueError, match="replicas must not be empty"):
         LeastQueueLengthBalancer().select({}, [])
 
 
 def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
@@ -214,21 +214,21 @@ def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
         "vllm_omni.distributed.omni_coordinator.load_balancer.random.choice",
         return_value=2,
     )
-    assert balancer.select({}, instances) == 2
+    assert balancer.select({}, replicas) == 2
 
 
 def test_least_queue_length_balancer_negative_queue_raises():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=-1,
             last_heartbeat=now,
             registered_at=now,
         ),
     ]
     with pytest.raises(ValueError, match="queue_length must be non-negative"):
-        LeastQueueLengthBalancer().select({}, instances)
+        LeastQueueLengthBalancer().select({}, replicas)

tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py

@@ -8,8 +8,8 @@ import pytest
 import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
-    InstanceList,
     OmniCoordClientForHub,
+    ReplicaList,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,8 +32,8 @@ def _wait_for_condition(cond, timeout: float = 2.0, interval: float = 0.01) -> b
     return False
 
 
-def test_hub_client_caches_instance_list_from_pub():
-    """Verify OmniCoordClientForHub receives instance list updates from OmniCoordinator and caches for get_instance_list()."""
+def test_hub_client_caches_replica_list_from_pub():
+    """Verify OmniCoordClientForHub receives replica list updates from OmniCoordinator and caches for get_replica_list()."""
     ctx, pub, endpoint = _bind_pub()
 
     client = OmniCoordClientForHub(endpoint)
@@ -41,7 +41,7 @@ def test_hub_client_caches_instance_list_from_pub():
     time.sleep(0.2)
 
     now = time.time()
-    instances_payload = [
+    replicas_payload = [
         {
             "input_addr": "tcp://stage:10001",
             "output_addr": "tcp://stage:10001-out",
@@ -71,37 +71,37 @@ def test_hub_client_caches_instance_list_from_pub():
         },
     ]
 
-    payload = {"instances": instances_payload, "timestamp": now}
+    payload = {"replicas": replicas_payload, "timestamp": now}
     pub.send(json.dumps(payload).encode("utf-8"))
 
-    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 3)
+    assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 3)
 
-    inst_list = client.get_instance_list()
-    assert isinstance(inst_list, InstanceList)
-    assert len(inst_list.instances) == 3
+    rep_list = client.get_replica_list()
+    assert isinstance(rep_list, ReplicaList)
+    assert len(rep_list.replicas) == 3
 
-    for src, inst in zip(instances_payload, inst_list.instances, strict=True):
-        assert inst.input_addr == src["input_addr"]
-        assert inst.output_addr == src["output_addr"]
-        assert inst.stage_id == src["stage_id"]
-        assert inst.status.value == src["status"]
+    for src, rep in zip(replicas_payload, rep_list.replicas, strict=True):
+        assert rep.input_addr == src["input_addr"]
+        assert rep.output_addr == src["output_addr"]
+        assert rep.stage_id == src["stage_id"]
+        assert rep.status.value == src["status"]
 
-    stage0 = client.get_instances_for_stage(0)
-    stage1 = client.get_instances_for_stage(1)
+    stage0 = client.get_replicas_for_stage(0)
+    stage1 = client.get_replicas_for_stage(1)
 
-    assert all(inst.stage_id == 0 for inst in stage0.instances)
-    assert all(inst.stage_id == 1 for inst in stage1.instances)
+    assert all(rep.stage_id == 0 for rep in stage0.replicas)
+    assert all(rep.stage_id == 1 for rep in stage1.replicas)
 
-    # Send an updated list with fewer instances and verify cache refresh.
+    # Send an updated list with fewer replicas and verify cache refresh.
     updated_payload = {
-        "instances": instances_payload[:2],
+        "replicas": replicas_payload[:2],
         "timestamp": now + 1.0,
     }
     pub.send(json.dumps(updated_payload).encode("utf-8"))
 
-    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 2)
-    updated_list = client.get_instance_list()
-    assert len(updated_list.instances) == 2
+    assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 2)
+    updated_list = client.get_replica_list()
+    assert len(updated_list.replicas) == 2
 
     client.close()
     pub.close(0)

tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py

@@ -9,7 +9,7 @@ import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
-    StageStatus,
+    ReplicaStatus,
 )
 from vllm_omni.distributed.omni_coordinator import (
     omni_coord_client_for_stage as stage_client_module,
@@ -47,7 +47,7 @@ def test_stage_client_auto_register_on_init():
     event = _recv_event(router)
 
     assert event["event_type"] == "update"
-    assert event["status"] == StageStatus.UP.value
+    assert event["status"] == ReplicaStatus.UP.value
     assert event["stage_id"] == stage_id
     assert event["input_addr"] == input_addr
     assert event["output_addr"] == output_addr
@@ -70,13 +70,13 @@ def test_stage_client_update_info_sends_correct_event():
     # Discard initial registration event.
     _recv_event(router)
 
-    client.update_info(status=StageStatus.ERROR)
+    client.update_info(status=ReplicaStatus.ERROR)
     client.update_info(queue_length=10)
 
     first = _recv_event(router)
     second = _recv_event(router)
 
-    assert first["status"] == StageStatus.ERROR.value
+    assert first["status"] == ReplicaStatus.ERROR.value
     assert first["stage_id"] == stage_id
     assert first["input_addr"] == input_addr
     assert first["output_addr"] == output_addr
@@ -107,7 +107,7 @@ def test_stage_client_close_sends_down_status():
     client.close()
 
     event = _recv_event(router)
-    assert event["status"] == StageStatus.DOWN.value
+    assert event["status"] == ReplicaStatus.DOWN.value
     assert event["stage_id"] == stage_id
     assert event["input_addr"] == input_addr
     assert event["output_addr"] == output_addr
@@ -274,6 +274,7 @@ def test_heartbeat_loop_retries_after_transient_send_failure():
     client._closed = False
     client._heartbeat_interval = 0.0
     client._stop_event = _FakeStopEvent()
+    client._on_heartbeat = None
 
     calls = {"count": 0}
 

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -11,14 +11,14 @@ from vllm.v1.utils import get_engine_client_zmq_addr
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
     OmniCoordinator,
-    StageStatus,
+    ReplicaStatus,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
-    """Receive InstanceList JSON from SUB socket. Returns None on timeout."""
+def _recv_replica_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
+    """Receive ReplicaList JSON from SUB socket. Returns None on timeout."""
     sub.setsockopt(zmq.RCVTIMEO, timeout_ms)
     try:
         data = sub.recv()
@@ -27,16 +27,16 @@ def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
         return None
 
 
-def _wait_for_instance_list(
+def _wait_for_replica_list(
     sub: zmq.Socket,
     expected_count: int,
     timeout: float = 3.0,
 ) -> dict | None:
-    """Wait until received InstanceList with expected_count active instances."""
+    """Wait until received ReplicaList with expected_count active replicas."""
     start = time.time()
     while time.time() - start < timeout:
-        msg = _recv_instance_list(sub, timeout_ms=500)
-        if msg is not None and len(msg.get("instances", [])) == expected_count:
+        msg = _recv_replica_list(sub, timeout_ms=500)
+        if msg is not None and len(msg.get("replicas", [])) == expected_count:
             return msg
     return None
 
@@ -45,7 +45,7 @@ def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None:
     """Drain queued SUB messages for a short window."""
     deadline = time.time() + max_seconds
     while time.time() < deadline:
-        _recv_instance_list(sub, timeout_ms=50)
+        _recv_replica_list(sub, timeout_ms=50)
 
 
 def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
@@ -81,7 +81,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
     )
 
     # Wait for initial registration broadcast and clear any queued messages.
-    msg = _wait_for_instance_list(sub, expected_count=1)
+    msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None
     _drain_sub_messages(sub)
 
@@ -96,7 +96,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
     deadline = time.time() + window_s
     recv_count = 0
     while time.time() < deadline:
-        if _recv_instance_list(sub, timeout_ms=100) is not None:
+        if _recv_replica_list(sub, timeout_ms=100) is not None:
             recv_count += 1
 
     assert recv_count < update_count // 2, (
@@ -110,8 +110,8 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
 
 
 def test_omni_coordinator_registration_broadcast():
-    """Verify that after multiple OmniCoordClientForStage instances register,
-    OmniCoordinator publishes an InstanceList containing all registered instances.
+    """Verify that after multiple OmniCoordClientForStage replicas register,
+    OmniCoordinator publishes a ReplicaList containing all registered replicas.
     """
     router_addr = get_engine_client_zmq_addr(
         local_only=False,
@@ -144,12 +144,12 @@ def test_omni_coordinator_registration_broadcast():
         OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
     ]
 
-    msg = _wait_for_instance_list(sub, expected_count=3)
-    assert msg is not None, "Expected InstanceList with 3 instances"
-    assert len(msg["instances"]) == 3
+    msg = _wait_for_replica_list(sub, expected_count=3)
+    assert msg is not None, "Expected ReplicaList with 3 replicas"
+    assert len(msg["replicas"]) == 3
     assert isinstance(msg["timestamp"], (int, float))
 
-    input_addrs = {inst["input_addr"] for inst in msg["instances"]}
+    input_addrs = {rep["input_addr"] for rep in msg["replicas"]}
     assert "tcp://stage:10001" in input_addrs
     assert "tcp://stage:10002" in input_addrs
     assert "tcp://stage:10003" in input_addrs
@@ -162,7 +162,7 @@ def test_omni_coordinator_registration_broadcast():
 
 
 def test_omni_coordinator_heartbeat_timeout_handling():
-    """Verify that when a stage instance stops sending heartbeats,
+    """Verify that when a stage replica stops sending heartbeats,
     OmniCoordinator marks it as unhealthy and excludes it from the active list.
     """
     router_addr = get_engine_client_zmq_addr(
@@ -201,23 +201,23 @@ def test_omni_coordinator_heartbeat_timeout_handling():
         "output_addr": "tcp://stage:c-out",
         "stage_id": 0,
         "event_type": "update",
-        "status": StageStatus.UP.value,
+        "status": ReplicaStatus.UP.value,
         "queue_length": 0,
     }
     dealer_c.send(json.dumps(reg_event).encode("utf-8"))
 
-    msg = _wait_for_instance_list(sub, expected_count=3)
-    assert msg is not None, "Expected initial 3 instances"
-    assert len(msg["instances"]) == 3
+    msg = _wait_for_replica_list(sub, expected_count=3)
+    assert msg is not None, "Expected initial 3 replicas"
+    assert len(msg["replicas"]) == 3
 
     # Wait for heartbeat timeout (timeout=5s, check interval ~2.5s).
     time.sleep(8.0)
 
     # Receive the update (C should be ERROR and excluded from active list).
-    msg_after_timeout = _wait_for_instance_list(sub, expected_count=2, timeout=5.0)
-    assert msg_after_timeout is not None, "Expected InstanceList with 2 instances after timeout"
-    instances = msg_after_timeout.get("instances", [])
-    input_addrs = {inst["input_addr"] for inst in instances}
+    msg_after_timeout = _wait_for_replica_list(sub, expected_count=2, timeout=5.0)
+    assert msg_after_timeout is not None, "Expected ReplicaList with 2 replicas after timeout"
+    replicas = msg_after_timeout.get("replicas", [])
+    input_addrs = {rep["input_addr"] for rep in replicas}
 
     assert "tcp://stage:a" in input_addrs
     assert "tcp://stage:b" in input_addrs
@@ -232,8 +232,8 @@ def test_omni_coordinator_heartbeat_timeout_handling():
     sub_ctx.term()
 
 
-def test_omni_coordinator_instance_shutdown_handling():
-    """Verify that when a stage instance sends status='down',
+def test_omni_coordinator_replica_shutdown_handling():
+    """Verify that when a stage replica sends status='down',
     OmniCoordinator removes it from the active list and broadcasts an updated list.
     """
     router_addr = get_engine_client_zmq_addr(
@@ -261,18 +261,18 @@ def test_omni_coordinator_instance_shutdown_handling():
 
     client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
 
-    msg = _wait_for_instance_list(sub, expected_count=1)
+    msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None
-    assert len(msg["instances"]) == 1
-    assert msg["instances"][0]["input_addr"] == "tcp://stage:shutdown"
+    assert len(msg["replicas"]) == 1
+    assert msg["replicas"][0]["input_addr"] == "tcp://stage:shutdown"
 
     # Send down status (simulating graceful shutdown).
-    client.update_info(status=StageStatus.DOWN)
+    client.update_info(status=ReplicaStatus.DOWN)
 
-    # Receive updated list (should have 0 active instances).
-    msg = _wait_for_instance_list(sub, expected_count=0)
+    # Receive updated list (should have 0 active replicas).
+    msg = _wait_for_replica_list(sub, expected_count=0)
     assert msg is not None
-    assert len(msg["instances"]) == 0
+    assert len(msg["replicas"]) == 0
 
     client.close()
     coordinator.close()

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -23,18 +23,31 @@ from tests.e2e.accuracy.helpers import (
     download_images,
     model_output_dir,
 )
-from tests.helpers.runtime import OmniRunner
+from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
 
 os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
 
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, *hardware_marks(res={"cuda": "H100"}, num_cards=8)]
 
 # ============================================================================
 # Configurable Parameters
 # ============================================================================
-AR_DEVICES = "0,1"
-DIT_DEVICES = "2,3"
+# Comma-separated logical CUDA device ids per stage: split visible GPUs (0..n-1), first half -> AR, second -> DiT.
+
+
+def _default_ar_dit_devices() -> tuple[str, str]:
+    """First floor(n/2) logical devices -> AR, rest -> DiT. ``device_count`` respects ``CUDA_VISIBLE_DEVICES``."""
+    n = torch.accelerator.device_count()
+    if n < 2:
+        return "0,1", "2,3"
+    split = n // 2
+    ar = ",".join(str(i) for i in range(split))
+    dit = ",".join(str(i) for i in range(split, n))
+    return ar, dit
+
+
+AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices()
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
 NUM_INFERENCE_STEPS = 50
 GUIDANCE_SCALE = 2.5
@@ -75,7 +88,7 @@ _BASE_CONFIG = {
                 "scheduler_cls": "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler",
                 "gpu_memory_utilization": 0.95, "enforce_eager": True, "trust_remote_code": True,
                 "engine_output_type": "latent", "enable_prefix_caching": False,
-                "max_num_batched_tokens": 32768, "tensor_parallel_size": 2, "pipeline_parallel_size": 1,
+                "max_num_batched_tokens": 32768, "tensor_parallel_size": AR_TP_SIZE, "pipeline_parallel_size": 1,
                 "hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}},
             },
             "is_comprehension": False, "final_output": True, "final_output_type": "text",
@@ -91,7 +104,7 @@ _BASE_CONFIG = {
             "engine_args": {
                 "model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM",
                 "enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp",
-                "parallel_config": {"tensor_parallel_size": 2, "enable_expert_parallel": True},
+                "parallel_config": {"tensor_parallel_size": DIT_TP_SIZE, "enable_expert_parallel": True},
             },
             "engine_input_source": [0],
             "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
@@ -134,12 +147,13 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None:
 def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
     from transformers import AutoTokenizer
 
+    from tests.helpers.runtime import OmniRunner
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
     from vllm_omni.platforms import current_omni_platform
 
     build_kwargs: dict = {"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified", "num_images": 2}
 
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
     result = build_prompt_tokens(
         PROMPT,
         tokenizer,
@@ -149,7 +163,7 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
     system_prompt_type = result.system_prompt_type
 
     ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer)
-    with OmniRunner(MODEL_NAME, stage_configs_path=stage_configs_path) as runner:
+    with OmniRunner(MODEL_PATH, stage_configs_path=stage_configs_path) as runner:
         params_list = list(runner.omni.default_sampling_params_list)
         for sp in params_list:
             if isinstance(sp, OmniDiffusionSamplingParams):
@@ -206,7 +220,10 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
     return image, cot_text, elapsed
 
 
-@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)")
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
+    reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
+)
 def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
     if importlib.util.find_spec("FlagEmbedding") is None:
         raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")

tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py

@@ -70,7 +70,7 @@ def test_text_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -93,7 +93,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -114,7 +114,7 @@ def test_image_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -135,7 +135,7 @@ def test_video_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -162,4 +162,4 @@ def test_mix_to_audio(omni_runner, omni_runner_handler):
         "modalities": ["audio"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"

tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py

@@ -87,7 +87,7 @@ def test_text_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -111,7 +111,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -133,7 +133,7 @@ def test_image_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -155,7 +155,7 @@ def test_video_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -177,7 +177,7 @@ def test_video_to_audio(omni_runner, omni_runner_handler):
         "modalities": ["audio"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
 
 
 # ------------------------------------------------------------------
@@ -204,4 +204,4 @@ def test_mix_to_audio(omni_runner, omni_runner_handler):
         "modalities": ["audio"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"

tests/e2e/offline_inference/test_voxcpm_expansion.py

@@ -1,99 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""E2E test for VoxCPM offline inference."""
-
-from typing import Any
-
-import numpy as np
-import pytest
-import torch
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path
-
-VOXCPM_MODEL = "OpenBMB/VoxCPM1.5"
-STAGE_CONFIG = get_deploy_config_path("voxcpm.yaml")
-SAMPLE_RATE = 24000
-
-# (model, stage_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
-_OMNI_RUNNER_PARAM = (VOXCPM_MODEL, STAGE_CONFIG)
-
-pytestmark = [
-    pytest.mark.full_model,
-    pytest.mark.tts,
-    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
-]
-
-
-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")
-
-
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_voxcpm_zero_shot_001(omni_runner: OmniRunner) -> None:
-    outputs = list(omni_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_wan22_autoround_w4a16_expansion.py

@@ -0,0 +1,311 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E tests for Wan2.2 AutoRound W4A16 quantized inference.
+
+These tests cover I2V (image-to-video) and T2V (text-to-video) generation
+with quantized weights.
+
+Requirements:
+  - CUDA GPU (H100 or equivalent, ~36 GiB for quantized model)
+  - The quantized model checkpoint (Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound,
+    Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound)
+"""
+
+import gc
+import os as _os
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+_os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+QUANTIZED_MODEL_I2V = "Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound"
+BASELINE_MODEL_I2V = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+QUANTIZED_MODEL_T2V = "Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound"
+BASELINE_MODEL_T2V = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+
+QUANTIZED_MODEL_I2V = _os.environ.get("WAN22_I2V_AUTOROUND_MODEL", QUANTIZED_MODEL_I2V)
+BASELINE_MODEL_I2V = _os.environ.get("WAN22_I2V_BASELINE_MODEL", BASELINE_MODEL_I2V)
+QUANTIZED_MODEL_T2V = _os.environ.get("WAN22_T2V_AUTOROUND_MODEL", QUANTIZED_MODEL_T2V)
+BASELINE_MODEL_T2V = _os.environ.get("WAN22_T2V_BASELINE_MODEL", BASELINE_MODEL_T2V)
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.diffusion,
+]
+
+# Small resolution to keep GPU memory & time manageable
+HEIGHT = 480
+WIDTH = 640
+NUM_FRAMES = 5  # must satisfy num_frames % 4 == 1 for Wan2.2
+NUM_STEPS = 2  # minimal for smoke-test
+
+# Parametrise: (model, stage_config_path=None, extra_omni_kwargs)
+# When stage_config_path is None, the engine auto-resolves from the model's own config.
+quant_i2v_params = [(QUANTIZED_MODEL_I2V, None, {"enforce_eager": True})]
+baseline_i2v_params = [(BASELINE_MODEL_I2V, None, {"enforce_eager": True})]
+quant_t2v_params = [(QUANTIZED_MODEL_T2V, None, {"enforce_eager": True})]
+baseline_t2v_params = [(BASELINE_MODEL_T2V, None, {"enforce_eager": True})]
+
+# Module-level storage for peak memory results across tests
+_memory_results: dict[str, float] = {}
+
+
+def _sampling_params_i2v() -> OmniDiffusionSamplingParams:
+    """Create sampling params for I2V generation."""
+    return OmniDiffusionSamplingParams(
+        height=HEIGHT,
+        width=WIDTH,
+        num_frames=NUM_FRAMES,
+        num_inference_steps=NUM_STEPS,
+        guidance_scale=5.0,
+        guidance_scale_2=6.0,
+        boundary_ratio=0.875,
+        generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+    )
+
+
+def _sampling_params_t2v() -> OmniDiffusionSamplingParams:
+    """Create sampling params for T2V generation."""
+    return OmniDiffusionSamplingParams(
+        height=HEIGHT,
+        width=WIDTH,
+        num_frames=NUM_FRAMES,
+        num_inference_steps=NUM_STEPS,
+        guidance_scale=4.0,
+        generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+    )
+
+
+def _create_test_image(width: int = WIDTH, height: int = HEIGHT) -> Image.Image:
+    """Create a deterministic test image for I2V tests."""
+    rng = np.random.RandomState(42)
+    arr = rng.randint(0, 256, (height, width, 3), dtype=np.uint8)
+    return Image.fromarray(arr)
+
+
+def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table, smooth motion") -> tuple:
+    """Generate one I2V video, return (frames, peak_memory_mb)."""
+    gc.collect()
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    current_omni_platform.reset_peak_memory_stats()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+
+    image = _create_test_image()
+    response = omni_runner_handler.send_diffusion_request(
+        {
+            "prompt": prompt,
+            "images": image,
+            "sampling_params": _sampling_params_i2v(),
+        },
+    )
+
+    peak = monitor.peak_used_mb
+    monitor.stop()
+
+    assert response.success, f"Request failed: {response.error_message}"
+    assert response.images is not None and len(response.images) > 0, "Expected image output"
+    frames = response.images[0]
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+
+    return frames, peak
+
+
+def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table") -> tuple:
+    """Generate one T2V video, return (frames, peak_memory_mb)."""
+    gc.collect()
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    current_omni_platform.reset_peak_memory_stats()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+
+    response = omni_runner_handler.send_diffusion_request(
+        {
+            "prompt": prompt,
+            "sampling_params": _sampling_params_t2v(),
+        },
+    )
+
+    peak = monitor.peak_used_mb
+    monitor.stop()
+
+    assert response.success, f"Request failed: {response.error_message}"
+    assert response.images is not None and len(response.images) > 0, "Expected image output"
+    frames = response.images[0]
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+
+    return frames, peak
+
+
+# ------------------------------------------------------------------
+# Test: I2V quantized model generates valid video
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True)
+def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler):
+    """Load the W4A16 quantized Wan2.2 I2V model and verify it produces a valid video."""
+    frames, _ = _generate_i2v_video(omni_runner_handler)
+
+    assert frames is not None, "Expected video frames output"
+    assert hasattr(frames, "shape"), "Expected frames to have a shape attribute"
+
+    # frames shape: (batch, num_frames, height, width, channels)
+    assert frames.shape[1] == NUM_FRAMES, f"Expected {NUM_FRAMES} frames, got {frames.shape[1]}"
+    assert frames.shape[2] == HEIGHT, f"Expected height {HEIGHT}, got {frames.shape[2]}"
+    assert frames.shape[3] == WIDTH, f"Expected width {WIDTH}, got {frames.shape[3]}"
+
+    # Sanity: video should not be blank (frames are [0, 1] floats)
+    arr = np.asarray(frames)
+    assert arr.std() > 0.01, "Generated video appears blank (std ≈ 0)"
+
+
+# ------------------------------------------------------------------
+# Test: T2V quantized model generates valid video
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True)
+def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler):
+    """Load the W4A16 quantized Wan2.2 T2V model and verify it produces a valid video."""
+    frames, _ = _generate_t2v_video(omni_runner_handler)
+
+    assert frames is not None, "Expected video frames output"
+    assert hasattr(frames, "shape"), "Expected frames to have a shape attribute"
+
+    assert frames.shape[1] == NUM_FRAMES, f"Expected {NUM_FRAMES} frames, got {frames.shape[1]}"
+    assert frames.shape[2] == HEIGHT, f"Expected height {HEIGHT}, got {frames.shape[2]}"
+    assert frames.shape[3] == WIDTH, f"Expected width {WIDTH}, got {frames.shape[3]}"
+
+    arr = np.asarray(frames)
+    assert arr.std() > 0.01, "Generated video appears blank (std ≈ 0)"
+
+
+# ------------------------------------------------------------------
+# Test: I2V quantized peak memory
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True)
+def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler):
+    """Measure peak GPU memory of W4A16 quantized I2V model."""
+    frames, peak = _generate_i2v_video(omni_runner_handler)
+
+    assert frames is not None, "Expected video frames output"
+    _memory_results["quant_i2v"] = peak
+    print(f"\nQuantized I2V (W4A16) peak memory: {peak:.0f} MB")
+
+
+# ------------------------------------------------------------------
+# Test: I2V baseline peak memory
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", baseline_i2v_params, indirect=True)
+def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler):
+    """Measure peak GPU memory of BF16 baseline I2V model."""
+    frames, peak = _generate_i2v_video(omni_runner_handler)
+
+    assert frames is not None, "Expected video frames output"
+    _memory_results["baseline_i2v"] = peak
+    print(f"\nBaseline I2V (BF16) peak memory: {peak:.0f} MB")
+
+
+# ------------------------------------------------------------------
+# Test: I2V memory savings
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_wan22_i2v_autoround_w4a16_memory_savings():
+    """Assert quantized I2V model uses meaningfully less memory than BF16 baseline."""
+    quant_peak = _memory_results["quant_i2v"]
+    baseline_peak = _memory_results["baseline_i2v"]
+
+    savings = baseline_peak - quant_peak
+    print(f"\nQuantized I2V (W4A16) peak memory: {quant_peak:.0f} MB")
+    print(f"Baseline I2V (BF16) peak memory:   {baseline_peak:.0f} MB")
+    print(f"Savings:                            {savings:.0f} MB")
+
+    # Wan2.2 I2V A14B transformer is ~28 GB in BF16; W4A16 should save ~20 GB.
+    # Use a conservative threshold to account for activations and overhead.
+    min_savings_mb = 5000
+    assert quant_peak + min_savings_mb < baseline_peak, (
+        f"Quantized model ({quant_peak:.0f} MB) should use at least "
+        f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)"
+    )
+
+
+# ------------------------------------------------------------------
+# Test: T2V quantized peak memory
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True)
+def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler):
+    """Measure peak GPU memory of W4A16 quantized T2V model."""
+    frames, peak = _generate_t2v_video(omni_runner_handler)
+
+    assert frames is not None, "Expected video frames output"
+    _memory_results["quant_t2v"] = peak
+    print(f"\nQuantized T2V (W4A16) peak memory: {peak:.0f} MB")
+
+
+# ------------------------------------------------------------------
+# Test: T2V baseline peak memory
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", baseline_t2v_params, indirect=True)
+def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler):
+    """Measure peak GPU memory of BF16 baseline T2V model."""
+    frames, peak = _generate_t2v_video(omni_runner_handler)
+
+    assert frames is not None, "Expected video frames output"
+    _memory_results["baseline_t2v"] = peak
+    print(f"\nBaseline T2V (BF16) peak memory: {peak:.0f} MB")
+
+
+# ------------------------------------------------------------------
+# Test: T2V memory savings
+# ------------------------------------------------------------------
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_wan22_t2v_autoround_w4a16_memory_savings():
+    """Assert quantized T2V model uses meaningfully less memory than BF16 baseline."""
+    quant_peak = _memory_results["quant_t2v"]
+    baseline_peak = _memory_results["baseline_t2v"]
+
+    savings = baseline_peak - quant_peak
+    print(f"\nQuantized T2V (W4A16) peak memory: {quant_peak:.0f} MB")
+    print(f"Baseline T2V (BF16) peak memory:   {baseline_peak:.0f} MB")
+    print(f"Savings:                            {savings:.0f} MB")
+
+    # Wan2.2 T2V A14B transformer is ~28 GB in BF16; W4A16 should save ~20 GB.
+    # Use a conservative threshold to account for activations and overhead.
+    min_savings_mb = 5000
+    assert quant_peak + min_savings_mb < baseline_peak, (
+        f"Quantized model ({quant_peak:.0f} MB) should use at least "
+        f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)"
+    )

tests/e2e/online_serving/test_omnivoice_expansion.py

@@ -105,16 +105,3 @@ class TestOmniVoiceVoiceCloning:
             "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
         }
         openai_client.send_audio_speech_request(request_config)
-
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_voice_clone_invalid_ref_audio_format(self, omni_server, openai_client) -> None:
-        """Test that invalid ref_audio format returns a clear error."""
-        request_config = {
-            "model": omni_server.model,
-            "input": get_prompt("text"),
-            "ref_audio": "not_a_valid_uri",
-            "response_format": "wav",
-            "timeout": 180.0,
-            "status_code": (400, 422),
-        }
-        openai_client.send_audio_speech_request(request_config)

tests/e2e/online_serving/test_voxtral_tts.py

@@ -91,20 +91,3 @@ class TestVoxtralTTSFixedVoice:
                     "timeout": 120.0,
                 }
             )
-
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "H100"}, num_cards=1)
-    def test_speech_invalid_voice_rejected(self, omni_server, openai_client) -> None:
-        """Request with a non-existent voice should be rejected (HTTP 200 with JSON or 4xx, depending on server)."""
-        openai_client.send_audio_speech_request(
-            {
-                "model": omni_server.model,
-                "input": "This voice does not exist.",
-                "voice": "nonexistent_voice_xyz",
-                "language": "English",
-                "response_format": "wav",
-                "timeout": 120.0,
-                "err_message": "Invalid speaker",
-            }
-        )

tests/engine/test_arg_utils.py

@@ -232,24 +232,6 @@ def test_stage_configs_path_field():
     assert args.stage_configs_path == "/some/path.yaml"
 
 
-def test_voxcpm_model_arch_injects_model_type_override(mocker):
-    """Ensure VoxCPM model_arch injects hf_overrides for config resolution."""
-    mocker.patch.object(OmniEngineArgs, "_ensure_omni_models_registered", return_value=True)
-    mocker.patch.object(OmniEngineArgs, "_patch_empty_hf_config")
-    mocker.patch.object(EngineArgs, "create_model_config", return_value=Mock())
-    mocker.patch.object(OmniModelConfig, "from_vllm_model_config", return_value=Mock())
-
-    args = OmniEngineArgs(
-        model="OpenBMB/VoxCPM1.5",
-        model_arch="VoxCPMForConditionalGeneration",
-    )
-    args.create_model_config()
-
-    assert args.hf_overrides["architectures"] == ["VoxCPMForConditionalGeneration"]
-    assert args.hf_overrides["model_type"] == "voxcpm"
-    args._patch_empty_hf_config.assert_called_once_with("voxcpm")
-
-
 def test_strip_single_engine_args():
     """_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
     kwargs = {

tests/engine/test_async_omni_engine_input.py

@@ -3,6 +3,7 @@ from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
 
+from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
 from vllm_omni.engine.stage_pool import StagePool
@@ -95,6 +96,44 @@ class _FakeStageClient:
     stage_type = "llm"
     final_output = False
 
+    def __init__(self, input_address: str | None = None):
+        if input_address is not None:
+            self.client_addresses = {"input_address": input_address}
+
+
+class _FakeHub:
+    def __init__(self, replicas: list[ReplicaInfo]):
+        self._replicas = replicas
+
+    def get_replicas_for_stage(self, stage_id: int):
+        return type(
+            "ReplicaList",
+            (),
+            {"replicas": [rep for rep in self._replicas if rep.stage_id == stage_id]},
+        )()
+
+
+class _RoundRobinLB:
+    def __init__(self):
+        self._next = 0
+
+    def select(self, task, replicas):  # noqa: ARG002
+        idx = self._next % len(replicas)
+        self._next += 1
+        return idx
+
+
+def _replica(input_addr: str) -> ReplicaInfo:
+    return ReplicaInfo(
+        input_addr=input_addr,
+        output_addr=input_addr.replace("input", "output"),
+        stage_id=0,
+        status=ReplicaStatus.UP,
+        queue_length=0,
+        last_heartbeat=0.0,
+        registered_at=0.0,
+    )
+
 
 def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
@@ -132,6 +171,93 @@ def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mo
     assert seen_uuids[0].removeprefix("stage0:rep0:") == seen_uuids[1].removeprefix("stage0:rep1:")
 
 
+@pytest.mark.asyncio
+async def test_build_add_request_message_scopes_mm_uuids_to_distributed_stage0_replica(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+
+    addr0 = "tcp://host-a:1000/input"
+    addr1 = "tcp://host-b:1000/input"
+    stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)])
+    stage_pool.attach_hub(_FakeHub([_replica(addr0), _replica(addr1)]))
+    stage_pool.attach_load_balancer(_RoundRobinLB())
+    engine.stage_pools = [stage_pool]
+
+    seen_uuids: list[str] = []
+
+    def process_inputs(**kwargs):
+        prompt = kwargs["prompt"]
+        seen_uuids.append(prompt["multi_modal_uuids"]["image"][0])
+        return _make_engine_core_request(kwargs["request_id"])
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = process_inputs
+    engine.input_processor = input_processor
+
+    for request_id in ("req-1", "req-2"):
+        engine._build_add_request_message(
+            request_id=request_id,
+            prompt={
+                "prompt": "describe",
+                "multi_modal_data": {"image": "same-image"},
+            },
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+    assert seen_uuids[0].startswith("stage0:rep0:")
+    assert seen_uuids[1].startswith("stage0:rep1:")
+    assert stage_pool.get_bound_replica_id("req-1") == 0
+    assert stage_pool.get_bound_replica_id("req-2") == 1
+    assert await stage_pool.pick("req-1") == 0
+    assert await stage_pool.pick("req-2") == 1
+
+
+def test_build_add_request_message_skips_distributed_mm_scope_when_no_replica(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+
+    addr0 = "tcp://host-a:1000/input"
+    addr1 = "tcp://host-b:1000/input"
+    stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)])
+    stage_pool.attach_hub(_FakeHub([]))
+    stage_pool.attach_load_balancer(_RoundRobinLB())
+    engine.stage_pools = [stage_pool]
+
+    seen_prompt: dict | None = None
+
+    def process_inputs(**kwargs):
+        nonlocal seen_prompt
+        seen_prompt = kwargs["prompt"]
+        return _make_engine_core_request(kwargs["request_id"])
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = process_inputs
+    engine.input_processor = input_processor
+
+    engine._build_add_request_message(
+        request_id="req-no-replica",
+        prompt={
+            "prompt": "describe",
+            "multi_modal_data": {"image": "same-image"},
+        },
+        sampling_params_list=[params],
+        final_stage_id=0,
+    )
+
+    assert seen_prompt is not None
+    assert "multi_modal_uuids" not in seen_prompt
+    assert stage_pool.get_bound_replica_id("req-no-replica") is None
+
+
 def test_stage_pool_replica_count_falls_back_to_clients():
     class PoolWithoutLiveNumReplicas:
         clients = [object(), None, object()]

tests/engine/test_async_omni_engine_stage_init.py

@@ -171,18 +171,6 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
     assert cleanup_clients == [shared, extra]
 
 
-def test_initialize_stages_rejects_non_diffusion_replicas_in_single_stage_mode():
-    engine = object.__new__(AsyncOmniEngine)
-    engine.single_stage_mode = True
-    engine.stage_configs = [types.SimpleNamespace(stage_id=0, runtime={"num_replicas": 2})]
-
-    with pytest.raises(
-        ValueError,
-        match="single_stage_mode only supports num_replicas > 1 for diffusion stages",
-    ):
-        engine._validate_single_stage_mode_replica_constraints()
-
-
 def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
     from vllm_omni.platforms import current_omni_platform
@@ -427,6 +415,7 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep
     engine.single_stage_mode = False
     engine._single_stage_id_filter = None
     engine._omni_master_server = None
+    engine._coordinator_runtime = None
     engine.stage_configs = [types.SimpleNamespace()]
 
     cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
@@ -471,6 +460,7 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
     engine.single_stage_mode = False
     engine._single_stage_id_filter = None
     engine._omni_master_server = None
+    engine._coordinator_runtime = None
     engine.stage_configs = [types.SimpleNamespace()]
 
     cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))

tests/engine/test_orchestrator.py

@@ -1052,7 +1052,7 @@ async def test_stage_pool_abort_requests_logs_when_binding_is_missing(caplog) ->
         target_logger.setLevel(prev_level)
 
     assert not stage0.abort_calls
-    assert "abort: no binding for req=missing-req in stage-0" in caplog.text
+    assert "abort: no live binding for req=missing-req in stage-0" in caplog.text
 
 
 @pytest.mark.asyncio

tests/engine/test_single_stage_mode.py

@@ -203,6 +203,32 @@ class TestOmniMasterServerAllocation:
         server = OmniMasterServer(master_address="127.0.0.1", master_port=15004, stage_ids=[3])
         assert server.get_allocation(3) is server._stage_routes[(3, 0)]
 
+    def test_next_free_replica_id_skips_head_local_slot_until_filled(self):
+        # Head pre-allocates slot (0, 0) for its own register_stage_with_omni_master
+        # call. A same-host headless that registers with auto-assign BEFORE the
+        # head's own registration must NOT be handed slot 0 — it should land on
+        # slot 1 instead. Without the head_local_replicas reservation,
+        # _next_free_replica_id would see (0, 0) absent from _stage_configs and
+        # return 0, colliding with the head's own bound sockets.
+        server = OmniMasterServer(
+            master_address="127.0.0.1",
+            master_port=15010,
+            stage_ids=[0],
+            head_local_replicas={0: [0]},
+        )
+        assert server._next_free_replica_id(0) == 1
+
+    def test_next_free_replica_id_uses_remote_slot_when_unowned(self):
+        # When the head pre-allocates a remote-only slot (the head's _initialize_*
+        # path waits on get_stage_config), auto-assign SHOULD fill it so the
+        # head's wait unblocks. This is the original behavior, preserved.
+        server = OmniMasterServer(
+            master_address="127.0.0.1",
+            master_port=15011,
+            stage_ids=[0],
+        )
+        assert server._next_free_replica_id(0) == 0
+
 
 # ---------------------------------------------------------------------------
 # OmniMasterServer registration flow
@@ -233,36 +259,6 @@ class TestOmniMasterServerRegistration:
             ctx.term()
             server.stop()
 
-    def test_server_handles_unknown_stage_id_gracefully(self):
-        import msgspec
-        import zmq
-        from vllm.utils.network_utils import get_open_port
-
-        master_port = get_open_port()
-        server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0])
-        server.start()
-
-        ctx = zmq.Context()
-        bad_sock = None
-        good_sock = None
-        try:
-            bad_sock = ctx.socket(zmq.DEALER)
-            bad_sock.connect(f"tcp://127.0.0.1:{master_port}")
-            bad_sock.send(msgspec.msgpack.encode({"stage_id": 99}))
-            assert not bad_sock.poll(timeout=500)
-
-            good_sock = ctx.socket(zmq.DEALER)
-            good_sock.connect(f"tcp://127.0.0.1:{master_port}")
-            good_sock.send(msgspec.msgpack.encode({"stage_id": 0}))
-            assert good_sock.poll(timeout=2_000)
-            good_sock.recv()
-        finally:
-            for sock in (bad_sock, good_sock):
-                if sock is not None:
-                    sock.close(linger=0)
-            ctx.term()
-            server.stop()
-
     def test_registration_stores_stage_config(self):
         import msgspec
         import zmq
@@ -474,6 +470,8 @@ class TestSingleStageInitialization:
         engine._omni_master_address = "127.0.0.1"
         engine._omni_master_port = 26000
         engine._omni_master_server = None
+        engine._omni_heartbeat_timeout = 30.0
+        engine._coordinator_runtime = None
         engine.async_chunk = False
         engine.diffusion_batch_size = 2
         return engine
@@ -508,10 +506,16 @@ class TestSingleStageInitialization:
 
     def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture):
         import vllm_omni.engine.async_omni_engine as engine_mod
+        from vllm_omni.distributed import omni_coordinator as omni_coord_mod
 
         engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7)
         mock_oms = mocker.Mock(spec=OmniMasterServer)
         mocker.patch.object(engine_mod, "OmniMasterServer", return_value=mock_oms)
+        mocker.patch.object(
+            omni_coord_mod,
+            "OmniCoordinatorRuntime",
+            return_value=mocker.Mock(router_address="tcp://127.0.0.1:9999"),
+        )
 
         stage_plans = [
             _make_llm_plan(0, configured_stage_id=7, launch_mode="local"),
@@ -520,12 +524,16 @@ class TestSingleStageInitialization:
 
         engine._start_omni_master_server(stage_plans)
 
-        engine_mod.OmniMasterServer.assert_called_once_with(
-            master_address="127.0.0.1",
-            master_port=26000,
-            stage_ids=[7, 11],
-            stage_replica_counts={7: 1, 11: 1},
-        )
+        call_kwargs = engine_mod.OmniMasterServer.call_args.kwargs
+        assert call_kwargs["master_address"] == "127.0.0.1"
+        assert call_kwargs["master_port"] == 26000
+        assert call_kwargs["stage_ids"] == [7, 11]
+        assert call_kwargs["stage_replica_counts"] == {7: 1, 11: 1}
+        # head_local_replicas reserves slots that the head will register
+        # itself (launch_mode == "local"). Stage 11 is remote, so it must
+        # NOT appear in the head-owned set — that slot is for the headless
+        # to fill via auto-assign.
+        assert call_kwargs["head_local_replicas"] == {7: [0]}
         mock_oms.start.assert_called_once()
 
     def test_start_omni_master_server_duplicate_stage_ids_raise(self):
@@ -545,7 +553,9 @@ class TestSingleStageInitialization:
         with pytest.raises(ValueError, match="requires both"):
             engine._start_omni_master_server([_make_llm_plan(0, configured_stage_id=7, launch_mode="local")])
 
-    def test_build_logical_stage_init_plans_clears_runtime_cfg_in_single_stage_mode(self, mocker: MockerFixture):
+    def test_build_logical_stage_init_plans_preserves_runtime_cfg_for_local_llm_in_single_stage_mode(
+        self, mocker: MockerFixture
+    ):
         import vllm_omni.engine.async_omni_engine as engine_mod
 
         engine = self._build_engine([_make_stage_cfg(7)], single_stage_mode=True, stage_id_filter=7)
@@ -574,7 +584,7 @@ class TestSingleStageInitialization:
         finally:
             monkeypatch.undo()
 
-        assert stage_plans[0].replicas[0].metadata.runtime_cfg is None
+        assert stage_plans[0].replicas[0].metadata.runtime_cfg == {"devices": "0"}
 
     def test_validate_single_stage_mode_allows_diffusion_replicas(self):
         stage_cfg = _make_stage_cfg(0, stage_type="diffusion")
@@ -583,14 +593,6 @@ class TestSingleStageInitialization:
 
         engine._validate_single_stage_mode_replica_constraints()
 
-    def test_validate_single_stage_mode_rejects_llm_replicas(self):
-        stage_cfg = _make_stage_cfg(0, stage_type="llm")
-        stage_cfg.runtime.num_replicas = 2
-        engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0)
-
-        with pytest.raises(ValueError, match="only supports num_replicas > 1 for diffusion"):
-            engine._validate_single_stage_mode_replica_constraints()
-
     def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_single_stage_mode(
         self, mocker: MockerFixture
     ):
@@ -809,6 +811,7 @@ class TestSingleStageReplicaInitialization:
         engine.model = "fake-model"
         engine.single_stage_mode = True
         engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
+        engine._coordinator_runtime = None
         engine.stage_configs = []
 
         fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace())
@@ -896,6 +899,7 @@ class TestSingleStageReplicaInitialization:
         engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
         engine._omni_master_server.address = "127.0.0.1"
         engine._omni_master_server.port = 25000
+        engine._coordinator_runtime = None
 
         plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0]
         sentinel_client = SimpleNamespace()
@@ -950,6 +954,9 @@ class TestSingleStageReplicaInitialization:
             handshake_address="tcp://hs",
             request_address="tcp://req",
             response_address="tcp://resp",
+            omni_coordinator_address=None,
+            omni_stage_id=5,
+            omni_replica_id=0,
         )
         mock_handshake.assert_called_once_with(proc, "tcp://hs", 60)
         mock_from_addresses.assert_called_once_with(
@@ -972,6 +979,7 @@ class TestSingleStageReplicaInitialization:
         engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
         engine._omni_master_server.address = "127.0.0.1"
         engine._omni_master_server.port = 25000
+        engine._coordinator_runtime = None
 
         plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0]
         proc = mocker.Mock()

tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py

@@ -5,6 +5,10 @@ E2E tests for speaker_embedding API parameter.
 
 Tests Base models (0.6B-Base and 1.7B-Base) which support voice cloning
 via pre-computed ECAPA-TDNN embeddings, bypassing ref_audio extraction.
+
+Invalid combinations (mutually exclusive with ref_audio, empty embedding,
+wrong embedding dimensions) live in
+``tests/dfx/reliability/http_invalid/test_invalid_audio_speech.py``.
 """
 
 import os
@@ -147,45 +151,6 @@ class TestSpeakerEmbeddingBase:
         assert len(response.content) > MIN_AUDIO_BYTES
         assert_not_silence(response.content)
 
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_speaker_embedding_mutually_exclusive_with_ref_audio(self, omni_server) -> None:
-        """Sending both speaker_embedding and ref_audio returns 400."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
-        payload = {
-            "model": omni_server.model,
-            "input": SYN_TEXT,
-            "task_type": "Base",
-            "speaker_embedding": DUMMY_EMBEDDING_1024,
-            "ref_audio": "https://example.com/audio.wav",
-            "response_format": "wav",
-        }
-        with httpx.Client(timeout=30.0) as client:
-            response = client.post(url, json=payload)
-
-        assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}"
-        assert "mutually exclusive" in response.text
-
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_speaker_embedding_empty_rejected(self, omni_server) -> None:
-        """Empty speaker_embedding list returns 400."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
-        payload = {
-            "model": omni_server.model,
-            "input": SYN_TEXT,
-            "task_type": "Base",
-            "speaker_embedding": [],
-            "response_format": "wav",
-        }
-        with httpx.Client(timeout=30.0) as client:
-            response = client.post(url, json=payload)
-
-        assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}"
-        assert "non-empty" in response.text
-
 
 # ── 1.7B-Base model tests (2048-dim embeddings) ──
 
@@ -216,26 +181,3 @@ class TestSpeakerEmbedding1_7B:
         assert response.headers.get("content-type") == "audio/wav"
         assert verify_wav_audio(response.content), "Response is not valid WAV"
         assert len(response.content) > MIN_AUDIO_BYTES, f"Audio too small: {len(response.content)} bytes"
-
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_1024_dim_on_1_7b_model_rejected(self, omni_server) -> None:
-        """1024-dim speaker_embedding on a 1.7B model (expects 2048) should fail gracefully."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
-        payload = {
-            "model": omni_server.model,
-            "input": SYN_TEXT,
-            "task_type": "Base",
-            "speaker_embedding": DUMMY_EMBEDDING_1024,
-            "x_vector_only_mode": True,
-            "response_format": "wav",
-            "max_new_tokens": MAX_NEW_TOKENS,
-        }
-        with httpx.Client(timeout=120.0) as client:
-            response = client.post(url, json=payload)
-
-        # Wrong dimensions should produce an error, not silent garbage.
-        assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}"
-        assert "speaker_embedding has 1024 dimensions" in response.text
-        assert "expected 2048" in response.text

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,5 +1,6 @@
 # tests/entrypoints/openai/test_serving_speech.py
 import asyncio
+import json
 import logging
 import os
 import struct
@@ -564,7 +565,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_success(self, client):
         """Upload a voice via speaker_embedding JSON."""
-        import json
 
         emb = [0.1] * 1024
         data = {
@@ -585,7 +585,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_appears_in_listing(self, client):
         """Embedding-uploaded voice appears in list with correct source."""
-        import json
 
         emb = [0.2] * 2048
         data = {
@@ -606,7 +605,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_and_audio_mutually_exclusive(self, client):
         """Providing both audio_sample and speaker_embedding returns 400."""
-        import json
 
         emb = [0.1] * 1024
         files = {"audio_sample": ("test.wav", b"fake", "audio/wav")}
@@ -632,7 +630,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_nan_rejected(self, client):
         """NaN values in speaker_embedding return 400."""
-        import json
 
         data = {
             "speaker_embedding": json.dumps([0.1] * 1023 + [float("nan")]),
@@ -930,7 +927,6 @@ class TestTTSMethods:
 
     def test_upload_voice_embedding_wrong_dims_rejected(self, speech_server):
         """Embedding uploads must match the loaded Qwen3-TTS model before being stored."""
-        import json
 
         speech_server._tts_model_type = "qwen3_tts"
         speech_server.engine_client.model_config = SimpleNamespace(

tests/entrypoints/openai_api/test_serving_speech_voxcpm.py

@@ -1,194 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""UTs for VoxCPM OpenAI speech serving behavior."""
-
-import asyncio
-from types import SimpleNamespace
-from unittest.mock import AsyncMock
-
-import pytest
-from pytest_mock import MockerFixture
-
-from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
-from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.fixture
-def voxcpm_server(mocker: MockerFixture):
-    mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
-    mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
-
-    mock_engine_client = mocker.MagicMock()
-    mock_engine_client.errored = False
-    mock_engine_client.model_config = mocker.MagicMock(model="OpenBMB/VoxCPM1.5")
-    mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
-    mock_engine_client.tts_batch_max_items = 32
-    mock_engine_client.generate = mocker.MagicMock(return_value="generator")
-    mock_engine_client.stage_configs = [
-        SimpleNamespace(
-            engine_args=SimpleNamespace(
-                model_stage="latent_generator",
-                model_arch="VoxCPMForConditionalGeneration",
-            ),
-            tts_args={},
-        ),
-        SimpleNamespace(
-            engine_args=SimpleNamespace(model_stage="vae"),
-            tts_args={},
-        ),
-    ]
-
-    mock_models = mocker.MagicMock()
-    mock_models.is_base_model.return_value = True
-
-    return OmniOpenAIServingSpeech(
-        engine_client=mock_engine_client,
-        models=mock_models,
-        request_logger=mocker.MagicMock(),
-    )
-
-
-@pytest.fixture
-def voxcpm2_server(mocker: MockerFixture):
-    mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
-
-    mock_engine_client = mocker.MagicMock()
-    mock_engine_client.errored = False
-    mock_engine_client.model_config = mocker.MagicMock(model="OpenBMB/VoxCPM2")
-    mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
-    mock_engine_client.tts_batch_max_items = 32
-    mock_engine_client.generate = mocker.MagicMock(return_value="generator")
-    mock_engine_client.stage_configs = [
-        SimpleNamespace(
-            engine_args=SimpleNamespace(
-                model_stage="latent_generator",
-                model_arch="VoxCPM2TalkerForConditionalGeneration",
-            ),
-            tts_args={},
-        ),
-        SimpleNamespace(
-            engine_args=SimpleNamespace(model_stage="vae"),
-            tts_args={},
-        ),
-    ]
-
-    mock_models = mocker.MagicMock()
-    mock_models.is_base_model.return_value = True
-
-    return OmniOpenAIServingSpeech(
-        engine_client=mock_engine_client,
-        models=mock_models,
-        request_logger=mocker.MagicMock(),
-    )
-
-
-class TestVoxCPMServing:
-    def test_voxcpm_model_type_detection(self, voxcpm_server):
-        assert voxcpm_server._tts_model_type == "voxcpm"
-        assert voxcpm_server._is_tts is True
-        assert voxcpm_server.supported_speakers == set()
-
-    @pytest.mark.parametrize(
-        ("request_kwargs", "expected_substring"),
-        [
-            ({"voice": "alice"}, "voice"),
-            ({"instructions": "whisper"}, "instructions"),
-            ({"language": "en"}, "language"),
-            ({"task_type": "CustomVoice"}, "plain tts"),
-            ({"x_vector_only_mode": True}, "x_vector_only_mode"),
-            ({"speaker_embedding": [0.1, 0.2]}, "speaker_embedding"),
-            ({"initial_codec_chunk_frames": 4}, "initial_codec_chunk_frames"),
-            ({"ref_text": "reference"}, "ref_audio"),
-        ],
-    )
-    def test_validate_voxcpm_rejects_unsupported_fields(self, voxcpm_server, request_kwargs, expected_substring):
-        request = OpenAICreateSpeechRequest(input="hello voxcpm", **request_kwargs)
-        error = voxcpm_server._validate_voxcpm_request(request)
-        assert error is not None
-        assert expected_substring in error.lower()
-
-    def test_validate_voxcpm_accepts_plain_tts_request(self, voxcpm_server):
-        request = OpenAICreateSpeechRequest(input="hello voxcpm", max_new_tokens=256)
-        assert voxcpm_server._validate_voxcpm_request(request) is None
-
-    def test_validate_voxcpm_accepts_voice_clone_request(self, voxcpm_server):
-        request = OpenAICreateSpeechRequest(
-            input="clone this voice",
-            ref_audio="data:audio/wav;base64,QUJD",
-            ref_text="reference transcript",
-            max_new_tokens=256,
-        )
-        assert voxcpm_server._validate_voxcpm_request(request) is None
-
-    def test_prepare_speech_generation_voxcpm_text_only(self, voxcpm_server):
-        request = OpenAICreateSpeechRequest(input="hello voxcpm", max_new_tokens=321)
-
-        request_id, generator, tts_params = asyncio.run(voxcpm_server._prepare_speech_generation(request))
-
-        assert request_id.startswith("speech-")
-        assert generator == "generator"
-        assert tts_params == {
-            "text": ["hello voxcpm"],
-            "cfg_value": [2.0],
-            "inference_timesteps": [10],
-            "min_len": [2],
-            "max_new_tokens": [321],
-        }
-
-        voxcpm_server.engine_client.generate.assert_called_once()
-        call = voxcpm_server.engine_client.generate.call_args
-        assert call.kwargs["prompt"] == {
-            "prompt_token_ids": [1],
-            "additional_information": tts_params,
-            "type": "token",
-        }
-        assert call.kwargs["output_modalities"] == ["audio"]
-
-    def test_prepare_speech_generation_voxcpm_voice_clone_resolves_ref_audio(self, voxcpm_server):
-        voxcpm_server._resolve_ref_audio = AsyncMock(return_value=([0.1, -0.1, 0.2], 16000))
-        request = OpenAICreateSpeechRequest(
-            input="clone this voice",
-            ref_audio="data:audio/wav;base64,QUJD",
-            ref_text="reference transcript",
-            max_new_tokens=512,
-        )
-
-        request_id, generator, tts_params = asyncio.run(voxcpm_server._prepare_speech_generation(request))
-
-        assert request_id.startswith("speech-")
-        assert generator == "generator"
-        assert tts_params == {
-            "text": ["clone this voice"],
-            "cfg_value": [2.0],
-            "inference_timesteps": [10],
-            "min_len": [2],
-            "max_new_tokens": [512],
-            "ref_text": ["reference transcript"],
-            "ref_audio": [[[0.1, -0.1, 0.2], 16000]],
-        }
-
-        voxcpm_server._resolve_ref_audio.assert_awaited_once_with("data:audio/wav;base64,QUJD")
-        call = voxcpm_server.engine_client.generate.call_args
-        assert call.kwargs["prompt"] == {
-            "prompt_token_ids": [1],
-            "additional_information": tts_params,
-            "type": "token",
-        }
-
-
-class TestVoxCPM2Serving:
-    """Regression tests for VoxCPM2 serving behavior."""
-
-    def test_voxcpm2_model_type_detection(self, voxcpm2_server):
-        assert voxcpm2_server._tts_model_type == "voxcpm2"
-        assert voxcpm2_server._is_tts is True
-
-    def test_voxcpm2_default_voice_is_supported(self, voxcpm2_server):
-        """VoxCPM2 should accept voice='default' (zero-shot mode, no voice cloning).
-
-        Regression: previously returned 400 Invalid voice 'default'.
-        """
-        assert "default" in voxcpm2_server.supported_speakers

tests/entrypoints/test_serve.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import argparse
+from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
@@ -31,143 +33,118 @@ def test_serve_parser_accepts_no_async_chunk() -> None:
     assert args.async_chunk is False
 
 
-def _make_headless_args() -> argparse.Namespace:
-    return argparse.Namespace(
+# ---------------------------------------------------------------------------
+# run_headless validation
+# ---------------------------------------------------------------------------
+
+
+def _make_headless_args(**overrides: Any) -> argparse.Namespace:
+    """Build an argparse.Namespace shaped like the headless CLI passes in.
+
+    Defaults pass every validation gate so individual tests can mutate just
+    the field they're exercising.
+    """
+    defaults = dict(
         model="fake-model",
-        stage_id=3,
+        stage_id=0,
         replica_id=0,
         omni_master_address="127.0.0.1",
         omni_master_port=26000,
+        omni_replica_address=None,
+        omni_dp_size_local=1,
         api_server_count=0,
         worker_backend="multi_process",
         stage_configs_path=None,
+        deploy_config=None,
         log_stats=False,
         disable_log_stats=False,
         stage_init_timeout=600,
+        tokenizer=None,
     )
+    defaults.update(overrides)
+    return argparse.Namespace(**defaults)
 
 
-def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    stage_cfg = mocker.Mock(stage_id=3)
-    stage_cfgs = [stage_cfg]
-    parallel_config = mocker.Mock(
-        data_parallel_size_local=2,
-        data_parallel_rank=4,
-        data_parallel_rank_local=1,
-        node_rank_within_dp=0,
-    )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    vllm_config.needs_dp_coordinator = False
-    executor_class = mocker.Mock()
-    engine_manager = mocker.Mock()
+def test_run_headless_requires_stage_id() -> None:
+    args = _make_headless_args(stage_id=None)
+    with pytest.raises(ValueError, match="--stage-id is required"):
+        run_headless(args)
 
+
+def test_run_headless_requires_master_address() -> None:
+    args = _make_headless_args(omni_master_address=None)
+    with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"):
+        run_headless(args)
+
+
+def test_run_headless_requires_master_port() -> None:
+    args = _make_headless_args(omni_master_port=None)
+    with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"):
+        run_headless(args)
+
+
+def test_run_headless_rejects_multi_api_server_count() -> None:
+    args = _make_headless_args(api_server_count=2)
+    with pytest.raises(ValueError, match="api_server_count can't be set"):
+        run_headless(args)
+
+
+def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
+    args = _make_headless_args(worker_backend="ray")
+    with pytest.raises(ValueError, match="worker_backend=multi_process"):
+        run_headless(args)
+
+
+def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None:
+    """Headless looks up its assigned stage_id in the loaded deploy YAML and
+    fails fast when the launcher's --stage-id doesn't match any entry."""
+    other_stage = SimpleNamespace(stage_id=99)
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
+        return_value=("/fake/stages.yaml", [other_stage]),
     )
-    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
-    mocker.patch(
-        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
-        return_value=(None, None, None),
-    )
-    mock_build_vllm_config = mocker.patch(
-        "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, executor_class),
-    )
-    mock_register = mocker.patch(
-        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
-    )
-    mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
-    mocker.patch("signal.signal")
-    run_headless(args)
 
-    mock_build_vllm_config.assert_called_once_with(
-        stage_cfg,
-        "fake-model",
-        stage_connector_spec={},
-        engine_args_dict={},
-        headless=True,
-    )
-    mock_register.assert_called_once_with(
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_stage_id=3,
-        omni_stage_config=stage_cfg,
-        coordinator=None,
-        replica_id=0,
-    )
-    mock_manager_cls.assert_called_once()
-    manager_kwargs = mock_manager_cls.call_args.kwargs
-    assert manager_kwargs["local_engine_count"] == 2
-    assert manager_kwargs["start_index"] == 4
-    assert manager_kwargs["local_start_index"] == 0
-    assert manager_kwargs["local_client"] is False
-    assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
-    assert manager_kwargs["log_stats"] is False
-    engine_manager.monitor_engine_liveness.assert_called_once_with()
-    engine_manager.shutdown.assert_called_once_with()
+    args = _make_headless_args(stage_id=0)
+    with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
+        run_headless(args)
 
 
-def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.log_stats = True
-    stage_cfg = mocker.Mock(stage_id=3)
-    stage_cfgs = [stage_cfg]
-    parallel_config = mocker.Mock(
-        data_parallel_size_local=2,
-        data_parallel_rank=4,
-        data_parallel_rank_local=1,
-        node_rank_within_dp=0,
-    )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    executor_class = mocker.Mock()
-    engine_manager = mocker.Mock()
-
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
-    )
-    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
-    mocker.patch(
-        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
-        return_value=(None, None, None),
-    )
-    mocker.patch(
-        "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, executor_class),
-    )
-    mocker.patch(
-        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
-    )
-    mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
-    mocker.patch("signal.signal")
-    run_headless(args)
-
-    manager_kwargs = mock_manager_cls.call_args.kwargs
-    assert manager_kwargs["log_stats"] is True
+# ---------------------------------------------------------------------------
+# run_headless happy paths
+# ---------------------------------------------------------------------------
 
 
-def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.replica_id = 2
-    stage_cfg = mocker.Mock(stage_id=3)
-    parallel_config = mocker.Mock(
+def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace:
+    """Build a stage config that satisfies every attribute run_headless reads.
+
+    Notably ``engine_args`` is a real dict (not a Mock) so
+    ``get_stage_devices_per_replica`` can call ``.get("tensor_parallel_size")``
+    and feed the result through ``int()`` without TypeError.
+    """
+    return SimpleNamespace(
+        stage_id=stage_id,
+        stage_type=stage_type,
+        # No "devices" key -> split_devices_for_replicas skipped, each replica
+        # inherits the launcher's CUDA_VISIBLE_DEVICES.
+        runtime=None,
+        engine_args={},
+    )
+
+
+def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: MockerFixture) -> None:
+    """LLM headless: each loop iteration registers with auto-assigned
+    replica_id (master picks a free slot) and spawns one
+    ``OmniCoreEngineProcManager`` per local replica."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(0, stage_type="llm")
+    parallel_config = SimpleNamespace(
         data_parallel_size_local=1,
         data_parallel_rank=0,
         data_parallel_rank_local=0,
         node_rank_within_dp=0,
     )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    vllm_config.needs_dp_coordinator = False
+    vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False)
     engine_manager = mocker.Mock()
 
     mocker.patch(
@@ -175,82 +152,250 @@ def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None:
         return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
     mocker.patch(
         "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
         return_value=(None, None, None),
     )
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
     mocker.patch(
         "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, mocker.Mock()),
+        return_value=(vllm_config, object),
     )
     mock_register = mocker.patch(
         "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address="tcp://127.0.0.1:26100",
+        ),
+    )
+    mock_manager_cls = mocker.patch(
+        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        return_value=engine_manager,
     )
-    mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
     mocker.patch("signal.signal")
 
-    run_headless(args)
+    run_headless(_make_headless_args(stage_id=0))
 
-    assert mock_register.call_args.kwargs["replica_id"] == 2
+    # The launcher must request auto-assignment (replica_id=None) and the
+    # full response so it can wire the master-allocated coordinator into the
+    # spawned subprocess. ``replica_binds_sockets=False`` is required for LLM
+    # because the head binds all three sockets (handshake, input, output).
+    assert mock_register.call_count == 1
+    kwargs = mock_register.call_args.kwargs
+    assert kwargs["omni_master_address"] == "127.0.0.1"
+    assert kwargs["omni_master_port"] == 26000
+    assert kwargs["omni_stage_id"] == 0
+    assert kwargs["omni_stage_config"] is stage_cfg
+    assert kwargs["replica_id"] is None
+    assert kwargs["return_full_response"] is True
+    assert kwargs["replica_binds_sockets"] is False
+
+    assert mock_manager_cls.call_count == 1
+    mgr_kwargs = mock_manager_cls.call_args.kwargs
+    assert mgr_kwargs["local_engine_count"] == 1
+    assert mgr_kwargs["local_client"] is False
+    assert mgr_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
+    assert mgr_kwargs["omni_stage_id"] == 0
+    assert mgr_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100"
+    assert mgr_kwargs["omni_replica_base_id"] == 0
+
+    engine_manager.monitor_engine_liveness.assert_called_once_with()
+    engine_manager.shutdown.assert_called_once_with()
 
 
-def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.replica_id = 1
-    stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion")
-    stage_cfg.engine_args = mocker.Mock()
-    stage_cfg.engine_input_source = []
-    stage_cfgs = [stage_cfg]
-    metadata = mocker.Mock(stage_id=3, stage_type="diffusion")
+def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: MockerFixture) -> None:
+    """``--omni-dp-size-local=N`` must spawn N managers, each with its own
+    master-assigned replica_id, and join all of them before returning."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(0, stage_type="llm")
+    parallel_config = SimpleNamespace(
+        data_parallel_size_local=1,
+        data_parallel_rank=0,
+        data_parallel_rank_local=0,
+        node_rank_within_dp=0,
+    )
+    vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False)
+    manager_a = mocker.Mock()
+    manager_b = mocker.Mock()
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        return_value=("/fake/stages.yaml", [stage_cfg]),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.build_vllm_config",
+        return_value=(vllm_config, object),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        side_effect=[
+            StageRegistrationResponse(
+                handshake_address=f"tcp://127.0.0.1:2700{idx}",
+                input_address=f"tcp://127.0.0.1:2710{idx}",
+                output_address=f"tcp://127.0.0.1:2720{idx}",
+                replica_id=idx,
+                coordinator_router_address=None,
+            )
+            for idx in (0, 1)
+        ],
+    )
+    mock_manager_cls = mocker.patch(
+        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        side_effect=[manager_a, manager_b],
+    )
+    mocker.patch("signal.signal")
+
+    run_headless(_make_headless_args(stage_id=0, omni_dp_size_local=2))
+
+    assert mock_manager_cls.call_count == 2
+    assigned_ids = [call.kwargs["omni_replica_base_id"] for call in mock_manager_cls.call_args_list]
+    assert assigned_ids == [0, 1]
+
+    # Multi-replica path joins the monitor threads instead of calling
+    # ``monitor_engine_liveness`` synchronously on the main thread, but every
+    # manager must still be shut down in the finally block.
+    manager_a.shutdown.assert_called_once_with()
+    manager_b.shutdown.assert_called_once_with()
+
+
+def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) -> None:
+    """Diffusion headless: registers as auto-assign, spawns a single
+    ``StageDiffusionProc`` per local replica, and waits for it via
+    ``multiprocessing.connection.wait``."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(1, stage_type="diffusion")
     od_config = mocker.Mock()
-    proc = mocker.Mock()
-    proc.exitcode = 0
+    proc = mocker.Mock(sentinel=object(), exitcode=0)
     proc.is_alive.return_value = False
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
     mocker.patch(
         "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
         return_value=(None, None, None),
     )
-    mocker.patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata)
-    mock_inject_stage_info = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
+    )
+    mock_inject = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
     mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config)
     mock_register = mocker.patch(
         "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address="tcp://127.0.0.1:26100",
+        ),
     )
     mock_spawn = mocker.patch(
         "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
-        return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        return_value=(proc, None, None, None),
     )
     mock_handshake = mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake")
+    # Replace the blocking wait with one that returns the only proc's sentinel
+    # immediately so the test does not hang.
+    mocker.patch(
+        "multiprocessing.connection.wait",
+        side_effect=lambda sentinels: [sentinels[0]],
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc")
     mocker.patch("signal.signal")
-    run_headless(args)
 
-    mock_inject_stage_info.assert_called_once_with(stage_cfg, 3)
-    mock_register.assert_called_once_with(
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_stage_id=3,
-        omni_stage_config=stage_cfg,
-        return_addresses=True,
-        replica_id=1,
-    )
-    mock_spawn.assert_called_once_with(
-        "fake-model",
-        od_config,
-        handshake_address="tcp://127.0.0.1:26001",
-        request_address="tcp://127.0.0.1:26002",
-        response_address="tcp://127.0.0.1:26003",
-    )
+    run_headless(_make_headless_args(stage_id=1))
+
+    mock_inject.assert_called_once()
+    assert mock_inject.call_args.args[0] is stage_cfg
+    assert mock_inject.call_args.args[1] == 1
+    assert mock_inject.call_args.args[2] == [stage_cfg]
+
+    reg_kwargs = mock_register.call_args.kwargs
+    assert reg_kwargs["omni_master_address"] == "127.0.0.1"
+    assert reg_kwargs["omni_master_port"] == 26000
+    assert reg_kwargs["omni_stage_id"] == 1
+    assert reg_kwargs["omni_stage_config"] is stage_cfg
+    assert reg_kwargs["replica_id"] is None
+    assert reg_kwargs["return_full_response"] is True
+
+    spawn_kwargs = mock_spawn.call_args.kwargs
+    assert spawn_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
+    assert spawn_kwargs["request_address"] == "tcp://127.0.0.1:26002"
+    assert spawn_kwargs["response_address"] == "tcp://127.0.0.1:26003"
+    assert spawn_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100"
+    assert spawn_kwargs["omni_stage_id"] == 1
+    assert spawn_kwargs["omni_replica_id"] == 0
+
     mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600)
-    proc.join.assert_called_once_with()
+
+
+def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None:
+    """A diffusion replica that exits with a non-zero code must surface as a
+    RuntimeError from ``run_headless`` (the head needs the signal to roll
+    back its own stage init)."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(1, stage_type="diffusion")
+    proc = mocker.Mock(sentinel=object(), exitcode=137, name="proc-stage1-rep0")
+    proc.is_alive.return_value = False
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        return_value=("/fake/stages.yaml", [stage_cfg]),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=mocker.Mock())
+    mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address=None,
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
+        return_value=(proc, None, None, None),
+    )
+    mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake")
+    mocker.patch(
+        "multiprocessing.connection.wait",
+        side_effect=lambda sentinels: [sentinels[0]],
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc")
+    mocker.patch("signal.signal")
+
+    with pytest.raises(RuntimeError, match=r"exited with code 137"):
+        run_headless(_make_headless_args(stage_id=1))

tests/entrypoints/test_utils.py

@@ -313,39 +313,6 @@ class TestResolveModelConfigPath:
         assert result is not None
         assert "glm_image.yaml" in result
 
-    def test_voxcpm_transformers_format_resolution(self, mocker: MockerFixture):
-        """Test VoxCPM transformers config resolves to the voxcpm stage config."""
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.get_config",
-            side_effect=ValueError("missing transformers config"),
-        )
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.file_or_path_exists",
-            side_effect=lambda _model, filename, revision=None: filename == "config.json",
-        )
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.get_hf_file_to_dict",
-            return_value={"model_type": "voxcpm"},
-        )
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path",
-            return_value="vllm_omni/deploy",
-        )
-
-        original_exists = os.path.exists
-
-        def mock_exists(path):
-            if "voxcpm.yaml" in str(path):
-                return True
-            return original_exists(path)
-
-        mocker.patch("os.path.exists", side_effect=mock_exists)
-
-        result = resolve_model_config_path("OpenBMB/VoxCPM1.5")
-
-        assert result is not None
-        assert "voxcpm.yaml" in result
-
 
 class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self):

tests/helpers/assertions.py

@@ -1,12 +1,13 @@
 """Assertion and response validation helpers for tests."""
 
 import io
+import json
 import tempfile
 import threading
 import wave
 from io import BytesIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 if TYPE_CHECKING:
     from tests.helpers.runtime import DiffusionResponse
@@ -495,62 +496,12 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                 )
 
 
-def _speech_combined_error_text_for_assert(response: Any) -> str:
-    """Join ``error_message`` and UTF-8 body (e.g. HTTP 200 + JSON) for substring checks."""
-    parts: list[str] = []
-    em = getattr(response, "error_message", None)
-    if em:
-        parts.append(str(em))
-    ab = getattr(response, "audio_bytes", None)
-    if ab:
-        parts.append(ab.decode("utf-8", errors="replace"))
-    return "\n".join(parts)
-
-
-def _assert_speech_error_substrings(haystack: str, err_expected: Any) -> None:
-    if isinstance(err_expected, (list, tuple)):
-        assert any(str(s) in haystack for s in err_expected), (
-            f"Expected error text to contain one of {err_expected!r}, got: {haystack!r}"
-        )
-    else:
-        assert str(err_expected) in haystack, f"Expected error text to contain {str(err_expected)!r}, got: {haystack!r}"
-
-
 def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None:
-    """Validate speech API results.
+    """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`.
 
-    If ``status_code`` is set: assert HTTP status, then optional ``err_message`` (4xx: ``error_message`` only;
-    HTTP 200: ``error_message`` and/or body, e.g. JSON with HTTP 200).
-
-    If only ``err_message`` is set: assert substring(s) on combined error text. Otherwise success/min_audio/etc.
+    Success path only. Use :meth:`OpenAIClientHandler.send_audio_speech_http_request` with ``err_code`` /
+    ``err_message`` for expected HTTP errors.
     """
-    # When present, only compare HTTP status and skip all other checks (e.g. expected 4xx).
-    # ``status_code`` may be a single int or a collection of allowed codes (e.g. ``(400, 422)``).
-    expected_status = request_config.get("status_code")
-    err_expected = request_config.get("err_message")
-    if expected_status is not None:
-        actual = getattr(response, "status_code", None)
-        assert actual is not None, (
-            "request_config has status_code but response has no .status_code "
-            "(set OmniResponse.status_code or pass a response object with status_code)"
-        )
-        if isinstance(expected_status, (list, tuple, set, frozenset)):
-            assert actual in expected_status, f"Expected HTTP status in {expected_status!r}, got {actual}"
-        else:
-            assert actual == int(expected_status), f"Expected HTTP status {int(expected_status)}, got {actual}"
-
-        if err_expected is not None:
-            if actual is not None and actual != 200:
-                msg = (getattr(response, "error_message", None) or "") + ""
-                _assert_speech_error_substrings(msg, err_expected)
-            else:
-                _assert_speech_error_substrings(_speech_combined_error_text_for_assert(response), err_expected)
-        return
-
-    if err_expected is not None:
-        _assert_speech_error_substrings(_speech_combined_error_text_for_assert(response), err_expected)
-        return
-
     assert response.success, "The request failed."
 
     # Optional floor on decoded audio size (models with very short clips may use a lower value).
@@ -599,10 +550,128 @@ def assert_diffusion_response(response: "DiffusionResponse", request_config: dic
         assert_audio_diffusion_response(response=response, request_config=request_config, run_level=run_level)
 
 
+def _http_response_body_materialize(resp: Any) -> tuple[bytes, dict[str, Any] | None]:
+    """Serialize ``HttpResponse``-like body to UTF-8 bytes and parse a JSON object when possible."""
+    jb = getattr(resp, "json_body", None)
+    if jb is not None:
+        raw = json.dumps(jb, ensure_ascii=False).encode("utf-8")
+        if isinstance(jb, dict):
+            return raw, jb
+        try:
+            parsed = json.loads(raw.decode("utf-8", errors="replace"))
+        except json.JSONDecodeError:
+            return raw, None
+        return raw, parsed if isinstance(parsed, dict) else None
+    err = getattr(resp, "error_message", None)
+    raw = (err or "").encode("utf-8", errors="replace")
+    try:
+        parsed = json.loads(raw.decode("utf-8", errors="replace"))
+    except json.JSONDecodeError:
+        return raw, None
+    return raw, parsed if isinstance(parsed, dict) else None
+
+
+def assert_err_message_in_text(
+    haystack: str,
+    err_message: str | tuple[str, ...] | list[str] | set[str] | frozenset[str],
+    *,
+    sequence_match: Literal["all", "any"] = "all",
+) -> None:
+    """Assert ``err_message`` appears in ``haystack`` (case-insensitive).
+
+    For non-string sequences: ``sequence_match='all'`` requires every substring (HTTP-style);
+    ``sequence_match='any'`` requires at least one (WebSocket first-frame JSON helpers).
+    """
+    hl = haystack.lower()
+    if isinstance(err_message, (list, tuple, set, frozenset)):
+        if sequence_match == "all":
+            missing = [s for s in err_message if str(s).lower() not in hl]
+            assert not missing, (
+                f"Expected error text to contain all of {err_message!r}; missing {missing!r}. haystack={haystack!r}"
+            )
+        else:
+            assert any(str(s).lower() in hl for s in err_message), (
+                f"Expected error text to contain one of {err_message!r}. haystack={haystack!r}"
+            )
+    else:
+        assert str(err_message).lower() in hl, f"Expected error text to contain {str(err_message)!r}, got: {haystack!r}"
+
+
+def assert_http_error(
+    resp: Any,
+    *,
+    err_code: int | tuple[int, ...] | list[int] | None = None,
+    err_message: str | tuple[str, ...] | list[str] | None = None,
+    websocket_json_message: bool = False,
+) -> dict[str, Any] | None:
+    """Validate a raw-HTTP :class:`~tests.helpers.runtime.HttpResponse`-like object.
+
+    Used by :class:`~tests.helpers.runtime.OpenAIClientHandler` ``send_*_http_request`` helpers when
+    ``request_config`` contains optional ``err_code`` and/or ``err_message``.
+
+    When ``websocket_json_message=True``, only ``json_body`` is checked (first JSON WebSocket text frame).
+    Tuple/list ``err_message`` then uses **any** substring match; HTTP mode still requires **all** pieces.
+
+    - ``err_code``: exact HTTP ``int``, or membership if a non-string sequence (e.g. ``(400, 422)``).
+      When ``err_code`` is set and the actual status is a client error in ``400..499`` other than ``404``,
+      the JSON body must include FastAPI ``detail`` and/or OpenAI-style ``error`` (2xx skips this).
+    - ``err_message``: substring match (case-insensitive) against serialized ``json_body`` and ``error_message``.
+      If ``err_message`` is a non-string sequence (``list`` / ``tuple`` / ``set`` / ``frozenset``), **every**
+      element must appear as a substring; a plain ``str`` still requires that single substring.
+    """
+    if websocket_json_message:
+        if err_code is not None:
+            raise ValueError("assert_http_error: err_code is incompatible with websocket_json_message=True")
+        jb_ws = getattr(resp, "json_body", None)
+        assert jb_ws is not None, resp
+        if err_message is None:
+            return jb_ws if isinstance(jb_ws, dict) else None
+        assert_err_message_in_text(
+            json.dumps(jb_ws, ensure_ascii=False),
+            err_message,
+            sequence_match="any",
+        )
+        return jb_ws if isinstance(jb_ws, dict) else None
+
+    if err_code is None and err_message is None:
+        return None
+
+    actual = getattr(resp, "status_code", None)
+    assert actual is not None, "response missing status_code"
+
+    if err_code is not None:
+        if isinstance(err_code, int):
+            assert actual == err_code, (resp, err_code)
+        else:
+            allowed = tuple(err_code)
+            assert actual in allowed, (resp, allowed)
+
+    body_bytes, payload = _http_response_body_materialize(resp)
+
+    if err_code is not None and actual is not None and 400 <= actual < 500 and actual != 404:
+        assert payload is not None, getattr(resp, "error_message", resp)
+        assert "detail" in payload or "error" in payload, payload
+
+    if err_message is not None:
+        pieces: list[str] = []
+        jb = getattr(resp, "json_body", None)
+        if jb is not None:
+            pieces.append(json.dumps(jb, ensure_ascii=False))
+        em = getattr(resp, "error_message", None)
+        if em:
+            pieces.append(str(em))
+        haystack = "\n".join(pieces) if pieces else body_bytes.decode("utf-8", errors="replace")
+        assert_err_message_in_text(haystack, err_message, sequence_match="all")
+
+    return payload
+
+
 __all__ = [
     "assert_audio_diffusion_response",
     "assert_audio_speech_response",
     "assert_diffusion_response",
+    "assert_err_message_in_text",
+    "assert_http_error",
     "assert_image_diffusion_response",
     "assert_image_valid",
     "assert_omni_response",

tests/helpers/runtime.py

@@ -1,5 +1,6 @@
 """Server/client/runner runtime primitives for tests."""
 
+import asyncio
 import base64
 import concurrent.futures
 import copy
@@ -19,6 +20,7 @@ from dataclasses import asdict, dataclass
 from io import BytesIO
 from pathlib import Path
 from typing import Any, NamedTuple
+from urllib.parse import quote
 
 import psutil
 import requests
@@ -37,6 +39,7 @@ from vllm.logger import init_logger
 from tests.helpers.assertions import (
     assert_audio_speech_response,
     assert_diffusion_response,
+    assert_http_error,
     assert_omni_response,
 )
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
@@ -86,6 +89,16 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
             logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
 
 
+def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None:
+    try:
+        data = r.json()
+        if isinstance(data, (dict, list)):
+            return data
+    except Exception:
+        pass
+    return None
+
+
 def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None:
     """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``)."""
     eb = cfg.get("extra_body")
@@ -647,6 +660,8 @@ class OmniServerStageCli(OmniServer):
 
 @dataclass
 class OmniResponse:
+    """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
+
     text_content: str | None = None
     audio_data: list[str] | None = None
     audio_content: str | None = None
@@ -656,16 +671,15 @@ class OmniResponse:
     #: OpenAI client call through response parsing and local post-process (e.g. audio decode).
     e2e_latency: float | None = None
     success: bool = False
-    error_message: str | None = None
     prompt_tokens: int | None = None
     cached_tokens: int | None = None
-    # Set for HTTP-level assertions (e.g. ``assert_audio_speech_response`` with ``status_code`` in config).
-    status_code: int | None = None
     logprobs: list | None = None
 
 
 @dataclass
 class DiffusionResponse:
+    """Decoded diffusion output from chat completions or offline runner (not raw ``requests``)."""
+
     text_content: str | None = None
     images: list[Image.Image] | None = None
     audios: list[Any] | None = None
@@ -674,7 +688,67 @@ class DiffusionResponse:
     #: ``chat.completions.create`` through local image / audio decode.
     e2e_latency: float | None = None
     success: bool = False
+
+
+@dataclass
+class HttpResponse:
+    """Normalized view of a ``requests`` response from :class:`OpenAIClientHandler` HTTP helpers."""
+
+    status_code: int
+    success: bool
     error_message: str | None = None
+    json_body: dict[str, Any] | list[Any] | None = None
+
+
+@dataclass
+class WebSocketJsonResponse:
+    """First JSON object delivered as a text WebSocket frame (streaming endpoints)."""
+
+    json_body: dict[str, Any]
+
+
+def _merge_http_expectation_kwargs(
+    base: dict[str, Any] | None,
+    *,
+    err_code: int | tuple[int, ...] | list[int] | None = None,
+    err_message: str | tuple[str, ...] | list[str] | None = None,
+) -> dict[str, Any]:
+    cfg = dict(base or {})
+    if err_code is not None:
+        cfg["err_code"] = err_code
+    if err_message is not None:
+        cfg["err_message"] = err_message
+    return cfg
+
+
+def _merge_ws_expectation_kwargs(
+    base: dict[str, Any] | None,
+    *,
+    err_message: str | tuple[str, ...] | list[str] | None = None,
+    ws_json_type: str | None = None,
+    ws_error_code: str | None = None,
+) -> dict[str, Any]:
+    cfg = dict(base or {})
+    if err_message is not None:
+        cfg["err_message"] = err_message
+    if ws_json_type is not None:
+        cfg["ws_json_type"] = ws_json_type
+    if ws_error_code is not None:
+        cfg["ws_error_code"] = ws_error_code
+    return cfg
+
+
+def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None:
+    jb = resp.json_body
+    want_type = cfg.get("ws_json_type")
+    if want_type is not None:
+        assert jb.get("type") == want_type, (jb, want_type)
+    want_code = cfg.get("ws_error_code")
+    if want_code is not None:
+        assert jb.get("code") == want_code, (jb, want_code)
+    err_message = cfg.get("err_message")
+    if err_message is not None:
+        assert_http_error(resp, err_message=err_message, websocket_json_message=True)
 
 
 def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse:
@@ -739,8 +813,8 @@ class OpenAIClientHandler:
             result.e2e_latency = time.perf_counter() - wall_start
             result.success = True
         except Exception as e:
-            result.error_message = f"Stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Stream processing error: {str(e)}"
+            print(f"Error: {msg}")
         return result
 
     def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
@@ -771,8 +845,8 @@ class OpenAIClientHandler:
                 result.logprobs = chat_completion.choices[0].logprobs.content
             result.success = True
         except Exception as e:
-            result.error_message = f"Non-stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Non-stream processing error: {str(e)}"
+            print(f"Error: {msg}")
         return result
 
     def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse:
@@ -811,11 +885,592 @@ class OpenAIClientHandler:
             result.e2e_latency = time.perf_counter() - wall_start
             result.success = True
         except Exception as e:
-            result.error_message = f"Diffusion response processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Diffusion response processing error: {str(e)}"
+            print(f"Error: {msg}")
         return result
 
+    def _http_response_from_requests(self, r: requests.Response) -> HttpResponse:
+        payload = _parse_response_json(r)
+        ok = 200 <= r.status_code < 300
+        return HttpResponse(
+            status_code=r.status_code,
+            success=ok,
+            error_message=None if ok else (r.text[:8000] if r.text else None),
+            json_body=payload,
+        )
+
+    def send_health_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/health`` (raw ``requests``).
+
+        ``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for
+        :func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args).
+        """
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0)))
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_models_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`)."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(
+            self._build_url("/v1/models"),
+            headers={"Accept": "application/json"},
+            timeout=float(cfg.get("timeout", 120.0)),
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_chat_completions_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests)."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_omni_sleep_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_omni_wakeup_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/omni/wakeup``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_voices_list_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/audio/voices``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(
+            self._build_url("/v1/audio/voices"),
+            headers={"Accept": "application/json"},
+            timeout=float(cfg.get("timeout", 120.0)),
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_voices_create_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_voices_delete_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        name = cfg["name"]
+        timeout = float(cfg.get("timeout", 120.0))
+        path = f"/v1/audio/voices/{quote(str(name), safe='')}"
+        r = requests.delete(
+            self._build_url(path),
+            headers={"Accept": "application/json"},
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_speech_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/speech`` with ``json`` or ``raw_body``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_speech_batch_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/speech/batch``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_generate_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/generate``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_images_generations_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_images_edits_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_videos_create_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_videos_sync_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/videos/sync``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_videos_list_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/videos`` — optional ``params``, ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(
+            self._build_url("/v1/videos"),
+            params=cfg.get("params"),
+            headers={"Accept": "application/json"},
+            timeout=float(cfg.get("timeout", 120.0)),
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_video_retrieve_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/videos/{video_id}``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        video_id = cfg["video_id"]
+        timeout = float(cfg.get("timeout", 120.0))
+        r = requests.get(
+            self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
+            headers={"Accept": "application/json"},
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_video_delete_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """DELETE ``/v1/videos/{video_id}``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        video_id = cfg["video_id"]
+        timeout = float(cfg.get("timeout", 120.0))
+        r = requests.delete(
+            self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
+            headers={"Accept": "application/json"},
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_video_content_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/videos/{video_id}/content``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        video_id = cfg["video_id"]
+        timeout = float(cfg.get("timeout", 120.0))
+        r = requests.get(
+            self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"),
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def _build_ws_url(self, path: str) -> str:
+        """Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers."""
+        base = self.base_url.rstrip("/")
+        suffix = "/" + path.lstrip("/")
+        if base.startswith("http://"):
+            return "ws://" + base.removeprefix("http://") + suffix
+        if base.startswith("https://"):
+            return "wss://" + base.removeprefix("https://") + suffix
+        raise ValueError(f"Unsupported base_url for WebSocket: {base!r}")
+
+    def _send_websocket_first_json_request(
+        self,
+        path: str,
+        cfg: dict[str, Any],
+    ) -> list[WebSocketJsonResponse]:
+        """Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`.
+
+        ``request_config`` keys:
+
+        - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server
+          speaks first, e.g. ``/v1/realtime`` rejection path).
+        - ``timeout``: seconds to wait for the first inbound text frame (default ``120``).
+        - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present.
+        """
+        send_frames_raw = cfg.get("send_frames")
+        if send_frames_raw is None:
+            frames: list[str] = []
+        elif isinstance(send_frames_raw, str):
+            frames = [send_frames_raw]
+        else:
+            frames = list(send_frames_raw)
+
+        timeout = float(cfg.get("timeout", 120.0))
+        uri = self._build_ws_url(path)
+
+        connect_kw: dict[str, Any] = {}
+        if "ws_max_size" in cfg:
+            connect_kw["max_size"] = cfg["ws_max_size"]
+
+        async def _recv_first_json_object() -> WebSocketJsonResponse:
+            import websockets
+
+            async with websockets.connect(uri, **connect_kw) as ws:
+                for frame in frames:
+                    await ws.send(frame)
+                raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
+                if not isinstance(raw, str):
+                    raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
+                try:
+                    data = json.loads(raw)
+                except json.JSONDecodeError as exc:
+                    raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
+                if not isinstance(data, dict):
+                    raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
+                return WebSocketJsonResponse(json_body=data)
+
+        resp = asyncio.run(_recv_first_json_object())
+        _run_ws_expectations_from_request_config(cfg, resp)
+        return [resp]
+
+    def send_audio_speech_stream_ws_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+        ws_json_type: str | None = None,
+        ws_error_code: str | None = None,
+    ) -> list[WebSocketJsonResponse]:
+        """WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame."""
+        cfg = _merge_ws_expectation_kwargs(
+            request_config,
+            err_message=err_message,
+            ws_json_type=ws_json_type,
+            ws_error_code=ws_error_code,
+        )
+        return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg)
+
+    def send_video_chat_stream_ws_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+        ws_json_type: str | None = None,
+        ws_error_code: str | None = None,
+    ) -> list[WebSocketJsonResponse]:
+        """WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame."""
+        cfg = _merge_ws_expectation_kwargs(
+            request_config,
+            err_message=err_message,
+            ws_json_type=ws_json_type,
+            ws_error_code=ws_error_code,
+        )
+        return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg)
+
+    def send_realtime_ws_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+        ws_json_type: str | None = None,
+        ws_error_code: str | None = None,
+    ) -> list[WebSocketJsonResponse]:
+        """WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated)."""
+        cfg = _merge_ws_expectation_kwargs(
+            request_config,
+            err_message=err_message,
+            ws_json_type=ws_json_type,
+            ws_error_code=ws_error_code,
+        )
+        return self._send_websocket_first_json_request("/v1/realtime", cfg)
+
     def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+        """Chat completions via the OpenAI Python SDK (not raw HTTP)."""
         responses: list[OmniResponse] = []
         stream = request_config.get("stream", False)
         modalities = request_config.get("modalities", ["text", "audio"])
@@ -939,8 +1594,8 @@ class OpenAIClientHandler:
                 result.audio_format = result.audio_format.headers.get("content-type", "")
 
         except Exception as e:
-            result.error_message = f"Audio speech stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Audio speech stream processing error: {str(e)}"
+            print(f"Error: {msg}")
 
         return result
 
@@ -978,8 +1633,8 @@ class OpenAIClientHandler:
                 result.audio_format = result.audio_format.headers.get("content-type", "")
 
         except Exception as e:
-            result.error_message = f"Audio speech non-stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Audio speech non-stream processing error: {str(e)}"
+            print(f"Error: {msg}")
 
         return result
 
@@ -1201,6 +1856,7 @@ class OpenAIClientHandler:
         Send OpenAI requests for diffusion models.
         If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel
         (scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order.
+
         Args:
             request_config: A single request configuration dict, or a list of
                 request configuration dicts (one request per element)
@@ -1300,10 +1956,13 @@ class OpenAIClientHandler:
         self, request_config: dict[str, Any], request_num: int = 1
     ) -> list[DiffusionResponse]:
         """
-        Send native /v1/videos requests.
+        Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content.
+
+        For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc.
         """
         if request_num != 1:
             raise NotImplementedError("Concurrent video diffusion requests are not currently implemented")
+
         form_data = request_config.get("form_data")
         if not isinstance(form_data, dict):
             raise ValueError("Video request_config must contain 'form_data'")
@@ -1346,6 +2005,54 @@ class OpenAIClientHandler:
             self._print_client_stat("[diffusion] request#1 completed")
         return [result]
 
+    def _post_json_endpoint(
+        self,
+        path: str,
+        request_config: dict[str, Any],
+        *,
+        default_timeout: float,
+    ) -> requests.Response:
+        url = self._build_url(path)
+        timeout = float(request_config.get("timeout", default_timeout))
+        if "raw_body" in request_config:
+            raw = request_config["raw_body"]
+            payload = raw.encode("utf-8") if isinstance(raw, str) else raw
+            return requests.post(
+                url,
+                data=payload,
+                headers={"Content-Type": "application/json", "Accept": "application/json"},
+                timeout=timeout,
+            )
+        if "json" not in request_config:
+            raise ValueError(f"{path} request_config must include 'json' or 'raw_body'")
+        return requests.post(
+            url,
+            json=request_config["json"],
+            headers={"Accept": "application/json", "Content-Type": "application/json"},
+            timeout=timeout,
+        )
+
+    def _post_form_endpoint(
+        self,
+        path: str,
+        request_config: dict[str, Any],
+        *,
+        default_timeout: float = 120.0,
+    ) -> requests.Response:
+        url = self._build_url(path)
+        timeout = float(request_config.get("timeout", default_timeout))
+        data = request_config.get("data")
+        files = request_config.get("files")
+        if data is None and not files:
+            data = {}
+        return requests.post(
+            url,
+            data=data,
+            files=files,
+            headers={"Accept": "application/json"} if not files else {"Accept": "application/json"},
+            timeout=timeout,
+        )
+
     def _wait_until_video_completed(
         self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300
     ) -> None:
@@ -1693,9 +2400,9 @@ class OmniRunnerHandler:
             result.text_content = text_content
             result.success = True
         except Exception as e:
-            result.error_message = f"Output processing error: {str(e)}"
+            msg = f"Output processing error: {str(e)}"
             result.success = False
-            print(f"Error: {result.error_message}")
+            print(f"Error: {msg}")
         return result
 
     def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse:
@@ -1831,15 +2538,11 @@ class OmniRunnerHandler:
                 mm_out = stage_out.request_output.outputs[0].multimodal_output
                 break
         if mm_out is None:
-            result = OmniResponse(success=False, error_message="No audio output from pipeline")
-            assert result.success, result.error_message
-            return result
+            raise AssertionError("No audio output from pipeline")
 
         audio_data = mm_out.get("audio")
         if audio_data is None:
-            result = OmniResponse(success=False, error_message="No audio tensor in multimodal output")
-            assert result.success, result.error_message
-            return result
+            raise AssertionError("No audio tensor in multimodal output")
 
         sr_raw = mm_out.get("sr")
         sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
@@ -2004,6 +2707,8 @@ def iter_omni_runner(
 
 __all__ = [
     "DiffusionResponse",
+    "HttpResponse",
+    "WebSocketJsonResponse",
     "OmniResponse",
     "OmniRunner",
     "OmniRunnerHandler",

tests/model_executor/models/qwen2_5_omni/test_weight_loading.py

@@ -0,0 +1,49 @@
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestQwen2_5OmniWeightLoading:
+    def test_qkv_weight_mapping(self):
+        mock_quant_description = {
+            "visual.blocks.0.attn.q.weight": {"quant_type": "FLOAT"},
+            "visual.blocks.0.attn.k.weight": {"quant_type": "FLOAT"},
+            "visual.blocks.0.attn.v.weight": {"quant_type": "FLOAT"},
+        }
+
+        prefix = "visual.blocks.0.attn.qkv"
+        packed_modules_mapping = {
+            "attn_akv_proj": [
+                "attn_q_proj",
+                "attn_k_proj",
+                "attn_v_proj",
+            ],
+            "qkv": [
+                "q",
+                "k",
+                "v",
+            ],
+        }
+
+        try:
+            proj_name = prefix.split(".")[-1]
+            if proj_name in packed_modules_mapping:
+                quant_type = None
+                shard_prefixes = [
+                    prefix.replace(proj_name, shard_proj_name) for shard_proj_name in packed_modules_mapping[proj_name]
+                ]
+                for shard_prefix in shard_prefixes:
+                    shard_quant_type = mock_quant_description[shard_prefix + ".weight"]
+
+                    if quant_type is None:
+                        quant_type = shard_quant_type
+                    elif shard_quant_type != quant_type:
+                        raise ValueError(
+                            f"Not all shards of {prefix} are quantized with same quant type."
+                            f"Shard {proj_name} uses {shard_quant_type}, but another shard"
+                            f"use {quant_type}. Please check quantization config."
+                        )
+            else:
+                quant_type = mock_quant_description[prefix + ".weight"]
+        except KeyError as e:
+            pytest.fail(f"KeyError was raised: {e}\n")

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py

@@ -74,6 +74,14 @@ class SyntheticDecoder(nn.Module):
         return self.out(x).clamp(min=-1, max=1)
 
 
+class ShortOutputDecoder(SyntheticDecoder):
+    """Decoder variant that returns fewer samples than seq_len * total_upsample."""
+
+    def forward(self, codes):
+        out = super().forward(codes)
+        return out[..., :-5] if out.shape[-1] > 5 else out
+
+
 @pytest.fixture(scope="module")
 def decoder():
     """Create a synthetic decoder on CUDA with fixed weights."""
@@ -228,6 +236,35 @@ def test_chunked_decode_output_survives_later_replay(wrapper):
     torch.testing.assert_close(graph_out, expected, atol=0, rtol=0)
 
 
+def test_chunked_decode_preserves_short_chunk_concat_semantics():
+    """Chunked decode must allow shorter-than-nominal chunk outputs.
+
+    Qwen3-Omni non-async startup can hit a decoder path where a chunk returns
+    fewer waveform samples than ``code_len * total_upsample``. The wrapper must
+    preserve the original eager concat behavior instead of copying into exact
+    nominal slices.
+    """
+    torch.manual_seed(123)
+    short_decoder = ShortOutputDecoder().to(DEVICE).eval()
+    short_wrapper = CUDAGraphDecoderWrapper(
+        decoder=short_decoder,
+        capture_sizes=[50],
+        capture_batch_sizes=[1],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    short_wrapper.warmup(DEVICE)
+    codes = _random_codes(100)
+
+    with torch.no_grad():
+        eager_out = _eager_chunked(short_decoder, codes, chunk_size=50, left_context_size=0)
+        graph_out = short_wrapper.chunked_decode_with_cudagraph(codes, chunk_size=50, left_context_size=0)
+
+    assert graph_out.shape == eager_out.shape
+    assert graph_out.shape[-1] < 100 * TOTAL_UPSAMPLE
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
 def test_batched_chunked_decode_variable_lengths_matches_per_request_eager(decoder, wrapper):
     """Variable-length chunk batching should match independent chunked decodes."""
     long_codes = _random_codes(100)

tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py

@@ -1,85 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""UTs for VoxCPM async-chunk stage input processing."""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.model_executor.stage_input_processors.voxcpm import (
-    _VOXCPM_LATENT_MAGIC,
-    _coerce_finished_flag,
-    latent2vae_async_chunk,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _request(*, finished):
-    return SimpleNamespace(is_finished=lambda: finished)
-
-
-def _decode_serialized_latent(codes: list[int]) -> torch.Tensor:
-    assert codes[0] == _VOXCPM_LATENT_MAGIC
-    latent_dim = codes[1]
-    time_dim = codes[2]
-    payload = torch.tensor(codes[3:], dtype=torch.int32).to(torch.uint16)
-    return payload.view(torch.bfloat16).to(torch.float32).reshape(1, latent_dim, time_dim)
-
-
-@pytest.mark.parametrize(
-    ("value", "expected"),
-    [
-        (None, False),
-        (False, False),
-        (True, True),
-        (torch.tensor(False), False),
-        (torch.tensor(True), True),
-        ([torch.tensor(True)], True),
-        (([True],), True),
-        ([], False),
-    ],
-)
-def test_coerce_finished_flag(value, expected):
-    assert _coerce_finished_flag(value) is expected
-
-
-def test_latent2vae_async_chunk_serializes_latent_payload():
-    latent = torch.arange(6, dtype=torch.float32).reshape(2, 3)
-
-    payload = latent2vae_async_chunk(
-        transfer_manager=None,
-        pooling_output={"latent_audio_feat": latent},
-        request=_request(finished=False),
-        is_finished=torch.tensor(False),
-    )
-
-    assert payload is not None
-    assert torch.equal(payload.meta.finished, torch.tensor(False, dtype=torch.bool))
-    recovered = _decode_serialized_latent(payload.codes.audio.tolist())
-    torch.testing.assert_close(recovered, latent.to(torch.bfloat16).to(torch.float32).unsqueeze(0))
-
-
-def test_latent2vae_async_chunk_returns_terminal_marker_without_latent():
-    payload = latent2vae_async_chunk(
-        transfer_manager=None,
-        pooling_output=None,
-        request=_request(finished=[torch.tensor(True)]),
-        is_finished=False,
-    )
-
-    assert payload.codes.audio.tolist() == []
-    assert torch.equal(payload.meta.finished, torch.tensor(True, dtype=torch.bool))
-
-
-def test_latent2vae_async_chunk_returns_none_for_nonterminal_empty_chunk():
-    payload = latent2vae_async_chunk(
-        transfer_manager=None,
-        pooling_output={"latent_audio_feat": torch.zeros((0,), dtype=torch.float32)},
-        request=_request(finished=False),
-        is_finished=False,
-    )
-
-    assert payload is None

tests/test_generate_nightly_perf_excel.py

@@ -35,6 +35,8 @@ def test_generate_excel_report_with_perf_templates(tmp_path: Path):
 
     omni_record = json.loads(omni_template_path.read_text(encoding="utf-8"))
     diffusion_records = json.loads(diffusion_template_path.read_text(encoding="utf-8"))
+    diffusion_records[0]["endpoint"] = "/v1/videos"
+    diffusion_records[0]["result"]["endpoint"] = "/v1/videos"
 
     input_dir = tmp_path / "input"
     diffusion_input_dir = tmp_path / "diffusion_input"
@@ -69,3 +71,6 @@ def test_generate_excel_report_with_perf_templates(tmp_path: Path):
     ws_omni_summary = wb["omni_summary"]
     assert _cell_value_by_header(ws_omni_summary, "commit_sha") == "test_commit_sha"
     assert _cell_value_by_header(ws_omni_summary, "build_id") == "test_build_id"
+
+    ws_diff_summary = wb["diffusion_summary"]
+    assert _cell_value_by_header(ws_diff_summary, "endpoint") == "/v1/videos"

tests/test_generate_nightly_perf_html.py

@@ -28,6 +28,8 @@ def test_generate_html_report_with_perf_templates(tmp_path: Path):
 
     omni_record = json.loads(omni_template_path.read_text(encoding="utf-8"))
     diffusion_records = json.loads(diffusion_template_path.read_text(encoding="utf-8"))
+    diffusion_records[0]["endpoint"] = "/v1/videos"
+    diffusion_records[0]["result"]["endpoint"] = "/v1/videos"
 
     input_dir = tmp_path / "input"
     diffusion_input_dir = tmp_path / "diffusion_input"
@@ -52,3 +54,4 @@ def test_generate_html_report_with_perf_templates(tmp_path: Path):
     assert "Omni records <strong>1</strong>" in html
     assert f"Diffusion records <strong>{len(diffusion_records)}</strong>" in html
     assert "const DIFF_DATA =" in html
+    assert '"endpoint": "/v1/videos"' in html