Update CUDA Docker base image to vLLM v0.21.0

Align Dockerfile.cuda with the v0.21.0 rebase on main and other
GPU Dockerfiles/docs. Related to #3858.

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/test-merge.yml

@@ -248,6 +248,47 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "TTS · GLM-TTS Test"
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              pytest -s -v tests/e2e/online_serving/test_glm_tts.py tests/e2e/offline_inference/test_glm_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
       - label: "Omni · Qwen3-Omni Test"
         timeout_in_minutes: 30

.buildkite/test-nightly.yml

@@ -959,7 +959,7 @@ steps:
       - label: ":full_moon: Diffusion X2V · Function Test"
         timeout_in_minutes: 90
         commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py  -m "full_model and cuda" --run-level "full_model"
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -1086,9 +1086,8 @@ steps:
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
     steps:
       - label: ":full_moon: Quantization Test with H100"
-        timeout_in_minutes: 90
         commands:
-          - timeout 30m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
+          - timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -1196,7 +1195,7 @@ steps:
       - nightly-diffusion-x2iat-performance-qwen-image-layered
       - nightly-diffusion-x2v-performance
       - nightly-testcase-statistics
-    if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
     commands:
       - pip install openpyxl
       - export DEFAULT_INPUT_DIR=tests/dfx/perf/results

benchmarks/diffusion/backends.py

@@ -10,6 +10,8 @@ from typing import Any
 import aiohttp
 from tqdm import tqdm
 
+DEFAULT_EDITS_BOT_TASK = "think"
+
 
 @dataclass
 class RequestFuncInput:
@@ -27,6 +29,7 @@ class RequestFuncInput:
     extra_body: dict[str, Any] = field(default_factory=dict)
     image_paths: list[str] | None = None
     request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
 
 
 @dataclass
@@ -53,6 +56,95 @@ def _encode_image_as_data_url(path: str) -> str:
     return f"data:{mime};base64,{encoded}"
 
 
+async def async_request_image_edits(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+    enable_diffusion_pipeline_profiler: bool = False,
+) -> RequestFuncOutput:
+    """POST /v1/images/edits (multipart)."""
+    del enable_diffusion_pipeline_profiler
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    extra_body = dict(input.extra_body)
+    width = input.width or extra_body.get("width") or 1024
+    height = input.height or extra_body.get("height") or 1024
+    edits_url = input.api_url
+
+    form = aiohttp.FormData()
+    form.add_field("model", input.model)
+    form.add_field("prompt", input.prompt)
+    form.add_field("size", f"{width}x{height}")
+    form.add_field("response_format", "b64_json")
+
+    if input.num_inference_steps is not None:
+        form.add_field("num_inference_steps", str(input.num_inference_steps))
+    elif extra_body.get("num_inference_steps") is not None:
+        form.add_field("num_inference_steps", str(extra_body["num_inference_steps"]))
+
+    if input.seed is not None:
+        form.add_field("seed", str(input.seed))
+    elif extra_body.get("seed") is not None:
+        form.add_field("seed", str(extra_body["seed"]))
+
+    if extra_body.get("guidance_scale") is not None:
+        form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
+    if extra_body.get("negative_prompt") is not None:
+        form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
+    if extra_body.get("true_cfg_scale") is not None:
+        form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"]))
+    if extra_body.get("sys_type") is not None:
+        form.add_field("sys_type", str(extra_body["sys_type"]))
+    if extra_body.get("system_prompt") is not None:
+        form.add_field("system_prompt", str(extra_body["system_prompt"]))
+
+    bot_task = extra_body.get("bot_task")
+    if bot_task is None and input.default_bot_task is not None:
+        bot_task = input.default_bot_task
+    if bot_task is not None:
+        form.add_field("bot_task", str(bot_task))
+
+    assert input.image_paths is not None
+    for img_path in input.image_paths:
+        if not os.path.exists(img_path):
+            output.error = f"Image file not found: {img_path}"
+            output.success = False
+            if pbar:
+                pbar.update(1)
+            return output
+        with open(img_path, "rb") as img_f:
+            image_bytes = img_f.read()
+        form.add_field(
+            "image",
+            image_bytes,
+            filename=os.path.basename(img_path),
+            content_type=_guess_mime_type(img_path),
+        )
+
+    try:
+        async with session.post(edits_url, data=form) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
 async def async_request_chat_completions(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
@@ -378,6 +470,7 @@ backends_function_mapping = {
     "2i": {
         "/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
         "/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
+        "/v1/images/edits": (async_request_image_edits, "/v1/images/edits"),
     },
     "2v": {
         "/v1/videos": (async_request_v1_videos, "/v1/videos"),

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -7,10 +7,13 @@ Benchmark online serving for diffusion models (Image/Video Generation).
 If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 
 Supports multiple endpoints:
-    - /v1/chat/completions: OpenAI chat-compatible image requests
+    - /v1/chat/completions: OpenAI chat-compatible image requests (e.g. t2i, Qwen i2i)
+    - /v1/images/edits: OpenAI image edit / IT2I (multipart; e.g. Hunyuan --bot-task think)
     - /v1/images/generations: OpenAI image generation requests
     - /v1/videos: Async video jobs
 
+Legacy --backend vllm-omni and openai are aliases for chat/completions and images/generations.
+
 Usage:
     # Video (/v1/videos endpoint)
     t2v:
@@ -39,7 +42,12 @@ Usage:
             {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
         ]'
 
-    i2i:
+    ti2i (Hunyuan / OpenAI image edit API):
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
+        --bot-task think
+
+    i2i (chat-based models such as Qwen-Image-Edit):
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
 
@@ -132,7 +140,7 @@ class VBenchDataset(BaseDataset):
     def _load_data(self) -> list[dict[str, Any]]:
         if self.args.task == "t2v":
             return self._load_t2v_prompts()
-        elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
+        elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
             return self._load_i2v_data()
         else:
             return self._load_t2v_prompts()
@@ -524,7 +532,7 @@ class TraceDataset(BaseDataset):
             single = row.get("image_path")
             image_paths = [single] if single else None
 
-        if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
+        if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
             raise ValueError(
                 f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
             )
@@ -585,7 +593,7 @@ class RandomDataset(BaseDataset):
             self._sampled_requests = None
 
         # Random image generate
-        if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
+        if self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
             self._random_image_path = self._generate_random_image_paths()
         else:
             self._random_image_path = None
@@ -891,7 +899,9 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
 def _default_endpoint_for_task(task: str) -> str:
     if task in {"t2v", "i2v", "ti2v"}:
         return "/v1/videos"
-    if task in {"t2i", "i2i", "ti2i"}:
+    if task in {"i2i", "ti2i", "it2i"}:
+        return "/v1/images/edits"
+    if task == "t2i":
         return "/v1/chat/completions"
     raise ValueError(f"Unsupported task for endpoint resolution: {task}")
 
@@ -902,7 +912,7 @@ async def benchmark(args):
         args.base_url = f"http://{args.host}:{args.port}"
 
     VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
-    IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
+    IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
 
     if args.task in VIDEO_TASKS:
         task_type = "2v"
@@ -947,6 +957,10 @@ async def benchmark(args):
     requests_list = dataset.get_requests()
     print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
 
+    if args.endpoint == "/v1/images/edits":
+        for req in requests_list:
+            req.default_bot_task = args.bot_task
+
     # Limit concurrency
     if args.max_concurrency is not None:
         semaphore = asyncio.Semaphore(args.max_concurrency)
@@ -998,6 +1012,8 @@ async def benchmark(args):
     metrics["model"] = args.model
     metrics["dataset"] = args.dataset
     metrics["task"] = args.task
+    if args.endpoint == "/v1/images/edits":
+        metrics["bot_task"] = args.bot_task
 
     print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
 
@@ -1088,7 +1104,7 @@ if __name__ == "__main__":
         "--task",
         type=str,
         default="t2v",
-        choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "t2i"],
+        choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
         help="Task type.",
     )
     parser.add_argument(
@@ -1195,9 +1211,15 @@ if __name__ == "__main__":
         default=1,
         help=(
             "Number of synthetic input images to attach for image-conditioned tasks "
-            "(i2v, ti2v, ti2i, i2i) when using random dataset."
+            "(i2v, ti2v, ti2i, i2i, it2i) when using random dataset."
         ),
     )
+    parser.add_argument(
+        "--bot-task",
+        type=str,
+        default="think",
+        help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."),
+    )
 
     args = parser.parse_args()
     asyncio.run(benchmark(args))

benchmarks/diffusion/quantization_quality.py

@@ -34,6 +34,16 @@ Video example (text-to-video):
         --height 720 --width 1280 \
         --num-frames 81 --num-inference-steps 40 --seed 42
 
+LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Lightricks/LTX-Video-2 \
+        --task t2v \
+        --quantization fp8 int8 \
+        --prompts \
+            "A serene lakeside sunrise with mist over the water" \
+        --height 704 --width 1216 \
+        --num-frames 121 --num-inference-steps 40 --seed 42
+
 Multiple quantization methods:
     python benchmarks/diffusion/quantization_quality.py \
         --model Tongyi-MAI/Z-Image-Turbo \
@@ -204,10 +214,19 @@ def _generate_video(omni, args, prompt, seed):
         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 (and similar audio+video models) may surface a dict or (video, audio) tuple
+    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:
@@ -225,9 +244,11 @@ def _generate_video(omni, args, prompt, seed):
     return frames_array, elapsed, peak_mem
 
 
-def _unload_omni(omni):
-    """Delete Omni instance and free GPU memory."""
-    del omni
+def _free_gpu_memory():
+    """Force GC and release cached GPU memory.
+
+    Must be called AFTER the caller has dropped (i.e., via `del`)
+    every reference to the Omni instance"""
     gc.collect()
     if torch.cuda.is_available():
         torch.accelerator.empty_cache()
@@ -267,7 +288,9 @@ def run_benchmark(args):
 
     bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
     bl_mem = baseline_outputs[prompts[0]][2]  # use first prompt's memory
-    _unload_omni(omni_bl)
+    omni_bl.shutdown()
+    del omni_bl
+    _free_gpu_memory()
 
     # Save baseline outputs
     bl_dir = output_dir / "baseline"
@@ -307,7 +330,9 @@ def run_benchmark(args):
 
         qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
         qt_mem = qt_outputs[prompts[0]][2]
-        _unload_omni(omni_qt)
+        omni_qt.shutdown()
+        del omni_qt
+        _free_gpu_memory()
 
         # Save quantized outputs
         qt_dir = output_dir / config_label.replace(" ", "_")

benchmarks/tts/README.md

@@ -40,7 +40,7 @@ vllm bench serve --omni \
     --num-prompts 20 --num-warmups 2 \
     --extra-body '{"task_type":"Base"}' \
     --max-concurrency 1 --request-rate inf \
-    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
+    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
     --save-result --result-dir ./results
 ```
 
@@ -58,7 +58,7 @@ vllm bench serve --omni \
     --num-prompts 20 --num-warmups 2 \
     --extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \
     --max-concurrency 1 --request-rate inf \
-    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
+    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
     --save-result --result-dir ./results
 ```
 
@@ -76,10 +76,28 @@ vllm bench serve --omni \
     --num-prompts 20 --num-warmups 2 \
     --extra-body '{"task_type":"VoiceDesign","language":"English"}' \
     --max-concurrency 1 --request-rate inf \
-    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
+    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
     --save-result --result-dir ./results
 ```
 
+#### Streaming continuity (`audio_underrun`)
+
+`audio_underrun` (seconds) is the per-request worst-case buffer deficit
+under a realtime-rate player simulation. The audio-speech backend records
+chunk arrival times during the SSE stream and surfaces:
+
+- `Mean / Median / P{50,99} AUDIO_UNDERRUN (s)` per the standard percentile
+  output - any positive value means at least one inter-chunk gap was longer
+  than the chunk's audio duration.
+- `Streaming continuity OK rate` - fraction of requests whose worst-case
+  underrun stayed under the threshold (default 100 ms, the commonly cited
+  "audible gap" budget). Override the threshold with
+  `VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S=<float>`.
+
+This captures the failure mode where `RTF_p50 < 1` (server keeps up in
+aggregate) but per-stream chunk arrival is bursty enough that listeners
+still hear gaps - common at high concurrency on streaming TTS pipelines.
+
 #### Add WER / SIM / UTMOS to any of the above
 
 Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
@@ -215,7 +233,6 @@ sentinel for regressions in this area.
 benchmarks/tts/
 ├── README.md                  (this file)
 ├── bench_tts.py               CLI — serve-mode benchmark driver
-├── bench_voxcpm_offline.py    CLI — offline VoxCPM benchmark (sync + streaming)
 ├── plot_results.py            Generate per-task / per-concurrency curves
 └── model_configs.yaml         Model registry (supported tasks + extra body)
 ```

benchmarks/tts/bench_tts.py

@@ -118,7 +118,7 @@ def build_bench_args(
         "--num-warmups",
         "2",
         "--percentile-metrics",
-        "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun",
     ]
 
     if resolved_dataset_path:

benchmarks/tts/bench_voxcpm_offline.py

@@ -1,922 +0,0 @@
-"""Offline VoxCPM benchmark for vLLM Omni.
-
-Supports both:
-- sync one-shot (Omni.generate)
-- streaming (AsyncOmni.generate with async_chunk config)
-- text-only synthesis
-- voice cloning
-- text/clone batch inputs from txt or jsonl
-
-Usage::
-
-    # Sync (default voice)
-    python benchmarks/tts/bench_voxcpm_offline.py \\
-        --model /path/to/VoxCPM \\
-        --text "Hello world" \\
-        --output-dir results/audio/
-
-    # Streaming (async_chunk)
-    python benchmarks/tts/bench_voxcpm_offline.py \\
-        --model /path/to/VoxCPM \\
-        --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\
-        --txt-prompts prompts.txt \\
-        --output-dir results/audio/
-
-    # Voice cloning batch via JSONL
-    python benchmarks/tts/bench_voxcpm_offline.py \\
-        --model /path/to/VoxCPM \\
-        --jsonl-prompts prompts.jsonl \\
-        --output-dir results/audio/
-"""
-
-from __future__ import annotations
-
-import asyncio
-import json
-import logging
-import os
-import tempfile
-import time
-import uuid
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import torch
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-from vllm_omni import AsyncOmni, Omni
-
-
-def _find_repo_root(start: Path) -> Path:
-    """Walk up from ``start`` until a repo marker is found.
-
-    Falls back to ``parents[2]`` for backwards compatibility if no marker hits
-    (which can only happen in unusual checkouts — the tree should always have
-    pyproject.toml + vllm_omni/ at the top level).
-    """
-    for candidate in [start, *start.parents]:
-        if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir():
-            return candidate
-    return start.parents[2]
-
-
-REPO_ROOT = _find_repo_root(Path(__file__).resolve())
-DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
-DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm.yaml"
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass(frozen=True, slots=True)
-class PromptSpec:
-    text: str
-    label: str
-    ref_audio: str | None = None
-    ref_text: str | None = None
-
-
-def _require_soundfile():
-    try:
-        import soundfile as sf  # type: ignore
-    except ModuleNotFoundError as exc:
-        raise RuntimeError(
-            "soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile"
-        ) from exc
-    return sf
-
-
-def _build_prompt(
-    args,
-    *,
-    text: str,
-    ref_audio: str | None = None,
-    ref_text: str | None = None,
-    global_request_id: str | None = None,
-) -> dict[str, Any]:
-    additional_information: dict[str, list[Any]] = {
-        "text": [text],
-        "cfg_value": [args.cfg_value],
-        "inference_timesteps": [args.inference_timesteps],
-        "min_len": [args.min_len],
-        "max_new_tokens": [args.max_new_tokens],
-    }
-    if args.streaming_prefix_len is not None:
-        additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
-
-    if ref_audio:
-        additional_information["ref_audio"] = [ref_audio]
-    if ref_text:
-        additional_information["ref_text"] = [ref_text]
-    if global_request_id is not None:
-        additional_information["global_request_id"] = [global_request_id]
-
-    return {
-        "prompt_token_ids": [1],
-        "additional_information": additional_information,
-    }
-
-
-def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
-    audio = mm.get("audio", mm.get("model_outputs"))
-    if audio is None:
-        raise ValueError("No audio output found in multimodal output.")
-    if isinstance(audio, list):
-        parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio]
-        audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
-    if not isinstance(audio, torch.Tensor):
-        audio = torch.as_tensor(audio)
-    return audio.float().cpu().reshape(-1)
-
-
-def _extract_sample_rate(mm: dict[str, Any]) -> int:
-    sr_raw = mm.get("sr", 24000)
-    if isinstance(sr_raw, list) and sr_raw:
-        sr_raw = sr_raw[-1]
-    if hasattr(sr_raw, "item"):
-        return int(sr_raw.item())
-    return int(sr_raw)
-
-
-def _emit_offline_metrics(
-    *,
-    request_id: str,
-    elapsed_s: float,
-    first_audio_elapsed: float | None,
-    audio_duration_s: float,
-) -> None:
-    metrics = {
-        "request_id": request_id,
-        "ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None,
-        "audio_duration_s": round(audio_duration_s, 6),
-        "rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None,
-    }
-    print(f"[OfflineMetrics] {metrics}")
-
-
-def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None:
-    sf = _require_soundfile()
-    if isinstance(audio_tensor, torch.Tensor):
-        audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy()
-    else:
-        audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy()
-    sf.write(
-        output_path,
-        audio_np,
-        sample_rate,
-        format="WAV",
-        subtype="PCM_16",
-    )
-
-
-def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path:
-    output_dir.mkdir(parents=True, exist_ok=True)
-    output_path = output_dir / f"output_{request_id}.wav"
-    _write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm))
-    return output_path
-
-
-def _iter_request_multimodal_outputs(request_output: Any):
-    outputs = getattr(request_output, "outputs", None)
-    if outputs:
-        for output in outputs:
-            mm = getattr(output, "multimodal_output", None)
-            if isinstance(mm, dict):
-                yield mm
-
-    mm = getattr(request_output, "multimodal_output", None)
-    if isinstance(mm, dict):
-        yield mm
-
-
-def _read_non_empty_lines(path: str) -> list[str]:
-    with open(path, encoding="utf-8") as f:
-        return [line.strip() for line in f if line.strip()]
-
-
-def _load_prompt_specs(args) -> list[PromptSpec]:
-    specs: list[PromptSpec] = []
-
-    if args.txt_prompts is not None:
-        texts = _read_non_empty_lines(args.txt_prompts)
-        if not texts:
-            raise ValueError(f"No prompts found in {args.txt_prompts}")
-        for idx, text in enumerate(texts, start=1):
-            specs.append(
-                PromptSpec(
-                    text=text,
-                    label=f"item{idx:03d}",
-                    ref_audio=args.ref_audio,
-                    ref_text=args.ref_text,
-                )
-            )
-        return specs
-
-    if args.jsonl_prompts is not None:
-        with open(args.jsonl_prompts, encoding="utf-8") as f:
-            for line_no, raw_line in enumerate(f, start=1):
-                line = raw_line.strip()
-                if not line:
-                    continue
-                try:
-                    item = json.loads(line)
-                except json.JSONDecodeError as exc:
-                    raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc
-                if not isinstance(item, dict):
-                    raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object")
-
-                text = item.get("text")
-                if not isinstance(text, str) or not text.strip():
-                    raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'")
-
-                ref_audio = item.get("ref_audio", args.ref_audio)
-                ref_text = item.get("ref_text", args.ref_text)
-                if (ref_audio is None) != (ref_text is None):
-                    raise ValueError(
-                        f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together"
-                    )
-
-                specs.append(
-                    PromptSpec(
-                        text=text.strip(),
-                        label=f"item{len(specs) + 1:03d}",
-                        ref_audio=ref_audio,
-                        ref_text=ref_text,
-                    )
-                )
-
-        if not specs:
-            raise ValueError(f"No prompts found in {args.jsonl_prompts}")
-        return specs
-
-    specs.append(
-        PromptSpec(
-            text=args.text,
-            label="item001",
-            ref_audio=args.ref_audio,
-            ref_text=args.ref_text,
-        )
-    )
-    return specs
-
-
-def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]:
-    return _build_prompt(
-        args,
-        text=spec.text,
-        ref_audio=spec.ref_audio,
-        ref_text=spec.ref_text,
-        global_request_id=global_request_id,
-    )
-
-
-def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int:
-    return sum(1 for spec in prompt_specs if spec.ref_audio is not None)
-
-
-def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]:
-    return prompt_specs[:1]
-
-
-def _extract_stream_finished(stage_output: Any) -> bool:
-    request_output = getattr(stage_output, "request_output", None)
-    request_finished = getattr(request_output, "finished", None)
-    if request_finished is not None:
-        return bool(request_finished)
-    return bool(getattr(stage_output, "finished", False))
-
-
-def _build_profiled_stage_config(
-    stage_configs_path: str,
-    profiler_dir: str,
-) -> str:
-    stage_config_path = Path(stage_configs_path)
-    yaml_text = stage_config_path.read_text(encoding="utf-8")
-    injected_lines: list[str] = []
-    injected_count = 0
-
-    for line in yaml_text.splitlines():
-        injected_lines.append(line)
-        if line.strip() != "engine_args:":
-            continue
-        indent = line[: len(line) - len(line.lstrip())]
-        child_indent = indent + "  "
-        grandchild_indent = child_indent + "  "
-        injected_lines.extend(
-            [
-                f"{child_indent}profiler_config:",
-                f'{grandchild_indent}profiler: "torch"',
-                f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"',
-                f"{grandchild_indent}torch_profiler_with_stack: true",
-            ]
-        )
-        injected_count += 1
-
-    if injected_count == 0:
-        raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}")
-
-    tmp = tempfile.NamedTemporaryFile(
-        mode="w",
-        encoding="utf-8",
-        delete=False,
-        suffix=".yaml",
-        prefix=f"{stage_config_path.stem}_profile_",
-    )
-    tmp.write("\n".join(injected_lines) + "\n")
-    tmp.close()
-    return tmp.name
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(
-        description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)"
-    )
-    parser.add_argument(
-        "--model",
-        type=str,
-        default=os.environ.get("VOXCPM_MODEL"),
-        help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.",
-    )
-    parser.add_argument(
-        "--text",
-        type=str,
-        default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.",
-        help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.",
-    )
-    parser.add_argument(
-        "--txt-prompts",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one synthesis text per line.",
-    )
-    parser.add_argument(
-        "--jsonl-prompts",
-        type=str,
-        default=None,
-        help=(
-            "Path to a .jsonl file. Each line must contain at least {'text': ...}; "
-            "clone rows can also set ref_audio/ref_text, and ref_text must be the "
-            "real transcript of ref_audio."
-        ),
-    )
-    parser.add_argument(
-        "--ref-audio",
-        type=str,
-        default=None,
-        help=(
-            "Optional reference audio path for voice cloning. With --txt-prompts, "
-            "the same reference is applied to every line."
-        ),
-    )
-    parser.add_argument(
-        "--ref-text",
-        type=str,
-        default=None,
-        help=(
-            "Real transcript of the reference audio. Placeholder text or mismatched "
-            "text will usually produce noisy/electronic clone audio."
-        ),
-    )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=str(DEFAULT_STAGE_SYNC),
-        help="Stage config YAML path. Routing is selected only from this path.",
-    )
-    parser.add_argument(
-        "--cfg-value",
-        type=float,
-        default=2.0,
-        help="Classifier-free guidance value for VoxCPM.",
-    )
-    parser.add_argument(
-        "--inference-timesteps",
-        type=int,
-        default=10,
-        help="Number of inference timesteps.",
-    )
-    parser.add_argument(
-        "--min-len",
-        type=int,
-        default=2,
-        help="Minimum generated token length.",
-    )
-    parser.add_argument(
-        "--max-new-tokens",
-        type=int,
-        default=4096,
-        help="Maximum generated token length.",
-    )
-    parser.add_argument(
-        "--streaming-prefix-len",
-        type=int,
-        default=None,
-        help="VoxCPM streaming window (optional, streaming mode only).",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        default=None,
-        help="Directory for output WAV files.",
-    )
-    parser.add_argument(
-        "--stage-init-timeout",
-        type=int,
-        default=600,
-        help="Stage initialization timeout in seconds.",
-    )
-    parser.add_argument(
-        "--log-stats",
-        dest="log_stats",
-        action="store_true",
-        help="Enable vLLM Omni stats logging.",
-    )
-    parser.add_argument(
-        "--no-log-stats",
-        dest="log_stats",
-        action="store_false",
-        help="Disable vLLM Omni stats logging.",
-    )
-    parser.set_defaults(log_stats=True)
-    parser.add_argument(
-        "--num-runs",
-        type=int,
-        default=1,
-        help="Number of full inference runs (same prompt each time). Default 1.",
-    )
-    parser.add_argument(
-        "--warmup-runs",
-        type=int,
-        default=0,
-        help=(
-            "Optional number of warmup passes before measured runs. Warmup uses only "
-            "the first prompt and does not save outputs."
-        ),
-    )
-    parser.add_argument(
-        "--enable-profiler",
-        action="store_true",
-        help=(
-            "Enable torch profiler for the configured stages. A temporary profiled "
-            "stage config is generated automatically."
-        ),
-    )
-    parser.add_argument(
-        "--profiler-dir",
-        type=str,
-        default=None,
-        help="Directory for profiler traces. Defaults to <output-dir>/profiler when profiling is enabled.",
-    )
-    parser.add_argument(
-        "--profiler-stages",
-        type=int,
-        nargs="*",
-        default=None,
-        help="Optional stage ids to profile. Defaults to all stages that have profiler_config.",
-    )
-    parser.add_argument(
-        "--profiler-wait-seconds",
-        type=float,
-        default=30.0,
-        help="Seconds to wait after stop_profile for trace files to flush.",
-    )
-    args = parser.parse_args()
-
-    if not args.model:
-        parser.error("--model is required unless $VOXCPM_MODEL is set")
-    if args.txt_prompts is not None and args.jsonl_prompts is not None:
-        parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive")
-    if (args.ref_audio is None) != (args.ref_text is None):
-        parser.error("--ref-audio and --ref-text must be provided together")
-    if args.num_runs < 1:
-        parser.error("--num-runs must be >= 1")
-    if args.warmup_runs < 0:
-        parser.error("--warmup-runs must be >= 0")
-    if args.output_dir is None:
-        args.output_dir = (
-            "output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio"
-        )
-    if args.enable_profiler and args.profiler_dir is None:
-        args.profiler_dir = str(Path(args.output_dir) / "profiler")
-    try:
-        args.prompt_specs = _load_prompt_specs(args)
-    except ValueError as exc:
-        parser.error(str(exc))
-
-    return args
-
-
-def _is_streaming_stage_config(stage_configs_path: str) -> bool:
-    cfg_name = Path(stage_configs_path).name.lower()
-    return "async_chunk" in cfg_name
-
-
-async def _collect_streaming_audio(
-    omni: AsyncOmni,
-    args: Any,
-    spec: PromptSpec,
-    request_id: str,
-    *,
-    phase_label: str,
-    prompt_index: int,
-    prompt_count: int,
-    print_prompt: bool = False,
-) -> tuple[torch.Tensor, int, float, float | None]:
-    prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id)
-    delta_chunks: list[torch.Tensor] = []
-    sample_rate = 24000
-    chunk_i = 0
-    prev_total_samples = 0
-    t_start = time.perf_counter()
-    first_audio_elapsed: float | None = None
-
-    if print_prompt:
-        print(f"---prompt---:{prompt}")
-
-    async for stage_output in omni.generate(prompt, request_id=request_id):
-        mm = getattr(stage_output, "multimodal_output", None)
-        if not isinstance(mm, dict):
-            ro = getattr(stage_output, "request_output", None)
-            if ro is None:
-                continue
-            mm = getattr(ro, "multimodal_output", None)
-            if not isinstance(mm, dict) and getattr(ro, "outputs", None):
-                seq = ro.outputs[0]
-                mm = getattr(seq, "multimodal_output", None)
-        if not isinstance(mm, dict):
-            continue
-        sample_rate = _extract_sample_rate(mm)
-        try:
-            w = _extract_audio_tensor(mm)
-            n = int(w.numel())
-            if n == 0:
-                continue
-            finished = _extract_stream_finished(stage_output)
-            if n > prev_total_samples:
-                delta = w.reshape(-1)[prev_total_samples:]
-                prev_total_samples = n
-            elif finished and n == prev_total_samples:
-                delta = w.reshape(-1)[:0]
-            else:
-                delta = w.reshape(-1)
-                prev_total_samples += int(delta.numel())
-            if int(delta.numel()) > 0:
-                delta_chunks.append(delta)
-            if first_audio_elapsed is None and int(delta.numel()) > 0:
-                first_audio_elapsed = time.perf_counter() - t_start
-            logger.info(
-                "%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s",
-                phase_label,
-                prompt_index + 1,
-                prompt_count,
-                chunk_i,
-                int(delta.numel()),
-                n,
-                finished,
-            )
-            chunk_i += 1
-        except ValueError:
-            if not _extract_stream_finished(stage_output):
-                logger.debug("skip non-audio partial output chunk=%d", chunk_i)
-
-    if not delta_chunks:
-        raise RuntimeError("No audio chunks received; check stage config and logs.")
-
-    audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0)
-    elapsed = time.perf_counter() - t_start
-    return audio_cat, sample_rate, elapsed, first_audio_elapsed
-
-
-async def _abort_streaming_residual_work(
-    omni: AsyncOmni,
-    request_id: str,
-    *,
-    settle_seconds: float = 0.1,
-) -> None:
-    """Stop any late stage-0 work once the final audio has been collected."""
-    await omni.engine.abort_async([request_id])
-    if settle_seconds > 0:
-        await asyncio.sleep(settle_seconds)
-
-
-async def _run_streaming_single(
-    omni: AsyncOmni,
-    args: Any,
-    spec: PromptSpec,
-    output_dir: Path,
-    request_id: str,
-    *,
-    run_index: int,
-    num_runs: int,
-    prompt_index: int,
-    prompt_count: int,
-) -> Path:
-    audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio(
-        omni,
-        args,
-        spec,
-        request_id,
-        phase_label=f"run={run_index + 1}/{num_runs}",
-        prompt_index=prompt_index,
-        prompt_count=prompt_count,
-        print_prompt=(run_index == 0 and prompt_index == 0),
-    )
-    await _abort_streaming_residual_work(omni, request_id)
-    output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav"
-    _write_audio_tensor(output_path, audio_cat, sample_rate)
-    audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0
-    ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
-    rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
-    print(
-        f"Saved (streaming) run {run_index + 1}/{num_runs}, "
-        f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})"
-    )
-    _emit_offline_metrics(
-        request_id=request_id,
-        elapsed_s=elapsed,
-        first_audio_elapsed=first_audio_elapsed,
-        audio_duration_s=audio_duration_s,
-    )
-    return output_path
-
-
-async def _run_streaming_warmup(args, omni: AsyncOmni) -> None:
-    if args.warmup_runs == 0:
-        return
-
-    warmup_specs = _get_warmup_specs(args.prompt_specs)
-    print(
-        f"Warmup: {args.warmup_runs} run(s) using the first prompt "
-        f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
-    )
-    for warmup_index in range(args.warmup_runs):
-        t_warmup = time.perf_counter()
-        tasks = []
-        request_ids: list[str] = []
-        for prompt_index, spec in enumerate(warmup_specs):
-            request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
-            request_ids.append(request_id)
-            tasks.append(
-                _collect_streaming_audio(
-                    omni,
-                    args,
-                    spec,
-                    request_id,
-                    phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}",
-                    prompt_index=prompt_index,
-                    prompt_count=len(warmup_specs),
-                )
-            )
-        results = await asyncio.gather(*tasks)
-        for request_id in request_ids:
-            await _abort_streaming_residual_work(omni, request_id)
-        total_samples = sum(int(audio.numel()) for audio, _, _, _ in results)
-        warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None]
-        ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else ""
-        print(
-            f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: "
-            f"{len(results)} prompt(s), {total_samples} sample(s) "
-            f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})"
-        )
-
-
-async def _run_streaming(args) -> list[Path]:
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    omni = AsyncOmni(
-        model=args.model,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
-
-    await _run_streaming_warmup(args, omni)
-    profiler_started = False
-    if args.enable_profiler:
-        profile_prefix = f"voxcpm_streaming_{int(time.time())}"
-        stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
-        print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}")
-        await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
-        profiler_started = True
-    t_total = time.perf_counter()
-    total_elapsed = 0.0
-    paths: list[Path] = []
-    prompt_specs: list[PromptSpec] = args.prompt_specs
-    try:
-        for run in range(args.num_runs):
-            for prompt_index, spec in enumerate(prompt_specs):
-                request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
-                paths.append(
-                    await _run_streaming_single(
-                        omni,
-                        args,
-                        spec,
-                        output_dir,
-                        request_id,
-                        run_index=run,
-                        num_runs=args.num_runs,
-                        prompt_index=prompt_index,
-                        prompt_count=len(prompt_specs),
-                    )
-                )
-        total_elapsed = time.perf_counter() - t_total
-    finally:
-        if profiler_started:
-            print("Stopping profiler (streaming)...")
-            await omni.stop_profile(stages=args.profiler_stages)
-            if args.profiler_wait_seconds > 0:
-                print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
-                await asyncio.sleep(args.profiler_wait_seconds)
-
-    print(
-        f"All streaming runs finished: {args.num_runs} run(s), "
-        f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total"
-    )
-    return paths
-
-
-def _run_sync(args) -> list[Path]:
-    output_dir = Path(args.output_dir)
-
-    omni = Omni(
-        model=args.model,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
-
-    def _run_sync_single(
-        spec: PromptSpec,
-        *,
-        request_prefix: str,
-        save_outputs: bool,
-        run_index: int | None = None,
-    ) -> tuple[list[Path], int, float | None, float, float, str]:
-        global_request_id = f"{request_prefix}_{spec.label}"
-        prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id)
-        if save_outputs and run_index == 0 and spec.label == "item001":
-            print(f"---prompt---:{prompt}")
-
-        saved_paths: list[Path] = []
-        output_count = 0
-        first_audio_elapsed: float | None = None
-        total_audio_duration_s = 0.0
-        metrics_request_id = global_request_id
-        t_start = time.perf_counter()
-        for stage_outputs in omni.generate(prompt):
-            request_output = stage_outputs.request_output
-            if request_output is None:
-                continue
-            request_output_id = getattr(request_output, "request_id", None)
-            if isinstance(request_output_id, str) and request_output_id:
-                metrics_request_id = request_output_id
-            for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)):
-                output_count += 1
-                if first_audio_elapsed is None:
-                    try:
-                        audio_tensor = _extract_audio_tensor(mm)
-                        if int(audio_tensor.numel()) > 0:
-                            first_audio_elapsed = time.perf_counter() - t_start
-                        total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
-                    except ValueError:
-                        pass
-                else:
-                    try:
-                        audio_tensor = _extract_audio_tensor(mm)
-                        total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
-                    except ValueError:
-                        pass
-                if not save_outputs:
-                    continue
-                save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}"
-                saved_paths.append(_save_wav(mm, output_dir, save_stem))
-
-        if output_count == 0:
-            raise RuntimeError("No output from Omni.generate")
-        elapsed_s = time.perf_counter() - t_start
-        return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id
-
-    if args.warmup_runs:
-        warmup_specs = _get_warmup_specs(args.prompt_specs)
-        print(
-            f"Warmup: {args.warmup_runs} run(s) using the first prompt "
-            f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
-        )
-        for warmup_index in range(args.warmup_runs):
-            t_warmup = time.perf_counter()
-            _, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single(
-                warmup_specs[0],
-                request_prefix=f"warmup_sync{warmup_index + 1}",
-                save_outputs=False,
-            )
-            ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
-            rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
-            print(
-                f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: "
-                f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})"
-            )
-
-    profiler_started = False
-    if args.enable_profiler:
-        profile_prefix = f"voxcpm_sync_{int(time.time())}"
-        stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
-        print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}")
-        omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
-        profiler_started = True
-
-    t_total = time.perf_counter()
-    total_elapsed = 0.0
-    saved_paths: list[Path] = []
-    prompt_specs: list[PromptSpec] = args.prompt_specs
-    try:
-        for run in range(args.num_runs):
-            t_run = time.perf_counter()
-            run_paths: list[Path] = []
-            for prompt_index, spec in enumerate(prompt_specs):
-                prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = (
-                    _run_sync_single(
-                        spec,
-                        request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}",
-                        save_outputs=True,
-                        run_index=run,
-                    )
-                )
-                run_paths.extend(prompt_paths)
-                ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
-                rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
-                print(
-                    f"Saved (sync) run {run + 1}/{args.num_runs}, "
-                    f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}"
-                )
-                _emit_offline_metrics(
-                    request_id=metrics_request_id,
-                    elapsed_s=elapsed_s,
-                    first_audio_elapsed=first_audio_elapsed,
-                    audio_duration_s=audio_duration_s,
-                )
-
-            saved_paths.extend(run_paths)
-            print(
-                f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)"
-            )
-            for path in run_paths:
-                print(f"  {path}")
-
-        total_elapsed = time.perf_counter() - t_total
-    finally:
-        if profiler_started:
-            print("Stopping profiler (sync)...")
-            omni.stop_profile(stages=args.profiler_stages)
-            if args.profiler_wait_seconds > 0:
-                print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
-                time.sleep(args.profiler_wait_seconds)
-
-    print(
-        f"All sync runs finished: {args.num_runs} run(s), "
-        f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total"
-    )
-    return saved_paths
-
-
-def main(args) -> int:
-    logging.basicConfig(level=logging.INFO)
-    profiled_stage_config_path: str | None = None
-    original_stage_config_path = args.stage_configs_path
-    if args.enable_profiler:
-        Path(args.profiler_dir).mkdir(parents=True, exist_ok=True)
-        profiled_stage_config_path = _build_profiled_stage_config(
-            args.stage_configs_path,
-            str(Path(args.profiler_dir).resolve()),
-        )
-        args.stage_configs_path = profiled_stage_config_path
-
-    is_streaming = _is_streaming_stage_config(args.stage_configs_path)
-    voice_clone_count = _count_voice_clone_prompts(args.prompt_specs)
-    print(f"Model: {args.model}")
-    print(f"Stage config: {original_stage_config_path}")
-    print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)")
-    print(f"Prompt count: {len(args.prompt_specs)}")
-    print("Batch mode: sequential (aligned with native VoxCPM)")
-    print(f"Warmup runs: {args.warmup_runs}")
-    print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}")
-    if args.enable_profiler:
-        print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})")
-        print(f"Profiled stage config: {args.stage_configs_path}")
-    if voice_clone_count:
-        print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.")
-    print(f"Num runs: {args.num_runs}")
-    try:
-        if is_streaming:
-            asyncio.run(_run_streaming(args))
-        else:
-            _run_sync(args)
-    finally:
-        if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path):
-            os.unlink(profiled_stage_config_path)
-    return 0
-
-
-if __name__ == "__main__":
-    os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-    raise SystemExit(main(parse_args()))

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.xpu

@@ -133,6 +133,11 @@ ENTRYPOINT []
 
 FROM ${VLLM_BASE} AS vllm-omni
 
+RUN apt-get update && \
+    apt-get install -y espeak-ng ffmpeg git jq && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
 WORKDIR /workspace/vllm-omni
 COPY . .
 

docs/.nav.yml

@@ -80,6 +80,8 @@ nav:
       - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
       - FP8 W8A8: user_guide/quantization/fp8.md
       - Int8 W8A8: user_guide/quantization/int8.md
+      - MXFP8 W8A8: user_guide/quantization/mxfp8.md
+      - MXFP4 W4A4: user_guide/quantization/mxfp4.md
       - ModelOpt: user_guide/quantization/modelopt.md
       - GGUF: user_guide/quantization/gguf.md
       - AutoRound: user_guide/quantization/autoround.md

docs/api/README.md

@@ -121,7 +121,6 @@ Configuration classes.
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
 - [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
-- [vllm_omni.transformers_utils.configs.voxcpm.VoxCPMConfig][]
 - [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
 
 ## Workers

docs/community/governance.md

@@ -39,11 +39,34 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
 - [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
 - [@princepride](https://github.com/princepride): Diffusion and Omni Support
-- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
+- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
+- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
 - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
-- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
+- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
 - [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
-- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
+- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
+- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
+
+## Diffusion Workload Division
+
+This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
+
+### Runtime-related
+
+- Diffusion models: @RuixiangMa, @wtomin
+
+### Optimization
+
+- Parallel: @RuixiangMa, @wtomin
+- Attention backend: @ZJY0516
+- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
+- Offload: @yuanheng-zhao
+- Quantization: @lishunyang12, @david6666666
+- Kernel fusion / communication-computation: @ZJY0516
+
+### Docs & Test
+
+- Diffusion docs: @RuixiangMa, @wtomin
 
 ## Meetings
 

docs/design/feature/diffusion_continuous_batching.md

@@ -83,7 +83,10 @@ There are two important details:
 
 Today that compatibility rule is still shape-sensitive. `height`, `width`,
 `num_frames`, and CFG-related fields remain part of the key, so different
-resolutions or incompatible guidance settings do **not** co-batch yet.
+resolutions or incompatible guidance settings do **not** co-batch yet. The
+key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
+targeting different adapters or scales run in separate batches and the
+worker can activate exactly one adapter per step.
 
 The current batching unit is one `OmniDiffusionRequest`. Requests with
 multiple prompts do not participate in batching today.

docs/design/feature/disaggregated_inference.md

@@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs:
 - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
 - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
 - [YuanrongConnector](omni_connectors/yuanrong_connector.md)
+- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
 
 ## Overview
 
@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
 | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
 | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
 | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
+| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
 
 ## Core API
 

docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md

@@ -0,0 +1,211 @@
+# YuanrongTransferEngineConnector
+
+## When to Use
+
+Use `YuanrongTransferEngineConnector` for high-performance peer-to-peer KV cache
+transfer on Ascend NPU. It is different from `YuanrongConnector`:
+
+- `YuanrongConnector` uses Yuanrong Datasystem as a distributed KV store and
+  requires Datasystem workers plus etcd.
+- `YuanrongTransferEngineConnector` uses Yuanrong TransferEngine directly. It
+  registers local memory in each stage worker and lets the receiver pull data
+  from the sender with TransferEngine.
+
+The implementation currently lives under
+`vllm_omni/platforms/npu/omni_connectors/` because the supported backend is
+Ascend NPU only. The connector is still registered with the generic
+OmniConnector factory under the name `YuanrongTransferEngineConnector`.
+
+For Ascend P2P transfer, the connector should use NPU memory as its transfer
+pool. The Mooncake-style CPU pool is not supported for this connector.
+
+## Prerequisites
+
+Install Yuanrong Datasystem/TransferEngine Python bindings in the runtime
+environment. The packaged install is:
+
+```bash
+pip install openyuanrong-datasystem
+```
+
+To build from source, follow the upstream project instructions:
+
+- Source repository: [openeuler/yuanrong-datasystem](https://atomgit.com/openeuler/yuanrong-datasystem)
+- Installation guide: [Yuanrong Datasystem Linux installation](https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/installation/installation_linux.html)
+
+The Ascend driver, CANN runtime, HCCL/HCCP, and NPU device network must also be
+configured on every machine participating in the transfer.
+
+## Check NPU Device IPv4
+
+Yuanrong TransferEngine with `protocol: "ascend"` uses the NPU device network,
+not only the host management IP. Each participating NPU must have an IPv4 HCCN
+device IP.
+
+Check the device IP, link, and health state:
+
+```bash
+for i in {0..7}; do
+  echo "===== NPU $i ====="
+  hccn_tool -i "$i" -ip -g
+  hccn_tool -i "$i" -link -g
+  hccn_tool -i "$i" -net_health -g
+done
+```
+
+Expected output for each NPU should include an IPv4 address and healthy link:
+
+```text
+ipaddr:10.86.1.80
+netmask:255.255.255.0
+link status: UP
+net health status: Success
+```
+
+If the output says `Get ipconf failed, because no ip was preset there!`, the NPU
+device IP is not configured and Ascend P2P will fail. If the address is IPv6,
+configure IPv4 for this connector path or validate that the underlying Yuanrong
+TransferEngine build supports the IPv6 HCCN path before using it.
+
+## Check Device-to-Device Connectivity
+
+After collecting the target-stage NPU IPv4 addresses, ping them from each
+source-stage NPU:
+
+```bash
+# Example: AR uses NPU 0..3 and DiT uses NPU 4..7.
+# Replace the destination IPs with the values reported by hccn_tool -i 4..7 -ip -g.
+for src in 0 1 2 3; do
+  for dst_ip in 10.86.1.78 10.86.1.79 10.86.1.80 10.86.1.81; do
+    echo "===== NPU $src -> $dst_ip ====="
+    hccn_tool -i "$src" -ping -g address "$dst_ip" || true
+  done
+done
+```
+
+The expected result is `0.00% packet loss` for every required source/target pair.
+This verifies basic HCCN IPv4 connectivity. TransferEngine can still fail later
+if RDMA/QP/HCCL setup is broken, but failed ping means the device network must be
+fixed first.
+
+## YAML Configuration
+
+For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU
+memory pool:
+
+```yaml
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  connectors:
+    yuanrong_te_connector:
+      name: YuanrongTransferEngineConnector
+      extra:
+        host: "auto"
+        zmq_port: 50051
+        rpc_port: "auto"
+        protocol: "ascend"
+        device_name: "auto"
+        memory_pool_size: 1073741824
+        memory_pool_device: "npu"
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+```
+
+Important fields:
+
+| Parameter | Recommended Value | Notes |
+|---|---|---|
+| `host` | `"auto"` | Advertises a routable local host IP for ZMQ metadata exchange. |
+| `zmq_port` | `50051` | Base ZMQ port. The runtime applies stage/rank offsets. |
+| `rpc_port` | `"auto"` | Lets TransferEngine choose an available RPC port. |
+| `protocol` | `"ascend"` | Uses the Ascend TransferEngine backend. |
+| `device_name` | `"auto"` | Each TP worker resolves to its local logical NPU, for example `npu:0`, `npu:1`, etc. |
+| `memory_pool_device` | `"npu"` | Required for the Ascend fast path. This registers NPU memory with TransferEngine. |
+| `memory_pool_size` | `1073741824` | 1 GiB per worker is a practical starting point for KV-cache transfer validation. Increase if the transfer pool is exhausted. |
+
+`memory_pool_device: "cpu"` is rejected for
+`YuanrongTransferEngineConnector` with `protocol: "ascend"`. CPU pool registration
+can lead to Ascend P2P failures such as HCCL RA init errors, QP timeouts, or
+FFTS/SDMA runtime errors. `memory_pool_device: "cpu"` is common for
+`MooncakeTransferEngineConnector`, but Yuanrong TE on Ascend only supports
+`"npu"`.
+
+## 4-Card AR-to-DiT TP Example
+
+For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in
+both stage `omni_kv_config` blocks:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1,2,3"
+    engine_args:
+      tensor_parallel_size: 4
+      omni_kv_config:
+        need_send_cache: true
+        rank_mapping:
+          from_tp: 4
+          to_tp: 4
+    output_connectors:
+      to_stage_1: yuanrong_te_connector
+
+  - stage_id: 1
+    runtime:
+      devices: "4,5,6,7"
+    engine_args:
+      parallel_config:
+        tensor_parallel_size: 4
+      omni_kv_config:
+        need_recv_cache: true
+        rank_mapping:
+          from_tp: 4
+          to_tp: 4
+    input_connectors:
+      from_stage_0: yuanrong_te_connector
+```
+
+In this topology each DiT rank receives the corresponding AR rank shard:
+
+```text
+DiT rank 0 <- AR rank 0
+DiT rank 1 <- AR rank 1
+DiT rank 2 <- AR rank 2
+DiT rank 3 <- AR rank 3
+```
+
+## Validate a Successful Run
+
+Search the inference log for Yuanrong TE receive-side messages:
+
+```bash
+grep -E "\[YR GET\]|Successfully received KV cache|transfer_engine get failed" inference.log
+```
+
+A successful Ascend fast-path transfer contains lines like:
+
+```text
+[YR GET] ... fast_path, zero-copy
+Successfully received KV cache ... across 1 key(s)
+Applied CFG KV caches
+[KV Reuse] ...
+```
+
+Sender-side `KV transfer OK` only means that the sender registered the payload
+and metadata. The actual P2P data transfer happens on the receiver side during
+`[YR GET]`.
+
+## Troubleshooting
+
+| Symptom | Likely Cause | Action |
+|---|---|---|
+| `IPv4 device IP not found` | NPU HCCN IPv4 is missing. | Run `hccn_tool -i <id> -ip -g` and configure IPv4 for every participating NPU. |
+| `P2PCommInitRootInfo failed`, `ra init failed`, `RA_QP_STATUS_TIMEOUT` | Ascend P2P/RDMA/QP setup failed. | First verify device IPv4, link health, and device-to-device ping. Then test a minimal Yuanrong TE pair outside vLLM. |
+| `fftsplus sdma error`, `context is abort` | Runtime SDMA/FFTS failure, often after an invalid or failed P2P transfer. | Use `memory_pool_device: "npu"` for Ascend TE, reduce `memory_pool_size` if OOM occurs, and isolate the failing NPU pair with a Yuanrong TE demo. |
+| Sender logs `KV transfer OK` but receiver never logs `[YR GET]` | Sender `put()` succeeded, receiver P2P read failed or timed out. | Check receiver logs and search for `transfer_engine get failed`. |
+| Pool allocation fails or OOM | Transfer pool is too small or too large for available NPU memory. | Start with 1 GiB (`1073741824`), reduce to 512 MiB for validation, or increase for larger KV payloads. |

docs/getting_started/installation/gpu/cuda.inc.md

@@ -103,10 +103,10 @@ docker run --runtime nvidia --gpus 2 \
     -p 8091:8091 \
     --ipc=host \
     vllm/vllm-omni:v0.21.0 \
-    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
 !!! tip
-    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+    The CUDA image does not define a default entrypoint, so include `vllm serve ... --omni` after the image name.
 
 # --8<-- [end:pre-built-images]

docs/models/supported_models.md

@@ -58,6 +58,7 @@ th {
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
@@ -71,4 +72,6 @@ th {
 | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
+
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

docs/serving/speech_api.md

@@ -5,6 +5,9 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
 - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
 - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
+- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
+
+See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
 
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
@@ -26,6 +29,10 @@ vllm serve fishaudio/s2-pro --omni --port 8091
 
 # Voxtral TTS
 vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
+
+# CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request)
+vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \
+    --omni --port 8091 --trust-remote-code
 ```
 
 ### Generate Speech

docs/user_guide/diffusion/step_execution.md

@@ -60,6 +60,9 @@ batching only becomes relevant when `--max-num-seqs > 1`.
 - `cache_backend` is not supported together with step execution.
 - Unsupported pipelines fail early during model loading.
 - Request-mode extras such as KV transfer are not wired into step mode yet.
+- LoRA is supported in step mode, but each batch must use a single adapter:
+  requests with different `lora_request` or `lora_scale` are scheduled into
+  separate batches.
 
 ## When To Use It
 

docs/user_guide/diffusion_features.md

@@ -115,6 +115,7 @@ The following tables show which models support each feature:
 | **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -184,7 +185,7 @@ The following tables show which models support each feature:
 | **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
 | **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
 | **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
-| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
+| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ✅ | |
 
 !!! info
 
@@ -193,7 +194,7 @@ The following tables show which models support each feature:
     3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
     4. CPU Offloading (Layerwise) supports single-card for now.
     5. Using FP8-Quant as an example of qunatization methods.
-    6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
+    6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
 
 
 ## Multi-Thread Weight Loading

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -17,9 +17,9 @@ list of supported architectures across all modalities, see
 | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
 |---|---|---|---|---|---|---|
 | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
-| VoxCPM | local model dir | split | ✓ | ✓ (`voxcpm_async_chunk.yaml`) | — | 24 kHz |
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
+| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
@@ -72,61 +72,6 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
 
 ---
 
-## VoxCPM
-
-Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
-
-### Prerequisites
-```bash
-pip install voxcpm soundfile
-# or use a local source tree:
-export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
-```
-
-If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
-
-```bash
-export VOXCPM_MODEL=/path/to/voxcpm-model
-export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
-mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
-cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
-cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
-python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
-```
-
-### Quick start
-```bash
-python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
-```
-
-### Voice cloning
-```bash
-python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This sentence is synthesized with a cloned voice." \
-    --ref-audio /path/to/reference.wav \
-    --ref-text  "The exact transcript spoken in reference.wav."
-```
-
-### Streaming
-Pass the async-chunk stage config:
-```bash
-python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model "$VOXCPM_MODEL" \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
-    --text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
-```
-
-### Notes
-- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
-- Streaming is currently single-request oriented.
-- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
-- For online serving, see [`examples/online_serving/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm/README.md). For benchmark reporting, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
-
----
-
 ## CosyVoice3
 
 2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
@@ -214,6 +159,35 @@ Streaming requires `async_chunk: true` in the stage config.
 
 ---
 
+## GLM-TTS
+
+2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+    --model zai-org/GLM-TTS \
+    --text "你好，这是语音合成测试。" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "这是参考音频的文本内容。" \
+    --output-dir ./output
+```
+
+### Architecture
+```
+Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
+        (Llama-based)    (32k vocab)      (Flow Matching)
+```
+
+### Notes
+- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
+- Reference audio should be 3-10 seconds.
+- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
+- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method).
+- The `--model` path should point to the repository root (not `llm/` subdirectory).
+
+---
+
 ## OmniVoice
 
 Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
@@ -369,6 +343,10 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
     ``````
+??? abstract "glm_tts/end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"
+    ``````
 ??? abstract "omnivoice/end2end.py"
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/omnivoice/end2end.py"
@@ -381,10 +359,6 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/qwen3_tts/end2end.py"
     ``````
-??? abstract "voxcpm/end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/text_to_speech/voxcpm/end2end.py"
-    ``````
 ??? abstract "voxcpm2/end2end.py"
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/voxcpm2/end2end.py"

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -17,14 +17,15 @@ For the full list of supported architectures across all modalities, see
 
 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
 |---|---|---|---|---|---|
+| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
+| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
-| VoxCPM | local model dir | ✓ | ✓ (PCM stream) | — | — |
 | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
 
-CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3) instead.
+For offline inference of any of these models, see the [offline TTS hub](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md).
 
 ## Common Quick Start
 
@@ -96,6 +97,54 @@ For full request-shape documentation (all parameters, response formats, error co
 
 ---
 
+## CosyVoice3
+
+2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.
+
+### Prerequisites
+```bash
+huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
+```
+
+If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class).
+
+### Launch
+```bash
+vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
+# or:
+./cosyvoice3/run_server.sh
+```
+
+Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
+
+### CLI client
+The client defaults to the official upstream zero-shot prompt, so it runs without extra flags:
+```bash
+cd examples/online_serving/text_to_speech/cosyvoice3
+python speech_client.py --text "收到好友从远方寄来的生日礼物。"
+```
+
+Pass your own reference clip and transcript for a different voice:
+```bash
+python speech_client.py --text "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Transcript of the reference audio."
+```
+
+Stream PCM instead of WAV:
+```bash
+python speech_client.py --text "Hello world" --stream --output output.pcm
+```
+
+The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`.
+
+### Notes
+- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
+- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
+
+---
+
 ## Fish Speech S2 Pro
 
 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
@@ -144,6 +193,45 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server
 
 ---
 
+## GLM-TTS
+
+2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
+
+### Launch
+```bash
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+# or:
+bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS
+```
+
+### Sending requests
+```bash
+# Voice cloning (required)
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "你好，这是语音克隆测试。" \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "这是参考音频的文本内容。"
+
+# Custom format
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "Hello, this is a voice cloning test." \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "Transcript of the reference audio." \
+    --response-format mp3 -o output.mp3
+```
+
+### Gradio demo
+```bash
+bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
+```
+
+### Notes
+- Output: 24 kHz mono WAV via HiFT vocoder.
+- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
+- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
+
+---
+
 ## OmniVoice
 
 Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
@@ -196,10 +284,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
 ./qwen3_tts/run_server.sh Base
 ```
 
-### Choosing an executor backend (uniproc vs mp)
-Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation.
+### Executor backend
+Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
 
-To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor.
+To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
 
 ### Sending requests
 ```bash
@@ -281,71 +369,11 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
 
 ### Notes
-- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
+- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
 
 ---
 
-## VoxCPM
-
-Split-stage TTS at 24 kHz.
-
-### Prerequisites
-```bash
-pip install voxcpm
-# or use a local source tree:
-export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
-```
-
-If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir:
-```bash
-export VOXCPM_MODEL=/path/to/voxcpm-model
-export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
-mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
-cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
-cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
-python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
-```
-
-### Launch
-```bash
-export VOXCPM_MODEL=/path/to/voxcpm-model
-./voxcpm/run_server.sh                # async-chunk streaming (default)
-./voxcpm/run_server.sh sync           # non-streaming
-```
-Or directly:
-```bash
-vllm serve "$VOXCPM_MODEL" \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
-    --trust-remote-code --enforce-eager --omni --port 8091
-```
-
-### Sending requests
-```bash
-# Basic TTS
-python voxcpm/openai_speech_client.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This is a VoxCPM online text-to-speech example."
-
-# Voice cloning
-python voxcpm/openai_speech_client.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This sentence is synthesized with a cloned voice." \
-    --ref-audio /path/to/reference.wav \
-    --ref-text "The exact transcript spoken in reference.wav."
-
-# Streaming PCM
-python voxcpm/openai_speech_client.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This is a streaming VoxCPM request." \
-    --stream --output voxcpm_stream.pcm
-```
-
-### Notes
-- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`.
-- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`.
-- For benchmark measurement, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
-
 ---
 
 ## VoxCPM2
@@ -407,6 +435,14 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
 
 ## Example materials
 
+??? abstract "cosyvoice3/run_server.sh"
+
+    --8<-- "examples/online_serving/text_to_speech/cosyvoice3/run_server.sh"
+
+??? abstract "cosyvoice3/speech_client.py"
+
+    --8<-- "examples/online_serving/text_to_speech/cosyvoice3/speech_client.py"
+
 ??? abstract "fish_speech/gradio_demo.py"
     ``````py
     --8<-- "examples/online_serving/text_to_speech/fish_speech/gradio_demo.py"
@@ -423,6 +459,22 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
     ``````py
     --8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
     ``````
+??? abstract "glm_tts/gradio_demo.py"
+    ``````py
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"
+    ``````
+??? abstract "glm_tts/openai_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py"
+    ``````
+??? abstract "glm_tts/run_gradio_demo.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh"
+    ``````
+??? abstract "glm_tts/run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/run_server.sh"
+    ``````
 ??? abstract "omnivoice/run_server.sh"
     ``````sh
     --8<-- "examples/online_serving/text_to_speech/omnivoice/run_server.sh"
@@ -463,14 +515,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
     ``````py
     --8<-- "examples/online_serving/text_to_speech/qwen3_tts/tts_common.py"
     ``````
-??? abstract "voxcpm/openai_speech_client.py"
-    ``````py
-    --8<-- "examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py"
-    ``````
-??? abstract "voxcpm/run_server.sh"
-    ``````sh
-    --8<-- "examples/online_serving/text_to_speech/voxcpm/run_server.sh"
-    ``````
 ??? abstract "voxcpm2/gradio_demo.py"
     ``````py
     --8<-- "examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py"

docs/user_guide/quantization/autoround.md

@@ -32,7 +32,9 @@ guide. AutoRound is Intel-supported.
 |-------|------------|-------|--------|---------|
 | FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 | Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
-| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated |
+| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 

docs/user_guide/quantization/modelopt.md

@@ -6,34 +6,46 @@ ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The
 quantized weights and scale tensors are generated before serving, so inference
 does not run online calibration or convert a BF16 checkpoint at startup.
 
-vLLM-Omni currently validates the ModelOpt FP8 checkpoint path for diffusion
-transformers. The loader auto-detects supported ModelOpt FP8 checkpoint configs
-and keeps non-transformer components, such as the tokenizer, scheduler, text
+vLLM-Omni validates ModelOpt FP8, ModelOpt NVFP4, and ModelOpt mixed
+FP8/NVFP4 checkpoint loading for diffusion transformer stages. The loader
+auto-detects supported ModelOpt checkpoint configs and keeps non-transformer
+components, such as the tokenizer, scheduler, text encoder, vision/audio
 encoder, and VAE, on the base checkpoint unless a model-specific guide says
 otherwise.
 
 !!! note
-    `--force-cutlass-fp8` is an explicit runtime override for diffusion
-    checkpoints that already carry a supported ModelOpt FP8 config. It does not
-    quantize BF16 checkpoints and it does not apply to online `--quantization
-    fp8`. The flag only takes effect for ModelOpt FP8 diffusion stages on CUDA
-    SM89+ devices; other platforms and non-ModelOpt FP8 paths fall back to the
-    normal vLLM kernel selection.
+    ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
+    `--quantization fp8` for these checkpoints. The checkpoint
+    `quantization_config` selects the ModelOpt path.
+
+!!! note
+    `--force-cutlass-fp8`, `--linear-backend cutlass`, and
+    `--moe-backend cutlass` are runtime backend selections for checkpoints that
+    already carry supported ModelOpt quantized weights and scales. They do not
+    quantize BF16 checkpoints at startup.
 
 ## Supported ModelOpt Checkpoint Formats
 
 vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The
 checkpoint config must identify ModelOpt as the quantization method or producer,
-and the quantization algorithm must be one of the validated FP8 algorithms.
+and the quantization algorithm must be one of the validated algorithms below.
 
 | Checkpoint field | Supported value |
 |------------------|-----------------|
-| `method` / `quant_method` | `modelopt` |
+| `method` / `quant_method` | `modelopt`, `modelopt_fp4`, `modelopt_mixed` |
 | `producer.name` | `modelopt` |
-| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` |
+| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `NVFP4`, `MIXED_PRECISION` |
 
-Other ModelOpt algorithms, such as NVFP4, are not enabled by this diffusion
-FP8 path until they have separate model and quality validation.
+| `quant_algo` | Runtime method | Typical use |
+|--------------|----------------|-------------|
+| `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` | `modelopt` | FP8 diffusion transformer checkpoints |
+| `NVFP4` | `modelopt_fp4` | NVFP4 diffusion transformer checkpoints |
+| `MIXED_PRECISION` | `modelopt_mixed` | Mixed FP8/NVFP4 checkpoints with a ModelOpt per-layer policy |
+
+For multi-component diffusion or omni models, only the checkpoint component
+that contains ModelOpt quantized weights should use the ModelOpt quantization
+method. Encoders, decoders, tokenizers, schedulers, and other BF16 components
+stay unquantized unless the model-specific recipe validates otherwise.
 
 ## Hardware Support
 
@@ -47,7 +59,10 @@ FP8 path until they have separate model and quality validation.
 | Ascend NPU | ❌ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
-The optional CUTLASS FP8 runtime override requires CUDA SM89+.
+The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
+and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
+CUDA systems in the recipes below; other CUDA generations require separate
+backend and quality validation.
 
 ## Model Type Support
 
@@ -56,12 +71,19 @@ The optional CUTLASS FP8 runtime override requires CUDA SM89+.
 | Model | HF checkpoint | Scope | Status |
 |-------|---------------|-------|--------|
 | Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy` | Diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
 | Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
 | Wan2.2 | Not available | Diffusion transformer | Not validated |
 
+For full serving commands and benchmark context, see
+[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
+and
+[`recipes/Tencent/HunyuanImage-3.0-Instruct.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md).
+
 ### Multi-Stage Omni/TTS Model
 
 | Model | Scope | Status |
@@ -80,15 +102,28 @@ ModelOpt targets yet.
 
 ## Configuration
 
-For pre-quantized ModelOpt FP8 checkpoints, no `--quantization fp8` flag is
-needed. The checkpoint config selects the ModelOpt path.
+For pre-quantized ModelOpt checkpoints, no `--quantization fp8` flag is needed.
+The checkpoint config selects the ModelOpt path.
 
 Online serving:
 
 ```bash
-vllm serve <modelopt-fp8-checkpoint> \
+vllm serve <modelopt-checkpoint> \
   --omni \
   --tensor-parallel-size <N> \
+  --linear-backend cutlass \
+  --force-cutlass-fp8
+```
+
+For mixed FP8/NVFP4 MoE checkpoints, also select the validated MoE backend:
+
+```bash
+vllm serve <modelopt-mixed-moe-checkpoint> \
+  --omni \
+  --tensor-parallel-size <N> \
+  --enable-expert-parallel \
+  --linear-backend cutlass \
+  --moe-backend cutlass \
   --force-cutlass-fp8
 ```
 
@@ -96,14 +131,14 @@ Offline inference:
 
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
-  --model <modelopt-fp8-checkpoint> \
+  --model <modelopt-checkpoint> \
   --tensor-parallel-size <N> \
   --prompt "a red ceramic teapot on a wooden table" \
   --height 1024 \
   --width 1024 \
   --num-inference-steps 20 \
   --seed 42 \
-  --output outputs/modelopt_fp8.png
+  --output outputs/modelopt.png
 ```
 
 Python API:
@@ -112,7 +147,7 @@ Python API:
 from vllm_omni import Omni
 
 omni = Omni(
-    model="<modelopt-fp8-checkpoint>",
+    model="<modelopt-checkpoint>",
     tensor_parallel_size=2,
     force_cutlass_fp8=True,
 )
@@ -123,16 +158,34 @@ omni = Omni(
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ |
+| `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages |
+| `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints |
 
 ## Validation and Notes
 
-1. Compare the ModelOpt FP8 checkpoint against the BF16 baseline with the same
+1. Compare the ModelOpt checkpoint against the BF16 baseline with the same
    prompt, resolution, seed, and inference steps.
 2. Use `tests/diffusion/quantization/test_quantization_quality.py` with
    `VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model
    paths.
-3. Report LPIPS, PSNR, MAE, throughput, latency, and peak memory when adding a
-   new validated ModelOpt diffusion checkpoint.
-4. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
-   ModelOpt path only when the checkpoint already contains ModelOpt FP8 weights
-   and scales.
+3. For HunyuanImage-3.0 quantized DiT checkpoints, the opt-in accuracy check is:
+
+   ```bash
+   CUDA_VISIBLE_DEVICES=2,3 \
+   HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY=1 \
+   HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
+   HUNYUAN_IMAGE3_QUANT_TP=2 \
+   HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
+   HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
+   HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+   PYTHONPATH=/path/to/vllm-omni:${PYTHONPATH:-} \
+   python -m pytest -s -v \
+     tests/e2e/accuracy/test_hunyuan_image3.py \
+     -k quantized_dit_matches_bf16_accuracy
+   ```
+
+4. Report CLIP score deltas, SSIM, PSNR, throughput, latency, and peak memory
+   when adding a new validated ModelOpt diffusion checkpoint.
+5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
+   ModelOpt path only when the checkpoint already contains ModelOpt quantized
+   weights and scales.

docs/user_guide/quantization/mxfp4.md

@@ -0,0 +1,479 @@
+# W4A4 MXFP4 Quantization
+
+## Overview
+
+W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4
+(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups
+of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale.
+
+vLLM-Omni provides two quantization methods with different scale structures:
+
+| Method | Scale structure | Mode | Use case |
+|--------|----------------|------|----------|
+| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
+| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
+
+!!! tip "Recommended: `mxfp4_dualscale` offline"
+    For production deployments, use the `mxfp4_dualscale` offline mode with a
+    pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
+    calibrated `mul_scale` tensors from disk, providing measurably better accuracy
+    than any online method. The one-time preprocessing cost amortises across all
+    subsequent inference runs.
+
+    Use `mxfp4` online only for quick experimentation where preprocessing time
+    is not acceptable and accuracy loss is tolerable.
+
+!!! warning "Online single-scale ≠ Offline dual-scale"
+    `mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`:
+    fine scale (per-32 K), coarse scale (per-512 K), and per-input-channel
+    `mul_scale` from calibration — all loaded from the checkpoint.
+    `mxfp4_dualscale` online mode uses `NPUMxfp4DualScaleOnlineLinearMethod`:
+    dual-level scales computed on the fly from BF16 weights; no calibration
+    `mul_scale` is available. Loading an offline checkpoint with the online
+    method (or vice versa) will produce incorrect results or shape errors.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
+| NVIDIA Ampere GPU (SM 80+) | ⭕ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU (Atlas 950 A5) | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
+
+## Model Type Support
+
+### Diffusion Model (Wan2.2)
+
+| Model | Online | Offline | Notes |
+|-------|--------|---------|-------|
+| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
+| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
+| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
+
+The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
+quantization quality, not model compatibility — both work on cascade (A14B) and
+single-transformer models alike, the same as `mxfp8` online:
+
+- `mxfp4`: single-scale, lower overhead, simpler compute, online only
+- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline
+
+**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced
+by the merge script); there is no offline `mxfp4` single-scale format.
+
+!!! note "Per-layer BF16 fallback in offline cascade models"
+    The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most
+    linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers
+    retain their original BF16 weights and are listed in `ignored_layers` inside
+    each transformer's `config.json`. The two transformers may have different
+    `ignored_layers` sets — the pipeline reads each transformer's own `config.json`
+    and rebuilds the config locally when they differ, so routing is always
+    per-transformer-accurate.
+
+!!! warning "TI2V-5B not supported"
+    Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter
+    count makes it significantly more sensitive to 4-bit quantization noise,
+    resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B.
+
+## Configuration
+
+### `mxfp4` — Single-Scale Online Mode
+
+Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
+MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
+(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
+calibration `mul_scale` is available. Applies equally to single-transformer
+and cascade (A14B) models — both transformers in a cascade receive the same
+quantization config automatically.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="<your-model>", quantization="mxfp4")
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+```bash
+# Single-transformer or cascade model — same command
+python text_to_video.py --model <your-model> --quantization mxfp4
+
+# Online serving
+vllm serve <your-model> --omni --quantization mxfp4
+```
+
+### `mxfp4_dualscale` — DualScale Online Mode
+
+Online DualScale mode computes both fine and coarse scales on the fly from BF16
+weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
+single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
+DualScale provides better quantization accuracy at higher compute cost.
+
+The default configuration keeps the leading 5 transformer blocks in BF16
+(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
+is sufficient to meet quality requirements and is the recommended setting.
+
+```python
+omni = Omni(model="<your-model>", quantization="mxfp4_dualscale")
+```
+
+```bash
+python text_to_video.py --model <your-model> --quantization mxfp4_dualscale
+```
+
+If accuracy debugging identifies additional precision-sensitive layers, they
+can be pinned to BF16 via the Python API:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization_config={
+        "method": "mxfp4_dualscale",
+        "ignored_layers": ["blocks.10.attn1.to_q"],   # explicit per-layer override
+    },
+)
+```
+
+BF16 fallback routing in online mode applies two rules in priority order:
+
+1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix
+   matches is kept in BF16 regardless of block index.
+2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N
+   transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults
+   to `5` (recommended). Layers outside `blocks.N.*`
+   (e.g. `condition_embedder`) are always quantized.
+
+### `mxfp4_dualscale` — DualScale Offline Mode (Recommended)
+
+Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A
+preprocessing step converts the raw quantized output to the diffusers format
+expected by vLLM-Omni and injects the quantization config into each
+`transformer/config.json` so that vLLM-Omni auto-detects the offline path
+without a `--quantization` flag.
+
+BF16 fallback layers may be interleaved anywhere in the transformer — they are
+not restricted to leading blocks. The merge script detects them from
+`quant_model_description.json` and writes their prefixes into `ignored_layers`
+inside `config.json`. At runtime, each layer's prefix is matched against
+`ignored_layers` to decide BF16 vs. MXFP4 DualScale.
+
+#### Checkpoint tensor layout
+
+Each quantized linear layer stores four tensors:
+
+| Tensor | Shape | dtype | Description |
+|--------|-------|-------|-------------|
+| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) |
+| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) |
+| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale |
+| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) |
+
+BF16 fallback layers have no quantization tensors; only the original `weight`
+(and optional `bias`) are present, loaded directly from the base checkpoint.
+
+#### Step 1 — Quantize with msModelSlim
+
+```bash
+msmodelslim quant \
+  --model_path /path/to/Wan2.2-T2V-A14B-Diffusers \
+  --save_path  /path/to/wan2_2_t2v_quantized_raw \
+  --device npu \
+  --model_type Wan2_2 \
+  --config_path /path/to/wan2_2_w4a4_mxfp4_dualscale.yaml \
+  --trust_remote_code True
+```
+
+After this step, `--save_path` contains raw quantized safetensors files,
+scale files, and a metadata JSON (`quant_model_description*.json`).
+
+For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two
+subdirectories: `high_noise_model/` (transformer) and `low_noise_model/`
+(transformer_2).
+
+#### Step 2 — Preprocess with merge_mxfp4_dualscale_checkpoint.py
+
+The script (`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`):
+
+1. Copies the original diffusers model to `--output-path` (VAE, text encoder,
+   scheduler, etc. are preserved).
+2. Remaps tensor names from msModelSlim convention to diffusers convention and
+   strips `.linear.` / `.div.` wrappers added by the quantization tool.
+3. Overlays MXFP4 tensors (weight, fine/coarse scales, `mul_scale`) onto the
+   BF16 base checkpoint. Non-quantized layers keep their original BF16 weights.
+4. Detects all linear layers that remain in BF16 and writes their prefixes into
+   `ignored_layers` in `config.json`.
+5. Injects `quantization_config` so vLLM-Omni auto-detects offline MXFP4
+   DualScale.
+
+For cascade MoE models, steps 2–5 run separately for each transformer.
+
+```bash
+python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \
+  --model-type     Wan2.2-T2V-A14B \
+  --original-model /path/to/Wan2.2-T2V-A14B-Diffusers \
+  --quant-path     /path/to/wan2_2_t2v_quantized_raw \
+  --output-path    /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
+```
+
+| Argument | Description |
+|----------|-------------|
+| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` |
+| `--original-model` | Root directory of the original BF16 diffusers model |
+| `--quant-path` | Root directory of the msModelSlim quantized output |
+| `--output-path` | Output directory for the merged model (created by the script) |
+
+The script outputs a complete diffusers model directory at `--output-path`,
+with each transformer subfolder containing:
+
+- `diffusion_pytorch_model.safetensors` — MXFP4 weights + scale tensors, with BF16 fallback layers from the base checkpoint
+- `config.json` — original transformer config with `quantization_config` injected
+- `quant_model_description.json` — quantization metadata (reference only)
+
+The `quantization_config` injected into `config.json` for each transformer:
+
+```json
+{
+  "quant_method": "mxfp4_dualscale",
+  "is_checkpoint_serialized": true,
+  "ignored_layers": [
+    "blocks.0.attn1.to_qkv",
+    "blocks.0.attn1.to_out",
+    "proj_out"
+  ]
+}
+```
+
+`ignored_layers` lists every linear layer that retains its original BF16 weight,
+using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out`
+unindexed). The exact entries are determined by the quantization tool (msModelSlim)
+and may differ between `transformer` and `transformer_2` in a cascade model.
+
+#### Step 3 — Serve
+
+```bash
+python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
+
+# Online serving
+vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni
+```
+
+```python
+omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale")
+```
+
+!!! note
+    No `--quantization` flag is needed for offline mode. The preprocessing
+    script injects `quantization_config` into each `transformer/config.json`,
+    which vLLM-Omni reads automatically to activate the correct offline path.
+
+## Parameters
+
+### `mxfp4` (single-scale, online only)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | — | `"mxfp4"` |
+| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 |
+
+### `mxfp4_dualscale` (dual-scale, online + offline)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | — | `"mxfp4_dualscale"` |
+| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script |
+| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
+| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B |
+
+#### BF16 fallback priority (online mode)
+
+```
+for each linear layer:
+    if prefix in ignored_layers               → BF16  (explicit override, highest priority)
+    elif block_idx < num_bf16_fallback_layers → BF16  (coarse leading-block rule)
+    else                                      → MXFP4 DualScale online
+```
+
+Layers outside `blocks.N.*` (e.g. `condition_embedder.*`) are always quantized
+unless they appear in `ignored_layers`.
+
+## Validation and Notes
+
+1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using
+   `npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is
+   available — all output partitions receive an identity pre-scale. No offline
+   checkpoint format exists for this method.
+
+2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)**
+   quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse
+   scales computed on the fly). No calibration `mul_scale`; leading blocks or
+   explicit `ignored_layers` stay in BF16 for accuracy.
+
+3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** —
+   **recommended for production** — loads four tensors per quantized layer: FP4
+   weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale
+   (`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback
+   layers have no quantization tensors and are routed via `ignored_layers`.
+
+4. **Scale dtype**: fine scales are stored as `uint8` in safetensors (same bit
+   layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy
+   float32 round-trip.
+
+5. **Cascade model config propagation**: in a cascade model (transformer +
+   transformer_2), vLLM-Omni reads each transformer's own `config.json` and
+   rebuilds the quant config locally when `ignored_layers` differs between
+   transformers, ensuring per-layer routing is accurate for each. The first
+   transformer's config is propagated to `od_config` so the second transformer
+   can reuse it as a starting point.
+
+6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a
+   single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the
+   fused name (`attn1.to_qkv`), written automatically by the merge script.
+
+7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The
+   DualScale offline method mitigates this with calibrated `mul_scale` smooth
+   quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
+   off compression vs. accuracy for precision-sensitive layers.
+
+## Adapting MXFP4 for a New Model
+
+This section is aimed at developers who want to add MXFP4 support to a model
+other than Wan2.2. The three integration points are: (1) discovering the correct
+runtime layer names, (2) wiring `ignored_layers` into the model, and (3) writing
+a merge script for offline checkpoints.
+
+### Step 1 — Discover runtime layer names
+
+`ignored_layers` entries must match the **runtime parameter names** used inside
+vllm-omni, which may differ from the names stored in the diffusers checkpoint.
+The canonical source of truth is the model's own `named_parameters()`.
+
+```python
+from vllm_omni import Omni
+
+# Load the model without quantization to inspect parameter names.
+omni = Omni(model="/path/to/your-model")  # no --quantization flag
+for name, _ in omni.pipeline.transformer.named_parameters():
+    if "weight" in name and "scale" not in name:
+        print(name)
+```
+
+Compare the printed names against the diffusers checkpoint keys
+(`safetensors.safe_open` or `torch.load`) to identify any renames your model
+applies. Common patterns that differ in Wan2.2 (and may appear in other
+models):
+
+| Diffusers checkpoint name | vllm-omni runtime name | Reason |
+|---------------------------|------------------------|--------|
+| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` |
+| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores |
+| `ffn.net.2` | `ffn.net_2` | Same underscore rule |
+| `to_out.0` | `to_out` | Sequential index stripped |
+
+If your model has different fusion patterns, inspect `packed_modules_mapping`
+on the model class — this dict records how checkpoint keys are mapped to
+fused runtime parameters.
+
+!!! warning "Partial QKV fallback is not allowed"
+    If your model fuses Q, K, V into a single layer, `ignored_layers` must
+    include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but
+    `to_k`, `to_v` quantized) cannot be expressed at runtime because they share
+    one `QKVParallelLinear`. The merge script enforces this and raises an error
+    if only some of the trio appear as non-quantized.
+
+### Step 2 — Add ignored_layers to the model
+
+#### Online mode
+
+Pass `ignored_layers` directly in the quantization config using the **runtime
+names** discovered in Step 1. No code changes to the model are required.
+
+```python
+omni = Omni(
+    model="/path/to/your-model",
+    quantization={
+        "method": "mxfp4_dualscale",
+        "ignored_layers": [
+            "blocks.0.attn1.to_qkv",   # runtime name, not diffusers name
+            "blocks.0.attn1.to_out",
+            "blocks.0.ffn.net_0.proj",
+        ],
+    },
+)
+```
+
+```bash
+# CLI does not support list-typed ignored_layers directly.
+# Use the Python API or set ignored_layers in config.json (offline).
+python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale
+```
+
+The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers
+individually: set it to N to keep all linear layers in blocks 0 … N-1 in BF16.
+The right value depends on the model's sensitivity; evaluate on a validation
+set and pick the smallest N that meets your accuracy target.
+
+#### Offline mode
+
+For offline checkpoints, `ignored_layers` is written into each transformer's
+`config.json` by the merge script (see Step 3). No manual editing is needed if
+the merge script is correct. The injected block:
+
+```json
+{
+  "quant_method": "mxfp4_dualscale",
+  "is_checkpoint_serialized": true,
+  "ignored_layers": [
+    "blocks.0.attn1.to_qkv",
+    "blocks.0.attn1.to_out"
+  ]
+}
+```
+
+To add a layer manually (e.g. to pin an additional layer to BF16 without
+re-running the merge script), edit `config.json` inside the transformer
+subfolder. Use runtime names, not diffusers checkpoint names.
+
+### Step 3 — Write a merge script for offline mode
+
+The merge script for a new model mirrors
+`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`. The four
+things it must do:
+
+1. **Remap tensor names** from the quantization tool convention to diffusers
+   convention (strip wrappers like `.linear.`, `.div.`; fix any prefix
+   differences).
+
+2. **Collect ignored_layers**: after loading, enumerate all `*.weight` keys that
+   have no corresponding `*.weight_scale` (i.e. layers the tool left in BF16).
+   Convert diffusers names to vllm-omni runtime names (fuse QKV, rename FFN
+   sub-modules, etc.). Write the result to `config.json`.
+
+3. **Inject `quantization_config`** into `config.json`:
+   ```python
+   config["quantization_config"] = {
+       "quant_method":              "mxfp4_dualscale",
+       "is_checkpoint_serialized":  True,
+       "ignored_layers":            ignored_layers,   # runtime names
+   }
+   ```
+
+4. **Save** the merged safetensors and the updated `config.json`.
+
+The key helper to implement is the diffusers-to-runtime name translator
+(equivalent to `_diffusers_to_vllm_ignored` in the Wan2.2 merge script).
+For each non-quantized diffusers weight key, apply your model's specific
+renaming rules and collect the results.
+
+!!! tip "Validate before serving"
+    After producing the offline checkpoint, load it without a `--quantization`
+    flag and verify that vLLM-Omni auto-detects the correct method. Check that
+    the layer count reported in the startup log matches expectations: quantized
+    layer count + `ignored_layers` count should equal total linear layer count.
+    Any mismatch indicates a name-mapping bug in the merge script.

docs/user_guide/quantization/online.md

@@ -9,20 +9,24 @@ separate quantized checkpoint.
 This mode is different from pre-quantized checkpoint formats such as GGUF,
 AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
 prepared before serving and are documented in their method-specific guides.
+For MXFP8 and MXFP4, use this page for load-time quantization from BF16
+checkpoints, and use the method-specific pages for offline checkpoints produced
+by msModelSlim and the merge tools.
 
 ## Hardware Support
 
-| Device | FP8 W8A8 | Int8 W8A8 |
-|--------|----------|-----------|
-| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ |
-| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ |
-| AMD ROCm | ⭕ | ⭕ |
-| Intel XPU | ⭕ | ⭕ |
-| Ascend NPU | ❌ | ✅ |
+| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 |
+|--------|----------|-----------|------------|------------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
+| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ |
+| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
-guide. FP8 on Ampere may use a weight-only path where available.
+guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and
+MXFP4 are documented for the Ascend NPU path.
 
 ## Model Type Support
 
@@ -32,6 +36,8 @@ guide. FP8 on Ampere may use a weight-only path where available.
 |--------|-------|----------------|--------|
 | FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
 | Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
+| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only |
+| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 
@@ -53,6 +59,9 @@ from vllm_omni import Omni
 
 omni_fp8 = Omni(model="<your-model>", quantization="fp8")
 omni_int8 = Omni(model="<your-model>", quantization="int8")
+omni_mxfp8 = Omni(model="<your-model>", quantization="mxfp8")
+omni_mxfp4 = Omni(model="<your-model>", quantization="mxfp4")
+omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale")
 ```
 
 CLI:
@@ -60,6 +69,9 @@ CLI:
 ```bash
 vllm serve <your-model> --omni --quantization fp8
 vllm serve <your-model> --omni --quantization int8
+vllm serve <your-model> --omni --quantization mxfp8
+vllm serve <your-model> --omni --quantization mxfp4
+vllm serve <your-model> --omni --quantization mxfp4_dualscale
 ```
 
 Per-component routing:
@@ -77,10 +89,11 @@ config = build_quant_config({
 
 | Parameter | Methods | Description |
 |-----------|---------|-------------|
-| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` |
-| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 |
+| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
+| `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
 | `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
 | `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
+| `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` |
 
 ## Validation and Notes
 
@@ -91,3 +104,5 @@ config = build_quant_config({
    model as supported.
 4. If a model already ships quantized weights, use the matching pre-quantized
    method guide instead of online quantization.
+5. For Ascend MXFP4 deployments, prefer offline `mxfp4_dualscale` checkpoints
+   when production quality is more important than avoiding preprocessing.

docs/user_guide/quantization/overview.md

@@ -9,20 +9,20 @@ type has a different quantization scope.
 
 | Mode | Guide | Description | Methods |
 |------|-------|-------------|---------|
-| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8 |
+| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 |
 | Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
-| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8 |
+| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
 
 ## Hardware Support
 
-| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | GGUF | AutoRound | msModelSlim |
-|--------|----------|-----------|----------|------------|------|-----------|-------------|
-| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
-| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
-| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
-| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
-| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ |
+| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim |
+|--------|----------|-----------|----------|------------|------------|------|-----------|-------------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
+| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
+| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
 guide. FP8 on Ampere may use a weight-only path where available.
@@ -42,6 +42,7 @@ otherwise.
 | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
 | ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
 | MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
+| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
 | GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
 | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
@@ -58,6 +59,7 @@ in BF16 unless the model guide explicitly adds support.
 | ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
 | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
 | MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
+| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
 | GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
 | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
@@ -73,6 +75,7 @@ attached to the intended stage rather than applied globally.
 | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
 | ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
 | MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
+| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
 | GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
 | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
@@ -100,7 +103,7 @@ config = build_quant_config({
 
 | Component | Default quantized? | Notes |
 |-----------|--------------------|-------|
-| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim |
+| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim |
 | Text encoder | No | Keep BF16 unless a method-specific guide documents support |
 | VAE | No | Keep BF16; storage-only paths are method-specific |
 | Scheduler/tokenizer | No | Loaded from the base model repository |

examples/offline_inference/image_to_video/image_to_video.py

@@ -173,8 +173,8 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
+        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
+        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
     )
 
     # Distributed and parallel execution

examples/offline_inference/text_to_image/README.md

@@ -35,6 +35,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
 | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
+| `HiDream-I1-Full` | 1024 x 1024 | 63.7  | 57.7 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -94,6 +95,7 @@ python text_to_image.py \
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
 | `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
 | `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
+| `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). |
 
 **NextStep-1.1 specific arguments:**
 
@@ -177,6 +179,24 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output flux2-dev.png
 ```
 
+
+### HiDream-I1-Full Models
+
+The `--auxiliary-text-encoder` parameter is required when running HiDream‑I1‑Full:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-I1-Full \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --seed 42 \
+  --guidance-scale 5.0 \
+  --tensor-parallel-size 1 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
+  --auxiliary-text-encoder meta-llama/Meta-Llama-3.1-8B-Instruct \
+  --output /output.png
+```
+
 ### Batch Requests (Multiple Prompts)
 
 You can pass multiple prompts in a single `generate` call.

examples/offline_inference/text_to_image/text_to_image.py

@@ -300,6 +300,12 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help=("Custom system prompt. Used when --use-system-prompt is custom. "),
     )
+    parser.add_argument(
+        "--auxiliary-text-encoder",
+        type=str,
+        default=None,
+        help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
+    )
     current_omni_platform.pre_register_and_update(parser)
     from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
@@ -392,6 +398,7 @@ def main():
         "profiler_config": args.profiler_config,
         "init_timeout": args.init_timeout,
         "stage_init_timeout": args.stage_init_timeout,
+        "auxiliary_text_encoder": args.auxiliary_text_encoder,
         **lora_args,
         **quant_kwargs,
     }

examples/offline_inference/text_to_speech/README.md

@@ -14,13 +14,13 @@ list of supported architectures across all modalities, see
 
 | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
 |---|---|---|---|---|---|---|
-| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 22.05 kHz |
+| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
+| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
 | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
-| VoxCPM | `openbmb/VoxCPM-0.5B` | split | ✓ | ✓ | — | 24 kHz |
 | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
 
@@ -41,7 +41,7 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
 
 ## CosyVoice3
 
-2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
+2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.
 
 ### Prerequisites
 ```bash
@@ -93,6 +93,35 @@ Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cos
 
 ---
 
+## GLM-TTS
+
+2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+    --model zai-org/GLM-TTS \
+    --text "你好，这是语音合成测试。" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "这是参考音频的文本内容。" \
+    --output-dir ./output
+```
+
+### Architecture
+```
+Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
+        (Llama-based)    (32k vocab)      (Flow Matching)
+```
+
+### Notes
+- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
+- Reference audio should be 3-10 seconds.
+- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
+- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method).
+- The `--model` path should point to the repository root (not `llm/` subdirectory).
+
+---
+
 ## Fish Speech S2 Pro
 
 4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
@@ -308,56 +337,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
 
 ---
 
-## VoxCPM
-
-Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
-
-### Prerequisites
-```bash
-pip install voxcpm soundfile
-# or use a local source tree:
-export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
-```
-
-### Quick start
-```bash
-python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
-    --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
-```
-
-### Voice cloning
-```bash
-python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
-    --text "This sentence is synthesized with a cloned voice." \
-    --ref-audio /path/to/reference.wav \
-    --ref-text  "The exact transcript spoken in reference.wav."
-```
-
-### Streaming
-Pass `--streaming` together with the legacy async-chunk stage config:
-```bash
-python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
-    --streaming \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
-    --text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
-```
-
-### Persistent HF config (optional)
-The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required.
-
-If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`.
-
-### Notes
-- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`.
-- Streaming is currently single-request oriented.
-- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
-- For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).
-
----
-
 ## VoxCPM2
 
 Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.

examples/offline_inference/text_to_speech/glm_tts/end2end.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS End-to-End Offline Inference Example.
+
+GLM-TTS is a two-stage TTS system:
+  - Stage 0 (AR): Llama-based model generates speech tokens from text
+  - Stage 1 (DiT): Flow matching model converts speech tokens to audio
+
+Usage:
+    # Sync two-stage (default)
+    python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+        --model /path/to/GLM-TTS \
+        --text "你好，这是一个语音合成测试。" \
+        --ref-audio /path/to/reference.wav \
+        --ref-text "参考音频的转录文本。" \
+        --output-dir ./output
+
+    # Async chunk mode (streaming DiT)
+    python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+        --model /path/to/GLM-TTS --async-chunk \
+        --text "你好，这是一个语音合成测试。" \
+        --ref-audio /path/to/reference.wav \
+        --ref-text "参考音频的转录文本。" \
+        --output-dir ./output
+"""
+
+import base64
+import io
+import logging
+import os
+import tempfile
+import time
+from typing import Any
+from urllib.request import urlopen
+
+import numpy as np
+import soundfile as sf
+import torch
+import yaml
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DEPLOY_CONFIG = os.path.join(
+    os.path.dirname(__file__),
+    "..",
+    "..",
+    "..",
+    "..",
+    "vllm_omni",
+    "deploy",
+    "glm_tts.yaml",
+)
+SAMPLE_RATE = 24000
+
+
+def _load_ref_audio(ref_audio: str) -> tuple[torch.Tensor, int]:
+    """Load reference audio from file path, URL, or data URI."""
+    if ref_audio.startswith(("http://", "https://")):
+        with urlopen(ref_audio, timeout=60) as response:
+            audio_obj: Any = io.BytesIO(response.read())
+    elif ref_audio.startswith("data:"):
+        _, _, encoded = ref_audio.partition(",")
+        audio_obj = io.BytesIO(base64.b64decode(encoded))
+    else:
+        audio_obj = ref_audio
+    wav_np, sr = sf.read(audio_obj, dtype="float32")
+    if wav_np.ndim > 1:
+        wav_np = wav_np.mean(axis=1)
+    return torch.from_numpy(wav_np), int(sr)
+
+
+def _concat_audio(audio_val: Any) -> np.ndarray:
+    """Concatenate audio tensors from multimodal output."""
+    if isinstance(audio_val, list):
+        tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None]
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False)
+    if isinstance(audio_val, torch.Tensor):
+        return audio_val.float().cpu().numpy().reshape(-1)
+    return np.asarray(audio_val, dtype=np.float32).reshape(-1)
+
+
+def _extract_sample_rate(audio_mm: dict) -> int:
+    """Extract sample rate from multimodal output dict."""
+    sr_raw = audio_mm.get("sr", SAMPLE_RATE)
+    if isinstance(sr_raw, list):
+        sr_raw = sr_raw[-1] if sr_raw else SAMPLE_RATE
+    if hasattr(sr_raw, "item"):
+        return int(sr_raw.item())
+    return int(sr_raw)
+
+
+def _modify_deploy_config(base_path: str, async_chunk: bool) -> str:
+    """Build deploy config with explicit sync/async mode and eager execution.
+
+    Mirrors the logic in ``tests/e2e/offline_inference/test_glm_tts.py``
+    (``_get_deploy_config``) so that example runs match CI behavior.
+    """
+    with open(base_path) as f:
+        cfg = yaml.safe_load(f)
+    cfg["async_chunk"] = async_chunk
+    for stage in cfg.get("stages", []):
+        stage["enforce_eager"] = True
+        if stage.get("stage_id") == 0:
+            stage["async_scheduling"] = bool(async_chunk)
+    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False, prefix="glm_tts_")
+    yaml.dump(cfg, tmp)
+    tmp.close()
+    return tmp.name
+
+
+def main(args):
+    """Run offline GLM-TTS inference."""
+    os.makedirs(args.output_dir, exist_ok=True)
+    base_deploy_config = args.deploy_config or DEFAULT_DEPLOY_CONFIG
+    deploy_config_path = _modify_deploy_config(base_deploy_config, args.async_chunk)
+
+    ref_audio_wav, ref_audio_sr = _load_ref_audio(args.ref_audio)
+    if not args.ref_text:
+        raise ValueError("GLM-TTS requires --ref-audio and --ref-text.")
+
+    inputs = [
+        {
+            "prompt": args.text,
+            "multi_modal_data": {
+                "audio": (ref_audio_wav.float().cpu().numpy(), ref_audio_sr),
+            },
+            "modalities": ["audio"],
+            "mm_processor_kwargs": {"prompt_text": args.ref_text},
+            "additional_information": build_glm_tts_prefill_metadata(
+                args.model,
+                args.text,
+                args.ref_text,
+            ),
+        }
+    ]
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=deploy_config_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    t_start = time.perf_counter()
+    outputs = omni.generate(inputs)
+    elapsed = (time.perf_counter() - t_start) * 1000
+
+    assert outputs, "No outputs returned"
+    audio_mm = outputs[0].multimodal_output
+    assert "audio" in audio_mm, "No audio output found"
+
+    audio = _concat_audio(audio_mm["audio"])
+    sr = _extract_sample_rate(audio_mm)
+    out_path = os.path.join(args.output_dir, "output.wav")
+    sf.write(out_path, audio, samplerate=sr, format="WAV")
+
+    logger.info("Saved %s (%.2fs @ %dHz)", out_path, len(audio) / sr, sr)
+    logger.info("Total inference: %.1f ms", elapsed)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="GLM-TTS Text-to-Speech Example")
+    parser.add_argument("--model", type=str, required=True, help="Model path")
+    parser.add_argument("--text", type=str, default="你好，这是一个语音合成测试。")
+    parser.add_argument("--output-dir", type=str, default="./output")
+    parser.add_argument("--ref-audio", type=str, required=True, help="Reference WAV path/URL")
+    parser.add_argument("--ref-text", type=str, required=True, help="Transcript of ref audio")
+    parser.add_argument("--deploy-config", type=str, default=None)
+    parser.add_argument(
+        "--async-chunk",
+        action="store_true",
+        default=False,
+        help="Enable async_chunk mode (streaming DiT). Default: sync two-stage.",
+    )
+    parser.add_argument("--log-stats", action="store_true")
+    parser.add_argument("--stage-init-timeout", type=int, default=600)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    main(parse_args())

examples/offline_inference/text_to_speech/voxcpm/end2end.py

@@ -1,220 +0,0 @@
-"""Minimal offline VoxCPM example for vLLM Omni."""
-
-from __future__ import annotations
-
-import asyncio
-import time
-from pathlib import Path
-from typing import Any
-
-import soundfile as sf
-import torch
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-from vllm_omni import AsyncOmni, Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-
-def _build_prompt(args) -> dict[str, Any]:
-    additional_information: dict[str, list[Any]] = {
-        "text": [args.text],
-        "cfg_value": [args.cfg_value],
-        "inference_timesteps": [args.inference_timesteps],
-        "min_len": [args.min_len],
-        "max_new_tokens": [args.max_new_tokens],
-    }
-    if args.streaming_prefix_len is not None:
-        additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
-    if args.ref_audio is not None:
-        additional_information["ref_audio"] = [args.ref_audio]
-    if args.ref_text is not None:
-        additional_information["ref_text"] = [args.ref_text]
-    return {
-        "prompt_token_ids": [1],
-        "additional_information": additional_information,
-    }
-
-
-def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
-    audio = mm.get("audio", mm.get("model_outputs"))
-    if audio is None:
-        raise ValueError("No audio output found in multimodal output.")
-    if isinstance(audio, list):
-        parts = [torch.as_tensor(item).float().cpu().reshape(-1) for item in audio]
-        audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
-    if not isinstance(audio, torch.Tensor):
-        audio = torch.as_tensor(audio)
-    return audio.float().cpu().reshape(-1)
-
-
-def _extract_sample_rate(mm: dict[str, Any]) -> int:
-    sr_raw = mm.get("sr", 24000)
-    if isinstance(sr_raw, list) and sr_raw:
-        sr_raw = sr_raw[-1]
-    if hasattr(sr_raw, "item"):
-        return int(sr_raw.item())
-    return int(sr_raw)
-
-
-def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path:
-    output_dir.mkdir(parents=True, exist_ok=True)
-    output_path = output_dir / f"output_{request_id}.wav"
-    sf.write(
-        output_path,
-        audio.float().cpu().clamp(-1.0, 1.0).numpy(),
-        sample_rate,
-        format="WAV",
-        subtype="PCM_16",
-    )
-    return output_path
-
-
-async def _run_streaming(args) -> Path:
-    prompt = _build_prompt(args)
-    output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio_streaming")
-    request_id = "streaming_example"
-    sample_rate = 24000
-    buffered_samples = 0
-    chunks: list[torch.Tensor] = []
-    started = time.perf_counter()
-    omni = AsyncOmni(
-        model=args.model,
-        stage_configs_path=args.stage_configs_path,
-        deploy_config=args.deploy_config,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
-    try:
-        async for stage_output in omni.generate(prompt, request_id=request_id):
-            mm = getattr(stage_output, "multimodal_output", None)
-            if not isinstance(mm, dict):
-                request_output = getattr(stage_output, "request_output", None)
-                if request_output is None:
-                    continue
-                mm = getattr(request_output, "multimodal_output", None)
-                if not isinstance(mm, dict) and getattr(request_output, "outputs", None):
-                    mm = getattr(request_output.outputs[0], "multimodal_output", None)
-            if not isinstance(mm, dict):
-                continue
-            audio = _extract_audio_tensor(mm)
-            if audio.numel() == 0:
-                continue
-            sample_rate = _extract_sample_rate(mm)
-            if audio.numel() > buffered_samples:
-                delta = audio[buffered_samples:]
-                buffered_samples = int(audio.numel())
-            else:
-                delta = audio
-                buffered_samples += int(delta.numel())
-            if delta.numel() > 0:
-                chunks.append(delta)
-        if not chunks:
-            raise RuntimeError("No streaming audio chunks received from VoxCPM.")
-        output_audio = torch.cat(chunks, dim=0)
-        output_path = _save_audio(output_audio, sample_rate, output_dir, request_id)
-        print(f"Saved streaming audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
-        return output_path
-    finally:
-        omni.shutdown()
-
-
-def _run_sync(args) -> Path:
-    prompt = _build_prompt(args)
-    output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio")
-    request_id = "sync_example"
-    started = time.perf_counter()
-    last_mm: dict[str, Any] | None = None
-    omni = Omni(
-        model=args.model,
-        stage_configs_path=args.stage_configs_path,
-        deploy_config=args.deploy_config,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
-    for stage_outputs in omni.generate(prompt):
-        request_output = getattr(stage_outputs, "request_output", None)
-        if request_output is None:
-            continue
-        outputs = getattr(request_output, "outputs", None)
-        if outputs:
-            for output in outputs:
-                mm = getattr(output, "multimodal_output", None)
-                if isinstance(mm, dict):
-                    last_mm = mm
-        mm = getattr(request_output, "multimodal_output", None)
-        if isinstance(mm, dict):
-            last_mm = mm
-    if last_mm is None:
-        raise RuntimeError("No audio output received from VoxCPM.")
-    output_path = _save_audio(
-        _extract_audio_tensor(last_mm),
-        _extract_sample_rate(last_mm),
-        output_dir,
-        request_id,
-    )
-    print(f"Saved audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
-    return output_path
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.")
-    parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help=("Override the deploy config path."),
-    )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help=(
-            "Legacy stage_args yaml path. Required for streaming "
-            "(vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml); "
-            "leave unset for non-streaming to use the auto-resolved deploy config."
-        ),
-    )
-    parser.add_argument(
-        "--streaming",
-        action="store_true",
-        default=False,
-        help="Stream audio chunks as they arrive via AsyncOmni.",
-    )
-    parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.")
-    parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
-    parser.add_argument("--ref-text", type=str, default=None, help="Transcript of the reference audio.")
-    parser.add_argument("--output-dir", type=str, default=None, help="Output directory for generated wav files.")
-    parser.add_argument("--cfg-value", type=float, default=2.0, help="Guidance value passed to VoxCPM.")
-    parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of diffusion timesteps.")
-    parser.add_argument("--min-len", type=int, default=2, help="Minimum latent length.")
-    parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum latent length.")
-    parser.add_argument(
-        "--streaming-prefix-len",
-        type=int,
-        default=3,
-        help="Streaming prefix length used by voxcpm_async_chunk.yaml.",
-    )
-    parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.")
-    parser.add_argument("--log-stats", action="store_true", help="Enable vLLM Omni stats logging.")
-    nullify_stage_engine_defaults(parser)
-    args = parser.parse_args()
-    if (args.ref_audio is None) != (args.ref_text is None):
-        raise ValueError("Voice cloning requires --ref-audio and --ref-text together.")
-    return args
-
-
-def main(args) -> None:
-    route = "streaming" if args.streaming else "sync"
-    print(f"Model: {args.model}")
-    print(f"Deploy config: {args.deploy_config or '<auto from HF model_type>'}")
-    print(f"Stage configs path: {args.stage_configs_path or '<unused>'}")
-    print(f"Route: {route}")
-    if route == "streaming":
-        asyncio.run(_run_streaming(args))
-    else:
-        _run_sync(args)
-
-
-if __name__ == "__main__":
-    main(parse_args())

examples/offline_inference/text_to_video/text_to_video.py

@@ -163,8 +163,8 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
+        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
+        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
     )
 
     # Distributed and parallel execution
@@ -237,6 +237,27 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
+def _extract_peak_memory_mb(result: Any) -> float:
+    """Pull worker-reported peak VRAM (MiB) generation result.
+
+    Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
+    """
+    if isinstance(result, list):
+        result = result[0] if result else None
+    if result is None:
+        return 0.0
+    val = getattr(result, "peak_memory_mb", 0.0)
+    if not val:
+        inner = getattr(result, "request_output", None)
+        if isinstance(inner, list):
+            inner = inner[0] if inner else None
+        val = getattr(inner, "peak_memory_mb", 0.0)
+    try:
+        return float(val or 0.0)
+    except (TypeError, ValueError):
+        return 0.0
+
+
 def main():
     args = parse_args()
     model_class_name = args.model_class_name
@@ -348,6 +369,10 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    peak_mb = _extract_peak_memory_mb(frames)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
+
     audio = None
     if isinstance(frames, list):
         frames = frames[0] if frames else None

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -101,3 +101,4 @@ Key arguments:
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
+- `--cache-backend`: enable `cache_dit` for acceleration.

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -63,6 +63,7 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -146,12 +147,29 @@ def main() -> None:
         hsdp_replicate_size=args.hsdp_replicate_size,
     )
 
+    cache_config = None
+    if args.cache_backend == "cache_dit":
+        cache_config = {
+            "Fn_compute_blocks": 1,
+            "Bn_compute_blocks": 0,
+            "max_warmup_steps": 4,
+            "max_cached_steps": 20,
+            "residual_diff_threshold": 0.24,
+            "max_continuous_cached_steps": 3,
+            "enable_taylorseer": False,
+            "taylorseer_order": 1,
+            "scm_steps_mask_policy": None,
+            "scm_steps_policy": "dynamic",
+        }
+
     omni = Omni(
         model=args.model,
         parallel_config=parallel_config,
         model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
     )
     start = time.perf_counter()
     outputs = omni.generate(prompt, sampling_params)

examples/online_serving/text_to_speech/README.md

@@ -15,11 +15,11 @@ For the full list of supported architectures across all modalities, see
 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
 |---|---|---|---|---|---|
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
+| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
 | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
 | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
-| VoxCPM | `openbmb/VoxCPM-0.5B` | ✓ | ✓ (PCM stream) | — | — |
 | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
 
@@ -95,6 +95,45 @@ For full request-shape documentation (all parameters, response formats, error co
 
 ---
 
+## GLM-TTS
+
+2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
+
+### Launch
+```bash
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+# or:
+bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS
+```
+
+### Sending requests
+```bash
+# Voice cloning (required)
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "你好，这是语音克隆测试。" \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "这是参考音频的文本内容。"
+
+# Custom format
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "Hello, this is a voice cloning test." \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "Transcript of the reference audio." \
+    --response-format mp3 -o output.mp3
+```
+
+### Gradio demo
+```bash
+bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
+```
+
+### Notes
+- Output: 24 kHz mono WAV via HiFT vocoder.
+- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
+- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
+
+---
+
 ## Fish Speech S2 Pro
 
 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
@@ -294,10 +333,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
 ./qwen3_tts/run_server.sh Base
 ```
 
-### Choosing an executor backend (uniproc vs mp)
-Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation.
+### Executor backend
+Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
 
-To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor.
+To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
 
 ### Sending requests
 ```bash
@@ -375,73 +414,11 @@ python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-de
 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
 
 ### Notes
-- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
+- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
 
 ---
 
-## VoxCPM
-
-Split-stage TTS at 24 kHz.
-
-### Prerequisites
-```bash
-pip install voxcpm
-# or use a local source tree:
-export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
-```
-
-If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir:
-```bash
-export VOXCPM_MODEL=/path/to/voxcpm-model
-export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
-mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
-cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
-cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
-python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
-```
-
-### Launch
-```bash
-export VOXCPM_MODEL=/path/to/voxcpm-model
-./voxcpm/run_server.sh                # async-chunk streaming (default)
-./voxcpm/run_server.sh sync           # non-streaming
-```
-Or directly:
-```bash
-vllm serve "$VOXCPM_MODEL" \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
-    --trust-remote-code --enforce-eager --omni --port 8091
-```
-
-### Sending requests
-```bash
-# Basic TTS
-python voxcpm/openai_speech_client.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This is a VoxCPM online text-to-speech example."
-
-# Voice cloning
-python voxcpm/openai_speech_client.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This sentence is synthesized with a cloned voice." \
-    --ref-audio /path/to/reference.wav \
-    --ref-text "The exact transcript spoken in reference.wav."
-
-# Streaming PCM
-python voxcpm/openai_speech_client.py \
-    --model "$VOXCPM_MODEL" \
-    --text "This is a streaming VoxCPM request." \
-    --stream --output voxcpm_stream.pcm
-```
-
-### Notes
-- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`.
-- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`.
-- For benchmark measurement, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).
-
----
-
 ## VoxCPM2
 
 Single-stage native AR TTS at 48 kHz.

examples/online_serving/text_to_speech/cosyvoice3/run_server.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Launch vLLM-Omni server for CosyVoice3 TTS
+#
+# Usage:
+#   ./run_server.sh
+#   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
+#
+# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
+# Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
+
+set -e
+
+MODEL="${MODEL:-FunAudioLLM/Fun-CosyVoice3-0.5B-2512}"
+PORT="${PORT:-8091}"
+
+EXTRA_ARGS=()
+if [[ -n "${NO_ASYNC_CHUNK:-}" ]]; then
+    EXTRA_ARGS+=(--no-async-chunk)
+fi
+
+echo "Starting CosyVoice3 server with model: $MODEL"
+
+vllm serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --trust-remote-code \
+    --omni \
+    "${EXTRA_ARGS[@]}"

examples/online_serving/text_to_speech/cosyvoice3/speech_client.py

@@ -0,0 +1,142 @@
+"""Client for CosyVoice3 TTS via /v1/audio/speech endpoint.
+
+CosyVoice3 has no built-in voice presets: every request is voice cloning
+driven by ``ref_audio`` + ``ref_text``. The defaults below point at the
+official upstream zero-shot prompt so the script runs out of the box.
+
+Examples:
+    # Voice cloning with the default upstream prompt
+    python speech_client.py --text "收到好友从远方寄来的生日礼物。"
+
+    # Custom reference clip + transcript
+    python speech_client.py --text "Hello, this is a cloned voice." \
+        --ref-audio /path/to/reference.wav \
+        --ref-text "Transcript of the reference audio."
+
+    # Streaming PCM output
+    python speech_client.py --text "Hello world" --stream --output output.pcm
+"""
+
+import argparse
+import base64
+import os
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8091"
+DEFAULT_API_KEY = "EMPTY"
+DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
+
+# Official CosyVoice zero-shot prompt and its transcript.
+DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。"
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file to a base64 data URL."""
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+    ext = audio_path.lower().rsplit(".", 1)[-1]
+    mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}
+    mime_type = mime_map.get(ext, "audio/wav")
+    with open(audio_path, "rb") as f:
+        audio_b64 = base64.b64encode(f.read()).decode("utf-8")
+    return f"data:{mime_type};base64,{audio_b64}"
+
+
+def run_tts(args) -> None:
+    """Generate speech via the /v1/audio/speech API."""
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "response_format": args.response_format,
+    }
+
+    if args.ref_audio.startswith(("http://", "https://")):
+        payload["ref_audio"] = args.ref_audio
+    else:
+        payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
+    payload["ref_text"] = args.ref_text
+
+    if args.stream:
+        payload["stream"] = True
+        payload["response_format"] = "pcm"
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}")
+    print("Generating audio...")
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    if args.stream:
+        output_path = args.output or "output.pcm"
+        with httpx.Client(timeout=300.0) as client:
+            with client.stream("POST", api_url, json=payload, headers=headers) as resp:
+                if resp.status_code != 200:
+                    print(f"Error: {resp.status_code}")
+                    print(resp.read().decode())
+                    return
+                total_bytes = 0
+                with open(output_path, "wb") as f:
+                    for chunk in resp.iter_bytes():
+                        f.write(chunk)
+                        total_bytes += len(chunk)
+                print(f"Streamed {total_bytes} bytes to: {output_path}")
+    else:
+        with httpx.Client(timeout=300.0) as client:
+            response = client.post(api_url, json=payload, headers=headers)
+
+        if response.status_code != 200:
+            print(f"Error: {response.status_code}")
+            print(response.text)
+            return
+
+        try:
+            text = response.content.decode("utf-8")
+            if text.startswith('{"error"'):
+                print(f"Error: {text}")
+                return
+        except UnicodeDecodeError:
+            pass
+
+        output_path = args.output or "output.wav"
+        with open(output_path, "wb") as f:
+            f.write(response.content)
+        print(f"Audio saved to: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="CosyVoice3 TTS client")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
+    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
+    parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name")
+    parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument(
+        "--ref-audio",
+        default=DEFAULT_REF_AUDIO,
+        help="Reference audio for voice cloning (path or URL)",
+    )
+    parser.add_argument(
+        "--ref-text",
+        default=DEFAULT_REF_TEXT,
+        help="Transcript of the reference audio",
+    )
+    parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)")
+    parser.add_argument(
+        "--response-format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio format (default: wav)",
+    )
+    parser.add_argument("--output", "-o", default=None, help="Output file path")
+    args = parser.parse_args()
+    run_tts(args)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_speech/glm_tts/gradio_demo.py

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gradio demo for GLM-TTS online serving via /v1/audio/speech API.
+
+Supports:
+  - Voice-cloned text-to-speech synthesis from reference audio (upload or URL)
+  - Streaming (progressive PCM chunks) and non-streaming modes
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_demo.py --api-base http://localhost:8091
+
+    # Or use run_gradio_demo.sh to start both server and demo together.
+"""
+
+import argparse
+import base64
+import io
+
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import httpx
+import numpy as np
+import soundfile as sf
+
+PCM_SAMPLE_RATE = 24000
+DEFAULT_API_BASE = "http://localhost:8091"
+
+
+def encode_audio_to_base64(audio_data: tuple) -> str:
+    """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
+    sample_rate, audio_np = audio_data
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = np.clip(audio_np, -1.0, 1.0)
+            audio_np = (audio_np * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:audio/wav;base64,{wav_b64}"
+
+
+def build_payload(
+    text: str,
+    ref_audio: tuple | None,
+    ref_audio_url: str,
+    ref_text: str,
+    response_format: str = "wav",
+    stream: bool = False,
+    max_new_tokens: int = 4096,
+) -> dict:
+    """Build the /v1/audio/speech request payload."""
+    if not text or not text.strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    payload: dict = {
+        "input": text.strip(),
+        "voice": "default",
+        "response_format": "pcm" if stream else response_format,
+        "stream": stream,
+        "max_new_tokens": max_new_tokens,
+    }
+
+    # Voice cloning: ref_audio takes priority over URL.
+    ref_url_stripped = ref_audio_url.strip() if ref_audio_url else ""
+    if ref_audio is not None:
+        payload["ref_audio"] = encode_audio_to_base64(ref_audio)
+    elif ref_url_stripped:
+        payload["ref_audio"] = ref_url_stripped
+    else:
+        raise gr.Error("GLM-TTS requires a reference audio file or URL.")
+
+    if not ref_text or not ref_text.strip():
+        raise gr.Error("GLM-TTS requires a transcript of the reference audio.")
+    payload["ref_text"] = ref_text.strip()
+
+    return payload
+
+
+def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens):
+    """Non-streaming: call /v1/audio/speech and return full audio."""
+    payload = build_payload(
+        text, ref_audio, ref_audio_url, ref_text, response_format, stream=False, max_new_tokens=max_new_tokens
+    )
+
+    try:
+        with httpx.Client(timeout=300.0) as client:
+            resp = client.post(
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            )
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+    if resp.status_code != 200:
+        raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+
+    content_type = resp.headers.get("content-type", "")
+    if "application/json" in content_type:
+        try:
+            raise gr.Error(f"Server error: {resp.json()}")
+        except ValueError:
+            pass
+
+    try:
+        if response_format == "pcm":
+            audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
+            return (PCM_SAMPLE_RATE, audio_np)
+        audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
+        if audio_np.ndim > 1:
+            audio_np = audio_np[:, 0]
+        return (sample_rate, audio_np.astype(np.float32))
+    except Exception as e:
+        raise gr.Error(f"Failed to decode audio: {e}")
+
+
+def generate_speech_stream(
+    api_base, text, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens, chunk_seconds=0.5
+):
+    """Streaming: yield PCM chunks for progressive playback."""
+    payload = build_payload(
+        text, ref_audio, ref_audio_url, ref_text, response_format, stream=True, max_new_tokens=max_new_tokens
+    )
+
+    min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
+    prebuffer_chunks = 2
+    pending: list[np.ndarray] = []
+    pending_len = 0
+    chunks_yielded = 0
+    prebuffer: list[np.ndarray] = []
+
+    try:
+        leftover = b""
+        with httpx.Client(timeout=300.0) as client:
+            with client.stream(
+                "POST",
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            ) as resp:
+                if resp.status_code != 200:
+                    resp.read()
+                    raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+                for chunk in resp.iter_bytes():
+                    if not chunk:
+                        continue
+                    raw = leftover + chunk
+                    usable = len(raw) - (len(raw) % 2)
+                    leftover = raw[usable:]
+                    if usable == 0:
+                        continue
+                    float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
+                    pending.append(float_samples)
+                    pending_len += len(float_samples)
+                    if pending_len >= min_chunk_samples:
+                        audio_chunk = np.concatenate(pending)
+                        pending.clear()
+                        pending_len = 0
+                        if chunks_yielded < prebuffer_chunks:
+                            prebuffer.append(audio_chunk)
+                            chunks_yielded += 1
+                            if chunks_yielded == prebuffer_chunks:
+                                yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
+                                prebuffer.clear()
+                        else:
+                            yield (PCM_SAMPLE_RATE, audio_chunk)
+        remaining = prebuffer + pending
+        if remaining:
+            yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+
+def on_stream_change(stream: bool):
+    """Lock format to PCM when streaming."""
+    if stream:
+        return gr.update(value="pcm", interactive=False)
+    return gr.update(value="wav", interactive=True)
+
+
+def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
+    """Build the Gradio interface."""
+    with gr.Blocks(title="GLM-TTS Demo") as demo:
+        gr.Markdown("# GLM-TTS - Voice Cloning")
+        gr.Markdown(f"**Server:** `{api_base}` | **Model:** GLM-TTS | **Output:** 24kHz")
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Enter text here (supports Chinese and English)...",
+                    lines=4,
+                )
+
+                with gr.Accordion("Reference Voice (required)", open=True):
+                    gr.Markdown("Upload or link a short reference audio (3-10s) and provide its transcript.")
+                    ref_audio = gr.Audio(
+                        label="Reference Audio",
+                        type="numpy",
+                        sources=["upload", "microphone"],
+                    )
+                    ref_audio_url = gr.Textbox(
+                        label="Reference Audio URL (alternative to upload)",
+                        placeholder="https://example.com/reference.wav",
+                        lines=1,
+                    )
+                    ref_text = gr.Textbox(
+                        label="Reference Audio Transcript (required for cloning)",
+                        placeholder="Exact transcript of the reference audio...",
+                        lines=2,
+                    )
+
+                with gr.Row():
+                    response_format = gr.Dropdown(
+                        choices=["wav", "mp3", "flac", "pcm"],
+                        value="wav",
+                        label="Audio Format",
+                        scale=1,
+                    )
+                    stream_checkbox = gr.Checkbox(
+                        label="Stream output",
+                        value=False,
+                        info="Progressive PCM streaming",
+                        scale=1,
+                    )
+                    max_new_tokens = gr.Slider(
+                        minimum=256,
+                        maximum=8192,
+                        value=4096,
+                        step=256,
+                        label="Max New Tokens",
+                        scale=1,
+                    )
+
+                generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    interactive=False,
+                    streaming=True,
+                    autoplay=True,
+                )
+                gr.Markdown(
+                    "### About\n"
+                    "- **GLM-TTS** by Zhipu AI: Llama-based AR + Flow Matching DiT\n"
+                    "- **Voice cloning**: upload 3-10s reference + transcript\n"
+                    "- **Streaming**: real-time PCM output\n"
+                    "- **24kHz** output via Vocos vocoder"
+                )
+
+        stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
+
+        all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens]
+
+        def dispatch(stream_enabled, *args):
+            if stream_enabled:
+                yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
+            else:
+                yield generate_speech(api_base, *args)
+
+        generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
+        demo.queue()
+    return demo
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Gradio demo for GLM-TTS")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
+    parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
+    parser.add_argument("--share", action="store_true", help="Share publicly")
+    parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk")
+    args = parser.parse_args()
+
+    print(f"Connecting to vLLM server at: {args.api_base}")
+    demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds)
+    demo.launch(server_name=args.host, server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OpenAI-compatible client for GLM-TTS via /v1/audio/speech endpoint.
+
+GLM-TTS is a two-stage TTS system (AR + DiT) that generates audio from text
+conditioned on reference speech. Each request requires ref_audio + ref_text.
+
+Usage:
+    # Voice cloning
+    python openai_speech_client.py --text "你好" --ref-audio file:///path/to/ref.wav --ref-text "参考文本"
+
+    # Streaming response, for async_chunk server mode
+    python openai_speech_client.py --text "你好" --stream --ref-audio file:///path/to/ref.wav --ref-text "参考文本"
+
+    # Specify output format
+    python openai_speech_client.py --text "你好" --ref-audio file:///path/to/ref.wav \
+        --ref-text "参考文本" --response-format mp3 -o output.mp3
+"""
+
+import argparse
+
+import httpx
+
+# Default server configuration
+DEFAULT_API_BASE = "http://localhost:8091"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def run_tts_generation(args) -> None:
+    """Run TTS generation via OpenAI-compatible /v1/audio/speech API."""
+    if not args.ref_audio or not args.ref_text:
+        raise ValueError("GLM-TTS requires --ref-audio and --ref-text for voice cloning.")
+
+    payload = {
+        "model": args.model,
+        "voice": "default",
+        "input": args.text,
+        "response_format": args.response_format,
+        "stream": bool(args.stream),
+        "ref_audio": args.ref_audio,
+        "ref_text": args.ref_text,
+    }
+    if args.max_new_tokens:
+        payload["max_new_tokens"] = args.max_new_tokens
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}")
+    print(f"Stream: {args.stream}")
+    print("Generating audio...")
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    if args.stream:
+        output_path = args.output or "tts_output.pcm"
+        with httpx.Client(timeout=300.0) as client, open(output_path, "wb") as f:
+            with client.stream("POST", api_url, json=payload, headers=headers) as response:
+                if response.status_code != 200:
+                    print(f"Error: {response.status_code}")
+                    response.read()
+                    print(response.text)
+                    return
+                for chunk in response.iter_bytes():
+                    f.write(chunk)
+        print(f"Streaming audio saved to: {output_path}")
+    else:
+        with httpx.Client(timeout=300.0) as client:
+            response = client.post(api_url, json=payload, headers=headers)
+        if response.status_code != 200:
+            print(f"Error: {response.status_code}")
+            print(response.text)
+            return
+        try:
+            text = response.content.decode("utf-8")
+            if text.startswith('{"error"'):
+                print(f"Error: {text}")
+                return
+        except UnicodeDecodeError:
+            pass
+        output_path = args.output or f"tts_output.{args.response_format}"
+        with open(output_path, "wb") as f:
+            f.write(response.content)
+        print(f"Audio saved to: {output_path}")
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="OpenAI-compatible client for GLM-TTS via /v1/audio/speech",
+    )
+
+    # Server configuration
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default=DEFAULT_API_BASE,
+        help=f"API base URL (default: {DEFAULT_API_BASE})",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_API_KEY,
+        help="API key (default: EMPTY)",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="glm-tts",
+        help="Model name/path",
+    )
+
+    # Input text
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Text to synthesize",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=None,
+        help="Maximum new tokens to generate (default: model default)",
+    )
+
+    # Output
+    parser.add_argument(
+        "--response-format",
+        type=str,
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio output format (default: wav)",
+    )
+    parser.add_argument(
+        "--stream",
+        action="store_true",
+        help="Request a streaming audio response (use with async_chunk server mode).",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=None,
+        help="Output audio file path (default: tts_output.<format>)",
+    )
+
+    # Voice cloning parameters
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio URL, file:// URI, or base64 data URL for voice cloning",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference audio (required with --ref-audio)",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_tts_generation(args)

examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+# Launch GLM-TTS server + Gradio demo together.
+#
+# Usage:
+#   ./run_gradio_demo.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
+
+set -e
+
+MODEL="${MODEL:-zai-org/GLM-TTS}"
+PORT="${PORT:-8091}"
+GRADIO_PORT="${GRADIO_PORT:-7860}"
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
+
+echo "Starting GLM-TTS server (port $PORT)..."
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm-omni serve "$MODEL" \
+    --deploy-config "$REPO_ROOT/vllm_omni/deploy/glm_tts.yaml" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni &
+SERVER_PID=$!
+
+cleanup() {
+    echo "Stopping server (PID $SERVER_PID)..."
+    kill $SERVER_PID 2>/dev/null
+    wait $SERVER_PID 2>/dev/null
+}
+trap cleanup EXIT
+
+# Wait for server to be ready.
+echo "Waiting for server to start..."
+for i in $(seq 1 120); do
+    if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
+        echo "Server ready."
+        break
+    fi
+    sleep 2
+done
+
+echo "Starting Gradio demo (port $GRADIO_PORT)..."
+python "$SCRIPT_DIR/gradio_demo.py" \
+    --api-base "http://localhost:$PORT" \
+    --port "$GRADIO_PORT"

examples/online_serving/text_to_speech/glm_tts/run_server.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+# Launch vLLM-Omni server for GLM-TTS models
+#
+# Usage:
+#   ./run_server.sh                           # Default model path, async_chunk mode
+#   ./run_server.sh /path/to/GLM-TTS          # Custom model path, async_chunk mode
+#   ./run_server.sh /path/to/GLM-TTS sync     # Sync two-stage mode
+#
+# NOTE: The model path should point to the repo ROOT (not llm/ subdirectory).
+# model_subdir/tokenizer_subdir in the pipeline config resolve subdirectories.
+
+set -e
+
+MODEL="${1:-zai-org/GLM-TTS}"
+MODE="${2:-async}"
+
+EXTRA_ARGS=()
+case "$MODE" in
+    async|async_chunk)
+        ;;
+    sync|no_async_chunk)
+        EXTRA_ARGS+=("--no-async-chunk")
+        ;;
+    *)
+        echo "Unknown mode: $MODE (expected async or sync)" >&2
+        exit 1
+        ;;
+esac
+
+echo "Starting GLM-TTS server with model: $MODEL (mode: $MODE)"
+
+vllm-omni serve "$MODEL" \
+    --deploy-config vllm_omni/deploy/glm_tts.yaml \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --trust-remote-code \
+    --omni \
+    "${EXTRA_ARGS[@]}"

examples/online_serving/text_to_speech/omnivoice/speech_client.py

@@ -6,6 +6,9 @@ Examples:
 
     # Specify language
     python speech_client.py --text "Bonjour, comment allez-vous?" --language French
+
+    # Use a specific uploaded/supported voice
+    python speech_client.py --text "Hello" --voice my_uploaded_voice
 """
 
 import argparse
@@ -21,15 +24,18 @@ def run_tts(args) -> None:
     payload = {
         "model": args.model,
         "input": args.text,
-        "voice": "default",
         "response_format": args.response_format,
     }
 
+    if args.voice:
+        payload["voice"] = args.voice
     if args.language:
         payload["language"] = args.language
 
     print(f"Model: {args.model}")
     print(f"Text: {args.text}")
+    if args.voice:
+        print(f"Voice: {args.voice}")
     if args.language:
         print(f"Language: {args.language}")
     print("Generating audio...")
@@ -68,6 +74,11 @@ def main():
     parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
     parser.add_argument("--model", "-m", default="k2-fsa/OmniVoice", help="Model name")
     parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument(
+        "--voice",
+        default=None,
+        help="Voice name (omit for auto voice; must match a supported or uploaded speaker if set)",
+    )
     parser.add_argument("--language", default=None, help="Language hint (e.g., English, Chinese, French)")
     parser.add_argument(
         "--response-format",

examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py

@@ -1,155 +0,0 @@
-"""OpenAI-compatible client for VoxCPM via /v1/audio/speech.
-
-Examples:
-    # Basic text-to-speech
-    python openai_speech_client.py --text "Hello from VoxCPM"
-
-    # Voice cloning
-    python openai_speech_client.py \
-        --text "This sentence uses the cloned voice." \
-        --ref-audio /path/to/reference.wav \
-        --ref-text "The exact transcript spoken in the reference audio."
-
-    # Streaming PCM output
-    python openai_speech_client.py \
-        --text "This is a streaming VoxCPM request." \
-        --stream \
-        --output output.pcm
-"""
-
-import argparse
-import base64
-import os
-
-import httpx
-
-DEFAULT_API_BASE = "http://localhost:8091"
-DEFAULT_API_KEY = "EMPTY"
-DEFAULT_MODEL = "OpenBMB/VoxCPM1.5"
-
-
-def encode_audio_to_base64(audio_path: str) -> str:
-    """Encode a local audio file to base64 data URL."""
-    if not os.path.exists(audio_path):
-        raise FileNotFoundError(f"Audio file not found: {audio_path}")
-
-    ext = audio_path.lower().rsplit(".", 1)[-1]
-    mime_map = {
-        "wav": "audio/wav",
-        "mp3": "audio/mpeg",
-        "flac": "audio/flac",
-        "ogg": "audio/ogg",
-    }
-    mime_type = mime_map.get(ext, "audio/wav")
-
-    with open(audio_path, "rb") as f:
-        audio_b64 = base64.b64encode(f.read()).decode("utf-8")
-    return f"data:{mime_type};base64,{audio_b64}"
-
-
-def build_payload(args) -> dict[str, object]:
-    payload: dict[str, object] = {
-        "model": args.model,
-        "input": args.text,
-        "response_format": "pcm" if args.stream else args.response_format,
-    }
-
-    if args.ref_audio:
-        if args.ref_audio.startswith(("http://", "https://", "data:")):
-            payload["ref_audio"] = args.ref_audio
-        else:
-            payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
-    if args.ref_text:
-        payload["ref_text"] = args.ref_text
-    if args.max_new_tokens is not None:
-        payload["max_new_tokens"] = args.max_new_tokens
-    if args.stream:
-        payload["stream"] = True
-
-    return payload
-
-
-def run_tts(args) -> None:
-    payload = build_payload(args)
-    api_url = f"{args.api_base}/v1/audio/speech"
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {args.api_key}",
-    }
-
-    print(f"Model: {args.model}")
-    print(f"Text: {args.text}")
-    if args.ref_audio:
-        print("Mode: voice cloning")
-        print(f"Reference audio: {args.ref_audio}")
-    else:
-        print("Mode: text-to-speech")
-
-    if args.stream:
-        output_path = args.output or "voxcpm_output.pcm"
-        with httpx.Client(timeout=300.0) as client:
-            with client.stream("POST", api_url, json=payload, headers=headers) as response:
-                if response.status_code != 200:
-                    print(f"Error: {response.status_code}")
-                    print(response.read().decode("utf-8", errors="ignore"))
-                    return
-
-                total_bytes = 0
-                with open(output_path, "wb") as f:
-                    for chunk in response.iter_bytes():
-                        if not chunk:
-                            continue
-                        f.write(chunk)
-                        total_bytes += len(chunk)
-        print(f"Streamed {total_bytes} bytes to: {output_path}")
-        return
-
-    with httpx.Client(timeout=300.0) as client:
-        response = client.post(api_url, json=payload, headers=headers)
-
-    if response.status_code != 200:
-        print(f"Error: {response.status_code}")
-        print(response.text)
-        return
-
-    try:
-        text = response.content.decode("utf-8")
-        if text.startswith('{"error"'):
-            print(f"Error: {text}")
-            return
-    except UnicodeDecodeError:
-        pass
-
-    output_path = args.output or "voxcpm_output.wav"
-    with open(output_path, "wb") as f:
-        f.write(response.content)
-    print(f"Audio saved to: {output_path}")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="VoxCPM OpenAI-compatible speech client")
-    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
-    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
-    parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path")
-    parser.add_argument("--text", required=True, help="Text to synthesize")
-    parser.add_argument("--ref-audio", default=None, help="Reference audio path, URL, or data URL")
-    parser.add_argument(
-        "--ref-text",
-        default=None,
-        help="The exact transcript spoken in the reference audio",
-    )
-    parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output")
-    parser.add_argument(
-        "--response-format",
-        default="wav",
-        choices=["wav", "pcm", "flac", "mp3", "aac", "opus"],
-        help="Audio format for non-streaming mode (default: wav)",
-    )
-    parser.add_argument("--max-new-tokens", type=int, default=None, help="Maximum tokens to generate")
-    parser.add_argument("--output", "-o", default=None, help="Output file path")
-    args = parser.parse_args()
-    run_tts(args)
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/text_to_speech/voxcpm/run_server.sh

@@ -1,38 +0,0 @@
-#!/bin/bash
-# Launch vLLM-Omni server for VoxCPM online speech serving.
-#
-# Usage:
-#   ./run_server.sh                 # default: async_chunk stage config
-#   ./run_server.sh async           # async_chunk stage config
-#   ./run_server.sh sync            # no-async-chunk stage config
-#   VOXCPM_MODEL=/path/to/model ./run_server.sh
-
-set -e
-
-MODE="${1:-async}"
-MODEL="${VOXCPM_MODEL:-OpenBMB/VoxCPM1.5}"
-
-case "$MODE" in
-    async)
-        STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml"
-        ;;
-    sync)
-        STAGE_CONFIG="vllm_omni/deploy/voxcpm.yaml"
-        ;;
-    *)
-        echo "Unknown mode: $MODE"
-        echo "Supported: async, sync"
-        exit 1
-        ;;
-esac
-
-echo "Starting VoxCPM server with model: $MODEL"
-echo "Stage config: $STAGE_CONFIG"
-
-vllm serve "$MODEL" \
-    --stage-configs-path "$STAGE_CONFIG" \
-    --host 0.0.0.0 \
-    --port 8091 \
-    --trust-remote-code \
-    --enforce-eager \
-    --omni

mkdocs.yml

@@ -120,7 +120,9 @@ plugins:
           inventories:
           - https://docs.python.org/3/objects.inv
           - https://typing-extensions.readthedocs.io/en/latest/objects.inv
-          - https://docs.aiohttp.org/en/stable/objects.inv
+          # Disabled because Read the Docs strict builds fail if this
+          # external inventory intermittently resets the connection.
+          # - https://docs.aiohttp.org/en/stable/objects.inv
           - https://pillow.readthedocs.io/en/stable/objects.inv
           - https://numpy.org/doc/stable/objects.inv
           # Temporarily disabled due to decompression errors

pyproject.toml

@@ -63,7 +63,7 @@ dev = [
     "funasr>=1.0.0",
     "lpips>=0.1.4", # for quantization fp8
     "orjson>=3.11.7", # to make v1/embeddings API fast
-    "voxcpm>=2.0.2", # for voxcpm model Test
+    "voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
     "torchdiffeq>=0.2.5" # for covo_audio Test
 ]
 
@@ -71,6 +71,11 @@ demo = [
     "gradio>=6.7.0",
 ]
 
+# Local-test-only extras
+local = [
+    "FlagEmbedding",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",
@@ -184,6 +189,7 @@ markers = [
     "core_model: L1&L2 tests (run in each PR)",
     "advanced_model: L3 level tests (run on each merge)",
     "full_model: L4 level tests (run nightly)",
+    "local_model: tests (run locally)",
     # function module markers
     "diffusion: Diffusion model tests",
     "omni: Omni model tests",
@@ -235,6 +241,8 @@ extend-ignore-identifiers-re = [
     ".*MoBA",
     ".*temperal_downsample",
     ".*_nd",
+    ".*nd$",
+    ".*EmbedND$",
     ".*nothink.*",
     ".*NOTHINK.*",
     ".*nin.*",

recipes/HiDream/HiDream-I1-Image.md

@@ -0,0 +1,106 @@
+# HiDream-I1-Full
+
+> Text-to-image serving and benchmark
+
+## Summary
+
+- Vendor: HiDream.ai
+- Model: `HiDream-ai/HiDream-I1-Full`
+- Task: Text-to-image generation
+- Mode: Online serving with the OpenAI-compatible API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`HiDream-ai/HiDream-I1-Full` on a single 80 GB A800, validate the normal online-serving path.
+
+## References
+
+- Model: <https://huggingface.co/HiDream-ai/HiDream-I1-Full>
+- Upstream or canonical docs:
+  [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
+- Related example under `examples/`:
+  [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
+
+
+## Hardware Support
+
+This recipe currently documents one CUDA GPU serving configuration.
+
+## GPU
+
+### 1x A800 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+Start the baseline server:
+
+```bash
+export MODEL_NAME_OR_PATH=HiDream-ai/HiDream-I1-Full
+vllm serve ${MODEL_NAME_OR_PATH} \
+   --omni \
+   --port 8092 \
+   --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \
+   --tensor-parallel-size 1 \
+   --vae_use_slicing \
+   --vae_use_tiling
+```
+
+
+You can also use the example launcher and pass the extra flags through:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-I1-Full \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --seed 42 \
+  --guidance-scale 5.0 \
+  --tensor-parallel-size 1 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
+  --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \
+  --output /workspace/mnt/cmss-wmd/result/output.png
+
+```
+
+#### Verification
+
+Run the existing client example after the server is ready:
+
+```bash
+python examples/online_serving/text_to_image/openai_chat_client.py \
+  --server http://localhost:8092 \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --output /tmp/hidream_i1_recipe.png
+```
+
+For a direct API smoke test:
+
+```bash
+curl -s http://localhost:8092/v1/chat/completions   -H "Content-Type: application/json"   -d '{
+    "messages": [
+      {"role": "user", "content": "The setting sun of late autumn dyes the riverside with a warm orange hue"}
+    ],
+    "extra_body": {
+      "num_inference_steps": 50,
+      "seed": 42,
+      "height": 1024,
+      "width": 1024
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output_online.png
+```
+
+
+
+#### Notes
+
+- Key flags: `--auxiliary-text-encoder` designates the path for the auxiliary text model meta-llama/Llama-3.1-8B-Instruct. For HiDream-I1-Full, unspecified use defaults meta-llama/Llama-3.1-8B-Instruct (downloaded from official Hugging Face), and custom paths are supported.

recipes/Qwen/Qwen-Image.md

@@ -130,3 +130,108 @@ and compare the output JSON or terminal metrics.
 - Key flags: `--step-execution` enables the step-wise runtime; `--max-num-seqs` controls how many compatible requests may stay active together.
 - Keep `--max-num-seqs 1` when you want the more conservative path, when traffic is mostly single-request, or when you are debugging correctness before measuring throughput.
 - Current batching is still shape-sensitive: different step progress can co-batch, but different resolutions do not yet co-batch.
+
+### 2x B200 ModelOpt mixed FP8/NVFP4
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA environment with two B200 GPUs
+- vLLM version: Match the vLLM-Omni checkout used for deployment
+- vLLM-Omni version or commit: Use the commit that contains ModelOpt mixed
+  FP8/NVFP4 checkpoint loading for diffusion models
+
+#### Command
+
+This configuration uses a pre-quantized ModelOpt mixed checkpoint for
+Qwen-Image-2512. The validated checkpoint keeps boundary, modulation,
+normalization, and output layers in BF16, uses FP8 for sensitive attention
+projections, and uses NVFP4 for heavier MLP / attention-output linear layers.
+
+Validated checkpoint:
+
+```text
+feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy
+```
+
+Start the server with the CUTLASS quantized linear backend:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+vllm serve feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091 \
+  --trust-remote-code \
+  --tensor-parallel-size 2 \
+  --linear-backend cutlass \
+  --force-cutlass-fp8
+```
+
+#### Verification
+
+Run a direct image generation request after the server is ready:
+
+```bash
+curl http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy",
+    "prompt": "A cinematic Warhammer 40,000 Ultramarine warrior in blue power armor, armor stained with alien ichor, swinging a chainsword at a Tyranid creature, gothic battlefield, dramatic lighting, high detail, no text, no watermark",
+    "negative_prompt": "blurry, low quality, text, watermark, logo, malformed armor, bad hands",
+    "size": "1024x1024",
+    "num_inference_steps": 20,
+    "true_cfg_scale": 4.0,
+    "guidance_scale": 4.0,
+    "seed": 48
+  }'
+```
+
+Check that:
+
+- The server responds on `http://localhost:8091/health`.
+- The response contains one generated image.
+- Logs show CUTLASS selection for ModelOpt FP8 / mixed-precision linear
+  layers instead of falling back to the default automatic backend choice.
+
+#### Benchmark
+
+The numbers below are loaded-once request-level measurements, not an HTTP
+concurrency benchmark. Each case loads the model once, runs warmup requests,
+then measures 100 sequential requests.
+
+```text
+Hardware: 2x B200
+Tensor parallel size: 2
+Requests: 100
+Concurrency: 1
+Resolution: 1024x1024
+Denoising steps: 20
+```
+
+| Config | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| BF16 | 2.696s | 2.700s | 2.707s | 99,000 MiB |
+| ModelOpt FP8 CUTLASS | 2.663s | 2.671s | 2.676s | 85,760 MiB |
+| Mixed FP8/NVFP4 CUTLASS | 2.584s | 2.590s | 2.594s | 84,412 MiB |
+
+The same 2-GPU setup was also used to sweep NVFP4 linear backends. Direct
+CUTLASS was the fastest validated backend for this checkpoint.
+
+| Backend | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| auto | 2.927s | 2.948s | 3.535s | 85,340 MiB |
+| flashinfer-cudnn | 3.523s | 3.528s | 3.766s | 84,500 MiB |
+| flashinfer-trtllm | 2.711s | 2.727s | 3.010s | 84,996 MiB |
+| flashinfer-cutlass | 2.935s | 2.950s | 3.011s | 85,340 MiB |
+| cutlass | 2.578s | 2.584s | 2.592s | 84,412 MiB |
+
+#### Notes
+
+- Use `--linear-backend cutlass` to select the validated quantized linear
+  backend from the CLI.
+- `--force-cutlass-fp8` keeps ModelOpt FP8 layers on the CUTLASS FP8 path.
+- The benchmark table intentionally reports Qwen-Image-2512 as
+  concurrency-1 request-level data; do not compare it directly with an online
+  HTTP concurrency benchmark.

recipes/README.md

@@ -33,15 +33,17 @@ recipes/
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
+| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
 | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
 | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
 | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
 | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
-| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark | 4x H100/H800 80GB |
+| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
+| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 
 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like

recipes/Tencent/HunyuanImage-3.0-Instruct.md

@@ -1,6 +1,8 @@
 # HunyuanImage-3.0-Instruct
 
-> DiT-only text-to-image serving and benchmark
+> DiT-only text-to-image serving and benchmark with FP8, tensor parallelism,
+> sequence parallelism, CFG parallelism, and ModelOpt mixed FP8/NVFP4
+> checkpoints.
 
 ## Summary
 
@@ -16,13 +18,15 @@ Use this recipe when you want to run or benchmark the HunyuanImage-3.0 DiT
 stage directly. This is the recommended first setup when validating DiT
 throughput, memory, FP8 kernels, sequence parallelism, or CFG parallelism.
 
-The recipe covers three 4-GPU configurations:
+The recipe covers three 4-GPU FP8 configurations and one 2-GPU ModelOpt mixed
+FP8/NVFP4 configuration:
 
 | Configuration | Parallelism | Notes |
 | --- | --- | --- |
 | `tp4_fp8` | TP=4 | Lowest per-GPU memory, higher communication overhead |
 | `tp2_fp8_sp2` | TP=2, SP=2, Ulysses=2 | Splits sequence work across two GPUs per TP group |
 | `tp2_fp8_cfgp2` | TP=2, CFG=2 | Runs CFG branches in parallel; fastest validated DiT setup |
+| `tp2_mixed_fp8_nvfp4` | TP=2, EP enabled | Uses FP8 dense layers and NVFP4 routed experts on B200 |
 
 ## References
 
@@ -222,3 +226,127 @@ of:
 - If you see OOM on 80GB GPUs, reduce image size, request concurrency, or use
   the TP=4 configuration before increasing batch size. GPU memory utilization
   is not a useful primary tuning knob for this DiT-only recipe.
+
+### 2x B200 ModelOpt mixed FP8/NVFP4
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA environment with two B200 GPUs
+- vLLM version: Match the vLLM-Omni checkout used for deployment
+- vLLM-Omni version or commit: Use the commit that contains ModelOpt mixed
+  FP8/NVFP4 checkpoint loading for diffusion models
+
+#### Command
+
+This configuration uses a pre-quantized ModelOpt mixed checkpoint for the
+HunyuanImage-3.0 DiT stage. Dense attention and shared dense projections use
+FP8, while routed MoE expert projections use NVFP4. Embeddings, norms,
+routers, output layers, and non-DiT components stay in BF16.
+
+Validated checkpoint:
+
+```text
+feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8
+```
+
+Start the server with CUTLASS selected from the CLI for quantized linear and
+MoE kernels:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+vllm serve feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091 \
+  --trust-remote-code \
+  --tensor-parallel-size 2 \
+  --enable-expert-parallel \
+  --linear-backend cutlass \
+  --moe-backend cutlass \
+  --force-cutlass-fp8
+```
+
+#### Verification
+
+Run a direct image generation request after the server is ready:
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8",
+    "messages": [
+      {"role": "user", "content": "A cinematic portrait of an Adeptus Custodes warrior standing on Terra before the Imperial Palace, ornate golden armor, guardian spear, crimson cloak, marble steps, gothic architecture, holy golden light, highly detailed, no text, no watermark"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 20,
+      "guidance_scale": 4.0,
+      "true_cfg_scale": 4.0,
+      "seed": 48
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' \
+     | cut -d',' -f2- \
+     | base64 -d > hunyuan_image3_mixed_nvfp4.png
+```
+
+Check that:
+
+- The server responds on `http://localhost:8091/health`.
+- The response contains one generated image.
+- Logs show CUTLASS selection for ModelOpt FP8 dense layers and ModelOpt
+  NVFP4 routed expert layers.
+
+#### Benchmark
+
+The numbers below are online HTTP measurements with 16 concurrent requests.
+Each case serves 100 image-generation requests on the same 2-GPU B200 setup.
+
+```text
+Hardware: 2x B200
+Tensor parallel size: 2
+Requests: 100
+Concurrency: 16
+Resolution: 1024x1024
+Denoising steps: 20
+```
+
+| Config | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| BF16 | 70.729s | 77.716s | 77.758s | 191,920 MiB |
+| ModelOpt FP8 | 58.603s | 64.435s | 64.469s | 115,816 MiB |
+| Mixed FP8/NVFP4 | 54.739s | 60.244s | 60.317s | 83,560 MiB |
+
+#### Accuracy
+
+Run the HunyuanImage-3.0 quantized DiT accuracy check with the BF16, FP8, and
+mixed FP8/NVFP4 checkpoints:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
+HUNYUAN_IMAGE3_QUANT_TP=2 \
+HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
+HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
+HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+PYTHONPATH=/path/to/vllm-omni:/path/to/vllm:${PYTHONPATH:-} \
+python -m pytest -s -v \
+  tests/e2e/accuracy/test_hunyuan_image3.py \
+  -k quantized_dit_matches_bf16_accuracy
+```
+
+The test compares quantized images with the BF16 image using CLIP-based image
+and text scores, plus structural SSIM / PSNR checks. The current structural
+thresholds are `SSIM >= 0.20` and `PSNR >= 10.0`.
+
+#### Notes
+
+- Use `--linear-backend cutlass` to select the validated quantized linear
+  backend from the CLI.
+- Use `--moe-backend cutlass` for the routed expert path.
+- `--force-cutlass-fp8` keeps the FP8 dense layers on the CUTLASS FP8 path.
+- The benchmark table is an online HTTP concurrency-16 result; do not compare
+  it directly with loaded-once request-level benchmarks.

recipes/fishaudio/Fish-Speech-S2-Pro.md

@@ -102,3 +102,30 @@ curl -X POST http://localhost:8091/v1/audio/speech \
   both `ref_audio` and `ref_text` parameters.
 - Memory usage: Model loads at ~48.3 GiB, peaks at ~48.9 GiB during inference
   headroom for video frames and audio caches.
+
+### 2x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Required package: `fish-speech` (for DAC codec)
+- CUDA 12.8
+- vLLM version: 0.19.0+
+
+#### Command
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve fishaudio/s2-pro --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
+```
+
+This profile pins Stage0 (Slow/Fast AR) to GPU 0 and Stage1 (DAC decoder)
+to GPU 1 to remove AR/DAC contention, sets `max_num_seqs=64` on both
+stages so concurrencies above 8 don't queue, and reduces
+`connector_get_sleep_s` from 10 ms to 1 ms so DAC chunks aren't held in
+the connector once they're ready.
+
+For benchmark numbers see the discussion in
+vllm-project/vllm-omni#2515 (initial H100x2 sweep by @Sy0307) and
+vllm-project/vllm-omni#3323 (H20x2 sweep with the shipped config).

recipes/zai-org/GLM-TTS.md

@@ -0,0 +1,105 @@
+# GLM-TTS for Chinese/English TTS on 1x GPU
+
+## Summary
+
+- Vendor: zai-org
+- Model: `zai-org/GLM-TTS`
+- Task: Zero-shot voice-cloned text-to-speech synthesis
+- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to serve GLM-TTS as a two-stage TTS system (AR + DiT
+flow-matching) for Chinese and English speech synthesis. Every request is
+conditioned on reference audio and its transcript.
+
+## References
+
+- Upstream or canonical docs:
+  [zai-org/GLM-TTS on HuggingFace](https://huggingface.co/zai-org/GLM-TTS)
+- GitHub repository:
+  [zai-org/GLM-TTS](https://github.com/zai-org/GLM-TTS)
+- Related example under `examples/`:
+  [`examples/online_serving/text_to_speech/README.md#glm-tts`](../../examples/online_serving/text_to_speech/README.md#glm-tts)
+- Offline inference example:
+  [`examples/offline_inference/text_to_speech/README.md#glm-tts`](../../examples/offline_inference/text_to_speech/README.md#glm-tts)
+
+## Hardware Support
+
+### GPU
+
+### 1x A40 48GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA CUDA environment with A40 48GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+Start the server from the repository root:
+
+```bash
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+```
+
+Async chunking is enabled by default in the bundled deployment config. For
+the sync (non-streaming) path:
+
+```bash
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
+```
+
+Use a custom deploy config for advanced cases:
+
+```bash
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
+  --deploy-config /path/to/your_glm_tts_overrides.yaml
+```
+
+#### Verification
+
+Run the bundled OpenAI-compatible client with reference audio:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+  --text "你好，这是一个语音合成测试。" \
+  --ref-audio file:///path/to/ref.wav \
+  --ref-text "这是参考音频的文本内容。"
+```
+
+For a quick API smoke test:
+
+```bash
+curl http://localhost:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "zai-org/GLM-TTS",
+    "input": "你好，这是一个语音合成测试。",
+    "response_format": "wav",
+    "ref_audio": "file:///path/to/ref.wav",
+    "ref_text": "这是参考音频的文本内容。"
+  }' --output test.wav
+```
+
+Voice cloning with reference audio:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+  --text "你好，这是语音克隆测试。" \
+  --ref-audio file:///path/to/ref.wav \
+  --ref-text "这是参考音频的文本内容。"
+```
+
+#### Notes
+
+- Hardware scope: the default bundled config is CUDA-only and verified on 1x A40 48GB (~16.6 GiB peak); fits 24GB cards. Split stages across GPUs for higher concurrency.
+- Memory usage: ~18-20GB total (AR ~10GB, DiT ~8GB); both stages share GPU 0 by default.
+- Audio output: 24kHz mono WAV via HiFT vocoder (Vocos2D 32kHz fallback with resampling).
+- Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback.
+- Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side.
+- Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40.

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/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/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

@@ -11,6 +11,8 @@ This runner separates two concepts:
 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, endpoint, benchmark_params,
@@ -652,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
@@ -699,16 +703,14 @@ def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, An
     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"]
@@ -744,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/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/cache/test_cache_backends.py

@@ -127,6 +127,10 @@ class TestCacheDiTBackend:
         """Test HunyuanImage3 custom cache-dit enabler is registered."""
         assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
 
+    def test_dreamid_custom_enabler_registered(self):
+        """Test DreamIDOmni custom cache-dit enabler is registered."""
+        assert "DreamIDOmniPipeline" in CUSTOM_DIT_ENABLERS
+
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
@@ -169,6 +173,47 @@ class TestCacheDiTBackend:
         with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
             backend.enable(mock_pipeline)
 
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.transformer = Mock()
+        mock_pipeline.transformer.fused_blocks = Mock()
+
+        mock_cache_dit.enable_cache = Mock()
+        mock_cache_dit.refresh_context = Mock()
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        assert backend._refresh_func is not None
+        mock_block_adapter.assert_called_once()
+        adapter_kwargs = mock_block_adapter.call_args.kwargs
+        assert adapter_kwargs["transformer"] is mock_pipeline.transformer
+        assert adapter_kwargs["blocks"] is mock_pipeline.transformer.fused_blocks
+        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_0
+        assert adapter_kwargs["has_separate_cfg"] is True
+        assert len(adapter_kwargs["params_modifiers"]) == 1
+        mock_cache_dit.enable_cache.assert_called_once()
+
+        backend.refresh(mock_pipeline, num_inference_steps=12)
+        mock_cache_dit.refresh_context.assert_called_once()
+        call_args = mock_cache_dit.refresh_context.call_args
+        assert call_args[0][0] is mock_pipeline.transformer
+        assert call_args[1]["num_inference_steps"] == 12
+
+    def test_enable_dreamid_pipeline_requires_fused_blocks(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.transformer = Mock(spec=[])
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+
+        with pytest.raises(ValueError, match="DreamIDOmniPipeline"):
+            backend.enable(mock_pipeline)
+
 
 class TestTeaCacheBackend:
     """Test TeaCacheBackend implementation."""

tests/diffusion/cache/test_cache_dit.py

@@ -8,14 +8,21 @@ Model specific tests for CacheDiT enablement.
 from unittest.mock import Mock, patch
 
 import pytest
+import torch
+from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 SEPARATE_CFG_ENABLERS = [
+    cd_backend.enable_cache_for_dreamid_omni,
     cd_backend.enable_cache_for_ltx2,
+    cd_backend.enable_cache_for_helios,
     cd_backend.enable_cache_for_wan22,
     cd_backend.enable_cache_for_longcat_image,
 ]
@@ -38,3 +45,85 @@ def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs
     assert adapter_kwargs["has_separate_cfg"] is True
+
+
+# This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM ROCm custom ops lack CPU fallback",
+)
+def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
+    """CacheDiT Pattern_0 treats the second positional arg as encoder_hidden_states,
+    which is a collision for one of the kwargs in LTX2 since we treat the audio
+    hidden states as encoder_hidden_states.
+
+    This test ensures that a tiny LTX2 transformer can be initialized and run
+    through the cache DiT backend without hitting a collision on the kwargs.
+    """
+    seq_len = 4
+    video_in = torch.full((1, seq_len, 16), 1.0)
+    audio_in = torch.full((1, seq_len, 16), 2.0)
+    text_in = torch.full((1, seq_len, 16), 3.0)
+    audio_text_in = torch.full((1, seq_len, 16), 4.0)
+
+    model = LTX2VideoTransformer3DModel(
+        in_channels=16,
+        out_channels=16,
+        patch_size=1,
+        patch_size_t=1,
+        num_attention_heads=2,
+        attention_head_dim=8,
+        cross_attention_dim=16,
+        audio_in_channels=16,
+        audio_out_channels=16,
+        audio_num_attention_heads=2,
+        audio_attention_head_dim=8,
+        audio_cross_attention_dim=16,
+        num_layers=2,
+        caption_channels=16,
+    )
+
+    # NOTE: This is currently using the LTX2 custom enabler, but the custom
+    # enablers will be consolidated after
+    # https://github.com/vllm-project/vllm-omni/pull/2527 lands.
+    LTX2Pipeline = type("LTX2Pipeline", (), {})
+    pipeline = LTX2Pipeline()
+    pipeline.transformer = model
+    backend = CacheDiTBackend(DiffusionCacheConfig())
+    backend.enable(pipeline)
+    backend.refresh(pipeline, num_inference_steps=5)
+
+    # Wrap call_Fn_blocks in CacheDiT so that we can verify the
+    # hidden/encoder states are what we expect them to be
+    captured = {}
+    original = CachedBlocks_Pattern_0_1_2.call_Fn_blocks
+
+    def call_Fn_blocks_and_capture(self, hidden_states, encoder_hidden_states, *a, **kw):
+        captured["hidden_states"] = hidden_states
+        captured["encoder_hidden_states"] = encoder_hidden_states
+        return original(self, hidden_states, encoder_hidden_states, *a, **kw)
+
+    # Also, map projections to identity so that we can just check
+    # the captured tensors directly instead of having to reproject
+    identity = torch.nn.Identity()
+    with (
+        patch.object(model, "proj_in", identity),
+        patch.object(model, "audio_proj_in", identity),
+        patch.object(CachedBlocks_Pattern_0_1_2, "call_Fn_blocks", call_Fn_blocks_and_capture),
+        torch.no_grad(),
+    ):
+        model(
+            hidden_states=video_in,
+            audio_hidden_states=audio_in,
+            encoder_hidden_states=text_in,
+            audio_encoder_hidden_states=audio_text_in,
+            timestep=torch.tensor([[1000.0] * seq_len]),
+            num_frames=1,
+            height=2,
+            width=2,
+            audio_num_frames=seq_len,
+        )
+
+    # Pattern_0 maps (hidden_states, encoder_hidden_states) to (video, audio)
+    assert torch.equal(captured["hidden_states"], video_in)
+    assert torch.equal(captured["encoder_hidden_states"], audio_in)

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/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/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/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/diffusion/test_omni_diffusion_config_master_port.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for OmniDiffusionConfig master_port resolution (issue #3794)."""
+
+import socket
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestOmniDiffusionConfigMasterPort:
+    def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("MASTER_PORT", "43210")
+        config = OmniDiffusionConfig(model="test")
+        assert config.master_port == 43210
+
+    def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        config = OmniDiffusionConfig(model="test", master_port=40000)
+        assert config.master_port == 40000
+
+    def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("MASTER_PORT", "43210")
+        config = OmniDiffusionConfig(model="test", master_port=40000)
+        assert config.master_port == 43210
+
+    def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        ports = {OmniDiffusionConfig(model="test", master_port=40123).master_port for _ in range(5)}
+        assert ports == {40123}
+
+    def test_default_master_port_is_available(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        config = OmniDiffusionConfig(model="test")
+        assert config.master_port is not None
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind(("", config.master_port))
+
+    def test_settle_port_steps_when_port_is_busy(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.bind(("", 0))
+            busy_port = sock.getsockname()[1]
+            config = OmniDiffusionConfig(model="test", master_port=busy_port)
+            assert config.master_port == busy_port + 37

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/e2e/accuracy/test_hunyuan_image3.py

@@ -6,12 +6,15 @@ from __future__ import annotations
 import copy
 import gc
 import importlib
+import json
 import os
 import tempfile
 import time
+from dataclasses import dataclass
 from pathlib import Path
 
 import pytest
+import requests
 import torch
 import yaml
 from PIL import Image
@@ -19,22 +22,37 @@ from PIL import Image
 from tests.e2e.accuracy.helpers import (
     CLIPScorer,
     SemanticSimilarityScorer,
+    assert_similarity,
     compute_image_ssim_psnr,
     download_images,
     model_output_dir,
 )
-from tests.helpers.runtime import OmniRunner
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner, OmniServer
 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]
 
 # ============================================================================
 # 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
@@ -63,54 +81,129 @@ THRESHOLDS = {
     "ssim": 0.20,  # Structural similarity
     "psnr": 11.0,  # Peak signal-to-noise ratio (dB)
 }
+
+QUANT_PROMPT = "A brown and white dog is running on the grass."
+QUANT_HEIGHT, QUANT_WIDTH = 1024, 1024
+QUANT_PSNR_THRESHOLD = 10.0
+QUANT_SSIM_THRESHOLD = 0.20
+QUANT_CLIP_SCORE_THRESHOLD = 20.0
+QUANT_CLIP_SCORE_DROP_THRESHOLD = float(os.environ.get("HUNYUAN_IMAGE3_QUANT_CLIP_SCORE_DROP_THRESHOLD", "5.0"))
+QUANT_RUN_ENV = "HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY"
+QUANT_BF16_ENV = "HUNYUAN_IMAGE3_BF16_MODEL"
+QUANT_FP8_ENV = "HUNYUAN_IMAGE3_FP8_MODEL"
+QUANT_NVFP4_ENV = "HUNYUAN_IMAGE3_NVFP4_MODEL"
+_TRUE_ENV_VALUES = {"1", "true", "yes", "on"}
 # fmt: off
-_BASE_CONFIG = {
-    "stage_args": [
+_DEPLOY_CONFIG = {
+    "pipeline": "hunyuan_image_3_moe",
+    "async_chunk": False,
+    "trust_remote_code": True,
+    "connectors": {
+        "shared_memory_connector": {
+            "name": "SharedMemoryConnector",
+            "extra": {"shm_threshold_bytes": 65536},
+        },
+    },
+    "stages": [
         {
-            "stage_id": 0, "stage_type": "llm",
-            "runtime": {"process": True, "devices": AR_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
-            "engine_args": {
-                "model_stage": "AR", "model_arch": "HunyuanImage3ForCausalMM",
-                "worker_cls": "vllm_omni.worker.gpu_ar_worker.GPUARWorker",
-                "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,
-                "hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}},
-            },
-            "is_comprehension": False, "final_output": True, "final_output_type": "text",
-            "default_sampling_params": {
-                "temperature": 0.0, "top_p": 1, "top_k": -1, "max_tokens": 8192,
-                "stop_token_ids": [128025], "detokenize": True, "skip_special_tokens": False,
+            "stage_id": 0,
+            "is_comprehension": False,
+            "final_output": True,
+            "final_output_type": "text",
+            "max_num_seqs": 1,
+            "gpu_memory_utilization": 0.95,
+            "enforce_eager": True,
+            "trust_remote_code": True,
+            "max_num_batched_tokens": 32768,
+            "devices": AR_DEVICES,
+            "tensor_parallel_size": AR_TP_SIZE,
+            "hf_overrides": {
+                "rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"},
             },
             "output_connectors": {"to_stage_1": "shared_memory_connector"},
+            "default_sampling_params": {
+                "temperature": 0.0,
+                "top_p": 1,
+                "top_k": -1,
+                "max_tokens": 8192,
+                "stop_token_ids": [128025],
+                "detokenize": True,
+                "skip_special_tokens": False,
+            },
         },
         {
-            "stage_id": 1, "stage_type": "diffusion",
-            "runtime": {"process": True, "devices": DIT_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
-            "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},
-            },
-            "engine_input_source": [0],
-            "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
-            "final_output": True, "final_output_type": "image",
-            "default_sampling_params": {"num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE},
+            "stage_id": 1,
+            "max_num_seqs": 1,
+            "enforce_eager": True,
+            "trust_remote_code": True,
+            "devices": DIT_DEVICES,
+            "distributed_executor_backend": "mp",
+            "parallel_config": {"tensor_parallel_size": DIT_TP_SIZE, "enable_expert_parallel": True},
             "input_connectors": {"from_stage_0": "shared_memory_connector"},
+            "default_sampling_params": {
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "guidance_scale": GUIDANCE_SCALE,
+            },
         },
     ],
-    "runtime": {
-        "enabled": True,
-        "connectors": {"shared_memory_connector": {
-            "name": "SharedMemoryConnector",
-            "extra": {"shm_threshold_bytes": 65536}
-        }},
-        "edges": [{"from": 0, "to": 1}],
-    },
+    "edges": [{"from": 0, "to": 1}],
 }
 # fmt: on
 
+_QUANT_DIT_CONFIG = {
+    "pipeline": "hunyuan_image_3_moe",
+    "async_chunk": False,
+    "trust_remote_code": True,
+    "stages": [
+        {
+            "stage_id": 0,
+            "model_stage": "dit",
+            "enforce_eager": True,
+            "trust_remote_code": True,
+            "devices": "0,1",
+            "distributed_executor_backend": "mp",
+            "force_cutlass_fp8": True,
+            "moe_backend": "cutlass",
+            "parallel_config": {
+                "tensor_parallel_size": 2,
+                "enable_expert_parallel": True,
+            },
+            "omni_kv_config": {"need_recv_cache": True},
+            "final_output": True,
+            "final_output_type": "image",
+            "is_comprehension": False,
+            "default_sampling_params": {"seed": SEED},
+        }
+    ],
+}
+
+
+@dataclass(frozen=True)
+class _QuantAccuracyCase:
+    name: str
+    model_env: str
+    nvfp4_backend: str | None = None
+
+
+def _quant_accuracy_cases() -> list[pytest.ParameterSet]:
+    cases = [
+        _QuantAccuracyCase(name="fp8", model_env=QUANT_FP8_ENV),
+        _QuantAccuracyCase(name="mixed_nvfp4", model_env=QUANT_NVFP4_ENV, nvfp4_backend="cutlass"),
+    ]
+    params: list[pytest.ParameterSet] = []
+    run_quant_accuracy = os.environ.get(QUANT_RUN_ENV, "").lower() in _TRUE_ENV_VALUES
+    for case in cases:
+        marks = []
+        if not run_quant_accuracy:
+            marks.append(pytest.mark.skip(reason=f"Set {QUANT_RUN_ENV}=1 to run HunyuanImage3 quant accuracy."))
+        if not os.environ.get(QUANT_BF16_ENV):
+            marks.append(pytest.mark.skip(reason=f"Set {QUANT_BF16_ENV} to run HunyuanImage3 quant accuracy."))
+        if not os.environ.get(case.model_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.model_env} to run this quant accuracy case."))
+        params.append(pytest.param(case, id=case.name, marks=marks))
+    return params
+
+
 # fmt: off
 COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，"
            "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质：它们具有光滑、高光的珐琅或烤漆质感，边缘有明显的金属包边，整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）"
@@ -125,13 +218,28 @@ COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo
 
 
 def _make_config(enable_kv_reuse: bool, path: Path) -> None:
-    config = copy.deepcopy(_BASE_CONFIG)
-    config["stage_args"][0]["engine_args"]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
-    config["stage_args"][1]["engine_args"]["omni_kv_config"] = {"need_recv_cache": enable_kv_reuse}
+    config = copy.deepcopy(_DEPLOY_CONFIG)
+    config["stages"][0]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
+    config["stages"][1]["omni_kv_config"] = {"need_recv_cache": enable_kv_reuse}
     path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
 
 
-def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+def _quant_devices() -> str:
+    return os.environ.get("HUNYUAN_IMAGE3_QUANT_DEVICES", "0,1")
+
+
+def _quant_tensor_parallel_size() -> int:
+    return int(os.environ.get("HUNYUAN_IMAGE3_QUANT_TP", str(len(_quant_devices().split(",")))))
+
+
+def _make_quant_dit_config(path: Path) -> None:
+    config = copy.deepcopy(_QUANT_DIT_CONFIG)
+    config["stages"][0]["devices"] = _quant_devices()
+    config["stages"][0]["parallel_config"]["tensor_parallel_size"] = _quant_tensor_parallel_size()
+    path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
+
+
+def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
     from transformers import AutoTokenizer
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -139,7 +247,7 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
 
     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 +257,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, deploy_config=deploy_config_path) as runner:
         params_list = list(runner.omni.default_sampling_params_list)
         for sp in params_list:
             if isinstance(sp, OmniDiffusionSamplingParams):
@@ -206,7 +314,170 @@ 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)")
+def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+    from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
+
+    server_args = [
+        "--stage-configs-path",
+        stage_configs_path,
+        "--stage-init-timeout",
+        "300",
+        "--init-timeout",
+        "900",
+    ]
+    try:
+        with OmniServer(MODEL_PATH, server_args, use_omni=True) as omni_server:
+            images = download_images(TEST_IMAGE_URLS)
+            t0 = time.perf_counter()
+            response = requests.post(
+                f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+                data={
+                    "model": omni_server.model,
+                    "prompt": PROMPT,
+                    "n": 1,
+                    "response_format": "b64_json",
+                    "num_inference_steps": NUM_INFERENCE_STEPS,
+                    "guidance_scale": GUIDANCE_SCALE,
+                    "seed": SEED,
+                    "sys_type": "en_unified",
+                    "bot_task": "think_recaption",
+                    "size": "1280x720",
+                },
+                files=[
+                    ("image", (f"image_{i}.png", pil_to_png_bytes(img), "image/png")) for i, img in enumerate(images)
+                ],
+                timeout=600,
+            )
+            elapsed = time.perf_counter() - t0
+            if not response.ok:
+                print(f"[ONLINE] HTTP {response.status_code} response body: {response.text}")
+            response.raise_for_status()
+            payload = response.json()
+            assert len(payload["data"]) == 1
+            image = decode_base64_image(payload["data"][0]["b64_json"])
+            image.load()
+            image.save(output_path / "image_online.png")
+            cot_text = payload.get("cot_output") or ""
+            (output_path / "cot_online.txt").write_text(cot_text, encoding="utf-8")
+            return image, cot_text, elapsed
+    finally:
+        gc.collect()
+        if torch.accelerator.is_available():
+            torch.accelerator.empty_cache()
+
+
+@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_online(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
+    """Online API test: same pipeline, same seed as offline → PSNR >= 10 dB."""
+    if importlib.util.find_spec("FlagEmbedding") is None:
+        raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
+    from tabulate import tabulate
+
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-online")
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp = Path(tmpdir)
+        _make_config(True, tmp / "online.yaml")
+        online_image, online_cot, _ = _run_online(str(tmp / "online.yaml"), output_dir)
+
+    online_cot = online_cot.lstrip("\n")
+    scorer = SemanticSimilarityScorer()
+    clip_scorer = CLIPScorer()
+    cot_results = scorer.text_similarity(online_cot, COT_REF)
+    image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
+    image_clip_score = clip_scorer.image_image_score(online_image, image_ref)
+    ssim_value, psnr_value = compute_image_ssim_psnr(prediction=online_image, reference=image_ref, compare_mode="RGB")
+
+    table = [
+        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
+        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
+        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
+        ["SSIM", f"{ssim_value:.4f}", 0.242],
+        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
+    ]
+    print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
+
+    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
+        f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}"
+    )
+    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
+        f"[ONLINE] COT prefix match {cot_results['text_prefix_match_count']} below threshold {THRESHOLDS['text_prefix_match']}"
+    )
+    assert image_clip_score >= THRESHOLDS["clip_score"], (
+        f"[ONLINE] Image-Image similarity {image_clip_score:.4f} below threshold {THRESHOLDS['clip_score']}"
+    )
+    assert ssim_value >= THRESHOLDS["ssim"], f"[ONLINE] SSIM {ssim_value:.4f} below threshold {THRESHOLDS['ssim']}"
+    assert psnr_value >= THRESHOLDS["psnr"], (
+        f"[ONLINE] PSNR {psnr_value:.2f} dB below threshold {THRESHOLDS['psnr']} dB"
+    )
+
+
+def _extract_image(outputs) -> Image.Image:
+    assert outputs, "Pipeline produced no outputs"
+    for output in outputs:
+        images = getattr(output, "images", None)
+        request_output = getattr(output, "request_output", None)
+        if not images and request_output is not None:
+            images = getattr(request_output, "images", None)
+        if images:
+            image = images[0].convert("RGB")
+            image.load()
+            return image
+    raise AssertionError("Pipeline output had no images")
+
+
+def _run_dit_model(
+    model: str,
+    deploy_config_path: str,
+    output_path: Path,
+    *,
+    nvfp4_backend: str | None = None,
+) -> tuple[Image.Image, float]:
+    from tests.helpers.runtime import OmniRunner
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    old_backend = os.environ.get("VLLM_NVFP4_GEMM_BACKEND")
+    if nvfp4_backend is not None:
+        os.environ["VLLM_NVFP4_GEMM_BACKEND"] = nvfp4_backend
+
+    try:
+        with OmniRunner(model, deploy_config=deploy_config_path, mode="text-to-image") as runner:
+            generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED)
+            params = OmniDiffusionSamplingParams(
+                height=QUANT_HEIGHT,
+                width=QUANT_WIDTH,
+                seed=SEED,
+                generator=generator,
+                num_inference_steps=20,
+                guidance_scale=4.0,
+                guidance_scale_provided=True,
+            )
+            t0 = time.perf_counter()
+            outputs = runner.omni.generate({"prompt": QUANT_PROMPT}, params)
+            elapsed = time.perf_counter() - t0
+            image = _extract_image(outputs)
+            image.save(output_path)
+            return image, elapsed
+    finally:
+        if nvfp4_backend is not None:
+            if old_backend is None:
+                os.environ.pop("VLLM_NVFP4_GEMM_BACKEND", None)
+            else:
+                os.environ["VLLM_NVFP4_GEMM_BACKEND"] = old_backend
+        gc.collect()
+        if torch.accelerator.is_available():
+            torch.accelerator.empty_cache()
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=8)
+@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")
@@ -248,3 +519,88 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
     )
     assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}"
     assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"
+
+
+@pytest.mark.parametrize("case", _quant_accuracy_cases())
+@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Needs 2+ GPUs for HunyuanImage3 DiT")
+def test_quantized_dit_matches_bf16_accuracy(
+    case: _QuantAccuracyCase,
+    accuracy_artifact_root: Path,
+) -> None:
+    """Quantized DiT checkpoints should preserve prompt-aligned image quality."""
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-quant")
+    bf16_model = os.environ[QUANT_BF16_ENV]
+    quant_model = os.environ[case.model_env]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        deploy_config_path = Path(tmpdir) / "hunyuan_image3_quant_dit.yaml"
+        _make_quant_dit_config(deploy_config_path)
+
+        bf16_image, bf16_time = _run_dit_model(
+            bf16_model,
+            str(deploy_config_path),
+            output_dir / "bf16.png",
+        )
+        quant_image, quant_time = _run_dit_model(
+            quant_model,
+            str(deploy_config_path),
+            output_dir / f"{case.name}.png",
+            nvfp4_backend=case.nvfp4_backend,
+        )
+
+    ssim_score, psnr_score = compute_image_ssim_psnr(
+        prediction=quant_image,
+        reference=bf16_image,
+    )
+    assert_similarity(
+        model_name=f"{MODEL_NAME} {case.name} vs bf16",
+        vllm_image=quant_image,
+        diffusers_image=bf16_image,
+        ssim_threshold=QUANT_SSIM_THRESHOLD,
+        psnr_threshold=QUANT_PSNR_THRESHOLD,
+        width=QUANT_WIDTH,
+        height=QUANT_HEIGHT,
+    )
+
+    clip_scorer = CLIPScorer()
+    bf16_clip_score = clip_scorer.score(bf16_image, QUANT_PROMPT)
+    quant_clip_score = clip_scorer.score(quant_image, QUANT_PROMPT)
+    clip_score_drop = bf16_clip_score - quant_clip_score
+
+    metrics = {
+        "case": case.name,
+        "bf16_model": bf16_model,
+        "quant_model": quant_model,
+        "prompt": QUANT_PROMPT,
+        "seed": SEED,
+        "height": QUANT_HEIGHT,
+        "width": QUANT_WIDTH,
+        "num_inference_steps": 20,
+        "guidance_scale": 4.0,
+        "bf16_elapsed_s": bf16_time,
+        "quant_elapsed_s": quant_time,
+        "ssim": ssim_score,
+        "psnr": psnr_score,
+        "bf16_clip_score": bf16_clip_score,
+        "quant_clip_score": quant_clip_score,
+        "clip_score_drop": clip_score_drop,
+    }
+    metrics_path = output_dir / f"{case.name}_metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
+
+    print(f"\nHunyuanImage3 quant accuracy ({case.name})")
+    print(f"  bf16 model:       {bf16_model}")
+    print(f"  quant model:      {quant_model}")
+    print(f"  BF16 CLIP score:  {bf16_clip_score:.4f}")
+    print(f"  quant CLIP score: {quant_clip_score:.4f} threshold>={QUANT_CLIP_SCORE_THRESHOLD:.4f}")
+    print(f"  CLIP score drop:  {clip_score_drop:.4f} threshold<={QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}")
+    print(f"  metrics:          {metrics_path}")
+
+    assert quant_clip_score >= QUANT_CLIP_SCORE_THRESHOLD, (
+        f"{case.name} CLIP score below threshold: got {quant_clip_score:.4f}, "
+        f"expected >= {QUANT_CLIP_SCORE_THRESHOLD:.4f}"
+    )
+    assert clip_score_drop <= QUANT_CLIP_SCORE_DROP_THRESHOLD, (
+        f"{case.name} CLIP score drop too large: got {clip_score_drop:.4f}, "
+        f"expected <= {QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}"
+    )

tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py

@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test for DreamID-Omni offline inference with Cache-Dit.
+
+Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of
+synthetic inputs so the coverage matches the upstream demo structure.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import re
+import time
+from functools import lru_cache
+from pathlib import Path
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import soundfile as sf
+from PIL import Image
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow]
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni"
+MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH"
+EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT"
+NUM_INFERENCE_STEPS = 45
+
+CACHE_CONFIG = {
+    "Fn_compute_blocks": 1,
+    "Bn_compute_blocks": 0,
+    "max_warmup_steps": 4,
+    "max_cached_steps": 20,
+    "residual_diff_threshold": 0.24,
+    "max_continuous_cached_steps": 3,
+    "enable_taylorseer": False,
+    "taylorseer_order": 1,
+    "scm_steps_mask_policy": None,
+    "scm_steps_policy": "dynamic",
+}
+
+
+class DreamIDGenerationResult(NamedTuple):
+    video: np.ndarray
+    audio: np.ndarray
+    fps: int
+    duration_s: float
+    elapsed_s: float
+
+
+class DreamIDExampleCase(NamedTuple):
+    prompt: str
+    image_paths: tuple[Path, ...]
+    audio_paths: tuple[Path, ...]
+
+
+def _require_model_path() -> str:
+    model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH))
+    if not Path(model_path).exists():
+        pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.")
+
+    if importlib.util.find_spec("dreamid_omni") is None:
+        pytest.skip(
+            "dreamid_omni dependency package is not importable. "
+            "Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first."
+        )
+
+    return model_path
+
+
+def _dreamid_repo_root_candidates() -> list[Path]:
+    candidates: list[Path] = []
+    example_root = os.environ.get(EXAMPLE_ROOT_ENV)
+    if example_root:
+        candidates.append(Path(example_root).expanduser().resolve())
+
+    spec = importlib.util.find_spec("dreamid_omni")
+    if spec is not None and spec.origin:
+        candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip")
+
+    return candidates
+
+
+@lru_cache(maxsize=1)
+def _load_example_case() -> DreamIDExampleCase:
+    for candidate_root in _dreamid_repo_root_candidates():
+        caption_path = candidate_root / "captions" / "20.json"
+        image_paths = (
+            candidate_root / "imgs" / "20" / "0.png",
+            candidate_root / "imgs" / "20" / "1.png",
+        )
+        audio_paths = (
+            candidate_root / "audios" / "20" / "0.wav",
+            candidate_root / "audios" / "20" / "1.wav",
+        )
+        if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths):
+            prompt = _load_prompt_from_json(caption_path)
+            return DreamIDExampleCase(
+                prompt=prompt,
+                image_paths=image_paths,
+                audio_paths=audio_paths,
+            )
+
+    searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or "<none>"
+    pytest.skip(
+        "DreamID-Omni official example assets not found. "
+        f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. "
+        f"Searched: {searched}"
+    )
+
+
+def _load_prompt_from_json(path: Path) -> str:
+    raw_text = path.read_text(encoding="utf-8").strip()
+    parsed = json.loads(raw_text)
+    if isinstance(parsed, str):
+        prompt = parsed
+    elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str):
+        prompt = parsed["content"]
+    else:
+        raise AssertionError(f"Unexpected DreamID caption format in {path}")
+
+    prompt = re.sub(
+        r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(
+        r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
+    prompt = re.sub(r"\n\s*\n", "\n", prompt).strip()
+    return prompt
+
+
+def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]:
+    images: list[Image.Image] = []
+    for path in example_case.image_paths:
+        with Image.open(path) as image:
+            images.append(image.convert("RGB"))
+    return images
+
+
+def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]:
+    audios: list[np.ndarray] = []
+    for path in example_case.audio_paths:
+        audio_array, sample_rate = sf.read(str(path), dtype="float32")
+        audio_np = np.asarray(audio_array, dtype=np.float32)
+        if audio_np.ndim == 2:
+            audio_np = np.mean(audio_np, axis=1)
+        audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)]
+        audios.append(audio_np.reshape(-1))
+    return audios
+
+
+def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        height=704,
+        width=1280,
+        num_inference_steps=num_inference_steps,
+        seed=42,
+        extra_args={
+            "solver_name": "unipc",
+            "shift": 5.0,
+        },
+    )
+
+
+def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]:
+    return {
+        "prompt": example_case.prompt,
+        "video_negative_prompt": "jitter, bad hands, blur, distortion",
+        "audio_negative_prompt": "robotic, muffled, echo, distorted",
+        "multi_modal_data": {
+            "image": _load_example_images(example_case),
+            "audio": _load_example_audios(example_case),
+        },
+    }
+
+
+def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult:
+    """Extract video and audio from raw Omni outputs"""
+    assert outputs, "DreamID-Omni returned no outputs"
+
+    result = outputs[0]
+    assert hasattr(result, "images") and result.images, "not return video frames"
+
+    video = result.images[0]
+    assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}"
+    assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}"
+
+    multimodal_output = getattr(result, "multimodal_output", None) or {}
+    assert "audio" in multimodal_output, "not return audio"
+    audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1)
+    assert audio.size > 0, "returned empty audio"
+
+    audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000))
+    fps = multimodal_output.get("fps", 24)
+    duration_s = float(audio.size) / float(audio_sample_rate)
+
+    generation = DreamIDGenerationResult(
+        video=video.astype(np.float32, copy=False),
+        audio=audio,
+        fps=fps,
+        duration_s=duration_s,
+        elapsed_s=elapsed_s,
+    )
+    return generation
+
+
+def _run_generation(
+    *,
+    model_path: str,
+    example_case: DreamIDExampleCase,
+    use_cache_dit: bool,
+) -> DreamIDGenerationResult:
+    runner_kwargs = {
+        "model_type": "dreamid-omni",
+        "parallel_config": DiffusionParallelConfig(cfg_parallel_size=1),
+    }
+    if use_cache_dit:
+        runner_kwargs["cache_backend"] = "cache_dit"
+        runner_kwargs["cache_config"] = CACHE_CONFIG
+
+    with OmniRunner(model_path, **runner_kwargs) as runner:
+        warmup_payload = _request_payload(example_case)
+        _ = runner.omni.generate(warmup_payload, _sampling_params(2))
+
+        elapsed_samples: list[float] = []
+        extracted_result: DreamIDGenerationResult | None = None
+
+        request_payload = _request_payload(example_case)
+        start = time.perf_counter()
+        outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS))
+        elapsed_s = time.perf_counter() - start
+        elapsed_samples.append(elapsed_s)
+        if extracted_result is None:
+            extracted_result = _extract_generation(outputs, elapsed_s)
+
+        assert extracted_result is not None
+        median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64)))
+        return extracted_result._replace(elapsed_s=median_elapsed_s)
+
+
+def _video_mae(a: np.ndarray, b: np.ndarray) -> float:
+    assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}"
+    return float(np.mean(np.abs(a - b)))
+
+
+def _audio_rms(audio: np.ndarray) -> float:
+    audio = np.asarray(audio, dtype=np.float32).reshape(-1)
+    return float(np.sqrt(np.mean(audio**2)))
+
+
+def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float:
+    rms_a = _audio_rms(a.audio)
+    rms_b = _audio_rms(b.audio)
+    return rms_b / max(rms_a, 1e-6)
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_dreamid_omni_cache_dit_e2e() -> None:
+    model_path = _require_model_path()
+    example_case = _load_example_case()
+
+    baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False)
+    cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True)
+
+    video_mae = _video_mae(baseline.video, cached.video)
+    audio_rms_ratio = _audio_rms_ratio(baseline, cached)
+    speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf")
+
+    assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}"
+    assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}"
+    assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x"

tests/e2e/offline_inference/test_glm_tts.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E Offline inference tests for GLM-TTS voice cloning.
+
+GLM-TTS is a zero-shot voice cloning model that always requires ref_audio.
+There is no text-only / non-clone inference path in the official implementation.
+"""
+
+import os
+import tempfile
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import numpy as np
+import pytest
+import soundfile as sf
+
+from tests.helpers.assertions import assert_audio_speech_response
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import convert_audio_file_to_text, get_asset_path
+from tests.helpers.runtime import OmniResponse
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata
+
+MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS")
+REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。"
+
+DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml")
+
+REFERENCE_PROMPT_WAV_PATH = get_asset_path("glm_tts/jiayan_zh.wav")
+
+ASYNC_CHUNK_MODES = [
+    pytest.param(False, id="sync"),
+    pytest.param(True, id="async_chunk"),
+]
+
+
+def _load_ref_audio() -> tuple[np.ndarray, int]:
+    audio, sr = sf.read(str(REFERENCE_PROMPT_WAV_PATH), dtype="float32", always_2d=False)
+    if isinstance(audio, np.ndarray) and audio.ndim > 1:
+        audio = np.mean(audio, axis=-1)
+    return np.asarray(audio, dtype=np.float32), int(sr)
+
+
+def _concat_audio(audio_val) -> np.ndarray:
+    import torch
+
+    if isinstance(audio_val, list):
+        tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None]
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False)
+    if isinstance(audio_val, torch.Tensor):
+        return audio_val.float().cpu().numpy().reshape(-1)
+    return np.asarray(audio_val, dtype=np.float32).reshape(-1)
+
+
+def _extract_sample_rate(audio_mm: dict) -> int:
+    sr_raw = audio_mm.get("sr", 24000)
+    if isinstance(sr_raw, list):
+        sr_raw = sr_raw[-1] if sr_raw else 24000
+    if hasattr(sr_raw, "item"):
+        return int(sr_raw.item())
+    return int(sr_raw)
+
+
+def _get_deploy_config(*, async_chunk: bool) -> str:
+    """Build deploy config with explicit sync/async mode and eager execution."""
+    return modify_stage_config(
+        DEPLOY_CONFIG,
+        updates={
+            "async_chunk": async_chunk,
+            "stages": {
+                0: {
+                    "enforce_eager": True,
+                    "async_scheduling": bool(async_chunk),
+                },
+                1: {
+                    "enforce_eager": True,
+                },
+            },
+        },
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("async_chunk", ASYNC_CHUNK_MODES)
+def test_offline_voice_clone_zh(async_chunk: bool) -> None:
+    """
+    Test voice cloning offline inference.
+    Deploy Setting: glm_tts.yaml with explicit sync/async mode, enforce_eager
+    Input Modal: text + ref_audio + ref_text
+    Output Modal: audio
+
+    Uses the official jiayan_zh.wav reference audio from the upstream
+    GLM-TTS repository (real human speech matching REF_TEXT).
+    """
+    synth_text = "今天天气真不错，适合出去散散步。"
+    prompt_audio = _load_ref_audio()
+    from tests.helpers.runtime import OmniRunner
+
+    with OmniRunner(
+        MODEL,
+        stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
+        stage_init_timeout=600,
+    ) as omni_runner:
+        outputs = omni_runner.omni.generate(
+            [
+                {
+                    "prompt": synth_text,
+                    "multi_modal_data": {"audio": prompt_audio},
+                    "modalities": ["audio"],
+                    "mm_processor_kwargs": {"prompt_text": REF_TEXT},
+                    "additional_information": build_glm_tts_prefill_metadata(
+                        MODEL,
+                        synth_text,
+                        REF_TEXT,
+                        trust_remote_code=True,
+                    ),
+                }
+            ],
+            omni_runner.get_default_sampling_params_list(),
+        )
+
+        assert outputs, "No outputs returned"
+        audio_mm = outputs[0].multimodal_output
+        assert "audio" in audio_mm, "No audio output found"
+        audio = _concat_audio(audio_mm["audio"])
+        assert audio.size > 0, "Generated audio is empty"
+
+        with tempfile.NamedTemporaryFile(suffix=".wav") as out_file:
+            sf.write(out_file.name, audio, samplerate=_extract_sample_rate(audio_mm), format="WAV")
+            transcript = convert_audio_file_to_text(out_file.name)
+            audio_bytes = open(out_file.name, "rb").read()
+
+        assert_audio_speech_response(
+            OmniResponse(
+                success=True,
+                audio_bytes=audio_bytes,
+                audio_content=transcript,
+                audio_format="audio/wav",
+            ),
+            {"input": synth_text, "response_format": "wav"},
+            run_level="advanced_model",
+        )

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_glm_tts.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for GLM-TTS model with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+the GLM-TTS two-stage pipeline (AR + DiT).
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS")
+REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。"
+
+DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml")
+
+# Official GLM-TTS reference audio and transcript. Vendored under tests/assets/
+# so Buildkite does not depend on raw.githubusercontent.com during E2E.
+REF_AUDIO_URL = load_test_audio_data_url("glm_tts/jiayan_zh.wav")
+
+
+SYNC_EXTRA_ARGS = [
+    "--trust-remote-code",
+    "--disable-log-stats",
+    "--no-async-chunk",
+]
+
+ASYNC_EXTRA_ARGS = [
+    "--trust-remote-code",
+    "--disable-log-stats",
+]
+
+tts_sync_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=SYNC_EXTRA_ARGS,
+            stage_init_timeout=600,
+        ),
+        id="glm_tts_sync",
+    )
+]
+
+tts_async_chunk_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=ASYNC_EXTRA_ARGS,
+            stage_init_timeout=600,
+        ),
+        id="glm_tts_async_chunk",
+    )
+]
+
+
+# GLM-TTS is a zero-shot voice cloning model (always requires ref_audio).
+# There is no text-only / non-clone path — the official implementation
+# asserts prompt_feat.shape[1] != 0 in token2wav_with_cache and every
+# example prompt includes a reference audio.  CosyVoice3 follows the same
+# pattern: all tests are voice_clone_*.
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True)
+def test_voice_clone_zh_sync(omni_server, openai_client) -> None:
+    """
+    Test voice cloning TTS with Chinese text.
+    Deploy Setting: glm_tts.yaml with ``--no-async-chunk`` (sync two-stage)
+    Input Modal: text + ref_audio + ref_text
+    Output Modal: audio
+    Input Setting: stream=False
+
+    Uses the official jiayan_zh.wav reference audio from the upstream
+    GLM-TTS repository (real human speech matching REF_TEXT).
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": "今天天气真不错，适合出去散散步。",
+        "stream": False,
+        "response_format": "wav",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
+def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None:
+    """
+    Test voice cloning TTS with Chinese text via async_chunk streaming.
+    Deploy Setting: glm_tts.yaml default ``async_chunk: true``
+    Input Modal: text + ref_audio + ref_text
+    Output Modal: audio
+    Input Setting: stream=True
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": "今天天气真不错，适合出去散散步。",
+        "stream": True,
+        "response_format": "wav",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True)
+def test_models_endpoint(omni_server, openai_client) -> None:
+    """Test the /v1/models endpoint returns loaded model."""
+    models = openai_client.client.models.list()
+    assert len(models.data) > 0

tests/e2e/online_serving/test_hidream_i1_full_expansion.py

@@ -0,0 +1,69 @@
+"""
+End-to-end diffusion coverage for HiDream-I1-Full in online serving mode.
+
+This test verifies that HiDream-I1-Full can be launched,
+accepts text-to-image requests through the OpenAI-compatible API, and returns
+valid generated images with the requested resolution.
+
+assert_diffusion_response validates successful generation and the expected
+image resolution.
+"""
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = "HiDream-ai/HiDream-I1-Full"
+PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography."
+NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_hidream_i1_image_feature_cases(model: str):
+    """Return HiDream-I1-Full diffusion feature cases."""
+
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--auxiliary-text-encoder",
+                    "meta-llama/Meta-Llama-3.1-8B-Instruct",
+                ],
+            ),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_hidream_i1_image_feature_cases(MODEL),
+    indirect=True,
+)
+def test_hidream_i1_image(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """Validate HiDream-I1-Full online serving with CPU offload."""
+
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

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/entrypoints/openai_api/test_serving_speech.py

@@ -2335,6 +2335,84 @@ class TestCosyVoice3Serving:
         cosyvoice3_server._build_cosyvoice3_prompt.assert_awaited_once()
 
 
+# ---- GLM-TTS Serving Tests ----
+
+
+@pytest.fixture
+def glm_tts_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="zai-org/GLM-TTS",
+        hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20),
+    )
+    mock_engine_client.default_sampling_params_list = [
+        SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
+    ]
+    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="glm_tts"),
+            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 TestGLMTTSServing:
+    def test_validate_glm_tts_requires_ref_audio(self, glm_tts_server):
+        request = OpenAICreateSpeechRequest(input="Hello", ref_text="Reference transcript")
+        error = glm_tts_server._validate_glm_tts_request(request)
+        assert error is not None
+        assert "ref_audio" in error
+
+    def test_validate_glm_tts_requires_ref_text(self, glm_tts_server):
+        request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc")
+        error = glm_tts_server._validate_glm_tts_request(request)
+        assert error is not None
+        assert "ref_text" in error
+
+    def test_estimate_glm_tts_target_text_len_uses_tokenizer_tokens(
+        self,
+        glm_tts_server,
+        mocker: MockerFixture,
+    ):
+        class FakeTokenizer:
+            def encode(self, text):
+                assert text == "normalized target"
+                return [10, 20, 30]
+
+        mocker.patch(
+            "vllm_omni.model_executor.models.glm_tts.glm_tts.resolve_glm_tts_tokenizer_path",
+            return_value="resolved-tokenizer",
+        )
+        load_tokenizer = mocker.patch(
+            "vllm_omni.model_executor.models.glm_tts.glm_tts.load_glm_tts_tokenizer",
+            return_value=FakeTokenizer(),
+        )
+        mocker.patch(
+            "vllm_omni.model_executor.models.glm_tts.text_frontend.GLMTTSTextFrontend.text_normalize",
+            return_value="normalized target",
+        )
+
+        text_token_len = glm_tts_server._estimate_glm_tts_text_token_len("abcdef")
+
+        assert text_token_len == 3
+        load_tokenizer.assert_called_once()
+
+
 class TestTTSAsyncOffloading:
     """Tests for event-loop-safe offloading of blocking TTS operations."""
 

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_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/fixtures/env.py

@@ -2,6 +2,7 @@ import os
 
 import pytest
 import torch
+from vllm.distributed import parallel_state
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -59,3 +60,15 @@ def default_vllm_config():
 
     with set_current_vllm_config(VllmConfig(device_config=device_config)):
         yield
+
+
+@pytest.fixture()
+def init_fake_tp_group(mocker):
+    """Provide a fake TP group so vllm linear layers can be instantiated."""
+    mock_tp = mocker.MagicMock()
+    mock_tp.world_size = 1
+    mock_tp.rank_in_group = 0
+    old = parallel_state._TP
+    parallel_state._TP = mock_tp
+    yield
+    parallel_state._TP = old