Merge branch 'main' into codex/unskip-fp8-z-image-quality

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -52,6 +52,7 @@ time timeout -k 30 30m docker run \
     pip install tblib==3.1.0
     cd /workspace/vllm-omni
     pytest -v -s -m "core_model and xpu and B60"
+    pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
     export VLLM_XPU_USE_SAMPLER_KERNEL=0    # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
     pytest -v -s -m "advanced_model and xpu and B60"
 '

.buildkite/test-merge.yml

@@ -389,6 +389,31 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: "Omni · Higgs-Audio V2 Test"
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_USE_DEEP_GEMM=0
+              export VLLM_MOE_USE_DEEP_GEMM=0
+              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "Diffusion · Qwen Image Edit Test"
         timeout_in_minutes: 20
         commands:

.buildkite/test-nightly.yml

@@ -305,6 +305,51 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Omni · Perf Test · Multi-Replica"
+        key: nightly-omni-performance-multi-replicas
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        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: 3
+                    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: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
         timeout_in_minutes: 45
@@ -525,6 +570,7 @@ steps:
             tests/e2e/online_serving/test_longcat_image_expansion.py
             tests/e2e/online_serving/test_longcat_image_edit_expansion.py
             tests/e2e/online_serving/test_flux_2_dev_expansion.py
+            tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
             -m "full_model and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
         agents:
@@ -797,8 +843,8 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
-        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
         timeout_in_minutes: 180
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
@@ -806,7 +852,7 @@ steps:
           - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
@@ -893,6 +939,54 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
+        key: nightly-diffusion-x2iat-performance-bagel
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - export CACHE_DIT_VERSION=1.3.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        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: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
         key: nightly-hunyuan-image3-performance
         soft_fail: true
@@ -998,6 +1092,8 @@ steps:
         timeout_in_minutes: 180
         commands:
           - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
         agents:
@@ -1188,11 +1284,13 @@ steps:
       - nightly-omni-performance-no-async-chunk
       - nightly-omni-performance-async-chunk
       - nightly-omni-performance-vllm-text
+      - nightly-omni-performance-multi-replicas
       - nightly-tts-performance
       - nightly-diffusion-x2iat-performance-qwen-image
       - nightly-diffusion-x2iat-performance-qwen-image-edit
-      - nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - nightly-diffusion-x2iat-performance-qwen-image-edit-2511
       - nightly-diffusion-x2iat-performance-qwen-image-layered
+      - nightly-diffusion-x2iat-performance-bagel
       - nightly-diffusion-x2v-performance
       - nightly-testcase-statistics
     if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
@@ -1204,10 +1302,12 @@ steps:
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
       - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
       - python tools/nightly/generate_nightly_perf_excel.py

.claude/skills/diffusion-perf-opt/SKILL.md

@@ -21,6 +21,9 @@ Before proposing changes, ask for the optimization scene if it is not already kn
 - Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
 - Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
 - Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
+- Quality validation method: exact-output comparison when feasible, or metrics
+  such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
+  checks, baseline self-run variance, and reviewable output artifacts.
 
 ## Workflow
 
@@ -79,7 +82,19 @@ Before proposing changes, ask for the optimization scene if it is not already kn
    - Record latency, stage timings, memory, output quality, and logs.
    - Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
 
-6. **Collect diagnostic trace only after narrowing hypotheses.**
+6. **Enforce the quality and precision gate for every optimization.**
+   - Do not mark an optimization ready until quality is checked on the same
+     workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
+     backend, parallelism, and output encoding.
+   - For math-preserving changes, require exact/near-exact agreement or show
+     that differences are within baseline self-run variance.
+   - For precision, quantization, approximate attention, sparse/custom kernels,
+     or backend changes, include reviewable artifacts and metrics such as SSIM,
+     PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
+   - If quality tolerance is not stated, default to quality-preserving behavior.
+     Failed or inconclusive quality validation blocks a ready-to-merge claim.
+
+7. **Collect diagnostic trace only after narrowing hypotheses.**
    - Use torch profiler for a small number of requests.
    - Run two separate diagnostic traces instead of mixing concerns:
      - **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
@@ -91,15 +106,16 @@ Before proposing changes, ask for the optimization scene if it is not already kn
    - Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
    - Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
 
-7. **Analyze host, communication, and operators.**
+8. **Analyze host, communication, and operators.**
    - Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
    - Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
    - Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
    - Rank operator work by total CUDA time and by repeated small-kernel launch count.
 
-8. **Produce an optimization plan.**
+9. **Produce an optimization plan.**
    - Classify candidates as P0/P1/P2.
    - For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
+   - Include a concrete quality/precision validation gate for each candidate.
    - Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
    - End the plan with a user-facing candidate selection table. The assistant
      should not automatically choose a risky optimization just because it is
@@ -124,7 +140,13 @@ Before proposing changes, ask for the optimization scene if it is not already kn
 - **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
 - **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
 
-Every implemented optimization needs A/B validation and quality regression. A/B means same workload and hardware before/after. Quality regression means checking generated image/video stability, artifacts, temporal flicker, and seed behavior when precision or approximate kernels change.
+Every implemented optimization needs A/B validation and a passing
+quality/precision gate, including math-preserving changes such as padding trim,
+layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
+when available, keep reviewable artifacts, and compare against baseline
+self-run variance for precision, quantization, approximate kernels, or backend
+changes. If quality validation fails or is inconclusive, do not present the
+optimization as ready to merge.
 
 ## Optimization Layers
 

.readthedocs.yml

@@ -10,6 +10,11 @@ build:
   jobs:
     post_checkout:
       - git fetch --unshallow || true
+    pre_install:
+      # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
+      # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
+      # 1 GB memory cap and time out the mkdocs build.
+      - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
 
 mkdocs:
   configuration: mkdocs.yml

benchmarks/README.md

@@ -1,44 +1,60 @@
-# Benchmarks Overview and Architecture
+# Benchmarks
 
-This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
+This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
 
-## Scope and goals
-- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
-- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
-- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
+## Benchmark families
 
-## Dataset and inputs
-- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
-- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
-- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
+### [TTS](tts/README.md) — Text-to-Speech
 
-## Directory layout
-- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
-- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
-- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
-  benchmark suites to vLLM-Omni serving and evaluation flows.
-- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
+Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
 
-## Reference workflows
-- **HF Transformers (offline, single process)**  
-  Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`  
-  Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
+- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
+- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
+- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
 
-- **vLLM-Omni end-to-end pipeline**  
-  Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`  
-  Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
+### [Diffusion](diffusion/README.md) — Image and Video Generation
 
-- **Adding a new task/model**  
-  1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.  
-  2) Add a task README describing dataset, configs, and expected outputs.  
-  3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
+Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
 
-## Metrics to watch
-- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
-- **Latency distribution**: look for long tails in `*.stats.jsonl`.
-- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
+- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
+- **Datasets**: `vbench`, `trace`, `random`
+- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
 
-## Troubleshooting
-- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
-- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
-- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.
+### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
+
+Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
+
+- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
+- **Tasks**: text-to-image and image-to-image
+- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
+
+### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
+
+RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
+
+- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
+- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
+
+### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
+
+Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
+
+- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
+- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
+
+### Common serving metrics framework
+
+`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
+
+- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
+- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
+- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
+
+See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
+
+## Adding a new benchmark
+
+1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
+2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
+3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
+4. Update this README with a link to the new benchmark family.

benchmarks/diffusion/bench_attention_backends.py

@@ -0,0 +1,328 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Standalone diagnostic for the diffusion attention backends added in #3079.
+
+Exercises the same synthetic attention shape through each backend and each
+torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
+CUDNN_ATTN underperforms on some SKUs. Two typical causes:
+
+  (a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
+      and silently falls back to MATH.
+  (b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
+      walking CUDNN -> FLASH -> MATH.
+
+Run:
+    python benchmarks/diffusion/bench_attention_backends.py --preset hv15
+    python benchmarks/diffusion/bench_attention_backends.py --preset wan22
+    python benchmarks/diffusion/bench_attention_backends.py \
+        --batch 1 --heads 24 --seq 14336 --head-dim 128
+
+Optional extras:
+    pip install --pre flash-attn-4        # FA4 is currently pre-release only
+    pip install -U flashinfer              # latest FlashInfer (0.6.9)
+
+The table at the end is the data we want on the PR — surface the row where a
+backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+# Presets mirror the DiT hot-path attention call (single-stream block) at the
+# shapes we run in the PR's validation configs.
+_PRESETS = {
+    # HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
+    "hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
+    # Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
+    "wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
+    # Smaller image-gen shape for quick smoke tests.
+    "flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
+}
+
+_SDPA_BACKENDS = [
+    ("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
+    ("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
+    ("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
+    ("MATH", [SDPBackend.MATH]),
+    # The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
+    # what torch actually picks when multiple are allowed.
+    ("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
+]
+
+
+def _env_report() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
+    cc = torch.cuda.get_device_capability()
+    print("=" * 72)
+    print(f"GPU           : {torch.cuda.get_device_name()}")
+    print(f"SM capability : sm_{cc[0]}{cc[1]}")
+    print(f"torch         : {torch.__version__}")
+    print(f"cuDNN version : {torch.backends.cudnn.version()}")
+    print("flashinfer    : ", end="")
+    try:
+        import flashinfer
+
+        print(getattr(flashinfer, "__version__", "present"))
+    except Exception as e:
+        print(f"not installed ({type(e).__name__})")
+    print("=" * 72)
+
+
+def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
+    # Layout: (B, S, H, D) — what the PR's backends accept as input.
+    q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    return q, k, v
+
+
+def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
+    # Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
+    mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
+    if pad_tokens > 0:
+        mask[..., -pad_tokens:] = float("-inf")
+    return mask
+
+
+def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
+    """Return (median ms, error string). Skips on any exception so that a
+    backend that rejects our args (wrong dtype, missing JIT module, unsupported
+    kwarg value) doesn't abort the whole sweep."""
+    try:
+        for _ in range(warmup):
+            fn(*args)
+        torch.accelerator.synchronize()
+        times = []
+        for _ in range(iters):
+            torch.accelerator.synchronize()
+            t0 = time.perf_counter()
+            fn(*args)
+            torch.accelerator.synchronize()
+            times.append((time.perf_counter() - t0) * 1000.0)
+        times.sort()
+        return times[len(times) // 2], ""
+    except Exception as e:  # noqa: BLE001 — probe script, keep the whole table going
+        msg = str(e).split("\n", 1)[0][:60]
+        return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
+
+
+def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
+    # q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
+    q_t = q.permute(0, 2, 1, 3)
+    k_t = k.permute(0, 2, 1, 3)
+    v_t = v.permute(0, 2, 1, 3)
+    rows: list[tuple[str, float, str]] = []
+    for name, backends in _SDPA_BACKENDS:
+
+        def _call(mask=attn_mask):
+            with sdpa_kernel(backends):
+                return torch.nn.functional.scaled_dot_product_attention(
+                    q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
+                )
+
+        ms, err = _time_call(_call)
+        rows.append((name, ms, err))
+    return rows
+
+
+def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
+    """Call FlashInfer's dense single-prefill.
+
+    ``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
+    that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
+    we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
+    Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
+    """
+    try:
+        import inspect
+
+        from flashinfer.prefill import single_prefill_with_kv_cache
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
+    if backend is not None:
+        sig = inspect.signature(single_prefill_with_kv_cache)
+        if "backend" not in sig.parameters:
+            return float("nan"), "no-backend-kwarg"
+        kwargs["backend"] = backend
+
+    if attn_mask is not None:
+        # Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
+        # FlashInfer expects True = keep, False = masked out.
+        mask_2d = attn_mask[0, 0]
+        kwargs["custom_mask"] = mask_2d != float("-inf")
+
+    def _call():
+        out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
+        return out.unsqueeze(0)
+
+    return _time_call(_call)
+
+
+def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
+    """Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
+    ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
+    cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
+    ``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
+    try:
+        from flash_attn.cute import flash_attn_func
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    def _call():
+        # FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
+        return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
+
+    return _time_call(_call)
+
+
+def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
+    """Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
+    (which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
+    straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
+    try:
+        from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    b, s, h, d = q.shape
+    qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
+    kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
+
+    def _call():
+        return cudnn_batch_prefill_with_kv_cache(
+            q.reshape(b * s, h, d),
+            k.reshape(b * s, h, d),
+            v.reshape(b * s, h, d),
+            qo_indptr=qo_indptr,
+            kv_indptr=kv_indptr,
+            max_qo_len=s,
+            max_kv_len=s,
+            sm_scale=scale,
+            causal=False,
+        )
+
+    return _time_call(_call)
+
+
+def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
+    print(f"\n{title}")
+    print("-" * 88)
+    header = f"{'backend':<24} {'median (ms)':>14}    {'vs baseline':>12}    status"
+    print(header)
+    print("-" * 88)
+    baseline_ms = None
+    if baseline_name is not None:
+        for name, ms, _ in rows:
+            if name == baseline_name and ms == ms:  # not NaN
+                baseline_ms = ms
+                break
+    for name, ms, err in rows:
+        ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
+        if baseline_ms is not None and ms == ms:
+            ratio = baseline_ms / ms
+            ratio_str = f"{ratio:>11.2f}x"
+        else:
+            ratio_str = f"{'—':>12}"
+        status = "ok" if not err else f"FAILED ({err})"
+        print(f"{name:<24} {ms_str}    {ratio_str}    {status}")
+
+
+def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
+    """Return (backend, ms) of the fastest non-failing row, or None."""
+    ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
+    if not ok_rows:
+        return None
+    return min(ok_rows, key=lambda x: x[1])
+
+
+def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
+    torch.manual_seed(0)
+    q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
+    scale = 1.0 / (shape["head_dim"] ** 0.5)
+
+    rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
+    rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
+    # `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
+    for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
+        rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
+    rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
+    rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
+
+    mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
+    rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
+    rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
+    rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
+
+    return rows_nomask, rows_mask
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
+    parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
+    parser.add_argument("--batch", type=int, default=None)
+    parser.add_argument("--heads", type=int, default=None)
+    parser.add_argument("--seq", type=int, default=None)
+    parser.add_argument("--head-dim", type=int, default=None)
+    parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
+    parser.add_argument("--device", default="cuda:0")
+    args = parser.parse_args()
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+    _env_report()
+    print(f"dtype         : {args.dtype}")
+
+    presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
+    summary: dict[str, tuple[str, float] | None] = {}
+
+    for preset_name in presets:
+        shape = dict(_PRESETS[preset_name])
+        if not args.sweep:
+            for k in ("batch", "heads", "seq", "head_dim"):
+                v = getattr(args, k if k != "head_dim" else "head_dim")
+                if v is not None:
+                    shape[k] = v
+
+        print("\n" + "=" * 88)
+        print(
+            f"Preset: {preset_name}  |  batch={shape['batch']} heads={shape['heads']} "
+            f"seq={shape['seq']} head_dim={shape['head_dim']}"
+        )
+        print("=" * 88)
+
+        rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
+        _print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
+        _print_table(
+            f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
+        )
+
+        summary[preset_name] = _pick_winner(rows_nomask)
+
+    if args.sweep or len(presets) > 1:
+        print("\n" + "=" * 88)
+        print("Winners per preset (no-mask path)")
+        print("=" * 88)
+        for preset_name, winner in summary.items():
+            if winner is None:
+                print(f"  {preset_name:<10} — no successful backend")
+            else:
+                name, ms = winner
+                print(f"  {preset_name:<10} {name:<24} {ms:>8.3f} ms")
+
+    print("\nNotes:")
+    print("  * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
+    print("  * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
+    print("  * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
+    print("  * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -47,6 +47,13 @@ Usage:
         --endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
         --bot-task think
 
+    python benchmarks/diffusion/diffusion_benchmark_serving.py \
+      --endpoint /v1/images/edits \
+      --dataset custom \
+      --dataset-path custom_requests.jsonl \
+      --task ti2i \
+      --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
@@ -566,6 +573,135 @@ class TraceDataset(BaseDataset):
         return [self[i] for i in range(len(self))]
 
 
+class CustomDataset(BaseDataset):
+    """
+    Custom dataset that loads requests from a JSONL file.
+
+    Each line in the JSONL file should be a JSON object with the following fields:
+    - prompt (required): The text prompt for the request
+    - width (optional): Image/video width
+    - height (optional): Image height
+    - num_inference_steps (optional): Number of diffusion steps
+    - seed (optional): Random seed
+    - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
+    - image_urls (optional): List of input image URLs (alternative to image_paths)
+
+    Example JSONL for ti2i:
+    {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
+    """
+
+    def __init__(self, args, api_url: str, model: str):
+        super().__init__(args, api_url, model)
+        self.dataset_path = args.dataset_path
+        if not self.dataset_path:
+            raise ValueError("--dataset-path must be provided when using 'custom' dataset")
+        self.load_data()
+
+    def load_data(self) -> None:
+        """Load data from JSONL file."""
+        import pandas as pd
+
+        if not self.dataset_path.endswith(".jsonl"):
+            raise ValueError("Custom dataset must be a JSONL file")
+
+        try:
+            df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
+        except Exception as e:
+            raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
+
+        if "prompt" not in df.columns:
+            raise ValueError("JSONL file must contain a 'prompt' column")
+
+        self.data = df.to_dict("records")
+        print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
+
+    def _resolve_image_paths(self, item: dict) -> list[str] | None:
+        # Handle local file paths
+        if "image_paths" in item and item["image_paths"]:
+            paths = item["image_paths"]
+            if isinstance(paths, str):
+                paths = [paths]
+
+            # Verify all paths exist
+            valid_paths = []
+            for path in paths:
+                if os.path.exists(path):
+                    valid_paths.append(path)
+                else:
+                    raise ValueError(f"Image file not found: {path}")
+
+            return valid_paths if valid_paths else None
+
+        # Handle URLs - download to temp files
+        if "image_urls" in item and item["image_urls"]:
+            urls = item["image_urls"]
+            if isinstance(urls, str):
+                urls = [urls]
+
+            downloaded_paths = []
+            for url in urls:
+                try:
+                    response = requests.get(url, timeout=30)
+                    response.raise_for_status()
+
+                    # Create temp file with appropriate extension
+                    suffix = os.path.splitext(url)[1] or ".png"
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                        tmp.write(response.content)
+                        downloaded_paths.append(tmp.name)
+                except Exception as e:
+                    raise ValueError(f"Failed to download image from {url}: {e}")
+
+            return downloaded_paths if downloaded_paths else None
+
+        return None
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> RequestFuncInput:
+        item = self.data[idx]
+
+        # Extract fields with fallback to args defaults
+        prompt = item["prompt"]
+        width = item.get("width", self.args.width)
+        height = item.get("height", self.args.height)
+        num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
+        seed = item.get("seed", self.args.seed)
+        reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
+        extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
+
+        # Handle image paths/URLs
+        image_paths = self._resolve_image_paths(item)
+
+        return RequestFuncInput(
+            prompt=prompt,
+            api_url=self.api_url,
+            model=self.model,
+            seed=seed,
+            image_paths=image_paths,
+            extra_body=extra_body,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+        )
+
+    def get_requests(self) -> list[RequestFuncInput]:
+        """Get all requests, cycling through data if num_prompts > dataset size."""
+        num_requests = getattr(self.args, "num_prompts", len(self.data))
+
+        if num_requests <= len(self.data):
+            return [self[i] for i in range(num_requests)]
+
+        # Cycle through the dataset to reach num_prompts
+        requests = []
+        for i in range(num_requests):
+            idx = i % len(self.data)
+            requests.append(self[idx])
+
+        return requests
+
+
 class RandomDataset(BaseDataset):
     def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
         super().__init__(args, api_url, model)
@@ -950,6 +1086,8 @@ async def benchmark(args):
         dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
         dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
+    elif args.dataset == "custom":
+        dataset = CustomDataset(args, api_url, args.model)
     else:
         raise ValueError(f"Unknown dataset: {args.dataset}")
 
@@ -1097,7 +1235,7 @@ if __name__ == "__main__":
         "--dataset",
         type=str,
         default="vbench",
-        choices=["vbench", "trace", "random"],
+        choices=["vbench", "trace", "random", "custom"],
         help="Dataset to use.",
     )
     parser.add_argument(

docs/design/module/dit_module.md

@@ -169,15 +169,13 @@ class DiffusionRequestStatus(enum.IntEnum):
 
 @dataclass
 class DiffusionRequestState:
-    sched_req_id: str
+    request_id: str
     req: OmniDiffusionRequest
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
 ```
 
 **Design Features**:
 
-- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
-
 - **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
 
 - **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
@@ -188,7 +186,6 @@ class DiffusionRequestState:
 class _BaseScheduler(SchedulerInterface):
     def __init__(self) -> None:
         self._request_states = {}
-        self._request_id_to_sched_req_id = {}
         self._waiting = deque()
         self._running = []
         self._finished_req_ids = set()
@@ -199,7 +196,8 @@ class _BaseScheduler(SchedulerInterface):
 
 - **Common state storage**: Shared request maps and waiting/running sets live in the base class.
 
-- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
+- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
+  duplicated in each policy.
 
 - **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
 
@@ -219,12 +217,12 @@ class RequestScheduler(_BaseScheduler):
 
 - **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
 
-- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
+- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
 
 #### 2.5 Engine-Driven Execution Loop
 
 ```python
-sched_req_id = scheduler.add_request(request)
+request_id = scheduler.add_request(request)
 while True:
     sched_output = scheduler.schedule()
     output = executor.add_req(req)

docs/mkdocs/hooks/generate_examples.py

@@ -16,6 +16,7 @@ ROOT_DIR_RELATIVE = "../../../../.."
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
 NAV_FILE = ROOT_DIR / "docs/.nav.yml"
+MAX_INLINE_MATERIAL_SIZE = 8 * 1024
 
 
 def fix_case(text: str) -> str:
@@ -191,18 +192,21 @@ class Example:
 
         return re.sub(link_pattern, replace_link, content)
 
+    def github_url(self, path: Path) -> str:
+        url = "https://github.com/vllm-project/vllm-omni/"
+        url += "tree/main" if path.is_dir() else "blob/main"
+        return f"{url}/{path.relative_to(ROOT_DIR).as_posix()}"
+
     def generate(self) -> str:
         content = f"# {self.title}\n\n"
-        url = "https://github.com/vllm-project/vllm-omni/"
-        url += "tree/main" if self.path.is_dir() else "blob/main"
-        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
+        content += f"Source <{self.github_url(self.path)}>.\n\n"
 
         # Use long code fence to avoid issues with
         # included files containing code fences too
         code_fence = "``````"
 
         if self.is_code:
-            main_file_rel = self.main_file.relative_to(ROOT_DIR)
+            main_file_rel = self.main_file.relative_to(ROOT_DIR).as_posix()
             content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
         else:
             with open(self.main_file, encoding="utf-8") as f:
@@ -216,10 +220,15 @@ class Example:
 
         content += "## Example materials\n\n"
         for file in sorted(self.other_files):
-            content += f'??? abstract "{file.relative_to(self.path)}"\n'
+            content += f'??? abstract "{file.relative_to(self.path).as_posix()}"\n'
+            if file.stat().st_size > MAX_INLINE_MATERIAL_SIZE:
+                content += (
+                    f"    Large file omitted from the rendered docs. View it on GitHub: <{self.github_url(file)}>.\n\n"
+                )
+                continue
             if file.suffix != ".md":
                 content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
+            content += f'    --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
             if file.suffix != ".md":
                 content += f"    {code_fence}\n"
 
@@ -291,7 +300,7 @@ def update_nav_file(examples: list[Example]):
         for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: str(rel_path)})
+            category_items.append({example.title: rel_path.as_posix()})
 
         if category_items:
             # Format category name (e.g., "offline_inference" -> "Offline Inference")

docs/mkdocs/hooks/search_index.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reduce browser search-index memory for generated API reference pages."""
+
+from material.plugins.search.plugin import SearchIndex
+
+_add_entry_from_context = SearchIndex.add_entry_from_context
+
+
+def add_entry_from_context(self: SearchIndex, page) -> None:
+    search = page.meta.get("search") or {}
+    if search.get("exclude"):
+        return
+
+    if str(page.url).startswith("api/"):
+        entry = {
+            "location": page.url,
+            "title": str(page.meta.get("title", page.title)),
+            "text": "",
+        }
+        if "boost" in search:
+            entry["boost"] = search["boost"]
+        self.entries.append(entry)
+        return
+
+    _add_entry_from_context(self, page)
+
+
+SearchIndex.add_entry_from_context = add_entry_from_context

docs/mkdocs/javascript/mathjax.js

@@ -1,20 +1,44 @@
-// Enables MathJax rendering
+// Enables MathJax rendering only on pages that contain math.
 window.MathJax = {
-    tex: {
-      inlineMath: [["\\(", "\\)"]],
-      displayMath: [["\\[", "\\]"]],
-      processEscapes: true,
-      processEnvironments: true
-    },
-    options: {
-      ignoreHtmlClass: ".*|",
-      processHtmlClass: "arithmatex"
-    }
-  };
+  tex: {
+    inlineMath: [["\\(", "\\)"]],
+    displayMath: [["\\[", "\\]"]],
+    processEscapes: true,
+    processEnvironments: true,
+  },
+  options: {
+    ignoreHtmlClass: ".*|",
+    processHtmlClass: "arithmatex",
+  },
+};
 
-  document$.subscribe(() => {
-    MathJax.startup.output.clearCache()
-    MathJax.typesetClear()
-    MathJax.texReset()
-    MathJax.typesetPromise()
-  })
+let mathJaxLoading;
+
+function loadMathJax() {
+  if (window.MathJax.typesetPromise) {
+    return Promise.resolve();
+  }
+  if (!mathJaxLoading) {
+    mathJaxLoading = new Promise((resolve, reject) => {
+      const script = document.createElement("script");
+      script.src = "https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js";
+      script.async = true;
+      script.onload = resolve;
+      script.onerror = reject;
+      document.head.appendChild(script);
+    });
+  }
+  return mathJaxLoading;
+}
+
+document$.subscribe(() => {
+  if (!document.querySelector(".arithmatex")) {
+    return;
+  }
+  loadMathJax().then(() => {
+    MathJax.startup.output.clearCache();
+    MathJax.typesetClear();
+    MathJax.texReset();
+    MathJax.typesetPromise();
+  });
+});

docs/mkdocs/javascript/memory.js

@@ -0,0 +1,2 @@
+// Opt out of bfcache so large generated API pages are released after navigation.
+window.addEventListener("unload", () => {});

docs/mkdocs/javascript/mermaid.js

@@ -1,21 +1,42 @@
-// Initialize Mermaid for diagram rendering
-mermaid.initialize({
-  startOnLoad: false,
-  theme: 'default',
-  securityLevel: 'loose',
-  flowchart: {
-    useMaxWidth: true,
-    htmlLabels: true
-  }
-});
+// Load Mermaid only on pages that contain diagrams.
+let mermaidLoading;
 
-// Render Mermaid diagrams when page content is ready
-document$.subscribe(() => {
-  const mermaidElements = document.querySelectorAll('.mermaid');
-  if (mermaidElements.length > 0) {
-    mermaid.run({
-      querySelector: '.mermaid',
-      nodes: mermaidElements
+function loadMermaid() {
+  if (window.mermaid) {
+    return Promise.resolve();
+  }
+  if (!mermaidLoading) {
+    mermaidLoading = new Promise((resolve, reject) => {
+      const script = document.createElement("script");
+      script.src = "https://unpkg.com/mermaid@10/dist/mermaid.min.js";
+      script.async = true;
+      script.onload = () => {
+        mermaid.initialize({
+          startOnLoad: false,
+          theme: "default",
+          securityLevel: "loose",
+          flowchart: {
+            useMaxWidth: true,
+            htmlLabels: true,
+          },
+        });
+        resolve();
+      };
+      script.onerror = reject;
+      document.head.appendChild(script);
     });
   }
+  return mermaidLoading;
+}
+
+document$.subscribe(() => {
+  if (!document.querySelector(".mermaid")) {
+    return;
+  }
+  loadMermaid().then(() => {
+    mermaid.run({
+      querySelector: ".mermaid",
+      nodes: document.querySelectorAll(".mermaid"),
+    });
+  });
 });

docs/models/supported_models.md

@@ -19,6 +19,7 @@ th {
 |--------------|--------|-------------------|------------|---------|-----|-----------|
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/serving/image_edit_api.md

@@ -13,7 +13,7 @@ For example...
 ```bash
 # Qwen-Image
 vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
-
+```
 
 ### Generate Images
 
@@ -53,7 +53,7 @@ result = client.images.edit(
     output_format='jpeg',
     # url格式
     extra_body={
-        "url": [input_image_url1,input_image_url],
+        "url": [input_image_url,input_image_url],
         "num_inference_steps": 50,
         "guidance_scale": 1,
         "seed": 777,
@@ -92,7 +92,8 @@ Content-Type: multipart/form-data
 | `user` | string | null | User identifier for tracking |
 | `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
 | `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
-| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s). |
+| `stream` | boolean | false | Return Server-Sent Events instead of a single JSON response. Currently supported only by multi-stage image edit pipelines that produce AR text before the final image, such as HunyuanImage3 IT2I AR+DiT. |
 
 #### vllm-omni Extension Parameters
 
@@ -109,6 +110,8 @@ Content-Type: multipart/form-data
 
 ### Response Format
 
+When `stream=false` or omitted, the endpoint returns the standard image edit response:
+
 ```json
 {
   "created": 1701234567,
@@ -120,10 +123,56 @@ Content-Type: multipart/form-data
     }
   ],
   "output_format": null,
-  "size": null,
+  "size": null
 }
 ```
 
+### Streaming Response Format
+
+Set `stream=true` when you want to receive the HunyuanImage3 IT2I AR
+recaption text before the final edited image is ready. Streaming is only
+available for multi-stage image edit pipelines where stage 0 produces AR text
+and a later diffusion stage produces the image. Single-stage diffusion image
+edit pipelines reject `stream=true` with HTTP 400.
+
+The response uses Server-Sent Events with `Content-Type: text/event-stream`.
+Each event is sent as one `data:` line. The stream order is:
+
+1. One or more AR text delta chunks.
+2. One final image chunk.
+3. `data: [DONE]`.
+
+AR text chunks expose only the generated text delta and the AR completion
+`index`. They do not expose token ids, ratio tokens, KV cache data, or internal
+prompt token ids.
+
+Example AR delta event:
+
+```text
+data: {"object":"image.edit.chunk","type":"ar_delta","delta":"A close-up product photo...","index":0,"created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
+```
+
+Example final image event:
+
+```text
+data: {"object":"image.edit.chunk","type":"image","data":[{"b64_json":"<base64-encoded PNG>","url":null,"revised_prompt":null}],"output_format":"png","size":"1024x1024","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
+```
+
+Terminal event:
+
+```text
+data: [DONE]
+```
+
+If the engine fails after the SSE response has started, the stream emits an
+error chunk followed by `[DONE]`:
+
+```text
+data: {"object":"error","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct","error":{"message":"<error message>","type":"server_error","code":500}}
+
+data: [DONE]
+```
+
 ## Examples
 
 ### Multiple Images input
@@ -134,12 +183,28 @@ curl -s -D >(grep -i x-request-id >&2) \
   -X POST "http://localhost:8000/v1/images/edits" \
   -F "model=xxx" \
   -F "image=@xx.png" \
-  -F "image=@xx.png"
+  -F "image=@xx.png" \
   -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
   -F "size=1024x1024" \
   -F "output_format=png"
 ```
 
+### Streaming HunyuanImage3 IT2I AR Text
+
+```bash
+curl -N -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=tencent/HunyuanImage-3.0-Instruct" \
+  -F "image=@./input.png" \
+  -F "prompt=Turn the product into a clean studio advertisement." \
+  -F "size=1024x1024" \
+  -F "output_format=png" \
+  -F "stream=true"
+```
+
+Use the `ar_delta` events for progressive display of the AR-generated
+recaption. Decode the `b64_json` field from the final `image` event to get the
+edited image.
+
 
 ## Parameter Handling
 

docs/user_guide/diffusion/attention_backends.md

@@ -1,28 +1,24 @@
 # Diffusion Attention Backends
 
-This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention.
+This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
 
 ## Overview
 
 Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
 
-This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
+This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
 
-On CUDA, the practical choices today are:
-
-- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
-- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
-- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
-
-If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
+The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
 
 ## Backend Options
 
 | Value | Notes |
 |---|---|
-| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
-| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
-| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
+| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
+| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
+| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
+| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
+| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
 
 ## Configuration
 
@@ -95,6 +91,80 @@ config = OmniDiffusionConfig(
 
 A plain dict is also accepted and normalized to `AttentionConfig`.
 
+## Platform Defaults
+
+### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
+
+Auto-route preference, in order:
+
+1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
+2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
+3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
+4. `TORCH_SDPA` — last resort
+
+The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
+
+**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
+
+### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
+
+Auto-route preference:
+
+1. `FLASH_ATTN` — when `flash-attn` is installed
+2. `TORCH_SDPA` — fallback
+
+`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
+
+## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
+
+Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
+
+| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
+|---|---|---|---|---|
+| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
+| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
+| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
+| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
+
+Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
+
+## Known Limitations
+
+### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
+
+LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
+
+**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 ...
+```
+
+### FA4 not yet integrated
+
+FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
+
+## Choosing a Backend Manually
+
+### When to override the default
+
+- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
+- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
+- **Workaround for known issues**: see Known Limitations above.
+
+### Verifying which backend is in use
+
+The startup log prints one of:
+
+```
+Using diffusion attention backend 'CUDNN_ATTN'           # explicit override
+Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
+Defaulting to diffusion attention backend SDPA           # nothing else available
+```
+
+If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
+
 ## SageAttention Installation
 
 vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
@@ -115,11 +185,33 @@ Quick check:
 python -c "import sageattention; print(sageattention.__file__)"
 ```
 
-## Usage
+## Usage Examples
 
-### Enable SageAttention
+### Default (auto-route)
 
-Example: HunyuanVideo-1.5 text-to-video
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
+    --output hv15.mp4
+```
+
+On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
+
+### Explicit backend selection
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
+    --output ltx2.mp4
+```
+
+### SageAttention (lossy)
 
 ```bash
 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
@@ -128,7 +220,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --height 480 --width 832 --num-frames 33 \
     --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
     --tensor-parallel-size 2 \
-    --output ../tmp/hv15_modelopt_sage.mp4
+    --output hv15_sage.mp4
 ```
 
 Example: Wan2.2 TI2V 5B
@@ -173,13 +265,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
 
 ## Validation Guidance
 
-Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
+Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
 
 Always compare:
 
 - End-to-end runtime
-- DiT / diffusion stage runtime
-- Output quality against a known-good baseline
+- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
+- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
 
 At minimum, keep the same:
 
@@ -187,6 +279,19 @@ At minimum, keep the same:
 - prompt
 - seed
 - resolution
-- frame count
-- inference steps
-- parallel config
+- frame count / step count
+- parallel config (TP / CFG-parallel / Ulysses degrees)
+
+## Reproducing the Benchmark Table
+
+The end-to-end numbers above were collected by running `text_to_video.py` /
+`text_to_image.py` with the same prompt and seed while varying
+`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
+backends without loading a model:
+
+```bash
+python benchmarks/diffusion/bench_attention_backends.py --preset hv15
+```
+
+It runs all three BF16 backends on representative DiT attention shapes and
+prints a ranking table at the end.

docs/user_guide/diffusion_features.md

@@ -113,7 +113,7 @@ The following tables show which models support each feature:
 | **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
@@ -145,7 +145,7 @@ The following tables show which models support each feature:
 | **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
 | 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 |
+| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz |
 
 ## Common Quick Start
 
@@ -333,6 +334,46 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
 - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
 - Run `--help` for the full argument surface.
 
+---
+
+## higgs-audio v2
+
+2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`.
+
+### Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+### Quick start (plain TTS)
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
+    --output-dir results/wavs
+```
+
+### Voice cloning
+
+Pass both `--ref-audio` and `--ref-text` together:
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --output-dir results/wavs
+```
+
+### Notes
+
+- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
+
 ## Example materials
 
 ??? abstract "cosyvoice3/end2end.py"
@@ -343,6 +384,9 @@ 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 "higgs_audio_v2/end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py"
 ??? abstract "glm_tts/end2end.py"
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -19,6 +19,7 @@ For the full list of supported architectures across all modalities, see
 |---|---|---|---|---|---|
 | 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) | — | ✓ |
+| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
 | 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) |
@@ -193,6 +194,45 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server
 
 ---
 
+## higgs-audio v2
+
+2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
+
+### Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec):
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+### Launch
+```bash
+GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
+```
+Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
+
+### Sending requests
+```bash
+# Plain TTS
+python higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --output-dir /tmp/higgs_out \
+    --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+
+# Voice cloning — pass a reference clip and its transcript together
+python higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --output-dir /tmp/higgs_clone \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --prompts "Hello, this is a cloned voice."
+```
+
+### Notes
+- Output: 24 kHz mono.
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
 ## GLM-TTS
 
 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
@@ -459,6 +499,13 @@ 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 "higgs_audio_v2/batch_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py"
+    ``````
+??? abstract "higgs_audio_v2/run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh"
 ??? abstract "glm_tts/gradio_demo.py"
     ``````py
     --8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"

docs/user_guide/quantization/autoround.md

@@ -52,8 +52,8 @@ vLLM-Omni's runtime module names.
 
 | Model | Scope | Status | Notes |
 |-------|-------|--------|-------|
+| GLM-Image | Diffusion transformer  | ✅ | `Intel/GLM-Image-int4-AutoRound` |
 | BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
-| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
 
 ## Configuration
 

docs/user_guide/quantization/mxfp8.md

@@ -7,26 +7,28 @@ using the OCP MX format: groups of 32 K-dimension elements share a single
 `float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise
 FP8 while keeping the same 8-bit weight footprint.
 
-This method supports two modes:
+This method supports three modes:
 
 | Mode | Description |
 |------|-------------|
 | **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
-| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
+| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
+| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
 
 ## 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) | ✅ |
+| Device | Online | Offline (Native) | Offline (AutoRound) |
+|--------|--------|------------------|---------------------|
+| 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.
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
+
+**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
 
 ## Model Type Support
 
@@ -82,9 +84,9 @@ python text_to_video.py --model <your-model> --quantization mxfp8
 vllm serve <your-model> --omni --quantization mxfp8
 ```
 
-### Offline Mode
+### Offline Mode (Native)
 
-Offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
+Native offline mode loads a pre-quantized 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 `transformer/config.json` so
 that vLLM-Omni auto-detects the offline path without a `--quantization` flag.
@@ -159,10 +161,41 @@ omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
 ```
 
 !!! note
-    No `--quantization` flag is needed for offline mode. The preprocessing
+    No `--quantization` flag is needed for native offline mode. The preprocessing
     script injects `quantization_config` into each `transformer/config.json`,
     which vLLM-Omni reads automatically to activate the offline MXFP8 method.
 
+### Offline Mode (AutoRound)
+
+AutoRound MXFP8 checkpoints declare `quant_method="auto-round"` with `data_type="mx_fp"` in their `config.json`. These are automatically detected and use the `IncMxfp8OfflineLinearMethod` backend.
+
+To use an AutoRound MXFP8 checkpoint:
+
+```bash
+python text_to_video.py --model <autoround-mxfp8-model>
+
+# Online serving
+vllm serve <autoround-mxfp8-model> --omni
+```
+
+Python API:
+
+```python
+omni = Omni(model="<autoround-mxfp8-model>")
+```
+
+!!! note
+    AutoRound MXFP8 checkpoints are auto-detected from `config.json` and do not require a `--quantization` flag. The config must include:
+    ```json
+    {
+      "quantization_config": {
+        "quant_method": "auto-round",
+        "data_type": "mx_fp",
+        ...
+      }
+    }
+    ```
+
 ## Parameters
 
 | Parameter | Type | Default | Description |

docs/user_guide/quantization/online.md

@@ -21,7 +21,7 @@ by msModelSlim and the merge tools.
 | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
 | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
 | AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
-| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ |
+| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
 | Ascend NPU | ❌ | ✅ | ✅ | ✅ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
@@ -36,7 +36,7 @@ MXFP4 are documented for the Ascend NPU path.
 |--------|-------|----------------|--------|
 | 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 |
+| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
 | 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)

examples/offline_inference/hunyuan_image3/end2end.py

@@ -17,7 +17,7 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniPromptType
 
 _REPO_ROOT = Path(__file__).resolve().parents[3]
-_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.yaml")
+_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml")
 _DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml")
 
 _MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = {

examples/offline_inference/internvla_a1/end2end.py

@@ -127,7 +127,12 @@ def run_one_path(
             ),
             args.device,
         )
-        pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+        pred = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-sample-{index}",
+        )
         pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
         results.append(
             {
@@ -150,6 +155,7 @@ def run_pipeline_forward(
     pipeline,
     batch_inputs: dict[str, torch.Tensor],
     noise: torch.Tensor,
+    request_id: str,
 ) -> torch.Tensor:
     output = pipeline.forward(
         OmniDiffusionRequest(
@@ -161,6 +167,7 @@ def run_pipeline_forward(
                     "decode_image": False,
                 }
             ),
+            request_id=request_id,
         )
     )
     if output.error:
@@ -212,23 +219,38 @@ def benchmark_forward(
 
     _synchronize(args.device)
     cold_start_begin = time.perf_counter()
-    pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+    pred = run_pipeline_forward(
+        pipeline,
+        batch_inputs,
+        noise,
+        request_id=f"internvla-a1-benchmark-{index}-cold",
+    )
     _synchronize(args.device)
     cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
 
     warmup_ms: list[float] = []
-    for _ in range(args.warmup_iters):
+    for iter_idx in range(args.warmup_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _ = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
+        )
         _synchronize(args.device)
         warmup_ms.append((time.perf_counter() - begin) * 1000.0)
 
     benchmark_ms: list[float] = []
-    for _ in range(args.benchmark_iters):
+    for iter_idx in range(args.benchmark_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _ = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
+        )
         _synchronize(args.device)
         benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
 
@@ -295,7 +317,12 @@ def main() -> None:
             dataset=dataset,
             train_meta=train_meta,
             collate_samples=collate_open_loop_samples,
-            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
+            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(
+                policy,
+                batch_inputs,
+                noise,
+                request_id="internvla-a1-open-loop",
+            ),
             num_episodes=args.num_episodes,
             seed=args.seed,
             device=args.device,

examples/offline_inference/ming_flash_omni/README.md

@@ -2,12 +2,13 @@
 
 [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
 
-vLLM-Omni supports two deployment modes:
+vLLM-Omni supports three deployment modes:
 
 | Mode | Deploy config | Output |
 |------|--------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image (online-serving only at the moment) |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -87,6 +88,17 @@ The default deploy YAML allocates thinker on GPUs 0–3 and talker on GPU 3 for
 
 Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to override the bundled deploy config.
 
+### Image generation (text-to-image / img2img)
+
+The diffusion-stage image-generation path is currently only wired through
+the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
+`"modalities": ["image"]`). For payloads, optional knobs
+(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
+reference-image flow, see the
+[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
+`end2end.py` does not yet exercise the imagegen stage; that is tracked as
+follow-up work.
+
 ## Online serving
 
 For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md).

examples/offline_inference/text_to_speech/higgs_audio_v2/README.md

@@ -0,0 +1,38 @@
+# higgs-audio v2 — offline example
+
+Drives Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) for `bosonai/higgs-audio-v2-generation-3B-base` end-to-end through the vLLM-Omni engine and writes a 24 kHz mono WAV per prompt.
+
+## Prerequisites
+
+Voice clone needs `transformers>=5.3.0` — vllm-omni loads the audio codec via HF's `HiggsAudioV2TokenizerModel`, instantiated from the `k2-fsa/OmniVoice/audio_tokenizer/` subdirectory (only that ~806 MB subdir is downloaded). The boson-ai standalone tokenizer repo's `model.safetensors` is actually a copy of the 3B talker LM, so HF can't load it directly; the k2 bundle ships the same codec weights repackaged with HF-compatible key naming.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+## Quick start
+
+Plain TTS:
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
+    --output-dir results/wavs
+```
+
+## Voice cloning
+
+Pass both `--ref-audio` and `--ref-text` together:
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --output-dir results/wavs
+```
+
+## Notes
+
+- Output: 24 kHz mono WAV.
+- Deploy config: `vllm_omni/deploy/higgs_audio_v2.yaml` (auto-loaded by HF `model_type`).
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- For online serving, see [`examples/online_serving/text_to_speech/higgs_audio_v2/`](../../../online_serving/text_to_speech/higgs_audio_v2/).

examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline higgs-audio v2 inference example.
+
+Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end
+through the vLLM-Omni engine without going through the HTTP server, and
+saves a 24 kHz mono WAV per prompt.
+
+Example:
+
+    python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\
+        --texts "Hello world." \\
+                "The quick brown fox jumps over the lazy dog." \\
+        --output-dir results/wavs
+"""
+
+from __future__ import annotations
+
+import os
+
+# DeepGEMM FP8 kernels require an optional backend that may not be installed.
+# Disable the warmup before importing vLLM so engine startup falls back to the
+# regular gemm path. Users with deep_gemm installed can override these.
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import time
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+SAMPLE_RATE = 24_000
+DEFAULT_TEXTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="bosonai/higgs-audio-v2-generation-3B-base",
+        help="Stage-0 talker model id or path.",
+    )
+    parser.add_argument(
+        "--texts",
+        type=str,
+        nargs="+",
+        default=list(DEFAULT_TEXTS),
+        help="One or more plain-text prompts to synthesize.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="results/wavs",
+        help="Directory to write per-prompt WAV files.",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Override the deploy config path. Auto-loads "
+        "vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=500,
+        help="Cap on Stage-0 codec frames per request.",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference clip. Required when --ref-audio is set.",
+    )
+    nullify_stage_engine_defaults(parser)
+    return parser.parse_args()
+
+
+def _slugify(text: str) -> str:
+    import re
+
+    slug = re.sub(r"\s+", "_", text.strip().lower())
+    slug = re.sub(r"[^a-z0-9_]+", "", slug)
+    return slug[:48] or "prompt"
+
+
+def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
+    """Pull the final concatenated PCM tensor out of a request's multimodal_output."""
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
+    if isinstance(audio, list):
+        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+        if not valid:
+            raise ValueError("audio list is empty")
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
+    arr = pcm.numpy()
+    if arr.dtype.kind == "f":
+        arr = np.clip(arr, -1.0, 1.0)
+        arr = (arr * 32767.0).astype(np.int16)
+    else:
+        arr = arr.astype(np.int16)
+    return arr
+
+
+def main():
+    args = parse_args()
+    if (args.ref_audio is None) != (args.ref_text is None):
+        raise SystemExit("--ref-audio and --ref-text must be supplied together")
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config)
+
+    # Build prompt_token_ids using the same path serving_speech.py takes online.
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_plain_text_prompt,
+        build_voice_clone_prompt,
+        input_ids_to_python_list,
+    )
+
+    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
+
+    # Voice-clone path: load the reference clip once. The HF processor will
+    # encode it via the bundled HiggsAudioV2TokenizerModel each time we build
+    # a prompt below. This is cheap on CPU for a few-second clip.
+    ref_wav: np.ndarray | None = None
+    ref_sr: int | None = None
+    if args.ref_audio is not None:
+        ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
+        if ref_wav.ndim == 2:
+            ref_wav = ref_wav.mean(axis=1)
+
+    print(f"Model       : {args.model}")
+    print(f"Prompts     : {len(args.texts)}")
+    print(f"Output dir  : {output_dir}")
+    print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}")
+
+    # Run one prompt at a time. The Stage-0 talker's per-slot audio state is
+    # request-scoped; submitting multiple prompts in the same engine.generate()
+    # call would batch them in the AR runner and exercise a code path that is
+    # not validated for this model yet.
+    total_elapsed = 0.0
+    total_dur = 0.0
+    for text in args.texts:
+        if ref_wav is not None:
+            out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text)
+            prompt = {
+                "prompt_token_ids": out["prompt_token_ids"],
+                # Bare tensors (NOT list-wrapped): the msgspec serializer in
+                # vllm_omni.data_entry_keys routes torch.Tensor → tensor_data
+                # and list[Tensor] → list_data (which silently strips tensors).
+                "additional_information": {
+                    "audio_input_ids": out["audio_input_ids"],
+                    "audio_input_ids_mask": out["audio_input_ids_mask"],
+                },
+            }
+        else:
+            inputs = build_plain_text_prompt(processor, text)
+            prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)}
+        t_start = time.perf_counter()
+        outputs = engine.generate([prompt])
+        elapsed = time.perf_counter() - t_start
+        total_elapsed += elapsed
+
+        mm = outputs[0].outputs[0].multimodal_output
+        pcm = _extract_pcm(mm)
+        slug = _slugify(text)
+        out_path = output_dir / f"{slug}.wav"
+        sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
+        dur = pcm.numel() / SAMPLE_RATE
+        total_dur += dur
+        print(f"  {slug:<50} dur={dur:5.2f}s  -> {out_path}")
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    main()

examples/offline_inference/vace/vace_video_generation.md

@@ -1,24 +1,105 @@
 # VACE Video Generation
 
-[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
+Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
 
-| Model | Architecture | Model Weights (bf16) | HuggingFace |
-|-------|-------------|----------------------|-------------|
-| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
-| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
+- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
 
-## Text-to-Video (T2V)
+## Table of Contents
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Key Arguments](#key-arguments)
+- [More CLI Examples](#more-cli-examples)
+
+## Overview
+
+[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
+
+### Supported Models
+
+| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
+|-------|-------------|-------------------|---------------------|-------------|
+| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
+| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
+
+!!! info
+*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
+
+Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
+
+## Quick Start
+
+### Python API
+
+Text-to-video generation:
+
+```python
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+if __name__ == "__main__":
+    omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
+    prompt = "A sleek robot stands in a vast warehouse filled with boxes"
+    outputs = omni.generate(
+        prompt,
+        OmniDiffusionSamplingParams(
+            height=480,
+            width=832,
+            num_frames=81,
+            num_inference_steps=30,
+            guidance_scale=5.0,
+        ),
+    )
+    video = outputs[0].images
+
+    from diffusers.utils import export_to_video
+    export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
+    omni.close()
+```
+
+### Local CLI Usage
 
 ```bash
 python vace_video_generation.py \
   --mode t2v \
   --prompt "A sleek robot stands in a vast warehouse filled with boxes" \
   --height 480 --width 832 --num-frames 81 \
-  --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
   --output t2v_output.mp4
 ```
 
-## Image-to-Video (I2V)
+## Key Arguments
+
+| Argument | Type | Default | Description |
+| -------- | ---- | ------- | ----------- |
+| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
+| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
+| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
+| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
+| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
+| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
+| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
+| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
+| `--num-frames` | int | `81` | Number of video frames to generate |
+| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
+| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
+| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
+| `--seed` | int | `42` | Random seed for deterministic sampling |
+| `--fps` | int | `16` | Frames per second for the saved MP4 |
+| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
+| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
+| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
+| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
+| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
+| `--tensor-parallel-size` | int | — | Tensor parallel size |
+| `--enforce-eager` | flag | off | Disable torch.compile |
+
+> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
+
+## More CLI Examples
+
+### Image-to-Video (I2V)
 
 First frame is kept, remaining frames are generated:
 
@@ -31,7 +112,7 @@ python vace_video_generation.py \
   --output i2v_output.mp4
 ```
 
-## First-Last-Frame Interpolation (FLF2V)
+### First-Last-Frame Interpolation (FLF2V)
 
 ```bash
 python vace_video_generation.py \
@@ -42,7 +123,7 @@ python vace_video_generation.py \
   --output flf2v_output.mp4
 ```
 
-## Inpainting
+### Inpainting
 
 Center vertical stripe is masked and regenerated:
 
@@ -55,7 +136,7 @@ python vace_video_generation.py \
   --output inpaint_output.mp4
 ```
 
-## Reference Image-guided (R2V)
+### Reference Image-guided (R2V)
 
 ```bash
 python vace_video_generation.py \
@@ -65,24 +146,3 @@ python vace_video_generation.py \
   --height 480 --width 832 --num-frames 81 \
   --output r2v_output.mp4
 ```
-
-## Key Arguments
-
-- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
-- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
-- `--image`: Input image for I2V, inpainting, and R2V modes.
-- `--last-image`: Last frame image for FLF2V mode.
-- `--prompt`: Text description of desired video.
-- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
-- `--num-frames`: Number of frames (default 81).
-- `--guidance-scale`: CFG scale (default 5.0).
-- `--flow-shift`: Scheduler flow shift (default 5.0).
-- `--num-inference-steps`: Number of denoising steps (default 30).
-- `--fps`: Frames per second for the saved MP4 (default 16).
-- `--output`: Path to save the generated video.
-- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
-- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
-- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
-- `--tensor-parallel-size`: Tensor parallel size.
-
-> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -102,3 +102,5 @@ Key arguments:
 - `--audio-negative-prompt`: negative prompt for audio generation.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
 - `--cache-backend`: enable `cache_dit` for acceleration.
+- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
+  (VAEs, the T5 text encoder, norms and modulation stay bf16.)

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -4,6 +4,7 @@
 import argparse
 import re
 import time
+from typing import Any
 
 import numpy as np
 from PIL import Image
@@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2, 3, 4],
         help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
     )
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=1,
+        help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
+    )
     parser.add_argument(
         "--video-negative-prompt",
         default="jitter, bad hands, blur, distortion",
@@ -52,6 +59,13 @@ def parse_args() -> argparse.Namespace:
         help="Negative prompt for audio.",
     )
     parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "int8"],
+        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -100,6 +114,27 @@ def load_image_and_audio(image_paths, audio_paths):
     return image, audio
 
 
+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() -> None:
     args = parse_args()
     if args.prompt is None and args.prompt_file is None:
@@ -142,6 +177,7 @@ def main() -> None:
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
@@ -162,7 +198,7 @@ def main() -> None:
             "scm_steps_policy": "dynamic",
         }
 
-    omni = Omni(
+    omni_kwargs = dict(
         model=args.model,
         parallel_config=parallel_config,
         model_type=args.model_type,
@@ -171,6 +207,9 @@ def main() -> None:
         cache_backend=args.cache_backend,
         cache_config=cache_config,
     )
+    if args.quantization is not None:
+        omni_kwargs["quantization"] = args.quantization
+    omni = Omni(**omni_kwargs)
     start = time.perf_counter()
     outputs = omni.generate(prompt, sampling_params)
     elapsed = time.perf_counter() - start
@@ -208,6 +247,9 @@ def main() -> None:
         f.write(video_bytes)
     print(f"Saved generated video to {output_path}")
     print(f"Total time: {elapsed:.2f}s")
+    peak_mb = _extract_peak_memory_mb(outputs)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
 
 
 if __name__ == "__main__":

examples/online_serving/ming_flash_omni/README.md

@@ -10,6 +10,7 @@ Please refer to [README.md](../../../README.md)
 |------|---------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm serve ... --omni` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -56,8 +57,98 @@ bash run_curl_multimodal_generation.sh use_image
 bash run_curl_multimodal_generation.sh use_audio
 bash run_curl_multimodal_generation.sh use_video
 bash run_curl_multimodal_generation.sh use_mixed_modalities
+bash run_curl_multimodal_generation.sh use_image_gen
 ```
 
+## Image generation (text-to-image)
+
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
+Launch with the dual-stage YAML, which adds an image-gen stage behind the
+thinker:
+
+```bash
+vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml \
+    --stage-init-timeout 1800 \
+    --port 8188
+```
+
+Then request image output by passing `"modalities": ["image"]`:
+
+```bash
+curl http://127.0.0.1:8188/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "Jonathan1909/Ming-flash-omni-2.0",
+      "messages": [{"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}],
+      "modalities": ["image"]
+    }' -o /tmp/ming_response.json
+
+python - <<'PY'
+import base64, json
+r = json.load(open('/tmp/ming_response.json'))
+url = r['choices'][0]['message']['content'][0]['image_url']['url']
+open('/tmp/ming_gen.png', 'wb').write(base64.b64decode(url.split(',', 1)[1]))
+print('saved /tmp/ming_gen.png')
+PY
+```
+
+### Optional knobs
+
+Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
+
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
+| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
+| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
+
+Example with all the knobs:
+
+```bash
+curl http://127.0.0.1:8188/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "Jonathan1909/Ming-flash-omni-2.0",
+      "modalities": ["image"],
+      "sampling_params_list": [
+        {"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
+         "extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
+        {"seed":42,"extra_args":{"image_gen":{
+            "steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
+            "byte5_text":["理解与生成统一"]}}}
+      ],
+      "messages": [{"role": "user", "content": "Draw a poster."}]
+    }' -o /tmp/ming_response.json
+```
+
+### img2img (reference image)
+
+Add an `image_url` content part to the user message; the parser routes it
+into the diffusion stage as `extra[reference_image]`:
+
+```jsonc
+"messages": [{
+  "role": "user",
+  "content": [
+    {"type": "text", "text": "Change the background to a sandy beach at sunset."},
+    {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
+  ]
+}]
+```
+
+### GPU layout
+
+The shipped `ming_flash_omni_dual.yaml` allocates the thinker on GPUs 0–3
+(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
+`runtime.devices` per stage to relocate; with fewer GPUs available, drop
+the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
+warmup takes roughly an extra 30–60 s on top of the thinker — set
+`--stage-init-timeout 1800` if the default 300 s is too tight.
+
 ## Modality control
 
 | `modalities` | Server config | Output |
@@ -65,6 +156,7 @@ bash run_curl_multimodal_generation.sh use_mixed_modalities
 | `["text"]` or omitted | Thinker only | Text |
 | `["audio"]` | Thinker + Talker | Audio (speech) |
 | `["text", "audio"]` | Thinker + Talker | Text + Audio |
+| `["image"]` | Thinker + Imagegen (dual YAML) | Image (PNG, base64 in `choices[0].message.content`) |
 
 For ready-to-copy curl examples (text / audio / multimodal input, SSE
 streaming, reasoning mode), see the recipe at

examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh

@@ -7,14 +7,15 @@ PORT="${PORT:-8091}"
 QUERY_TYPE="${1:-text}"
 
 # Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities)$ ]]; then
+if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then
     echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities]"
+    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]"
     echo "  text: Text-only query"
     echo "  use_audio: Audio + Text query"
     echo "  use_image: Image + Text query"
     echo "  use_video: Video + Text query"
     echo "  use_mixed_modalities: Audio + Image + Video + Text query"
+    echo "  use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally"
     exit 1
 fi
 
@@ -100,6 +101,43 @@ esac
 echo "Running query type: $QUERY_TYPE"
 echo ""
 
+if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
+    # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
+    # The shipped ming_flash_omni_dual.yaml runs the thinker on cards 0-3
+    # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
+    # section in README.md for knobs and img2img.
+    out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
+    request_body=$(cat <<'EOF'
+{
+  "model": "Jonathan1909/Ming-flash-omni-2.0",
+  "modalities": ["image"],
+  "messages": [
+    {"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}
+  ]
+}
+EOF
+)
+    output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+        -X POST "http://localhost:${PORT}/v1/chat/completions" \
+        -H "Content-Type: application/json" \
+        -d "$request_body")
+    # Server returns base64-encoded PNG either as a data URL in
+    # ``choices[0].message.content[0].image_url.url`` or as a raw string.
+    b64=$(echo "$output" | jq -r '.choices[0].message.content
+        | if type == "array"
+            then (map(select(.image_url? != null)) | .[0].image_url.url // "")
+            elif type == "string" then . else "" end')
+    if [[ -z "$b64" ]]; then
+        echo "Error: no image returned. Raw response:"
+        echo "$output" | jq '.'
+        exit 1
+    fi
+    if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi
+    printf '%s' "$b64" | base64 --decode > "$out_path"
+    echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)"
+    exit 0
+fi
+
 request_body=$(cat <<EOF
 {
   "model": "Jonathan1909/Ming-flash-omni-2.0",

examples/online_serving/text_to_speech/README.md

@@ -143,6 +143,26 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
 pip install fish-speech
 ```
 
+### Kvcache attention fast path
+
+Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
+default on CUDA builds. Set `VLLM_OMNI_FISH_KVCACHE_ATTN=0` to disable it, or
+`VLLM_OMNI_FISH_KVCACHE_ATTN=required` to fail fast if the fast path cannot be
+installed.
+
+```bash
+# Verify fast path availability.
+python - <<'PY'
+from vllm_omni.attention import fish_kvcache_attn
+
+print(fish_kvcache_attn.is_available())
+print(fish_kvcache_attn.load_error())
+PY
+
+# Optional: disable the runtime fast path.
+export VLLM_OMNI_FISH_KVCACHE_ATTN=0
+```
+
 ### Launch
 ```bash
 vllm serve fishaudio/s2-pro --omni --port 8091

examples/online_serving/text_to_speech/higgs_audio_v2/README.md

@@ -0,0 +1,63 @@
+# higgs-audio v2 online example
+
+This directory contains the online-serving entry points for boson-ai's higgs-audio v2 as integrated by vllm-omni: a 2-stage TTS pipeline (Llama-3.2-3B talker with DualFFN audio expert + HiggsAudio codec decoder) emitting 24 kHz mono speech.
+
+## Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec). Only that ~806 MB subdir is downloaded.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+## Files
+
+- `run_server.sh` — launch the vllm-omni server with the bundled `vllm_omni/deploy/higgs_audio_v2.yaml` deploy config.
+- `batch_speech_client.py` — send a list of prompts to `/v1/audio/speech` and save the returned WAV / PCM bytes to a directory; optionally passes `--ref-audio` + `--ref-text` for shallow voice clone.
+
+## Launching the server
+
+```bash
+GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
+```
+
+Environment overrides:
+
+- `MODEL` — HF id of the talker (default `bosonai/higgs-audio-v2-generation-3B-base`).
+- `PORT` — server port (default `8094`).
+- `GPUS` — `CUDA_VISIBLE_DEVICES` value (default `6,7`).
+- `GPU_UTIL` — `--gpu-memory-utilization` (default `0.4`).
+
+The script also exports `VLLM_USE_DEEP_GEMM=0` / `VLLM_MOE_USE_DEEP_GEMM=0` so the example works on images without the optional `deep_gemm` backend.
+
+The deploy YAML ships with `async_chunk: false` and `codec_streaming: true`, i.e. Stage 0 finishes its codec frames before Stage 1 starts decoding, and Stage 1 streams WAV/PCM bytes to the client chunk-by-chunk.
+
+## Driving the server
+
+Plain TTS:
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --model bosonai/higgs-audio-v2-generation-3B-base \
+    --output-dir /tmp/higgs_audio_v2_batch \
+    --prompts "Hello world." \
+              "The quick brown fox jumps over the lazy dog."
+```
+
+Voice clone — pass a reference clip and its transcript (both required together):
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --model bosonai/higgs-audio-v2-generation-3B-base \
+    --output-dir /tmp/higgs_audio_v2_clone \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --prompts "Hello, this is a cloned voice."
+```
+
+## Notes
+
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with explicit 4xx by the request validator): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.

examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Batch client for the higgs-audio v2 online server.
+
+Sends a fixed list of prompts to ``/v1/audio/speech`` and saves the returned
+WAV files (or raw PCM bytes when ``--format pcm``) into ``--output-dir``.
+
+Usage (plain text -> speech):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+      --base-url http://localhost:8094 \
+      --output-dir /tmp/higgs_audio_v2_batch \
+      --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+
+Usage (shallow voice clone — pass a reference clip + its transcript):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+      --base-url http://localhost:8094 \
+      --output-dir /tmp/higgs_audio_v2_clone \
+      --ref-audio path/to/reference.wav \
+      --ref-text "the transcript of the reference clip" \
+      --prompts "Hello world."
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import sys
+from pathlib import Path
+
+DEFAULT_PROMPTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "It was the night before my birthday.",
+    "Innovation distinguishes between a leader and a follower.",
+)
+
+
+def _slug(text: str) -> str:
+    import re
+
+    s = re.sub(r"\s+", "_", text.strip().lower())
+    return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--base-url", default="http://localhost:8094")
+    parser.add_argument("--model", default="higgs_audio_v2")
+    parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
+    parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_audio_v2_batch"))
+    parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
+    parser.add_argument("--max-new-tokens", type=int, default=300)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--timeout-s", type=float, default=120.0)
+    parser.add_argument(
+        "--ref-audio",
+        type=Path,
+        default=None,
+        help="Reference clip for voice clone (path to a WAV file). Must be paired with --ref-text.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference clip. Required when --ref-audio is set.",
+    )
+    args = parser.parse_args()
+
+    if (args.ref_audio is None) != (args.ref_text is None):
+        print("--ref-audio and --ref-text must be supplied together", file=sys.stderr)
+        return 2
+
+    ref_audio_data_url: str | None = None
+    if args.ref_audio is not None:
+        if not args.ref_audio.exists():
+            print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
+            return 2
+        mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
+        ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
+        ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
+
+    try:
+        import httpx
+    except ImportError:
+        print(
+            "this client needs `httpx`. Install with `pip install httpx`.",
+            file=sys.stderr,
+        )
+        return 2
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    url = args.base_url.rstrip("/") + "/v1/audio/speech"
+    failures = 0
+    with httpx.Client(timeout=args.timeout_s) as client:
+        for prompt in args.prompts:
+            payload = {
+                "model": args.model,
+                "input": prompt,
+                "response_format": args.format,
+                "max_new_tokens": args.max_new_tokens,
+                "seed": args.seed,
+            }
+            if ref_audio_data_url is not None:
+                payload["ref_audio"] = ref_audio_data_url
+                payload["ref_text"] = args.ref_text
+            resp = client.post(url, json=payload)
+            if resp.status_code != 200:
+                print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
+                failures += 1
+                continue
+            suffix = ".wav" if args.format == "wav" else ".pcm"
+            out = args.output_dir / f"{_slug(prompt)}{suffix}"
+            out.write_bytes(resp.content)
+            print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
+
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Launch vLLM-Omni server for higgs-audio v2.
+#
+# v1 scope: plain text -> 24 kHz speech only. Voice cloning, multi-speaker,
+# ChatML rich content, and language overrides are rejected by the validator
+# with explicit 4xx (see vllm_omni/entrypoints/openai/serving_speech.py).
+#
+# Usage:
+#   ./run_server.sh                 # default port 8094, GPUs 6 and 7
+#   PORT=8095 GPUS=6,7 ./run_server.sh
+#   MODEL=bosonai/higgs-audio-v2-generation-3B-base ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-bosonai/higgs-audio-v2-generation-3B-base}"
+PORT="${PORT:-8094}"
+GPUS="${GPUS:-6,7}"
+GPU_UTIL="${GPU_UTIL:-0.4}"
+
+echo "Starting higgs-audio v2 server"
+echo "  MODEL=$MODEL"
+echo "  PORT=$PORT"
+echo "  CUDA_VISIBLE_DEVICES=$GPUS"
+
+# DeepGEMM FP8 kernels are optional and trip warmup on builds without
+# the deep_gemm backend; disable them so the example works out of the box.
+# Users with deep_gemm installed can re-enable via the same env vars.
+CUDA_VISIBLE_DEVICES="$GPUS" \
+VLLM_USE_DEEP_GEMM=0 \
+VLLM_MOE_USE_DEEP_GEMM=0 \
+vllm-omni serve "$MODEL" \
+    --deploy-config vllm_omni/deploy/higgs_audio_v2.yaml \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization "$GPU_UTIL" \
+    --trust-remote-code \
+    --omni

mkdocs.yml

@@ -38,13 +38,12 @@ theme:
   features:
     - content.action.edit
     - content.code.copy
-    - navigation.instant
-    - navigation.instant.progress
     - navigation.tracking
     - navigation.tabs
     - navigation.tabs.sticky
     - navigation.sections
     - navigation.indexes
+    - navigation.prune
     - navigation.top
     - search.suggest
     - search.highlight
@@ -56,6 +55,7 @@ theme:
   custom_dir: docs/mkdocs/overrides
 
 hooks:
+  - docs/mkdocs/hooks/search_index.py
   - docs/mkdocs/hooks/generate_api_readme.py
   - docs/mkdocs/hooks/url_schemes.py
   - docs/mkdocs/hooks/generate_examples.py
@@ -117,6 +117,7 @@ plugins:
             separate_signature: true
             show_overloads: true
             signature_crossrefs: true
+            show_source: false
           inventories:
           - https://docs.python.org/3/objects.inv
           - https://typing-extensions.readthedocs.io/en/latest/objects.inv
@@ -167,9 +168,8 @@ extra_css:
   - mkdocs/stylesheets/extra.css
 
 extra_javascript:
+  - mkdocs/javascript/memory.js
   - mkdocs/javascript/mathjax.js
-  - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
-  - https://unpkg.com/mermaid@10/dist/mermaid.min.js
   - mkdocs/javascript/mermaid.js
   - mkdocs/javascript/edit_and_feedback.js
   - mkdocs/javascript/slack_and_forum.js

pyproject.toml

@@ -249,6 +249,7 @@ extend-ignore-identifiers-re = [
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
     ".*[Nn][Dd]3.*",
+    "byt5.*",
 ]
 extend-ignore-re = [
     "\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -261,3 +262,4 @@ Vai = "Vai"
 tockens = "tockens"
 CANN = "CANN"
 ND = "ND"
+byt5 = "byt5"

recipes/LTX/LTX-2.3.md

@@ -86,7 +86,6 @@ curl -X POST http://localhost:8000/v1/videos \
 - Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
 - Known limitations:
   - No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder)
-  - No CFG-parallel support (single-GPU only)
   - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
 
 ## Hardware Support

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

@@ -50,6 +50,26 @@ fi
 pip install fish-speech
 ```
 
+#### Kvcache attention fast path
+
+Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
+default on CUDA builds. Set `VLLM_OMNI_FISH_KVCACHE_ATTN=0` to disable it, or
+`VLLM_OMNI_FISH_KVCACHE_ATTN=required` to fail fast if the fast path cannot be
+installed.
+
+```bash
+# Verify fast path availability.
+python - <<'PY'
+from vllm_omni.attention import fish_kvcache_attn
+
+print(fish_kvcache_attn.is_available())
+print(fish_kvcache_attn.load_error())
+PY
+
+# Optional: disable the runtime fast path.
+export VLLM_OMNI_FISH_KVCACHE_ATTN=0
+```
+
 #### Command
 
 ```bash

setup.py

@@ -62,6 +62,12 @@ def detect_target_device() -> str:
         else:
             print(f"Warning: Invalid VLLM_OMNI_TARGET_DEVICE '{target_device}', falling back to auto-detection")
 
+    # Priority 1.5: ReadTheDocs builds use CPU requirements (no GPU, 1 GB RAM cap).
+    # Avoids pulling ~2 GB of CUDA libs that push the mkdocs build into swap.
+    if os.environ.get("READTHEDOCS"):
+        print("Detected READTHEDOCS=1; using cpu requirements for docs build")
+        return "cpu"
+
     # Priority 2: Torch backend detection
     # This is a code path for when user is using
     # --no-build-isolation flag

tests/core/sched/test_generation_scheduler_restore.py

@@ -28,7 +28,7 @@ class FakeAdapter:
             req = running.pop()
             self.waiting_for_chunk_running_requests.append(req)
 
-    def restore_queues(self, waiting, running):
+    def restore_queues(self, waiting, running, scheduler_requests=None):
         """Put requests back."""
         self.restore_called = True
         running.extend(self.waiting_for_chunk_running_requests)

tests/core/sched/test_omni_scheduler_mixin_timeouts.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit coverage for _process_pending_input_timeouts.
+
+Verifies that the mixin correctly *delegates* timed-out requests to the
+base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
+The end-to-end effect (queue removal + status set + per-request cleanup +
+client-facing FINISHED_ERROR emission) is the responsibility of upstream
+vLLM's ``finish_requests`` implementation and is covered by upstream tests;
+this file only asserts the wiring from the mixin to that API.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+
+class _FakeCoordinator:
+    def __init__(self, timed_out_ids):
+        self._timed_out_ids = set(timed_out_ids)
+        self.calls = []
+
+    def collect_timed_out_request_ids(self, timeout_s):
+        self.calls.append(timeout_s)
+        return set(self._timed_out_ids)
+
+
+class _FakeScheduler(OmniSchedulerMixin):
+    def __init__(self, requests, coordinator):
+        self.requests = requests
+        self.input_coordinator = coordinator
+        self.finish_calls = []
+
+    def finish_requests(self, req_ids, status):
+        self.finish_calls.append((set(req_ids), status))
+
+
+def test_process_pending_input_timeouts_delegates_to_finish_requests():
+    """Timed-out request present in self.requests is forwarded to finish_requests."""
+    req_id = "stuck-req"
+    requests = {req_id: SimpleNamespace(request_id=req_id)}
+    coord = _FakeCoordinator(timed_out_ids={req_id})
+    scheduler = _FakeScheduler(requests, coord)
+
+    scheduler._process_pending_input_timeouts()
+
+    assert len(coord.calls) == 1, "coordinator should be polled once"
+    assert coord.calls[0] > 0, "timeout must be positive when enabled"
+
+    assert len(scheduler.finish_calls) == 1
+    finished_ids, status = scheduler.finish_calls[0]
+    assert finished_ids == {req_id}
+    # RequestStatus is the upstream enum; the mixin imports it as
+    # RequestStatus.FINISHED_ERROR.  Check by name to avoid hard import here.
+    assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
+
+
+def test_process_pending_input_timeouts_skips_already_freed_request():
+    """Timed-out id no longer in self.requests must not be forwarded."""
+    coord = _FakeCoordinator(timed_out_ids={"already-freed"})
+    scheduler = _FakeScheduler(requests={}, coordinator=coord)
+
+    scheduler._process_pending_input_timeouts()
+
+    assert coord.calls == [coord.calls[0]] and coord.calls[0] > 0
+    assert scheduler.finish_calls == []
+
+
+def test_process_pending_input_timeouts_noop_without_coordinator():
+    """No coordinator => no finish_requests call, no crash."""
+
+    class _NoCoord(OmniSchedulerMixin):
+        def __init__(self):
+            self.requests = {}
+            self.input_coordinator = None
+            self.finish_calls = []
+
+        def finish_requests(self, req_ids, status):
+            self.finish_calls.append((set(req_ids), status))
+
+    scheduler = _NoCoord()
+    scheduler._process_pending_input_timeouts()
+    assert scheduler.finish_calls == []
+
+
+def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
+    """Setting DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0 disables the safety net."""
+    from vllm_omni.core.sched import omni_scheduler_mixin
+
+    monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
+
+    coord = _FakeCoordinator(timed_out_ids={"r1"})
+    scheduler = _FakeScheduler(requests={"r1": SimpleNamespace(request_id="r1")}, coordinator=coord)
+    scheduler._process_pending_input_timeouts()
+    assert coord.calls == [], "coordinator must not be polled when timeout is disabled"
+    assert scheduler.finish_calls == []
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -16,7 +16,7 @@ import torch
 import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
-    uses_qwen3_omni_full_payload_input_coordinator,
+    uses_full_payload_input_coordinator,
 )
 
 # ------------------------------------------------------------------ #
@@ -92,47 +92,96 @@ class MockQueue:
 
 
 class TestFullPayloadCoordinatorSelection(unittest.TestCase):
-    def test_qwen3_omni_talker_and_code2wav_use_full_payload_input_coordinator(self):
-        for model_stage in ("talker", "code2wav"):
+    """Tests for the (model_arch, model_stage) whitelist gate.
+
+    The init_omni_connectors arch allowlist is keyed by ``model_arch`` and
+    is a superset of the stages registered here -- consumer-wait stages
+    must be registered explicitly in ``_FULL_PAYLOAD_INPUT_STAGES``, while
+    the init allowlist covers both producer- and consumer-side runners.
+    These tests pin which ``(arch, stage)`` pairs the gate fires for today.
+    """
+
+    # Expected whitelist (model_arch, model_stage).  Hardcoded to avoid the
+    # tautology of importing _FULL_PAYLOAD_INPUT_STAGES and asserting it
+    # against itself; any drift between this matrix and the whitelist will
+    # fail loudly here.
+    EXPECTED_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
+        {
+            ("Qwen3OmniMoeForConditionalGeneration", "talker"),
+            ("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
+            ("Qwen2_5OmniForConditionalGeneration", "talker"),
+            ("Qwen2_5OmniForConditionalGeneration", "code2wav"),
+            ("CovoAudioForConditionalGeneration", "code2wav"),
+            ("MiMoAudioModel", "code2wav"),
+            ("Qwen3TTSCode2Wav", "code2wav"),
+            ("CosyVoice3Model", "cosyvoice3_code2wav"),
+            ("DyninOmniForConditionalGeneration", "token2image"),
+            ("DyninOmniForConditionalGeneration", "token2audio"),
+        }
+    )
+
+    def test_whitelist_matches_expected_matrix(self):
+        """_FULL_PAYLOAD_INPUT_STAGES must equal the hardcoded expected matrix.
+
+        Catches both accidental additions (which would silently enable the
+        consumer-wait gate for a new arch) and accidental removals (which
+        would silently disable an enabled arch).
+        """
+        from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
+
+        self.assertEqual(
+            frozenset(_FULL_PAYLOAD_INPUT_STAGES),
+            self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES,
+            msg="_FULL_PAYLOAD_INPUT_STAGES drifted from the expected matrix; "
+            "update EXPECTED_FULL_PAYLOAD_INPUT_STAGES if intentional.",
+        )
+
+    def test_all_whitelisted_arch_stage_pairs_fire_gate(self):
+        """Every (arch, stage) pair in the expected matrix must fire
+        the gate when stage_id > 0 and async_chunk=False.
+        """
+        for arch, stage in self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES:
             model_config = SimpleNamespace(
                 stage_id=1,
                 async_chunk=False,
-                model_arch="Qwen3OmniMoeForConditionalGeneration",
-                model_stage=model_stage,
+                model_arch=arch,
+                model_stage=stage,
+            )
+            self.assertTrue(
+                uses_full_payload_input_coordinator(model_config),
+                msg=f"expected gate to fire for {arch}/{stage}",
             )
 
-            self.assertTrue(uses_qwen3_omni_full_payload_input_coordinator(model_config))
-
-    def test_async_chunk_and_non_qwen3_omni_do_not_use_full_payload_input_coordinator(self):
+    def test_other_arch_or_stage_or_mode_does_not_fire(self):
         cases = [
             SimpleNamespace(
-                stage_id=1,
-                async_chunk=True,
-                model_arch="Qwen3OmniMoeForConditionalGeneration",
-                model_stage="talker",
+                stage_id=1, async_chunk=True, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="talker"
+            ),
+            SimpleNamespace(
+                stage_id=0, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="thinker"
             ),
             SimpleNamespace(
                 stage_id=1,
                 async_chunk=False,
-                model_arch="Qwen3TTSForConditionalGeneration",
-                model_stage="code2wav",
-            ),
-            SimpleNamespace(
-                stage_id=1,
-                async_chunk=False,
-                model_arch="Qwen2_5OmniForConditionalGeneration",
-                model_stage="talker",
-            ),
-            SimpleNamespace(
-                stage_id=0,
-                async_chunk=False,
                 model_arch="Qwen3OmniMoeForConditionalGeneration",
-                model_stage="thinker",
+                model_stage="some_future_stage",
+            ),
+            SimpleNamespace(
+                stage_id=1, async_chunk=False, model_arch="Qwen3TTSForConditionalGeneration", model_stage="code2wav"
+            ),
+            SimpleNamespace(
+                stage_id=1, async_chunk=False, model_arch="MingFlashOmniForConditionalGeneration", model_stage="talker"
+            ),
+            SimpleNamespace(stage_id=1, async_chunk=False, model_arch=None, model_stage="talker"),
+            SimpleNamespace(
+                stage_id=1, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage=None
             ),
         ]
-
         for model_config in cases:
-            self.assertFalse(uses_qwen3_omni_full_payload_input_coordinator(model_config))
+            self.assertFalse(
+                uses_full_payload_input_coordinator(model_config),
+                msg=f"expected gate OFF for {model_config}",
+            )
 
 
 class TestChunkCoordinatorStateTransition(unittest.TestCase):

tests/core/test_prefix_cache.py

@@ -13,7 +13,7 @@ DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE])
 
 
 class MockInputBatch:
-    def __init__(self, num_computed_tokens_cpu):
+    def __init__(self, num_computed_tokens_cpu, block_table=None):
         self.req_ids = ["req1", "req2"]
         self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)}
         self.num_computed_tokens_cpu = num_computed_tokens_cpu
@@ -21,11 +21,33 @@ class MockInputBatch:
         # Block table is only mocked for validation of length;
         # we don't actually need to add valid values here since
         # we patch the table when testing.
-        class _DummyBlockTable:
-            pass
+        if block_table is None:
 
-        self.block_table = _DummyBlockTable()
-        self.block_table.block_tables = [None]
+            class _DummyBlockTable:
+                pass
+
+            self.block_table = _DummyBlockTable()
+            self.block_table.block_tables = [None]
+        else:
+
+            class _TensorWrapper:
+                def __init__(self, tensor):
+                    self.cpu = tensor
+
+            class _BlockGroup:
+                def __init__(self, tensor):
+                    self.block_table = _TensorWrapper(tensor)
+
+            class _BlockTable:
+                def __init__(self, tensor):
+                    self._group = _BlockGroup(tensor)
+                    self.block_tables = [self._group.block_table]
+
+                def __getitem__(self, idx):
+                    assert idx == 0
+                    return self._group
+
+            self.block_table = _BlockTable(block_table)
 
 
 def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache:
@@ -171,6 +193,94 @@ def test_update_with_multimodal_outputs(feat_dims):
             assert torch.all(slot_states == new_output)
 
 
+def test_deferred_multimodal_cache_write_commits_on_completion():
+    cache = get_omni_pcache()
+    mm_key = "codes.audio"
+    block_table = torch.tensor(
+        [
+            [2, 3, 4, 5],
+            [6, 7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+    input_batch = MockInputBatch(
+        num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
+        block_table=block_table,
+    )
+
+    first_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=None,
+        multimodal_outputs={mm_key: first_chunk},
+        num_tokens_unpadded=4,
+        slot_mapping=torch.arange(8, 12),
+        skip_mm_cache_keys={mm_key},
+    )
+
+    rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
+    assert torch.all(rows[8:12] == 0)
+
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 4], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: first_chunk},
+        num_scheduled_tokens={"req1": 4, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+
+    second_chunk = torch.arange(8, 16, dtype=torch.long).reshape(4, 2)
+    input_batch.num_computed_tokens_cpu = torch.tensor([8, 0], dtype=torch.long)
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 4], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: second_chunk},
+        num_scheduled_tokens={"req1": 4, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+
+    cache.commit_deferred_mm_outputs({"req1"}, input_batch)
+
+    expected = torch.cat([first_chunk, second_chunk], dim=0)
+    rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
+    assert torch.equal(rows[8:16], expected)
+
+
+def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
+    cache = get_omni_pcache()
+    mm_key = "codes.audio"
+    block_table = torch.tensor(
+        [
+            [2, 3, 4, 5],
+            [6, 7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+    input_batch = MockInputBatch(
+        num_computed_tokens_cpu=torch.tensor([8, 0], dtype=torch.long),
+        block_table=block_table,
+    )
+    cached_codes = torch.arange(16, dtype=torch.long).reshape(8, 2)
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 8], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: cached_codes},
+        num_scheduled_tokens={"req1": 8, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+    cache.commit_deferred_mm_outputs({"req1"}, input_batch)
+
+    new_codes = torch.tensor([[99, 100]], dtype=torch.long)
+    cache.add_prefix_cached_new_req_id("req1")
+    merged = cache.get_merged_multimodal_states(
+        query_start_loc=torch.tensor([0, 1], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: new_codes},
+        num_scheduled_tokens={"req1": 1, "req2": 0},
+    )
+
+    assert torch.equal(merged[mm_key]["req1"], torch.cat([cached_codes, new_codes], dim=0))
+
+
 ### Tests for Merging
 def fake_get_cached_block_ids(self, req_idx, *args, **kwargs):
     """Fake block table lookup.

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -41,6 +41,72 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
 
+
+# ---------------------------------------------------------------------------
+# Inline field processing
+# ---------------------------------------------------------------------------
+def _process_inline_fields(obj: Any, parent_key: str = "") -> None:
+    """Recursively process '*-inline' fields into temp files."""
+
+    if isinstance(obj, list):
+        for item in obj:
+            _process_inline_fields(item, parent_key)
+        return
+
+    if not isinstance(obj, dict):
+        return
+
+    import atexit
+
+    import yaml
+
+    for key in list(obj.keys()):
+        value = obj[key]
+
+        if not key.endswith("-inline"):
+            _process_inline_fields(value, key)
+            continue
+
+        base_key = key[:-7]
+        full_key = f"{parent_key}.{key}" if parent_key else key
+
+        try:
+            if not isinstance(value, dict):
+                raise ValueError("must be a dict")
+
+            file_type = value.get("type")
+            content = value.get("content")
+
+            if file_type not in {"yaml", "jsonl"}:
+                raise ValueError(f"invalid type: {file_type}")
+
+            fd, path = tempfile.mkstemp(
+                suffix=f".{file_type}",
+                prefix=f"{base_key}_",
+            )
+
+            atexit.register(Path(path).unlink, missing_ok=True)
+
+            with os.fdopen(fd, "w", encoding="utf-8") as f:
+                if file_type == "jsonl":
+                    items = content if isinstance(content, list) else [content]
+                    f.writelines(json.dumps(x, ensure_ascii=False) + "\n" for x in items)
+                else:
+                    yaml.dump(
+                        content,
+                        f,
+                        allow_unicode=True,
+                        sort_keys=False,
+                        indent=2,
+                    )
+
+            obj[base_key] = path
+            del obj[key]
+
+        except Exception as e:
+            print(f"Warning: failed processing '{full_key}': {e}")
+
+
 # ---------------------------------------------------------------------------
 # Paths
 # ---------------------------------------------------------------------------
@@ -102,11 +168,14 @@ def _resolve_refs(configs: list[dict[str, Any]], config_dir: Path) -> list[dict[
 
 
 def load_configs(config_path: str) -> list[dict[str, Any]]:
+    """Load benchmark configs from JSON file and process inline fields."""
     try:
         abs_path = Path(config_path).resolve()
         with open(abs_path, encoding="utf-8") as f:
             configs = json.load(f)
-        return _resolve_refs(configs, abs_path.parent)
+        configs = _resolve_refs(configs, abs_path.parent)
+        _process_inline_fields(configs)
+        return configs
     except json.JSONDecodeError as e:
         raise ValueError(f"JSON parsing error: {str(e)}")
     except FileNotFoundError:
@@ -654,9 +723,47 @@ def run_benchmark(
 # ---------------------------------------------------------------------------
 
 
-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)
+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."""
+    if baseline_raw is None:
+        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)):
+        if sweep_index is None:
+            raise ValueError("sweep_index is required when baseline is a list or tuple")
+        return baseline_raw[sweep_index]
+    return baseline_raw
+
+
+def assert_result(
+    result: dict[str, Any],
+    params: dict[str, Any],
+    num_prompts: int,
+    *,
+    sweep_index: int | None = None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+    assert_baseline: bool = True,
+) -> None:
+    """Assert that benchmark metrics satisfy the configured baselines."""
     completed = result.get("completed_requests", result.get("completed", 0))
     assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
 
@@ -666,9 +773,15 @@ def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_base
         print("Skipping performance assertions.")
         return
 
-    for metric, threshold in params.get("baseline", {}).items():
+    for metric, baseline_raw in params.get("baseline", {}).items():
         current = result.get(metric)
         assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}"
+        threshold = _resolve_baseline_value(
+            baseline_raw,
+            sweep_index=sweep_index,
+            max_concurrency=max_concurrency,
+            request_rate=request_rate,
+        )
         if "throughput" in metric:
             assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}"
         else:
@@ -692,6 +805,111 @@ def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, An
     return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i")))
 
 
+def _to_list(value: Any) -> list[Any]:
+    if value is None:
+        return []
+    return [value] if not isinstance(value, (list, tuple)) else list(value)
+
+
+def _build_run_params(
+    params: dict[str, Any],
+    *,
+    num_prompts: int,
+    sweep_index: int | None = None,
+    request_rate: Any | None = None,
+    max_concurrency: Any | None = None,
+) -> dict[str, Any]:
+    run_params = {
+        key: value for key, value in params.items() if key not in {"request-rate", "max-concurrency", "num-prompts"}
+    }
+    run_params["num-prompts"] = num_prompts
+    if request_rate is not None:
+        run_params["request-rate"] = request_rate
+    if max_concurrency is not None:
+        run_params["max-concurrency"] = max_concurrency
+    if "baseline" in params:
+        run_params["baseline"] = {
+            metric: _resolve_baseline_value(
+                baseline_raw,
+                sweep_index=sweep_index,
+                max_concurrency=max_concurrency,
+                request_rate=request_rate,
+            )
+            for metric, baseline_raw in params["baseline"].items()
+        }
+    return run_params
+
+
+def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
+    request_rate_list = _to_list(params.get("request-rate"))
+    num_prompt_list = _to_list(params.get("num-prompts", 10))
+    max_concurrency_list = _to_list(params.get("max-concurrency"))
+
+    max_len = max(len(request_rate_list), len(max_concurrency_list))
+    if len(num_prompt_list) == 1 and max_len > 1:
+        num_prompt_list = num_prompt_list * max_len
+    elif max_len == 1 and len(num_prompt_list) > 1:
+        if len(request_rate_list) == 1:
+            request_rate_list = request_rate_list * len(num_prompt_list)
+        if len(max_concurrency_list) == 1:
+            max_concurrency_list = max_concurrency_list * len(num_prompt_list)
+        max_len = max(len(request_rate_list), len(max_concurrency_list))
+    elif len(num_prompt_list) != max_len and max_len > 0:
+        raise ValueError("The number of prompts does not match the request-rate or max-concurrency")
+
+    sweep_runs: list[dict[str, Any]] = []
+
+    for i, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
+        sweep_runs.append(
+            {
+                "params": _build_run_params(
+                    params,
+                    request_rate=request_rate,
+                    num_prompts=num_prompts,
+                    sweep_index=i,
+                ),
+                "num_prompts": num_prompts,
+                "sweep_index": i,
+                "request_rate": request_rate,
+                "max_concurrency": None,
+            }
+        )
+
+    for i, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
+        sweep_runs.append(
+            {
+                "params": _build_run_params(
+                    params,
+                    max_concurrency=max_concurrency,
+                    num_prompts=num_prompts,
+                    sweep_index=i,
+                    request_rate="inf",
+                ),
+                "num_prompts": num_prompts,
+                "sweep_index": i,
+                "request_rate": None,
+                "max_concurrency": max_concurrency,
+            }
+        )
+
+    if not sweep_runs:
+        default_num_prompts = num_prompt_list[0]
+        sweep_runs.append(
+            {
+                "params": _build_run_params(
+                    params,
+                    num_prompts=default_num_prompts,
+                ),
+                "num_prompts": default_num_prompts,
+                "sweep_index": None,
+                "request_rate": None,
+                "max_concurrency": None,
+            }
+        )
+
+    return sweep_runs
+
+
 # ---------------------------------------------------------------------------
 # Test entry point
 # ---------------------------------------------------------------------------
@@ -715,36 +933,47 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
     server_cfg = getattr(diffusion_server, "server_cfg", {})
-    endpoint = _resolve_benchmark_endpoint(server_cfg, params)
+    sweep_runs = _iter_sweep_runs(params)
 
-    result = run_benchmark(
-        host=diffusion_server.host,
-        port=diffusion_server.port,
-        model=diffusion_server.model,
-        params=params,
-        test_name=test_name,
-        endpoint=endpoint,
-        server_cfg=server_cfg,
-        source_file=cast(str, CONFIG_FILE_PATH),
-    )
+    for sweep_run in sweep_runs:
+        endpoint = _resolve_benchmark_endpoint(server_cfg, sweep_run["params"])
+        result = run_benchmark(
+            host=diffusion_server.host,
+            port=diffusion_server.port,
+            model=diffusion_server.model,
+            params=sweep_run["params"],
+            test_name=test_name,
+            endpoint=endpoint,
+            server_cfg=server_cfg,
+            source_file=cast(str, CONFIG_FILE_PATH),
+        )
 
-    print(f"\n{'=' * 60}")
-    print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
-    for key in (
-        "throughput_qps",
-        "latency_mean",
-        "latency_median",
-        "latency_p50",
-        "latency_p99",
-        "peak_memory_mb_max",
-        "peak_memory_mb_mean",
-        "peak_memory_mb_median",
-    ):
-        if key in result:
-            print(f"  {key}: {result[key]:.4f}")
+        print(f"\n{'=' * 60}")
+        print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
+        for key in (
+            "throughput_qps",
+            "latency_mean",
+            "latency_median",
+            "latency_p50",
+            "latency_p99",
+            "peak_memory_mb_max",
+            "peak_memory_mb_mean",
+            "peak_memory_mb_median",
+        ):
+            if key in result:
+                print(f"  {key}: {result[key]:.4f}")
 
-    print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
-    print("=" * 60)
+        print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
+        print("=" * 60)
 
-    assert_baseline = request.config.getoption("--assert-baseline", default=False)
-    assert_result(result, params, assert_baseline=assert_baseline)
+        assert_baseline = request.config.getoption("--assert-baseline", default=False)
+
+        assert_result(
+            result,
+            params,
+            sweep_run["num_prompts"],
+            sweep_index=sweep_run["sweep_index"],
+            max_concurrency=sweep_run["max_concurrency"],
+            request_rate=sweep_run["request_rate"],
+            assert_baseline=assert_baseline,
+        )

tests/dfx/perf/tests/test_bagel_vllm_omni.json

@@ -0,0 +1,152 @@
+[
+  {
+    "test_name": "test_bagel_single_device_single_stage_t2i",
+    "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel_single_stage.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.90,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_t2i_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.3678,
+          "latency_mean": 2.691,
+          "peak_memory_mb_max": 33083.6,
+          "peak_memory_mb_mean": 33082.5
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_bagel_single_device_single_stage_i2i",
+    "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel_single_stage.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.90,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_i2i_steps20",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "num-input-images": 1,
+        "baseline": {
+          "throughput_qps": 0.3124,
+          "latency_mean": 3.1686,
+          "peak_memory_mb_max": 33448.8,
+          "peak_memory_mb_mean": 33447.7
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_bagel_single_device_multi_stage_t2i",
+    "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.45,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_t2i_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.3746,
+          "latency_mean": 2.6425,
+          "peak_memory_mb_max": 33077.0,
+          "peak_memory_mb_mean": 33077.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_bagel_single_device_multi_stage_i2i",
+    "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.45,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_i2i_steps20",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "num-input-images": 1,
+        "baseline": {
+          "throughput_qps": 0.1349,
+          "latency_mean": 7.339,
+          "peak_memory_mb_max": 33321.2,
+          "peak_memory_mb_mean": 33321.2
+        }
+      }
+    ]
+  }
+]

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -0,0 +1,131 @@
+[
+    {
+        "test_name": "test_hunyuan_image3_ti2i_tp4_tp4_inline",
+        "description": "AR (TP=4) + DiT (TP=4) - using inline dataset and deploy-config",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/images/edits",
+        "server_params": {
+            "model": "tencent/HunyuanImage-3.0-Instruct",
+            "serve_args": {
+                "deploy-config-inline": {
+                    "type": "yaml",
+                    "content": {
+                        "pipeline": "hunyuan_image_3_moe",
+                        "async_chunk": false,
+                        "trust_remote_code": true,
+                        "connectors": {
+                            "rdma_connector": {
+                                "name": "MooncakeTransferEngineConnector",
+                                "extra": {
+                                    "host": "auto",
+                                    "zmq_port": 50051,
+                                    "protocol": "rdma",
+                                    "device_name": "",
+                                    "memory_pool_size": 4294967296,
+                                    "memory_pool_device": "cpu"
+                                }
+                            }
+                        },
+                        "stages": [
+                            {
+                                "stage_id": 0,
+                                "is_comprehension": true,
+                                "final_output": true,
+                                "final_output_type": "text",
+                                "max_num_seqs": 4,
+                                "gpu_memory_utilization": 0.95,
+                                "enforce_eager": false,
+                                "enable_prefix_caching": false,
+                                "max_num_batched_tokens": 32768,
+                                "devices": "0,1,2,3",
+                                "tensor_parallel_size": 4,
+                                "hf_overrides": {
+                                    "rope_parameters": {
+                                        "mrope_section": [0, 32, 32],
+                                        "rope_type": "default"
+                                    }
+                                },
+                                "omni_kv_config": {
+                                    "need_send_cache": true
+                                },
+                                "output_connectors": {
+                                    "to_stage_1": "rdma_connector"
+                                },
+                                "default_sampling_params": {
+                                    "temperature": 0.0,
+                                    "top_p": 1,
+                                    "top_k": -1,
+                                    "max_tokens": 8192,
+                                    "detokenize": true,
+                                    "skip_special_tokens": false,
+                                    "stop_token_ids": [128025]
+                                }
+                            },
+                            {
+                                "stage_id": 1,
+                                "max_num_seqs": 1,
+                                "enforce_eager": true,
+                                "devices": "4,5,6,7",
+                                "distributed_executor_backend": "mp",
+                                "omni_kv_config": {
+                                    "need_recv_cache": true
+                                },
+                                "parallel_config": {
+                                    "tensor_parallel_size": 4,
+                                    "enable_expert_parallel": true
+                                },
+                                "input_connectors": {
+                                    "from_stage_0": "rdma_connector"
+                                },
+                                "default_sampling_params": {
+                                    "num_inference_steps": 50,
+                                    "guidance_scale": 0
+                                }
+                            }
+                        ],
+                        "edges": [
+                            {
+                                "from": 0,
+                                "to": 1,
+                                "window_size": -1,
+                                "max_inflight": 1
+                            }
+                        ]
+                    }
+                },
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "inline_dataset_5_requests",
+                "dataset": "custom",
+                "dataset-path-inline": {
+                    "type": "jsonl",
+                    "content": [
+                        {
+                            "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
+                            "image_urls": ["https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"],
+                            "num_inference_steps": 8,
+                            "seed": 42,
+                            "guidance_scale": 1.0
+                        }
+                    ]
+                },
+                "task": "ti2i",
+                "bot-task": "think_recaption",
+                "max-concurrency": 4,
+                "num-prompts": 16,
+                "skip-performance-assertion": true,
+                "baseline": {
+                    "throughput_qps": 0.1,
+                    "latency_p99": 44.46,
+                    "latency_mean": 35.57,
+                    "latency_std_dev": 0.5,
+                    "peak_memory_mb_max": 70000,
+                    "peak_memory_mb_mean": 70000
+                }
+            }
+        ]
+    }
+]

tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json

@@ -0,0 +1,115 @@
+[
+    {
+        "test_name": "test_qwen3_omni_3gpu_replica2",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_overrides": {
+                "1": {"num_replicas": 2, "devices": "1,2"},
+                "2": {"num_replicas": 2, "devices": "1,2"}
+            },
+            "extra_cli_args": [
+                "--async-chunk"
+            ]
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [128, 128, 128, 128],
+                "max_concurrency": [8, 16, 24, 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, 2000, 3000, 4000],
+                    "mean_tpot_ms": [200, 200, 200, 200],
+                    "mean_audio_ttfp_ms": [10000, 15000, 20000, 25000],
+                    "mean_audio_rtf": [1.5, 1.5, 1.5, 1.5]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [10],
+                "request_rate": [0.1],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "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]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [40],
+                "request_rate": [0.5],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "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]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [100],
+                "request_rate": [1.0],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "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]
+                }
+            }
+        ]
+    }
+]

tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

@@ -1,10 +1,10 @@
 [
     {
-        "test_name": "test_qwen_image_edit_2509_single_device",
+        "test_name": "test_qwen_image_edit_2511_single_device",
         "description": "Single-device baseline (two input images)",
         "server_type": "vllm-omni",
         "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2509",
+            "model": "Qwen/Qwen-Image-Edit-2511",
             "serve_args": {
                 "enable-diffusion-pipeline-profiler": true
             }
@@ -49,11 +49,11 @@
         ]
     },
     {
-        "test_name": "test_qwen_image_edit_2509_ulysses2_cfg2_vae_patch4",
+        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
         "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
         "server_type": "vllm-omni",
         "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2509",
+            "model": "Qwen/Qwen-Image-Edit-2511",
             "serve_args": {
                 "ulysses-degree": 2,
                 "cfg-parallel-size": 2,
@@ -63,24 +63,6 @@
             }
         },
         "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.15,
-                    "latency_mean": 10,
-                    "peak_memory_mb_max": 55000,
-                    "peak_memory_mb_mean": 55000
-                }
-            },
             {
                 "name": "1536x1536_steps35_i2i_2img",
                 "dataset": "random",
@@ -102,11 +84,11 @@
         ]
     },
     {
-        "test_name": "test_qwen_image_edit_2509_ulysses2_cfg2_cache_dit",
+        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
         "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
         "server_type": "vllm-omni",
         "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2509",
+            "model": "Qwen/Qwen-Image-Edit-2511",
             "serve_args": {
                 "ulysses-degree": 2,
                 "cfg-parallel-size": 2,

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -61,23 +61,6 @@
             }
         },
         "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.21,
-                    "latency_mean": 10,
-                    "peak_memory_mb_max": 56000,
-                    "peak_memory_mb_mean": 56000
-                }
-            },
             {
                 "name": "1536x1536_steps35_i2i",
                 "dataset": "random",

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -181,5 +181,39 @@
                 }
             }
         ]
+    },
+    {
+        "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
+        "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true,
+                "step-execution": true,
+                "max-num-seqs": 8
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20_high_concurrency",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": [20, 40, 80, 160],
+                "max-concurrency": [1, 2, 4, 8],
+                "warmup-requests": 8,
+                "warmup-concurrency": 8,
+                "warmup-num-inference-steps": 20,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": [0.30, 0.50, 0.60, 0.70],
+                    "latency_mean": [3.50, 4.00, 7.00, 12.00],
+                    "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
+                }
+            }
+        ]
     }
 ]

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -60,6 +60,7 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
             output_type="pil",
             generator_device="cpu",
         ),
+        "request_id": "diffusers-backend",
     }
     defaults.update(overrides)
     return OmniDiffusionRequest(**defaults)

tests/diffusion/models/bagel/test_forward_cache_update_vae.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Bagel.forward_cache_update_vae packed-sequence assembly."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from transformers import PretrainedConfig
+from vllm.transformers_utils.configs.bagel import BagelConfig
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import (
+    Bagel,
+    BaseNavitOutputWithPast,
+    NaiveCache,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+HIDDEN_SIZE = 32
+NUM_LAYERS = 2
+
+
+def _make_bagel_config() -> BagelConfig:
+    llm_config = PretrainedConfig()
+    llm_config.hidden_size = HIDDEN_SIZE
+    llm_config.layer_module = "Qwen2MoTDecoderLayer"
+    llm_config.num_hidden_layers = NUM_LAYERS
+    llm_config.num_attention_heads = 4
+
+    vae_config = PretrainedConfig()
+    vae_config.downsample = 8
+    vae_config.z_channels = 4
+
+    return BagelConfig(
+        llm_config=llm_config,
+        vae_config=vae_config,
+        visual_gen=True,
+        visual_und=False,
+        latent_patch_size=2,
+        max_latent_size=32,
+    )
+
+
+def _make_tracking_language_model() -> MagicMock:
+    """Return a mock LM that records forward() kwargs."""
+    calls: list[dict] = []
+    hidden_size = HIDDEN_SIZE
+
+    def forward(*_args, **kwargs):
+        calls.append(dict(kwargs))
+        if kwargs.get("return_embeddings_only"):
+            text_ids = kwargs["packed_text_ids"]
+            emb = torch.randn(text_ids.shape[0], hidden_size)
+            return BaseNavitOutputWithPast(packed_query_sequence=emb)
+
+        packed_query_sequence = kwargs.get("packed_query_sequence")
+        if packed_query_sequence is None:
+            text_ids = kwargs["packed_text_ids"]
+            packed_query_sequence = torch.randn(text_ids.shape[0], hidden_size)
+        return BaseNavitOutputWithPast(
+            packed_query_sequence=packed_query_sequence,
+            past_key_values=kwargs.get("past_key_values"),
+        )
+
+    lm = MagicMock()
+    lm.forward = forward
+    lm._calls = calls
+    return lm
+
+
+def _vae_transform(image: Image.Image) -> torch.Tensor:
+    arr = torch.from_numpy(np.array(image.convert("RGB"), dtype=np.float32))
+    arr = arr / 127.5 - 1.0
+    return arr.permute(2, 0, 1)
+
+
+@pytest.fixture
+def bagel_and_vae_input():
+    language_model = _make_tracking_language_model()
+    bagel = Bagel(language_model=language_model, vit_model=None, config=_make_bagel_config())
+
+    image = Image.new("RGB", (64, 64), color=(128, 64, 32))
+    new_token_ids = {"start_of_image": 100, "end_of_image": 101}
+
+    gen_input, _newlens, _new_rope = bagel.prepare_vae_images(
+        curr_kvlens=[0],
+        curr_rope=[0],
+        images=[image],
+        transforms=_vae_transform,
+        new_token_ids=new_token_ids,
+    )
+
+    vae_model = MagicMock()
+    patch_shapes = gen_input["patchified_vae_latent_shapes"]
+    p = bagel.latent_patch_size
+    c = bagel.latent_channel
+
+    def encode(padded_images: torch.Tensor) -> torch.Tensor:
+        # Match AutoEncoder.encode layout: [B, z_channels, h*p, w*p] per image.
+        batch = padded_images.shape[0]
+        assert batch == len(patch_shapes)
+        latents = [torch.randn(c, h * p, w * p) for h, w in patch_shapes]
+        return torch.stack(latents, dim=0)
+
+    vae_model.encode = encode
+
+    expected_seq_len = int(gen_input["packed_seqlens"].sum().item())
+    num_text_tokens = int(gen_input["packed_text_ids"].numel())
+
+    return SimpleNamespace(
+        bagel=bagel,
+        language_model=language_model,
+        vae_model=vae_model,
+        gen_input=gen_input,
+        past_key_values=NaiveCache(NUM_LAYERS),
+        expected_seq_len=expected_seq_len,
+        num_text_tokens=num_text_tokens,
+    )
+
+
+class TestForwardCacheUpdateVae:
+    """Regression tests for single-stage img2img VAE cache prefill."""
+
+    def test_forwards_full_packed_query_sequence(self, bagel_and_vae_input):
+        """Update forward must use packed_query_sequence spanning text + VAE tokens."""
+        ctx = bagel_and_vae_input
+        assert ctx.expected_seq_len > ctx.num_text_tokens, "test setup should include many VAE tokens"
+
+        ctx.bagel.forward_cache_update_vae(
+            ctx.vae_model,
+            ctx.past_key_values,
+            **ctx.gen_input,
+        )
+
+        update_calls = [c for c in ctx.language_model._calls if not c.get("return_embeddings_only")]
+        assert len(update_calls) == 1
+
+        call = update_calls[0]
+        assert "packed_query_sequence" in call
+        assert "packed_text_ids" not in call
+
+        packed = call["packed_query_sequence"]
+        assert packed.shape[0] == ctx.expected_seq_len
+        assert call.get("mode") == "gen"
+
+    def test_packed_sequence_covers_text_and_vae_index_ranges(self, bagel_and_vae_input):
+        """Text and VAE index ranges must lie within packed_query_sequence."""
+        ctx = bagel_and_vae_input
+        ctx.bagel.forward_cache_update_vae(
+            ctx.vae_model,
+            ctx.past_key_values,
+            **ctx.gen_input,
+        )
+
+        update_call = next(c for c in ctx.language_model._calls if not c.get("return_embeddings_only"))
+        packed = update_call["packed_query_sequence"]
+
+        text_idx = ctx.gen_input["packed_text_indexes"]
+        vae_idx = ctx.gen_input["packed_vae_token_indexes"]
+        assert int(text_idx.max()) < packed.shape[0]
+        assert int(vae_idx.max()) < packed.shape[0]
+        assert int(text_idx.min()) >= 0
+        assert int(vae_idx.min()) >= 0

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py

@@ -48,6 +48,7 @@ def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=prompts or [{"prompt": "a cat dancing"}],
         sampling_params=sp,
+        request_id="dmd2-sanitize",
     )
 
 

tests/diffusion/models/dmd2/test_dmd2_scheduler.py

@@ -48,7 +48,7 @@ def _make_pipeline(cls):
 
 def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
     sp = OmniDiffusionSamplingParams(**sp_kwargs)
-    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
+    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp, request_id="dmd2-scheduler")
 
 
 @pytest.fixture(

tests/diffusion/models/ernie_image/test_ernie_image_pe.py

@@ -73,7 +73,7 @@ def test_should_apply_pe_respects_sampling_extra_args():
 
 
 def test_should_apply_pe_disables_dummy_warmup_request():
-    req = SimpleNamespace(request_ids=["dummy_req_id"], sampling_params=SimpleNamespace(extra_args={}))
+    req = SimpleNamespace(is_dummy_run=lambda: True, sampling_params=SimpleNamespace(extra_args={}))
 
     assert ErnieImagePipeline._should_apply_pe(req) is False
 

tests/diffusion/models/glm_image/test_glm_image_quantization.py

@@ -0,0 +1,741 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for GLM-Image quantization support (W4A16/AutoRound).
+
+These tests verify that the GLM-Image DiT transformer correctly accepts and uses
+quantization configs for W4A16/AutoRound quantization support.
+"""
+
+import os
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+    ColumnParallelGELU,
+    ColumnParallelSiLU,
+    GlmImageAdaLayerNormContinuous,
+    GlmImageAdaLayerNormZero,
+    GlmImageAttention,
+    GlmImageFeedForward,
+    GlmImageImageProjector,
+    GlmImagePrepare,
+    GlmImageRotaryPosEmbed,
+    GlmImageTransformer2DModel,
+    GlmImageTransformerBlock,
+    _positive_divisors,
+    validate_glm_image_tp_constraints,
+)
+from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(scope="function", autouse=True)
+def _init_distributed():
+    """Initialize the minimal distributed environment required by
+    ReplicatedLinear (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29502")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
+
+    vLLM's dispatch_unquantized_gemm() selects the backend by platform
+    (e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
+    CPU test tensors crash with NotImplementedError on ROCm.  Monkeypatch
+    the dispatcher to always return the default (torch.nn.functional.linear)
+    implementation which works on any device."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_mocks(mocker: MockerFixture):
+    """Set up common mocks for all tests."""
+    mocker.patch(
+        "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
+        return_value=1,
+    )
+    mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
+    mock_tp_group = mocker.MagicMock()
+    mock_tp_group.world_size = 1
+    mock_get_tp_group.return_value = mock_tp_group
+    yield
+
+
+class TestPositiveDivisors:
+    """Test _positive_divisors helper function."""
+
+    def test_divisors_of_1(self):
+        assert _positive_divisors(1) == {1}
+
+    def test_divisors_of_12(self):
+        assert _positive_divisors(12) == {1, 2, 3, 4, 6, 12}
+
+    def test_divisors_of_prime(self):
+        assert _positive_divisors(7) == {1, 7}
+
+    def test_divisors_of_0_returns_empty(self):
+        assert _positive_divisors(0) == set()
+
+    def test_divisors_of_negative_returns_empty(self):
+        assert _positive_divisors(-5) == set()
+
+
+class TestValidateGlmImageTpConstraints:
+    """Test TP constraint validation for GLM-Image."""
+
+    def test_valid_tp_size_1(self):
+        """TP=1 should always be valid."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=1,
+        )
+        assert 1 in result
+
+    def test_valid_tp_size_2_for_divisible_dim(self):
+        """TP=2 is valid when all dims are divisible by 2."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=2,
+        )
+        assert 1 in result
+        assert 2 in result
+
+    def test_valid_tp_size_4_for_divisible_dim(self):
+        """TP=4 is valid when all dims are divisible by 4."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=4,
+        )
+        assert 1 in result
+        assert 2 in result
+        assert 4 in result
+
+    def test_invalid_tp_size_3_for_divisible_dim(self):
+        """TP=3 is invalid when dim is not divisible by 3."""
+        with pytest.raises(ValueError, match="dim % tensor_parallel_size == 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,  # 2560 % 3 != 0
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=3,
+            )
+
+    def test_invalid_tp_size_zero(self):
+        """TP=0 should raise error."""
+        with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=0,
+            )
+
+    def test_invalid_tp_size_negative(self):
+        """Negative TP size should raise error."""
+        with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=-1,
+            )
+
+
+class TestGlmImageAdaLayerNormZeroQuantization:
+    """Test GlmImageAdaLayerNormZero with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the class accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=mock_quant_config,
+            prefix="test.norm1",
+        )
+        assert layer.linear.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=None,
+        )
+        assert layer.linear.quant_config is None
+
+    def test_forward_handles_tuple_return_from_linear(self):
+        """Verify forward handles tuple returns from ReplicatedLinear."""
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=None,
+        )
+
+        batch_size = 2
+        seq_len = 10
+        hidden_states = torch.randn(batch_size, seq_len, 2560)
+        encoder_hidden_states = torch.randn(batch_size, seq_len, 2560)
+        temb = torch.randn(batch_size, 512)
+
+        # This should work regardless of whether linear returns tuple or tensor
+        result = layer(hidden_states, encoder_hidden_states, temb)
+        assert len(result) == 10  # Should return 10 chunks
+
+
+class TestGlmImageAdaLayerNormContinuousQuantization:
+    """Test GlmImageAdaLayerNormContinuous with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the class accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=mock_quant_config,
+            prefix="test.norm_out",
+        )
+        assert layer.linear.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=None,
+        )
+        assert layer.linear.quant_config is None
+
+    def test_forward_handles_tuple_return_from_linear(self):
+        """Verify forward handles tuple returns from ReplicatedLinear."""
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=None,
+        )
+
+        batch_size = 2
+        seq_len = 10
+        x = torch.randn(batch_size, seq_len, 2560)
+        conditioning_embedding = torch.randn(batch_size, 512)
+
+        result = layer(x, conditioning_embedding)
+        assert result.shape == (batch_size, seq_len, 2560)
+
+
+class TestGlmImageAttentionQuantization:
+    """Test GlmImageAttention with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageAttention accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        attn = GlmImageAttention(
+            dim=2560,
+            num_heads=64,
+            head_dim=40,
+            quant_config=mock_quant_config,
+            prefix="test.attn1",
+        )
+        assert attn.to_qkv.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        attn = GlmImageAttention(
+            dim=2560,
+            num_heads=64,
+            head_dim=40,
+            quant_config=None,
+        )
+        assert attn.to_qkv.quant_config is None
+
+
+class TestColumnParallelModulesQuantization:
+    """Test ColumnParallelGELU and ColumnParallelSiLU with quantization config."""
+
+    def test_column_parallel_gelu_accepts_quant_config(self, mocker: MockerFixture):
+        """Verify ColumnParallelGELU accepts quant_config."""
+        mock_quant_config = mocker.MagicMock()
+        layer = ColumnParallelGELU(
+            dim_in=2560,
+            dim_out=10240,
+            quant_config=mock_quant_config,
+            prefix="test.gelu",
+        )
+        assert layer.proj.quant_config is mock_quant_config
+
+    def test_column_parallel_silu_accepts_quant_config(self, mocker: MockerFixture):
+        """Verify ColumnParallelSiLU accepts quant_config."""
+        mock_quant_config = mocker.MagicMock()
+        layer = ColumnParallelSiLU(
+            dim_in=2560,
+            dim_out=10240,
+            quant_config=mock_quant_config,
+            prefix="test.silu",
+        )
+        assert layer.proj.quant_config is mock_quant_config
+
+
+class TestGlmImageFeedForwardQuantization:
+    """Test GlmImageFeedForward with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageFeedForward accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="gelu-approximate",
+            quant_config=mock_quant_config,
+            prefix="test.ff",
+        )
+        # Check that the first layer (ColumnParallelGELU) has quant_config
+        gelu_layer = ff.net[0]
+        assert gelu_layer.proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="gelu-approximate",
+            quant_config=None,
+        )
+        gelu_layer = ff.net[0]
+        assert gelu_layer.proj.quant_config is None
+
+    def test_linear_silu_activation(self):
+        """Test linear-silu activation function initialization works."""
+        # This test verifies the module can be instantiated with linear-silu activation.
+        # Full forward testing requires proper TP group setup, so we just verify construction.
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="linear-silu",
+            quant_config=None,
+        )
+        # Verify the FFN has the correct structure
+        assert len(ff.net) == 3  # ColumnParallelSiLU, Identity, RowParallelLinear
+        assert isinstance(ff.net[0], ColumnParallelSiLU)
+
+
+class TestGlmImageTransformerBlockQuantization:
+    """Test GlmImageTransformerBlock with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageTransformerBlock accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+        block = GlmImageTransformerBlock(
+            dim=2560,
+            num_attention_heads=64,
+            attention_head_dim=40,
+            time_embed_dim=512,
+            parallel_config=parallel_config,
+            quant_config=mock_quant_config,
+            prefix="test.block",
+        )
+        # Check that inner modules have quant_config
+        assert block.norm1.linear.quant_config is mock_quant_config
+        assert block.attn1.to_qkv.quant_config is mock_quant_config
+        assert block.ff.net[0].proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        block = GlmImageTransformerBlock(
+            dim=2560,
+            num_attention_heads=64,
+            attention_head_dim=40,
+            time_embed_dim=512,
+            quant_config=None,
+        )
+        assert block.norm1.linear.quant_config is None
+
+
+class TestGlmImagePrepareModule:
+    """Test GlmImagePrepare module."""
+
+    def test_prepare_module_exists(self):
+        """Verify GlmImagePrepare module exists and works."""
+        projector = GlmImageImageProjector(in_channels=16, hidden_size=2560, patch_size=2)
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2)
+        prepare = GlmImagePrepare(
+            image_projector=projector,
+            rope=rope,
+            patch_size=2,
+        )
+
+        # Test forward pass
+        hidden_states = torch.randn(1, 16, 64, 64)  # [B, C, H, W]
+        result = prepare(hidden_states)
+
+        assert len(result) == 5
+        hidden_out, rope_cos, rope_sin, height, width = result
+        assert hidden_out.shape[0] == 1  # batch size
+        assert hidden_out.shape[1] == 1024  # seq_len (32 * 32)
+        assert hidden_out.shape[2] == 2560  # hidden_dim
+
+
+class TestGlmImagePipelineConfig:
+    """Test GLM_IMAGE_PIPELINE configuration."""
+
+    def test_glm_image_pipeline_config_exists(self):
+        """Verify GLM_IMAGE_PIPELINE config exists."""
+        assert GLM_IMAGE_PIPELINE is not None
+
+    def test_pipeline_has_correct_model_type(self):
+        """Verify pipeline has correct model_type."""
+        assert GLM_IMAGE_PIPELINE.model_type == "glm_image"
+
+    def test_pipeline_has_correct_model_arch(self):
+        """Verify pipeline has correct model_arch."""
+        assert GLM_IMAGE_PIPELINE.model_arch == "GlmImageForConditionalGeneration"
+
+    def test_pipeline_has_two_stages(self):
+        """Verify pipeline has two stages (AR + DiT)."""
+        assert len(GLM_IMAGE_PIPELINE.stages) == 2
+
+    def test_stage_0_is_ar_llm(self):
+        """Verify stage 0 is the AR LLM stage."""
+        stage_0 = GLM_IMAGE_PIPELINE.stages[0]
+        assert stage_0.stage_id == 0
+        assert stage_0.model_stage == "ar"
+        assert stage_0.execution_type.value == "llm_ar"
+        assert stage_0.owns_tokenizer is True
+
+    def test_stage_1_is_diffusion(self):
+        """Verify stage 1 is the DiT diffusion stage."""
+        stage_1 = GLM_IMAGE_PIPELINE.stages[1]
+        assert stage_1.stage_id == 1
+        assert stage_1.model_stage == "dit"
+        assert stage_1.execution_type.value == "diffusion"
+        assert stage_1.final_output is True
+        assert stage_1.final_output_type == "image"
+        assert stage_1.input_sources == (0,)  # Takes input from stage 0
+
+    def test_pipeline_has_diffusers_class_name(self):
+        """Verify pipeline has diffusers_class_name."""
+        assert GLM_IMAGE_PIPELINE.diffusers_class_name == "GlmImagePipeline"
+
+
+class TestGlmImageImageProjector:
+    """Test GlmImageImageProjector module."""
+
+    def test_projector_output_shape(self):
+        """Verify projector produces correct output shape."""
+        projector = GlmImageImageProjector(
+            in_channels=16,
+            hidden_size=2560,
+            patch_size=2,
+        )
+
+        # Input: [B, C, H, W] = [1, 16, 64, 64]
+        hidden_states = torch.randn(1, 16, 64, 64)
+        output = projector(hidden_states)
+
+        # After patchify: [B, H/2*W/2, D] = [1, 32*32, 2560] = [1, 1024, 2560]
+        assert output.shape == (1, 1024, 2560)
+
+    def test_projector_with_different_patch_sizes(self):
+        """Verify projector works with different patch sizes."""
+        for patch_size in [2, 4]:
+            projector = GlmImageImageProjector(
+                in_channels=16,
+                hidden_size=2560,
+                patch_size=patch_size,
+            )
+            h, w = 64, 64
+            hidden_states = torch.randn(1, 16, h, w)
+            output = projector(hidden_states)
+            expected_seq_len = (h // patch_size) * (w // patch_size)
+            assert output.shape[1] == expected_seq_len
+
+
+class TestGlmImageRotaryPosEmbed:
+    """Test GlmImageRotaryPosEmbed module."""
+
+    def test_rope_output_shape(self):
+        """Verify RoPE produces correct output shape."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(1, 16, 64, 64)  # [B, C, H, W]
+
+        cos, sin = rope(hidden_states)
+
+        # After patchify: height=32, width=32
+        # Output: [32*32, dim] = [1024, 40]
+        assert cos.shape[0] == 1024
+        assert cos.shape[1] == 40
+        assert sin.shape == cos.shape
+
+    def test_rope_value_range(self):
+        """Verify RoPE produces values in valid range."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(2, 16, 32, 32)
+
+        cos, sin = rope(hidden_states)
+
+        # cos and sin should be in [-1, 1]
+        assert cos.min() >= -1.0 and cos.max() <= 1.0
+        assert sin.min() >= -1.0 and sin.max() <= 1.0
+
+    def test_rope_consistency(self):
+        """Verify RoPE is consistent for same input."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(1, 16, 32, 32)
+
+        cos1, sin1 = rope(hidden_states)
+        cos2, sin2 = rope(hidden_states)
+
+        # Same input should produce same output
+        assert torch.allclose(cos1, cos2)
+        assert torch.allclose(sin1, sin2)
+
+
+class TestGlmImageTransformer2DModelQuantization:
+    """Test GlmImageTransformer2DModel with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the model accepts quant_config parameter."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        # Create a minimal mock od_config
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2  # Small number for testing
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=mock_quant_config,
+        )
+
+        # Check that quantization config was passed to transformer blocks
+        for block in model.transformer_blocks:
+            assert block.norm1.linear.quant_config is mock_quant_config
+            assert block.attn1.to_qkv.quant_config is mock_quant_config
+            assert block.ff.net[0].proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self, mocker: MockerFixture):
+        """Verify quant_config=None is accepted."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        # Check that quantization config is None
+        for block in model.transformer_blocks:
+            assert block.norm1.linear.quant_config is None
+            assert block.attn1.to_qkv.quant_config is None
+
+    def test_norm_out_has_no_quantization(self, mocker: MockerFixture):
+        """Verify norm_out (output layer) does NOT use quantization to preserve precision."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=mock_quant_config,
+        )
+
+        # norm_out.linear should NOT have quant_config to preserve output precision
+        assert model.norm_out.linear.quant_config is None
+
+    def test_model_has_sp_plan(self, mocker: MockerFixture):
+        """Verify model has _sp_plan defined for sequence parallelism."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        # Verify _sp_plan exists
+        assert hasattr(model, "_sp_plan")
+        assert "prepare" in model._sp_plan
+        assert "proj_out" in model._sp_plan
+
+    def test_model_has_hsdp_shard_conditions(self, mocker: MockerFixture):
+        """Verify model has _hsdp_shard_conditions for HSDP parallelism."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        assert hasattr(model, "_hsdp_shard_conditions")
+        assert len(model._hsdp_shard_conditions) > 0
+
+    def test_model_creates_kv_cache(self, mocker: MockerFixture):
+        """Verify model can create KV cache for image editing."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        kv_cache = model.create_kv_cache()
+        assert kv_cache is not None
+        assert len(kv_cache) == 2  # num_layers

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -23,6 +23,7 @@ NUM_HEADS = 4
 NUM_KV_HEADS = 2
 HEAD_DIM = 16
 IMAGE_TOKEN_LEN = 8
+SUFFIX_TOKEN_LEN = 3
 SCALING = 1.0 / math.sqrt(HEAD_DIM)
 
 
@@ -79,6 +80,10 @@ def _make_known_kv(num_tokens, base=0.0):
     return k, v
 
 
+def _gen_timestep_index(bs: int, current_start: int) -> torch.Tensor:
+    return torch.full((bs, 1), current_start, dtype=torch.long)
+
+
 def _call_mgr(
     mgr,
     bs,
@@ -90,6 +95,8 @@ def _call_mgr(
     uncond_cfg_prefill=False,
     num_image_tokens=IMAGE_TOKEN_LEN,
     shard_image_size=None,
+    gen_timestep_scatter_index=None,
+    position_ids=None,
 ):
     query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
     attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -104,6 +111,8 @@ def _call_mgr(
         uncond_cfg_prefill=uncond_cfg_prefill,
         num_image_tokens=num_image_tokens,
         shard_image_size=shard_image_size,
+        gen_timestep_scatter_index=gen_timestep_scatter_index,
+        position_ids=position_ids,
     )
 
 
@@ -117,54 +126,65 @@ def test_no_ar_kv(bs):
     """
     No AR KV injected. Tests the basic first_step cache → update reuse path.
 
-    Sequence layout per batch on first_step (q_len=12, IMAGE_TOKEN_LEN=8):
-        [prompt(3) | image(8) | eoi(1)]
-        image_size = IMAGE_TOKEN_LEN + 1 = 9
-        cached_prompt_len = seq_len - image_size = 12 - 9 = 3
+    Sequence layout per batch on first_step (q_len=14, IMAGE_TOKEN_LEN=8):
+        [prompt(3) | current timestep/image(8) | suffix(3)]
+        gen_timestep_scatter_index = 3
+        cached_prompt_len = 3
 
     After first_step:
-        image_kv_cache_map stores prompt tokens [0:3] for each batch (total 3*bs tokens flat).
+        image_kv_cache_map stores prompt tokens [0:3] for each batch.
 
-    Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + eoi(1) + image(8) = 12):
-        _reuse_prompt_kv produces [cached_prompt(3) | new_image(8) | zero_eoi(1)] per batch.
+    Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + current(8) = 11):
+        _reuse_prompt_kv produces [cached_prompt(3) | new_current(8)] per batch.
     """
     mgr = _make_cache_mgr()
     assert mgr.image_kv_cache_map is None
 
     # --- first_step ---
-    q_len = 12
+    prompt_len = 3
+    q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
 
-    _call_mgr(mgr, bs, q_len=q_len, seq_len=q_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=q_len,
+        seq_len=q_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
     # 3 prompt tokens cached per batch
-    assert cached_key.shape[0] == 3 * bs
+    assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
+    assert torch.equal(mgr.image_kv_cache_lens, torch.full((bs,), prompt_len))
     for b in range(bs):
         flat_offset = b * q_len
-        cache_offset = b * 3
-        assert torch.allclose(cached_key[cache_offset : cache_offset + 3], k_flat[flat_offset : flat_offset + 3])
-        assert torch.allclose(cached_value[cache_offset : cache_offset + 3], v_flat[flat_offset : flat_offset + 3])
+        assert torch.allclose(cached_key[b, :prompt_len], k_flat[flat_offset : flat_offset + prompt_len])
+        assert torch.allclose(cached_value[b, :prompt_len], v_flat[flat_offset : flat_offset + prompt_len])
 
     # --- update step ---
     img_q_len = IMAGE_TOKEN_LEN
-    update_seq_len = 3 + 1 + img_q_len  # cached_prompt(3) + eoi(1) + image(8) = 12
+    update_seq_len = prompt_len + img_q_len  # cached_prompt(3) + current(8) = 11
     new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
 
     key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
     val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
-    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
+    position_ids = torch.arange(prompt_len, prompt_len + img_q_len).repeat(bs, 1)
+    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs, position_ids=position_ids)
 
     assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
     for b in range(bs):
-        cache_offset = b * 3
         img_offset = b * img_q_len
         # Cached prompt preserved
-        assert torch.allclose(result_k[b, :3], cached_key[cache_offset : cache_offset + 3])
+        assert torch.allclose(result_k[b, :prompt_len], cached_key[b, :prompt_len])
         # New image tokens
-        assert torch.allclose(result_k[b, 3 : 3 + img_q_len], new_img_k[img_offset : img_offset + img_q_len])
-        # Zero eoi at the end
-        assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
+        assert torch.allclose(
+            result_k[b, prompt_len : prompt_len + img_q_len],
+            new_img_k[img_offset : img_offset + img_q_len],
+        )
 
 
 # ============================================================
@@ -178,16 +198,17 @@ def test_ar_kv_no_cfg(sp_size):
     AR KV injected, bs=1, no CFG. Tests AR prefix prepend + cache + reuse.
 
     sp_size=1:
-        first_step input: q_len=12, ar_len=5, seq_len=17
-        Sequence layout: [ar(5) | prompt(3) | image(8) | eoi(1)]
-        image_size = IMAGE_TOKEN_LEN + 1 = 9
-        cached_prompt_len = seq_len - image_size = 17 - 9 = 8 (= ar(5) + prompt(3))
+        first_step input: q_len=14, ar_len=5, seq_len=19
+        Sequence layout: [ar(5) | prompt(3) | current(8) | suffix(3)]
+        gen_timestep_scatter_index = 3
+        cached_prompt_len = ar(5) + prompt(3) = 8
 
-        Update step: q_len=8, seq_len = cached_prompt(8) + eoi(1) + image(8) = 17
-        Result: [cached(8) | new_image(8) | zero_eoi(1)]
+        Update step: q_len=8, seq_len = cached_prompt(8) + current(8) = 16
+        Result: [cached(8) | new_current(8)]
 
     sp_size=2:
         shard_image_size = 4 (passed externally, simulating SP sharding)
+        gen_timestep_scatter_index = 8
         cached_prompt_len = seq_len - shard_image_size = 17 - 4 = 13 (= ar(5) + prompt(8))
         Note: with SP, more of the sequence is treated as "prompt" because
         image tokens are sharded across ranks.
@@ -204,15 +225,23 @@ def test_ar_kv_no_cfg(sp_size):
     mgr._injected_ar_kv = [(ar_k.clone(), ar_v.clone())]
 
     # --- first_step: call _cache_prompt_kv directly ---
-    bs, q_len = 1, 12
-    seq_len = q_len + ar_len  # 17
+    bs = 1
+    q_len = 12 if sp_size > 1 else 14
+    seq_len = q_len + ar_len
     k_raw, v_raw = _make_known_kv(q_len, base=1.0)
     key_4d = k_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
     val_4d = v_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
 
     shard_image_size = 4 if sp_size > 1 else None
+    current_start = 8 if sp_size > 1 else 3
     mgr.image_kv_cache_map = None
-    mgr._cache_prompt_kv(key_4d, val_4d, seq_len, shard_image_size)
+    mgr._cache_prompt_kv(
+        key_4d,
+        val_4d,
+        seq_len,
+        shard_image_size,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
     if sp_size == 1:
@@ -222,25 +251,32 @@ def test_ar_kv_no_cfg(sp_size):
         # cached = seq_len - shard_image_size = 17 - 4 = 13
         expected_cached_len = seq_len - shard_image_size
 
-    assert cached_key.shape[0] == expected_cached_len
+    assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)
     # AR KV always at the front
-    assert torch.allclose(cached_key[:ar_len], ar_k)
-    assert torch.allclose(cached_value[:ar_len], ar_v)
+    assert torch.allclose(cached_key[0, :ar_len], ar_k)
+    assert torch.allclose(cached_value[0, :ar_len], ar_v)
     # Prompt tokens follow AR
     prompt_cached = expected_cached_len - ar_len
-    assert torch.allclose(cached_key[ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
+    assert torch.allclose(cached_key[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
     # AR KV consumed
     assert mgr._injected_ar_kv is None
 
     # --- update step: call _reuse_prompt_kv directly ---
     if sp_size == 1:
         img_q_len = IMAGE_TOKEN_LEN
-        update_seq_len = expected_cached_len + 1 + img_q_len  # 8 + 1 + 8 = 17
+        update_seq_len = expected_cached_len + img_q_len  # 8 + 8 = 16
         new_img_k, new_img_v = _make_known_kv(img_q_len, base=50.0)
 
         key_input = new_img_k.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
         val_input = new_img_v.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
-        result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=1)
+        position_ids = torch.arange(expected_cached_len, expected_cached_len + img_q_len).reshape(1, img_q_len)
+        result_k, result_v = mgr._reuse_prompt_kv(
+            key_input,
+            val_input,
+            update_seq_len,
+            bs=1,
+            position_ids=position_ids,
+        )
 
         assert result_k.shape == (1, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
         # AR + prompt preserved
@@ -248,8 +284,6 @@ def test_ar_kv_no_cfg(sp_size):
         assert torch.allclose(result_k[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
         # New image tokens
         assert torch.allclose(result_k[0, expected_cached_len : expected_cached_len + img_q_len], new_img_k)
-        # Zero eoi
-        assert torch.allclose(result_k[0, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
     else:
         # SP path: _reuse_prompt_kv returns only cached prompt (no image concat)
         img_q_len = shard_image_size  # 4
@@ -292,13 +326,12 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             → _injected_ar_kv becomes [(pos_ar(10), pos_av(10)), (neg_k(10), neg_v(10))]
 
         first_step (bs=2, q_len=12, seq_len=22):
-            Batch 0 (pos): [pos_ar(10) | prompt(3) | image(8) | eoi(1)]
-            Batch 1 (neg): [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
-            cached_prompt_len per batch = 22 - 9 = 13
-            Total cached = 13 * 2 = 26
+            Batch 0 (pos): [pos_ar(10) | prompt(3) | current(8) | suffix(3)]
+            Batch 1 (neg): [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
+            cached_prompt_len per batch = 10 + 3 = 13
 
-        Update (bs=2, seq_len = 13 + 1 + 8 = 22):
-            Result per batch: [cached(13) | new_image(8) | zero_eoi(1)]
+        Update (bs=2, seq_len = 13 + 8 = 21):
+            Result per batch: [cached(13) | new_current(8)]
 
     CFG Parallel (cfg_parallel=True, bs=1):
         This rank handles only the negative branch.
@@ -308,11 +341,11 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             → _injected_ar_kv = [(neg_k(10), neg_v(10))]
 
         first_step (bs=1, q_len=12, seq_len=22):
-            [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
-            cached_prompt_len = 22 - 9 = 13
+            [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
+            cached_prompt_len = 13
 
-        Update (bs=1, seq_len = 13 + 1 + 8 = 22):
-            Result: [cached(13) | new_image(8) | zero_eoi(1)]
+        Update (bs=1, seq_len = 13 + 8 = 21):
+            Result: [cached(13) | new_current(8)]
     """
     positive_reuse_len = 10
     negative_reuse_len = 6
@@ -351,46 +384,62 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
         mgr._injected_ar_kv = [mgr._injected_ar_kv[1]]
 
     # --- first_step ---
-    q_len = 12
+    prompt_len = 3
+    q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     seq_len = q_len + positive_reuse_len  # 22
     k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
 
-    _call_mgr(mgr, bs, q_len=q_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=q_len,
+        seq_len=seq_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
-    cached_prompt_len_per_batch = seq_len - (IMAGE_TOKEN_LEN + 1)  # 22 - 9 = 13
-    assert cached_key.shape[0] == cached_prompt_len_per_batch * bs
+    cached_prompt_len_per_batch = positive_reuse_len + prompt_len
+    assert cached_key.shape == (bs, cached_prompt_len_per_batch, NUM_KV_HEADS, HEAD_DIM)
     assert mgr._injected_ar_kv is None
 
     if not cfg_parallel:
         # bs=2: batch 0 = pos, batch 1 = neg
         # Batch 0: pos_ar(10) + prompt(3)
-        assert torch.allclose(cached_key[:positive_reuse_len], pos_ar_k)
-        assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
+        assert torch.allclose(cached_key[0, :positive_reuse_len], pos_ar_k)
+        assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
         # Batch 1: neg_ar(10) + prompt(3)
-        assert torch.allclose(cached_key[13:23], neg_ar_k)
-        assert torch.allclose(cached_key[23:26], k_flat[q_len : q_len + 3])
+        assert torch.allclose(cached_key[1, :positive_reuse_len], neg_ar_k)
+        assert torch.allclose(cached_key[1, positive_reuse_len:13], k_flat[q_len : q_len + 3])
     else:
         # bs=1: only neg branch
-        assert torch.allclose(cached_key[:positive_reuse_len], neg_ar_k)
-        assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
+        assert torch.allclose(cached_key[0, :positive_reuse_len], neg_ar_k)
+        assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
 
     # --- update step ---
     img_q_len = IMAGE_TOKEN_LEN
-    update_seq_len = cached_prompt_len_per_batch + 1 + img_q_len  # 13 + 1 + 8 = 22
+    update_seq_len = cached_prompt_len_per_batch + img_q_len  # 13 + 8 = 21
     new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
 
     key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
     val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
-    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
+    position_ids = torch.arange(cached_prompt_len_per_batch, cached_prompt_len_per_batch + img_q_len).repeat(bs, 1)
+    result_k, result_v = mgr._reuse_prompt_kv(
+        key_input,
+        val_input,
+        update_seq_len,
+        bs=bs,
+        position_ids=position_ids,
+    )
 
     assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
     for b in range(bs):
-        cache_offset = b * cached_prompt_len_per_batch
         # Cached prompt preserved
         assert torch.allclose(
             result_k[b, :cached_prompt_len_per_batch],
-            cached_key[cache_offset : cache_offset + cached_prompt_len_per_batch],
+            cached_key[b, :cached_prompt_len_per_batch],
         )
         # New image tokens
         img_offset = b * img_q_len
@@ -398,8 +447,6 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             result_k[b, cached_prompt_len_per_batch : cached_prompt_len_per_batch + img_q_len],
             new_img_k[img_offset : img_offset + img_q_len],
         )
-        # Zero eoi
-        assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
 
 
 # ============================================================
@@ -412,7 +459,7 @@ def test_cross_request_isolation():
     Verify leftover image_kv_cache_map from a previous request is NOT treated as AR KV.
 
     Setup: mgr has stale cache from a prior request (9 tokens, base=999).
-    New request: first_step with q_len=12, no AR KV.
+    New request: first_step with q_len=14, no AR KV.
 
     Expected: stale cache is overwritten. New cache = prompt tokens from current request.
     The stale values (999.x) must NOT appear in the new cache.
@@ -421,20 +468,33 @@ def test_cross_request_isolation():
 
     # Simulate leftover from previous request
     leftover_k, leftover_v = _make_known_kv(9, base=999.0)
-    mgr.image_kv_cache_map = (leftover_k, leftover_v)
+    mgr.image_kv_cache_map = (
+        leftover_k.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
+        leftover_v.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
+    )
     assert mgr._injected_ar_kv is None
 
     # New request first_step
-    bs, seq_len = 1, 12
+    bs = 1
+    prompt_len = 3
+    seq_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     k_flat, v_flat = _make_known_kv(bs * seq_len, base=1.0)
 
-    _call_mgr(mgr, bs, q_len=seq_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=seq_len,
+        seq_len=seq_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
-    # cached_prompt_len = 12 - 9 = 3
-    assert cached_key.shape[0] == 3
+    assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
     # Must be from current request, not stale
-    assert torch.allclose(cached_key[:3], k_flat[:3])
-    assert torch.allclose(cached_value[:3], v_flat[:3])
+    assert torch.allclose(cached_key[0, :prompt_len], k_flat[:prompt_len])
+    assert torch.allclose(cached_value[0, :prompt_len], v_flat[:prompt_len])
     # Stale values must not be present
     assert not torch.any(cached_key >= 999.0)

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -15,12 +15,29 @@ These tests verify:
 import json
 import os
 import tempfile
+from types import SimpleNamespace
 
 import pytest
+import torch
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _make_ltx23_pipeline(sequence_parallel_size: int = 1):
+    from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+    pipeline = object.__new__(LTX23Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae_temporal_compression_ratio = 4
+    pipeline.audio_vae_mel_compression_ratio = 4
+    pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
+    pipeline.audio_vae = SimpleNamespace(
+        latents_mean=torch.tensor(0.0),
+        latents_std=torch.tensor(1.0),
+    )
+    return pipeline
+
+
 class TestPipelineIndependence:
     """Verify LTX23Pipeline is fully independent from LTX2Pipeline."""
 
@@ -48,6 +65,313 @@ class TestPipelineIndependence:
 
         assert issubclass(LTX23Pipeline, ProgressBarMixin)
 
+    def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
+        """LTX23Pipeline must use the shared CFG parallel implementation."""
+        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        assert issubclass(LTX23Pipeline, CFGParallelMixin)
+
+
+class TestCFGParallelHelpers:
+    """Test LTX-2.3 CFG helper math without loading model weights."""
+
+    def test_combine_cfg_noise_matches_x0_space_formula(self):
+        import torch
+
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        pipe = object.__new__(LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        video_combined, audio_combined = pipe.combine_cfg_noise(
+            (video_pos, audio_pos),
+            (video_neg, audio_neg),
+            scale,
+            video_latents=video_sample,
+            audio_latents=audio_sample,
+            video_sigma=video_sigma,
+            audio_sigma=audio_sigma,
+        )
+
+        x0_video_cond = video_sample - video_pos * video_sigma
+        x0_video_uncond = video_sample - video_neg * video_sigma
+        x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond)
+        expected_video = (video_sample - x0_video_guided) / video_sigma
+
+        x0_audio_cond = audio_sample - audio_pos * audio_sigma
+        x0_audio_uncond = audio_sample - audio_neg * audio_sigma
+        x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond)
+        expected_audio = (audio_sample - x0_audio_guided) / audio_sigma
+        assert torch.allclose(video_combined, expected_video)
+        assert torch.allclose(audio_combined, expected_audio)
+
+    def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
+        audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if tensor.shape == video_pos.shape:
+                    return [video_pos, video_neg]
+                return [audio_pos, audio_neg]
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_sample,
+            video_pos,
+            video_neg,
+            video_sigma,
+            scale,
+        )
+        expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_sample,
+            audio_pos,
+            audio_neg,
+            audio_sigma,
+            scale,
+        )
+
+        for rank, expected_branch in ((0, "positive"), (1, "negative")):
+            calls = []
+            monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
+
+            def fake_predict_noise(**kwargs):
+                calls.append(kwargs["branch"])
+                if kwargs["branch"] == "positive":
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+            object.__setattr__(pipe, "predict_noise", fake_predict_noise)
+            video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
+                true_cfg_scale=scale,
+                positive_kwargs={"branch": "positive"},
+                negative_kwargs={"branch": "negative"},
+                cfg_normalize=False,
+                video_latents=video_sample,
+                audio_latents=audio_sample,
+                video_sigma=video_sigma,
+                audio_sigma=audio_sigma,
+            )
+
+            assert calls == [expected_branch]
+            torch.testing.assert_close(video_combined, expected_video)
+            torch.testing.assert_close(audio_combined, expected_audio)
+
+        assert "_cfg_video_latents" not in pipe.__dict__
+        assert "_cfg_audio_latents" not in pipe.__dict__
+
+
+class TestCFGParallelForwardPath:
+    """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
+
+    @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
+    def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
+        self,
+        monkeypatch,
+        cfg_rank,
+        expected_prompt_value,
+    ):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 1
+        pipe.vae_spatial_compression_ratio = 32
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.audio_sampling_rate = 1
+        pipe.audio_hop_length = 1
+        pipe.audio_vae_temporal_compression_ratio = 1
+        pipe.audio_vae_mel_compression_ratio = 1
+        pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+        pipe.tokenizer = SimpleNamespace(padding_side="left")
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.audio_vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(mel_bins=2, latent_channels=1),
+        )
+
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg):
+                    return [video_pos, video_neg]
+                if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg):
+                    return [audio_pos, audio_neg]
+                raise AssertionError(f"Unexpected gathered tensor: {tensor}")
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
+            scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
+            return torch.tensor([1.0, 0.5], device=device), 2
+
+        monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
+
+        class FakeScheduler:
+            def __init__(self, name="video", calls=None):
+                self.name = name
+                self.calls = [] if calls is None else calls
+                self.config = {
+                    "max_image_seq_len": 4096,
+                    "base_image_seq_len": 1024,
+                    "base_shift": 0.95,
+                    "max_shift": 2.05,
+                }
+                self.sigmas = torch.tensor([0.25, 0.25])
+
+            def __deepcopy__(self, memo):
+                return FakeScheduler("audio", self.calls)
+
+            def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+                self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone()))
+                return (latents - noise_pred,)
+
+        class FakeConnectors:
+            def to(self, device):
+                return self
+
+            def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
+                assert padding_side == "left"
+                assert prompt_embeds.shape[0] == 2
+                return prompt_embeds, prompt_embeds, prompt_attention_mask
+
+        class FakeRope:
+            def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
+                return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
+
+            def prepare_audio_coords(self, batch_size, num_frames, device):
+                return torch.zeros(batch_size, num_frames, 1, device=device)
+
+        class FakeTransformer:
+            def __init__(self):
+                self.config = SimpleNamespace(in_channels=2)
+                self.rope = FakeRope()
+                self.audio_rope = FakeRope()
+                self.calls = []
+
+            def __call__(self, **kwargs):
+                self.calls.append(kwargs)
+                expected_prompt = torch.full((1, 1, 1), expected_prompt_value)
+                torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt)
+                torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt)
+                assert kwargs["hidden_states"].shape == (1, 1, 2)
+                assert kwargs["audio_hidden_states"].shape == (1, 1, 2)
+                if cfg_rank == 0:
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+        class DummyProgress:
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                return False
+
+            def update(self):
+                pass
+
+        pipe.scheduler = FakeScheduler()
+        pipe.connectors = FakeConnectors()
+        pipe.transformer = FakeTransformer()
+        object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress())
+
+        def fake_encode_prompt(**kwargs):
+            return (
+                torch.ones(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+                torch.zeros(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+            )
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        video_latents = torch.tensor([[[1.0, -2.0]]])
+        audio_latents = torch.tensor([[[0.5, 3.0]]])
+        req = OmniDiffusionRequest(
+            prompts=[{"prompt": "prompt", "negative_prompt": "negative"}],
+            sampling_params=OmniDiffusionSamplingParams(
+                height=32,
+                width=32,
+                num_frames=1,
+                frame_rate=1.0,
+                num_inference_steps=2,
+                guidance_scale=4.0,
+                latents=video_latents,
+                audio_latents=audio_latents,
+                output_type="latent",
+            ),
+            request_id="ltx23-cfg-parallel-forward-test",
+        )
+
+        output = pipe.forward(req)
+
+        expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_latents,
+            video_pos,
+            video_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_latents,
+            audio_pos,
+            audio_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
+        assert scheduler_call_names == ["video", "audio", "video", "audio"]
+        assert len(pipe.transformer.calls) == 2
+        torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
+
+        video_out, audio_out = output.output
+        torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
+        torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
+
 
 class TestRegistryIntegration:
     """Verify all LTX-2.3 pipeline variants are registered."""
@@ -228,3 +552,85 @@ class TestInitExports:
         for name in expected_classes:
             assert hasattr(ltx2, name), f"{name} not exported from ltx2 package"
             assert name in ltx2.__all__, f"{name} not in ltx2.__all__"
+
+
+class TestAudioLatentSPPadding:
+    def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=2)
+
+        latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+            batch_size=1,
+            num_channels_latents=8,
+            num_mel_bins=64,
+            audio_latent_length=1,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+        )
+
+        assert original_num_frames == 1
+        assert padded_num_frames == 2
+        assert latents.shape == (1, 2, 128)
+
+    def test_prepare_audio_latents_pads_provided_packed_sequence_dim_for_sp(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+
+        padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+            batch_size=1,
+            num_channels_latents=2,
+            num_mel_bins=8,
+            audio_latent_length=10,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        assert original_num_frames == 10
+        assert padded_num_frames == 12
+        assert padded.shape == (1, 12, 4)
+        torch.testing.assert_close(padded[:, :10], latents)
+        torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
+
+    def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
+
+        audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
+        padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+            batch_size=1,
+            num_channels_latents=2,
+            num_mel_bins=16,
+            audio_latent_length=audio_latent_length,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        assert audio_latent_length == 10
+        assert original_num_frames == 10
+        assert padded_num_frames == 12
+        assert padded.shape == (1, 12, 8)
+        torch.testing.assert_close(padded, pipeline._pack_audio_latents(latents))
+
+    def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.zeros(1, 2, 13, 4)
+
+        audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
+
+        assert audio_latent_length == 13
+
+    def test_prepare_audio_latents_rejects_incompatible_provided_length(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.zeros(1, 11, 4)
+
+        with pytest.raises(ValueError, match="incompatible audio frame count"):
+            pipeline.prepare_audio_latents(
+                batch_size=1,
+                num_channels_latents=2,
+                num_mel_bins=8,
+                audio_latent_length=10,
+                dtype=torch.float32,
+                device=torch.device("cpu"),
+                latents=latents,
+            )

tests/diffusion/models/ovis_image/test_ovis_image.py

@@ -167,6 +167,7 @@ def test_basic_generation(ovis_pipeline):
     # Setup request
     req = OmniDiffusionRequest(
         prompts=["A photo of a cat"],
+        request_id="ovis-basic",
         sampling_params=OmniDiffusionSamplingParams(
             height=256,
             width=256,
@@ -206,6 +207,7 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch):
                 "negative_prompt": "bad quality",
             }
         ],
+        request_id="ovis-guidance",
         sampling_params=OmniDiffusionSamplingParams(
             height=256,
             width=256,
@@ -225,6 +227,7 @@ def test_resolution_check(ovis_pipeline):
     # Pass odd resolution
     req = OmniDiffusionRequest(
         prompts=["test"],
+        request_id="ovis-resolution",
         sampling_params=OmniDiffusionSamplingParams(
             height=250,  # Not divisible by 16 (8*2)
             width=250,

tests/diffusion/quantization/test_mxfp8_config.py

@@ -115,37 +115,64 @@ def test_mxfp8_quant_config_structure():
 # ---------------------------------------------------------------------------
 
 
-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
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native MXFP8 offline only supported on NPU")
+def test_get_quant_method_offline_npu(mocker: MockerFixture):
+    """Offline (serialized) path must return NPUMxfp8LinearMethod on NPU."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
 
     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)
+    assert type(method).__name__ == "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
+@pytest.mark.skipif(not current_omni_platform.is_xpu(), reason="XPU platform not available")
+def test_get_quant_method_offline_xpu_raises(mocker: MockerFixture):
+    """XPU offline mode must raise NotImplementedError (use AutoRound MXFP8 instead)."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
+    layer = mocker.Mock(spec=LinearBase)
+
+    with pytest.raises(NotImplementedError, match="Native MXFP8 offline mode is not supported on XPU"):
+        config.get_quant_method(layer, "blocks.0.attn1.to_q")
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_npu() or current_omni_platform.is_xpu()), reason="MXFP8 only supported on NPU and XPU"
+)
+def test_get_quant_method_online(mocker: MockerFixture):
+    """Online (BF16 checkpoint) path must return platform-specific method on current platform."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    # Mock the vLLM online method to avoid config dependency for XPU
+    if current_omni_platform.is_xpu():
+        mock_inner = mocker.Mock()
+        mocker.patch(
+            "vllm_omni.quantization.mxfp8_config.VllmMxfp8OnlineLinearMethod.__init__",
+            lambda self: setattr(self, "_inner", mock_inner),
+        )
 
     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)
+
+    if current_omni_platform.is_npu():
+        assert type(method).__name__ == "NPUMxfp8OnlineLinearMethod"
+    elif current_omni_platform.is_xpu():
+        assert type(method).__name__ == "VllmMxfp8OnlineLinearMethod"
 
 
-def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
-    """Non-NPU platform must raise NotImplementedError (CUDA not yet supported)."""
+def test_get_quant_method_unsupported_platform_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Unsupported platform (CUDA, ROCm) must raise NotImplementedError."""
     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_xpu", lambda: False)
     monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
 
     with pytest.raises(NotImplementedError):
@@ -159,6 +186,8 @@ def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pyte
     config = DiffusionMXFP8Config(ignored_layers=["proj_out"])
     layer = mocker.Mock(spec=LinearBase)
     monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+    monkeypatch.setattr(current_omni_platform, "is_xpu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
 
     method = config.get_quant_method(layer, "proj_out")
     assert isinstance(method, UnquantizedLinearMethod)

tests/diffusion/quantization/test_quantization_quality.py

@@ -71,7 +71,7 @@ class QualityTestConfig:
     prompt: str  # generation prompt
     max_lpips: float  # fail threshold — higher = more lenient
     model: str | None = None  # HF model name
-    quantization: str | None = None  # quantization method, e.g. "fp8"
+    quantization: str | dict[str, object] | None = None  # quantization method/config, e.g. "fp8"
     baseline_model: str | None = None  # explicit BF16/local baseline path
     quantized_model: str | None = None  # explicit quantized/local model path
     height: int = 1024
@@ -118,9 +118,57 @@ QUALITY_CONFIGS = [
     QualityTestConfig(
         id="fp8_z_image",
         model="Tongyi-MAI/Z-Image-Turbo",
-        quantization="fp8",
+        quantization={
+            "method": "fp8",
+            "ignored_layers": [
+                "img_mlp",
+                *[
+                    f"layers.{layer_id}.{suffix}"
+                    for layer_id in range(15, 30)
+                    for suffix in (
+                        "attention.to_qkv",
+                        "attention.to_out.0",
+                        "feed_forward.w13",
+                        "feed_forward.w2",
+                    )
+                ],
+                *[
+                    f"model.layers.{layer_id}.{suffix}"
+                    for layer_id in range(28, 36)
+                    for suffix in (
+                        "self_attn.q_proj",
+                        "self_attn.k_proj",
+                        "self_attn.v_proj",
+                        "self_attn.o_proj",
+                        "mlp.gate_proj",
+                        "mlp.up_proj",
+                        "mlp.down_proj",
+                    )
+                ],
+            ],
+        },
         task="t2i",
-        prompt="a cup of coffee on a wooden table, morning light",
+        prompt=(
+            "A breathtaking twilight scene atop a floating archipelago of crystalline islands suspended in an "
+            "endless nebula-drenched sky, where waterfalls of liquid starlight cascade from the edges of each "
+            "island into the cosmic abyss below. The largest island hosts an ancient, overgrown observatory "
+            "crafted from pearlescent white stone and living silverwood trees whose bioluminescent leaves pulse "
+            "with soft cerulean and violet light. At the observatory's center stands a solitary figure—a young "
+            "astronomer in flowing robes woven from woven moonlight and deep-space fabric, their face illuminated "
+            "by the glow of a holographic star chart hovering above an altar of polished obsidian. The sky swirls "
+            "with the birth of a new galaxy: ribbons of magenta and gold gas clouds twist around nascent stars, "
+            "while distant supernovae bloom like cosmic flowers in the far reaches of the void. Below, the abyss "
+            "reveals faint echoes of forgotten civilizations—ghostly silhouettes of submerged cities shimmering "
+            "in layers of atmospheric haze. Cinematic volumetric lighting cuts through the scene as twin moons—one "
+            "copper-hued, one opalescent—rise on opposite horizons, casting long, intersecting shadows across "
+            "moss-covered ruins and crystalline flora that refract light into prismatic halos. Hyper-detailed, "
+            "photorealistic rendering with the atmospheric depth of Roger Deakins' cinematography, the "
+            "architectural grandeur of Zaha Hadid, and the cosmic wonder of James Jean's illustrations. Shot on "
+            "a mythical 150mm lens with shallow depth of field, 8K resolution, Unreal Engine 5 realism, subsurface "
+            "scattering on organic elements, and ray-traced reflections dancing across every water droplet in the "
+            "starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation "
+            "at the edge of existence."
+        ),
         max_lpips=0.15,
         num_inference_steps=20,
     ),
@@ -391,12 +439,6 @@ _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUA
 
 def _quality_param(c: QualityTestConfig):
     marks = list(_marks)
-    if c.id == "fp8_z_image":
-        marks.append(
-            pytest.mark.skip(
-                reason="Z-Image FP8 quality gate temporarily disabled: https://github.com/vllm-project/vllm-omni/issues/3531"
-            )
-        )
     if c.id == "fp8_qwen_image":
         marks.append(
             pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")

tests/diffusion/test_diffusion_engine.py

@@ -22,15 +22,15 @@ class DiffusionSchedulerOutput:
     num_waiting_reqs: int = 0
 
     @property
-    def scheduled_req_ids(self):
+    def scheduled_request_ids(self):
         ids = [req.request_id for req in self.scheduled_new_reqs]
         if self.scheduled_cached_reqs:
-            ids.extend(self.scheduled_cached_reqs.sched_req_ids)
+            ids.extend(self.scheduled_cached_reqs.request_ids)
         return ids
 
     @property
     def is_empty(self):
-        return len(self.scheduled_req_ids) == 0
+        return len(self.scheduled_request_ids) == 0
 
 
 class MockScheduler:
@@ -71,31 +71,36 @@ async def test_async_add_req_and_wait_for_response():
     engine.scheduler = MockScheduler()
     engine._out_queue = {}
     engine.abort_queue: queue.Queue[str] = queue.Queue()
-    engine._rpc_lock = threading.Lock()
+    engine._rpc_queue = queue.Queue()
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine._init_lock = asyncio.Lock()
+    engine._closed = False
+    engine._loop_started = False
+    engine.main_loop = None
 
     engine._finalize_finished_request = lambda rid, out, err: out.result
 
     def mock_execute_batch(sched_output):
-        req_ids = sched_output.scheduled_req_ids
+        request_ids = sched_output.scheduled_request_ids
 
         time.sleep(1)
 
         class MockRunnerOutput:
             def __init__(self, ids):
-                self.req_id = ids
+                self.request_id = ids
                 self.step_index = 0
                 self.finished = True
                 self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
 
-            def get_req_output(self, rid):
+            def get_request_output(self, rid):
                 return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
 
-        return MockRunnerOutput(req_ids)
+        return MockRunnerOutput(request_ids)
 
     engine.execute_fn = mock_execute_batch
 
-    engine.stop_event = threading.Event()
-    engine.start_background_loop()
+    await engine._check_and_start_background_loop()
 
     async def run_task(rid):
         req = SimpleNamespace(request_id=rid)
@@ -105,10 +110,13 @@ async def test_async_add_req_and_wait_for_response():
 
     task_ids = [f"req_{i}" for i in range(5)]
     tasks = [run_task(rid) for rid in task_ids]
-    results = await asyncio.gather(*tasks)
-
-    engine.stop_event.set()
-    engine.worker_thread.join(timeout=5)
+    try:
+        results = await asyncio.gather(*tasks)
+    finally:
+        with engine._cv:
+            engine.stop_event.set()
+            engine._cv.notify_all()
+        engine.worker_thread.join(timeout=5)
     assert len(results) == 5
     for rid, res, elapsed in results:
         assert rid in res.result_data

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -17,15 +17,10 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-def _make_request(
-    request_id: str,
-    *,
-    request_ids: list[str] | None = None,
-) -> OmniDiffusionRequest:
+def _make_request(request_id: str) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[f"prompt_{request_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-        request_ids=request_ids or [request_id],
         request_id=request_id,
     )
 
@@ -79,15 +74,15 @@ def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
     engine._finalize_finished_request.assert_not_called()
 
 
-def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None:
+def test_abort_request_id_aborts_scheduler_request() -> None:
     engine = _make_engine()
-    request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"])
-    sched_req_id = engine.scheduler.add_request(request)
+    request = _make_request("batch-parent")
+    request_id = engine.scheduler.add_request(request)
 
     engine.abort("batch-parent")
     engine._process_aborts_queue()
 
-    state = engine.scheduler.get_request_state(sched_req_id)
+    state = engine.scheduler.get_request_state(request_id)
     assert state is not None
     assert state.status == DiffusionRequestStatus.FINISHED_ABORTED
 

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -89,7 +89,7 @@ class _ConcurrencyTrackingExecutor:
             if method in {"execute_model", "execute_stepwise", "generate"}:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
-                tag = req.request_ids[0] if req is not None and hasattr(req, "request_ids") else "unknown"
+                tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
                 return DiffusionOutput(error=f"result_for_{tag}")
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
@@ -108,7 +108,7 @@ class _ConcurrencyTrackingExecutor:
             exec_all_ranks=True,
         )
         return RunnerOutput(
-            req_id=new_req.sched_req_id,
+            request_id=new_req.request_id,
             step_index=None,
             finished=True,
             result=result,
@@ -120,7 +120,7 @@ class _ConcurrencyTrackingExecutor:
 
 def _make_request(tag: str):
     return SimpleNamespace(
-        request_ids=[tag],
+        request_id=tag,
         prompts=[f"prompt_{tag}"],
         sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
     )
@@ -148,6 +148,7 @@ def _make_engine_with_loop(
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
     engine._out_queue = {}
+    engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
 

tests/diffusion/test_diffusion_request.py

@@ -5,7 +5,7 @@ import random
 
 import pytest
 
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -15,9 +15,53 @@ def _make_request() -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[{"prompt": "a cup of coffee on a table"}],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id="request-test",
     )
 
 
+def test_request_id_is_required():
+    with pytest.raises(TypeError, match="request_id"):
+        OmniDiffusionRequest(
+            prompts=[{"prompt": "a cup of coffee on a table"}],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        )
+
+
+def test_request_ids_identity_list_is_removed():
+    req = _make_request()
+
+    assert req.request_id == "request-test"
+    assert not hasattr(req, "request_ids")
+
+
+def test_request_id_must_be_non_empty():
+    with pytest.raises(ValueError, match="request_id must be a non-empty string"):
+        OmniDiffusionRequest(
+            prompts=[{"prompt": "a cup of coffee on a table"}],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_id="",
+        )
+
+
+def test_dummy_run_request_is_identified_by_reserved_request_id():
+    req = OmniDiffusionRequest(
+        prompts=[{"prompt": "dummy run"}],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id=DUMMY_DIFFUSION_REQUEST_ID,
+    )
+
+    assert req.is_dummy_run()
+    assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
+
+
+def test_non_dummy_request_is_not_identified_as_dummy_run():
+    req = _make_request()
+
+    assert not req.is_dummy_run()
+    assert not OmniDiffusionRequest.is_dummy_run_request_id(req.request_id)
+    assert not OmniDiffusionRequest.is_dummy_run_request_id(None)
+
+
 def test_tp_seed_same_across_ranks_and_varies_across_requests():
     random.seed(0)
     n_requests = 5

tests/diffusion/test_diffusion_scheduler.py

@@ -31,13 +31,13 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[f"prompt_{req_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
 def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
     return RunnerOutput(
-        req_id=req_id,
+        request_id=req_id,
         step_index=None,
         finished=finished,
         result=DiffusionOutput(output=None, error=error),
@@ -52,7 +52,7 @@ def _make_step_output(
     error: str | None = None,
 ):
     return RunnerOutput(
-        req_id=req_id,
+        request_id=req_id,
         step_index=step_index,
         finished=finished,
         result=DiffusionOutput(output=None, error=error) if error is not None else None,
@@ -73,16 +73,16 @@ def _make_step_request(
             num_inference_steps=num_inference_steps,
             step_index=step_index,
         ),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
 def _new_ids(sched_output) -> list[str]:
-    return [req.sched_req_id for req in sched_output.scheduled_new_reqs]
+    return [req.request_id for req in sched_output.scheduled_new_reqs]
 
 
 def _cached_ids(sched_output) -> list[str]:
-    return list(sched_output.scheduled_cached_reqs.sched_req_ids)
+    return list(sched_output.scheduled_cached_reqs.request_ids)
 
 
 class _StubScheduler(SchedulerInterface):
@@ -90,7 +90,7 @@ class _StubScheduler(SchedulerInterface):
         self._request = request
         self._output = output
         self.initialized_with = None
-        self._sched_req_id = request.request_ids[0]
+        self._request_id = request.request_id
         self._state = None
         self._scheduled = False
         self.max_num_running_reqs = 1
@@ -100,22 +100,22 @@ class _StubScheduler(SchedulerInterface):
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         assert request is self._request
-        self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request)
-        return self._sched_req_id
+        self._state = SimpleNamespace(request_id=self._request_id, req=request)
+        return self._request_id
 
     def schedule(self):
         if self._scheduled or self._state is None:
             return SimpleNamespace(
                 scheduled_new_reqs=[],
                 scheduled_cached_reqs=CachedRequestData.make_empty(),
-                scheduled_req_ids=[],
+                scheduled_request_ids=[],
                 is_empty=True,
             )
         self._scheduled = True
         return SimpleNamespace(
             scheduled_new_reqs=[NewRequestData.from_state(self._state)],
             scheduled_cached_reqs=CachedRequestData.make_empty(),
-            scheduled_req_ids=[self._state.sched_req_id],
+            scheduled_request_ids=[self._state.request_id],
             is_empty=False,
         )
 
@@ -123,30 +123,25 @@ class _StubScheduler(SchedulerInterface):
         del sched_output
         assert output is self._output
         self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
-        return {self._sched_req_id}
+        return {self._request_id}
 
     def has_requests(self) -> bool:
         return not self._scheduled
 
-    def get_request_state(self, sched_req_id: str):
-        del sched_req_id
+    def get_request_state(self, request_id: str):
+        del request_id
         return self._state
 
-    def get_sched_req_id(self, request_id: str) -> str | None:
-        if request_id in self._request.request_ids:
-            return self._sched_req_id
-        return None
-
-    def pop_request_state(self, sched_req_id: str):
-        del sched_req_id
+    def pop_request_state(self, request_id: str):
+        del request_id
         return self._state
 
-    def preempt_request(self, sched_req_id: str) -> bool:
-        del sched_req_id
+    def preempt_request(self, request_id: str) -> bool:
+        del request_id
         return False
 
-    def finish_requests(self, sched_req_ids, status) -> None:
-        del sched_req_ids, status
+    def finish_requests(self, request_ids, status) -> None:
+        del request_ids, status
         return None
 
     def close(self) -> None:
@@ -171,7 +166,7 @@ class TestGetSamplingParamsKey:
         return OmniDiffusionRequest(
             prompts=["prompt"],
             sampling_params=sp,
-            request_ids=[f"req-{lora_int_id}-{lora_scale}"],
+            request_id=f"req-{lora_int_id}-{lora_scale}",
         )
 
     def test_distinguishes_lora_id(self) -> None:
@@ -291,7 +286,7 @@ class TestRequestScheduler:
             OmniDiffusionRequest(
                 prompts=["prompt_b"],
                 sampling_params=OmniDiffusionSamplingParams(width=768),
-                request_ids=["b"],
+                request_id="b",
             )
         )
         scheduler.add_request(_make_request("c"))
@@ -333,7 +328,7 @@ class TestRequestScheduler:
         assert state_a.status == DiffusionRequestStatus.FINISHED_ABORTED
 
         assert self.scheduler.has_requests() is False
-        assert self.scheduler.schedule().scheduled_req_ids == []
+        assert self.scheduler.schedule().scheduled_request_ids == []
 
     def test_has_requests_state_transition(self) -> None:
         assert self.scheduler.has_requests() is False
@@ -348,60 +343,28 @@ class TestRequestScheduler:
         assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
         assert self.scheduler.has_requests() is False
 
-    def test_request_id_mapping_lifecycle(self) -> None:
+    def test_request_id_is_scheduler_key(self) -> None:
         request = OmniDiffusionRequest(
             prompts=["prompt_map_a", "prompt_map_b"],
             sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["map-a", "map-b"],
             request_id="map-parent",
         )
 
-        sched_req_id = self.scheduler.add_request(request)
+        request_id = self.scheduler.add_request(request)
 
-        assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
-        assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
-        assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id
+        assert request_id == "map-parent"
+        state = self.scheduler.get_request_state("map-parent")
+        assert state.request_id == "map-parent"
 
-        self.scheduler.pop_request_state(sched_req_id)
+        self.scheduler.pop_request_state("map-parent")
 
-        assert self.scheduler.get_sched_req_id("map-a") is None
-        assert self.scheduler.get_sched_req_id("map-b") is None
-        assert self.scheduler.get_sched_req_id("map-parent") is None
+        assert self.scheduler.get_request_state("map-parent") is None
 
-    def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None:
-        self.scheduler.add_request(
-            OmniDiffusionRequest(
-                prompts=["prompt_existing"],
-                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-                request_ids=["existing-child"],
-                request_id="duplicate-parent",
-            )
-        )
+    def test_duplicate_request_id_is_rejected(self) -> None:
+        self.scheduler.add_request(_make_request("dup"))
 
-        colliding_request = OmniDiffusionRequest(
-            prompts=["prompt_unique"],
-            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["unique-child"],
-            request_id="duplicate-parent",
-        )
-
-        with pytest.raises(ValueError, match="duplicate-parent"):
-            self.scheduler.add_request(colliding_request)
-
-        assert self.scheduler.get_request_state("unique-child") is None
-        assert self.scheduler.get_sched_req_id("unique-child") is None
-
-        valid_request = OmniDiffusionRequest(
-            prompts=["prompt_unique"],
-            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["unique-child"],
-            request_id="unique-parent",
-        )
-        sched_req_id = self.scheduler.add_request(valid_request)
-
-        assert sched_req_id == "unique-child"
-        assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id
-        assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id
+        with pytest.raises(ValueError, match="request_id 'dup' is already active"):
+            self.scheduler.add_request(_make_request("dup"))
 
 
 class TestDiffusionEngine:
@@ -662,7 +625,7 @@ class TestStepScheduler:
         finished = self.scheduler.update_from_output(
             sched_output,
             RunnerOutput(
-                req_id=req_id,
+                request_id=req_id,
                 step_index=None,
                 finished=True,
                 result=None,
@@ -934,7 +897,7 @@ class TestStepScheduler:
 
         for _ in range(expected_steps - 1):
             sched_output = self.scheduler.schedule()
-            assert sched_output.scheduled_req_ids == [req_id]
+            assert sched_output.scheduled_request_ids == [req_id]
             next_step = request.sampling_params.step_index + 1
             assert (
                 self.scheduler.update_from_output(
@@ -945,7 +908,7 @@ class TestStepScheduler:
             )
 
         final_output = self.scheduler.schedule()
-        assert final_output.scheduled_req_ids == [req_id]
+        assert final_output.scheduled_request_ids == [req_id]
         assert self.scheduler.update_from_output(
             final_output,
             _make_step_output(req_id, step_index=expected_steps, finished=True),

tests/diffusion/test_diffusion_step_pipeline.py

@@ -206,7 +206,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
 def _make_step_request(num_inference_steps: int = 2):
     return SimpleNamespace(
         prompts=["a prompt"],
-        request_ids=["req-1"],
+        request_id="req-1",
         sampling_params=SimpleNamespace(
             generator=None,
             seed=None,
@@ -227,7 +227,7 @@ def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) ->
     return OmniDiffusionRequest(
         prompts=[f"prompt-{req_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
@@ -274,10 +274,10 @@ def _make_distributed_runner(mode: str, device: torch.device):
     return runner
 
 
-def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
+def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
-        scheduled_new_reqs=[NewRequestData(sched_req_id=sched_req_id, req=req)],
+        scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)],
         scheduled_cached_reqs=CachedRequestData.make_empty(),
         finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
         num_running_reqs=1,
@@ -287,7 +287,7 @@ 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]
+    new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs]
     return DiffusionSchedulerOutput(
         step_id=step_id,
         scheduled_new_reqs=new_reqs,
@@ -298,11 +298,11 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
     )
 
 
-def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
+def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
         scheduled_new_reqs=[],
-        scheduled_cached_reqs=CachedRequestData(sched_req_ids=[sched_req_id]),
+        scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
         finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
         num_running_reqs=1,
         num_waiting_reqs=0,
@@ -359,7 +359,7 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
             runner,
             _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
         )
-        output = result.get_req_output("req-1")
+        output = result.get_request_output("req-1")
 
         assert output.finished is True
         assert output.result is not None
@@ -384,16 +384,16 @@ class TestRunner:
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-        first = result.get_req_output("req-1")
-        assert first.req_id == "req-1"
+        first = result.get_request_output("req-1")
+        assert first.request_id == "req-1"
         assert first.step_index == 1
         assert first.finished is False
         assert first.result is None
         assert "req-1" in runner.state_cache
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
-        second = result.get_req_output("req-1")
-        assert second.req_id == "req-1"
+        second = result.get_request_output("req-1")
+        assert second.request_id == "req-1"
         assert second.step_index == 2
         assert second.finished is True
         assert second.result is not None
@@ -410,13 +410,13 @@ class TestRunner:
         runner = _make_runner()
         req_1 = _make_step_request()
         req_2 = _make_step_request()
-        req_2.request_ids = ["req-2"]
+        req_2.request_id = "req-2"
 
         scheduler_output = DiffusionSchedulerOutput(
             step_id=0,
             scheduled_new_reqs=[
-                NewRequestData(sched_req_id="req-1", req=req_1),
-                NewRequestData(sched_req_id="req-2", req=req_2),
+                NewRequestData(request_id="req-1", req=req_1),
+                NewRequestData(request_id="req-2", req=req_2),
             ],
             scheduled_cached_reqs=CachedRequestData.make_empty(),
             finished_req_ids=set(),
@@ -431,7 +431,7 @@ class TestRunner:
         runner = _make_runner()
 
         with pytest.raises(ValueError, match="Missing cached state"):
-            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
+            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing"))
 
     def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
         runner = _make_runner()
@@ -440,8 +440,8 @@ class TestRunner:
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-        output = result.get_req_output("req-1")
-        assert output.req_id == "req-1"
+        output = result.get_request_output("req-1")
+        assert output.request_id == "req-1"
         assert output.step_index == 0
         assert output.finished is True
         assert output.result is not None
@@ -515,7 +515,7 @@ def _make_step_worker(lora_manager=None, *, expected_output=None):
     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")
+    output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1")
     worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
     return worker
 
@@ -525,9 +525,9 @@ class TestWorker:
     """DiffusionWorker.execute_stepwise"""
 
     def test_delegates_to_model_runner(self):
-        expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
+        expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None)
         worker = _make_step_worker(expected_output=expected)
-        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
 
         output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -536,7 +536,7 @@ class TestWorker:
     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")
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -552,7 +552,7 @@ class TestWorker:
 
         manager = _RecordingLoRAManager()
         worker = _make_step_worker(lora_manager=manager)
-        scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(request, request_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -568,8 +568,8 @@ class TestWorker:
 
         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)
+        first = _make_scheduler_output(request, request_id="req-1")
+        second = _make_cached_scheduler_output(request_id="req-1", step_id=1)
 
         DiffusionWorker.execute_stepwise(worker, first)
         DiffusionWorker.execute_stepwise(worker, second)
@@ -610,10 +610,10 @@ class TestWorker:
         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")
+        first = _make_scheduler_output(finishing, request_id="req-1")
         next_batch = _make_scheduler_output(
             next_request,
-            sched_req_id="req-2",
+            request_id="req-2",
             step_id=1,
             finished_req_ids={"req-1"},
         )
@@ -633,11 +633,11 @@ class TestExecutor:
     def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
         executor = object.__new__(MultiprocDiffusionExecutor)
         executor._ensure_open = lambda: None
-        expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
+        expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
         executor.collective_rpc = mocker.Mock(return_value=expected)
 
         request = _make_engine_request("req-step", num_inference_steps=2)
-        scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
+        scheduler_output = _make_scheduler_output(request, request_id="req-step")
 
         output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
 
@@ -653,7 +653,7 @@ class TestEngine:
         [
             (
                 lambda _: RunnerOutput(
-                    req_id="req-error",
+                    request_id="req-error",
                     step_index=1,
                     finished=True,
                     result=DiffusionOutput(error="boom"),
@@ -688,7 +688,7 @@ class TestEngine:
             call_count["n"] += 1
             finished = call_count["n"] == 2
             return RunnerOutput(
-                req_id="req-step",
+                request_id="req-step",
                 step_index=call_count["n"],
                 finished=finished,
                 result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
@@ -714,7 +714,7 @@ class TestEngine:
             step["n"] += 1
             engine.abort("req-stop")
             return RunnerOutput(
-                req_id="req-stop",
+                request_id="req-stop",
                 step_index=1,
                 finished=False,
                 result=None,
@@ -741,7 +741,7 @@ class TestEngine:
                 assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
                 engine.abort("req-mid")
             return RunnerOutput(
-                req_id="req-mid",
+                request_id="req-mid",
                 step_index=step["n"],
                 finished=False,
                 result=None,
@@ -760,7 +760,7 @@ class TestEngine:
         engine = _make_engine(
             scheduler,
             execute_fn=lambda _: RunnerOutput(
-                req_id="req-missing",
+                request_id="req-missing",
                 step_index=1,
                 finished=True,
                 result=None,
@@ -777,7 +777,7 @@ class TestEngine:
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
         tensor = torch.zeros(300_000, dtype=torch.float32)
-        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+        output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
 
         packed = pack_diffusion_output_shm(output)
         assert isinstance(packed.result.output, dict)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -36,7 +36,7 @@ def _tagged_output(tag: str) -> DiffusionOutput:
 def _mock_request(tag: str):
     """Return a lightweight request object identifiable by *tag*."""
     return SimpleNamespace(
-        request_ids=[tag],
+        request_id=tag,
         prompts=[f"prompt_{tag}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
     )
@@ -97,8 +97,8 @@ def _start_worker(req_q, res_q, count=2):
             req = req_q.get(timeout=10)
             method = req.get("method", "")
             args = req.get("args", ())
-            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
-                tag = f"result_for_{args[0].request_ids[0]}"
+            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
+                tag = f"result_for_{args[0].request_id}"
             elif args:
                 tag = f"result_for_{args[0]}"
             else:

tests/diffusion/test_stage_diffusion_proc.py

@@ -68,7 +68,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
 
         request = captured["request"]
         assert request.request_id == "req-parent"
-        assert request.request_ids == ["req-parent-0", "req-parent-1"]
+        assert not hasattr(request, "request_ids")
         assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
         assert result.request_id == "req-parent"
         assert result.images == ["img-1", "img-2"]

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -33,6 +33,20 @@ def test_tensor_serialization():
     assert torch.equal(tensor, deserialized)
 
 
+def test_noncontiguous_tensor_serialization():
+    """Test torch.Tensor serialization for strided views."""
+    import torch
+
+    tensor = torch.arange(40, dtype=torch.long).reshape(4, 10).transpose(0, 1)
+    assert not tensor.is_contiguous()
+    assert tensor.stride(-1) != 1
+
+    serialized = OmniSerializer.serialize({"codes": tensor})
+    deserialized = OmniSerializer.deserialize(serialized)
+
+    assert torch.equal(tensor, deserialized["codes"])
+
+
 def test_ndarray_serialization():
     """Test numpy.ndarray serialization."""
     import numpy as np

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -37,6 +37,7 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
         status=status,
         prompt_token_ids=[],
         num_computed_tokens=0,
+        num_output_placeholders=0,
         additional_information=None,
         is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
     )
@@ -44,9 +45,16 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
 
 @pytest.fixture
 def build_adapter(monkeypatch, mocker: MockerFixture):
-    def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2):
+    def _build(
+        *,
+        stage_id: int = 1,
+        model_mode: str = "ar",
+        max_num_seqs: int = 2,
+        connector_extra: dict | None = None,
+    ):
         connector = mocker.MagicMock()
         connector.stage_id = stage_id
+        connector.config = {"extra": connector_extra or {}}
         connector.get.return_value = None
         connector.put.return_value = (True, 1, {})
 
@@ -127,6 +135,61 @@ def test_load_poll(build_adapter):
     assert "req-1" not in adapter._pending_load_reqs
 
 
+def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter):
+    adapter, connector = build_adapter(stage_id=1, model_mode="generation")
+    request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor")
+
+    codes = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)
+    payload: OmniPayload = {
+        "codes": {"audio": codes},
+        "meta": {
+            "left_context_size": 1,
+            "finished": torch.tensor(False, dtype=torch.bool),
+        },
+    }
+    connector.get.return_value = (payload, 16)
+
+    adapter._poll_single_request(request)
+
+    assert request.prompt_token_ids == [0]
+    assert request.num_computed_tokens == 0
+    assert torch.equal(request.additional_information["codes"]["audio"], codes)
+    assert request.additional_information["meta"]["left_context_size"] == 1
+    assert "finished" not in request.additional_information["meta"]
+    assert "req-tensor" in adapter._finished_load_reqs
+
+
+def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapter):
+    adapter, connector = build_adapter(stage_id=1, model_mode="generation")
+    request = _req("req-empty-tensor", RequestStatus.WAITING, external_req_id="external-empty")
+
+    empty_payload: OmniPayload = {
+        "codes": {"audio": torch.empty((4, 0), dtype=torch.long)},
+        "meta": {
+            "left_context_size": 0,
+            "finished": torch.tensor(False, dtype=torch.bool),
+        },
+    }
+    ready_payload: OmniPayload = {
+        "codes": {"audio": torch.tensor([[1, 2]], dtype=torch.long)},
+        "meta": {
+            "left_context_size": 0,
+            "finished": torch.tensor(False, dtype=torch.bool),
+        },
+    }
+    connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)]
+
+    assert adapter._poll_single_request(request) is False
+    assert request.request_id not in adapter._finished_load_reqs
+    assert request.request_id not in adapter.requests_with_ready_chunks
+    assert adapter.get_req_chunk[request.request_id] == 1
+
+    assert adapter._poll_single_request(request) is True
+    assert request.request_id in adapter._finished_load_reqs
+    assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"])
+    assert adapter.get_req_chunk[request.request_id] == 2
+
+
 def test_save_async(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
@@ -140,6 +203,18 @@ def test_save_async(build_adapter):
     assert task["is_finished"] is False
 
 
+def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-async", RequestStatus.WAITING, external_req_id="external-async")
+    request.num_computed_tokens = 10
+    request.num_output_placeholders = 2
+
+    adapter.save_async(pooling_output=None, request=request)
+
+    assert adapter.requests_num_chunks_sent["external-async"] == 8
+    assert len(adapter._pending_save_reqs) == 1
+
+
 def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
     """Producer may return a struct with ``meta=None`` (e.g. payload that
     carries only ``embed`` or ``codes``). The sender's ``meta is not None``
@@ -458,6 +533,52 @@ def test_finish_requests_restores_status(build_adapter):
     assert req_id not in adapter.requests_origin_status
 
 
+def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
+    other = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
+    adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
+    adapter.waiting_for_chunk_running_requests = deque([other, zombie])
+    adapter.requests_with_ready_chunks.add("req-zombie")
+    adapter.finished_requests.add("req-zombie")
+    requests_map = {
+        "req-zombie": zombie,
+        "req-live": other,
+    }
+
+    adapter.finish_requests(
+        ["req-zombie"],
+        RequestStatus.FINISHED_ABORTED,
+        requests_map,
+    )
+
+    assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
+    assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
+    assert "req-zombie" not in adapter.requests_with_ready_chunks
+    assert "req-zombie" not in adapter.finished_requests
+
+
+def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
+    live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
+    waiting_queue = DummyWaitingQueue()
+    running_queue = []
+    adapter.waiting_for_chunk_waiting_requests = deque([zombie, live])
+    adapter.waiting_for_chunk_running_requests = deque([zombie, live])
+
+    adapter.restore_queues(
+        waiting_queue,
+        running_queue,
+        scheduler_requests={"req-live": live},
+    )
+
+    assert [req.request_id for req in waiting_queue] == ["req-live"]
+    assert [req.request_id for req in running_queue] == ["req-live"]
+    assert not adapter.waiting_for_chunk_waiting_requests
+    assert not adapter.waiting_for_chunk_running_requests
+
+
 # ---------------------------------------------------------------
 # Scheduler trigger tests
 # ---------------------------------------------------------------

tests/distributed/omni_connectors/test_kv_flow.py

@@ -384,7 +384,7 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
     req = OmniDiffusionRequest(
         prompts=["test_recv"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
+        request_id=req_id,
     )
     # req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
     # But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
@@ -428,7 +428,6 @@ def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_conf
     req = OmniDiffusionRequest(
         prompts=["prompt-a", "prompt-b"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
         request_id=parent_req_id,
     )
 
@@ -454,7 +453,6 @@ def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_con
     req = OmniDiffusionRequest(
         prompts=["prompt-a", "prompt-b"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=["req-parent-0", "req-parent-1"],
         request_id="req-parent",
     )
     req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
@@ -516,7 +514,7 @@ def test_integration_flow(common_constants):
     req = OmniDiffusionRequest(
         prompts=["test_integ"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
     # Receive

tests/e2e/accuracy/conftest.py

@@ -62,6 +62,19 @@ def pytest_addoption(parser):
         default=1200,
         help="Online serving timeout in seconds for Wan2.2 I2V accuracy tests.",
     )
+    group.addoption(
+        "--hunyuanvideo15-i2v-image-source",
+        action="store",
+        default=None,
+        help="Image source for HunyuanVideo-1.5 I2V accuracy tests. Can be local path or remote URL",
+    )
+    group.addoption(
+        "--hunyuanvideo15-online-timeout-seconds",
+        action="store",
+        type=int,
+        default=3600,
+        help="Online serving timeout in seconds for HunyuanVideo-1.5 accuracy tests.",
+    )
 
 
 def _hf_cache_root() -> Path:
@@ -168,6 +181,17 @@ def wan22_i2v_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
     return int(request.config.getoption("wan22_i2v_online_timeout_seconds"))
 
 
+@pytest.fixture(scope="session")
+def hunyuanvideo15_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
+    value = request.config.getoption("hunyuanvideo15_i2v_image_source")
+    return str(value) if value else None
+
+
+@pytest.fixture(scope="session")
+def hunyuanvideo15_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
+    return int(request.config.getoption("hunyuanvideo15_online_timeout_seconds"))
+
+
 @pytest.fixture(scope="session")
 def gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
     return int(request.config.getoption("gebench_samples_per_type"))

tests/e2e/accuracy/helpers.py

@@ -1,7 +1,17 @@
+import json
 import os
+import re
+import shutil
+import subprocess
+import time
+from base64 import b64decode, b64encode
+from fractions import Fraction
+from hashlib import sha1
 from io import BytesIO
+from mimetypes import guess_type
 from pathlib import Path
 from typing import Any
+from urllib.parse import urlparse
 
 import numpy as np
 import pytest
@@ -63,6 +73,274 @@ def model_output_dir(parent_dir: Path, model: str) -> Path:
     return path
 
 
+_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
+_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
+
+
+def parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
+    streams = payload.get("streams")
+    if not isinstance(streams, list) or not streams:
+        raise ValueError(f"ffprobe payload did not include video streams: {payload}")
+
+    stream = streams[0]
+    width = int(stream["width"])
+    height = int(stream["height"])
+    fps = float(Fraction(stream["avg_frame_rate"]))
+    frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
+    if frame_count_value is None:
+        raise ValueError(f"ffprobe payload did not include frame count: {payload}")
+    frame_count = int(frame_count_value)
+    return {
+        "width": width,
+        "height": height,
+        "fps": fps,
+        "frame_count": frame_count,
+    }
+
+
+def parse_ssim_score(output: str) -> float:
+    match = _SSIM_RE.search(output)
+    if match is None:
+        raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
+    return float(match.group("score"))
+
+
+def parse_psnr_score(output: str) -> float:
+    match = _PSNR_RE.search(output)
+    if match is None:
+        raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
+    return float(match.group("score"))
+
+
+def probe_binary(binary: str, *, test_name: str = "video similarity e2e test") -> str:
+    resolved = shutil.which(binary)
+    if resolved is None:
+        pytest.skip(f"{binary} is required for {test_name}.")
+    return resolved
+
+
+def probe_video(path: Path) -> dict[str, int | float]:
+    probe_binary("ffprobe")
+    result = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "v:0",
+            "-count_frames",
+            "-show_entries",
+            "stream=width,height,avg_frame_rate,nb_read_frames",
+            "-of",
+            "json",
+            str(path),
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return parse_video_metadata(json.loads(result.stdout))
+
+
+def run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
+    probe_binary("ffmpeg")
+    result = subprocess.run(
+        [
+            "ffmpeg",
+            "-hide_banner",
+            "-i",
+            str(first),
+            "-i",
+            str(second),
+            "-lavfi",
+            f"[0:v][1:v]{filter_name}",
+            "-f",
+            "null",
+            "-",
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return result.stderr
+
+
+def online_timeout_seconds(configured: int | None, *, default: int = 1200) -> int:
+    return int(configured) if configured is not None else default
+
+
+def is_remote_image_source(source: str) -> bool:
+    return source.startswith("http://") or source.startswith("https://")
+
+
+def validate_image_source(source: str) -> None:
+    if is_remote_image_source(source):
+        response = requests.get(source, timeout=60)
+        response.raise_for_status()
+        return
+
+    image_path = Path(source)
+    if not image_path.exists():
+        raise FileNotFoundError(f"Local image source does not exist: {image_path}")
+
+
+def build_online_image_reference(source: str) -> str:
+    if source.startswith("data:image"):
+        return source
+
+    if is_remote_image_source(source):
+        response = requests.get(source, timeout=60)
+        response.raise_for_status()
+        mime_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
+        if not mime_type.startswith("image/"):
+            mime_type = guess_type(urlparse(source).path)[0] or "image/png"
+        encoded = b64encode(response.content).decode("ascii")
+        return f"data:{mime_type};base64,{encoded}"
+
+    image_path = Path(source)
+    mime_type = guess_type(image_path.name)[0] or "image/png"
+    encoded = b64encode(image_path.read_bytes()).decode("ascii")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def materialize_image_source(source: str, output_dir: Path) -> str:
+    if not is_remote_image_source(source):
+        return source
+
+    from diffusers.utils import load_image
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    source_name = Path(urlparse(source).path).name or "input.png"
+    safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
+    if "." not in safe_name:
+        safe_name = f"{safe_name}.png"
+    image_path = output_dir / safe_name
+    if not image_path.exists():
+        image = load_image(source).convert("RGB")
+        image.save(image_path)
+    return str(image_path)
+
+
+def video_artifact_dir(result_root: Path, source: str) -> Path:
+    if is_remote_image_source(source):
+        source_name = Path(urlparse(source).path).stem or "remote"
+    else:
+        source_name = Path(source).stem or "local"
+
+    safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
+    digest = sha1(source.encode("utf-8")).hexdigest()[:8]
+    return result_root / f"{safe_name}-{digest}"
+
+
+def send_video_request_with_timeout(
+    openai_client,
+    request_config: dict[str, Any],
+    *,
+    timeout_seconds: int,
+) -> bytes:
+    form_data = request_config.get("form_data")
+    if not isinstance(form_data, dict):
+        raise ValueError("Video request_config must contain 'form_data'")
+
+    if not form_data.get("prompt"):
+        raise ValueError("Video request_config['form_data'] must contain 'prompt'")
+
+    normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
+
+    files: dict[str, tuple[str, BytesIO, str]] = {}
+    image_reference = request_config.get("image_reference")
+    if image_reference:
+        if image_reference.startswith("data:image"):
+            header, encoded = image_reference.split(",", 1)
+            content_type = header.split(";")[0].removeprefix("data:")
+            extension = content_type.split("/")[-1]
+            file_data = b64decode(encoded)
+            files["input_reference"] = (
+                f"reference.{extension}",
+                BytesIO(file_data),
+                content_type,
+            )
+        else:
+            normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
+
+    start_time = time.perf_counter()
+    create_url = openai_client._build_url("/v1/videos")
+    response = requests.post(
+        create_url,
+        data=normalized_form_data,
+        files=files,
+        headers={"Accept": "application/json"},
+        timeout=60,
+    )
+    try:
+        response.raise_for_status()
+    except requests.HTTPError as exc:
+        raise requests.HTTPError(f"{exc}; response body: {response.text}", response=response) from exc
+    job_data = response.json()
+    video_id = job_data["id"]
+    openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
+    video_content = openai_client._download_video_content(video_id)
+    print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
+    return video_content
+
+
+def assert_video_metadata(
+    metadata: dict[str, int | float],
+    *,
+    width: int,
+    height: int,
+    fps: int | float,
+    frame_count: int,
+) -> None:
+    assert metadata["width"] == width
+    assert metadata["height"] == height
+    assert metadata["fps"] == float(fps)
+    assert metadata["frame_count"] == frame_count
+
+
+def assert_video_similarity_metrics(
+    *,
+    label: str,
+    online_path: Path,
+    offline_path: Path,
+    ssim_threshold: float,
+    psnr_threshold: float,
+) -> None:
+    ssim_output = run_ffmpeg_similarity("ssim", online_path, offline_path)
+    psnr_output = run_ffmpeg_similarity("psnr", online_path, offline_path)
+    ssim_score = parse_ssim_score(ssim_output)
+    psnr_score = parse_psnr_score(psnr_output)
+
+    print(f"{label} similarity metrics:")
+    print(
+        "  SSIM:"
+        f" value={ssim_score:.6f},"
+        f" threshold>={ssim_threshold:.6f},"
+        " range=[0, 1],"
+        " higher_is_better=True,"
+        " interpretation=structural_similarity"
+    )
+    print(
+        "  PSNR:"
+        f" value={psnr_score:.6f} dB,"
+        f" threshold>={psnr_threshold:.6f} dB,"
+        " range=[0, +inf),"
+        " higher_is_better=True,"
+        " interpretation=pixel_error_in_decibels"
+    )
+    print(f"online_video={online_path}")
+    print(f"offline_video={offline_path}")
+
+    assert ssim_score >= ssim_threshold, (
+        f"SSIM below threshold: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}. "
+        f"online={online_path} offline={offline_path}"
+    )
+    assert psnr_score >= psnr_threshold, (
+        f"PSNR below threshold: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}. "
+        f"online={online_path} offline={offline_path}"
+    )
+
+
 def assert_similarity(
     *,
     model_name: str,

tests/e2e/accuracy/hunyuanvideo15_i2v/__init__.py

@@ -0,0 +1 @@
+"""HunyuanVideo-1.5 I2V accuracy test cases."""

tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.e2e.accuracy.helpers import (
+    assert_video_metadata,
+    assert_video_similarity_metrics,
+    build_online_image_reference,
+    materialize_image_source,
+    probe_binary,
+    probe_video,
+    send_video_request_with_timeout,
+    validate_image_source,
+    video_artifact_dir,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL_NAME = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v"
+IMAGE_URL = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
+PROMPT = (
+    "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. "
+    "The fluffy-furred feline gazes directly at the camera with a relaxed expression. "
+    "Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, "
+    "and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, "
+    "as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's "
+    "intricate details and the refreshing atmosphere of the seaside."
+)
+WIDTH = 832
+HEIGHT = 480
+FPS = 24
+NUM_FRAMES = 33
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 6.0
+FLOW_SHIFT = 5.0
+SEED = 42
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+WORKSPACE_ROOT = REPO_ROOT.parent
+RUNNER_PATH = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
+RESULT_ROOT = Path(__file__).parent / "result"
+VIDEO_TIMEOUT_SECONDS = 60 * 60
+SERVER_CASES = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_NAME,
+            server_args=[
+                "--flow-shift",
+                str(FLOW_SHIFT),
+                "--enforce-eager",
+            ],
+            env_dict={"VLLM_OMNI_STORAGE_PATH": str(RESULT_ROOT / "storage")},
+            use_omni=True,
+        ),
+        id="hunyuanvideo15_i2v_default",
+    )
+]
+
+
+def _runner_env() -> dict[str, str]:
+    env = os.environ.copy()
+    pythonpath_parts = [
+        str(REPO_ROOT),
+        str(WORKSPACE_ROOT / "diffusers" / "src"),
+    ]
+    existing_pythonpath = env.get("PYTHONPATH")
+    if existing_pythonpath:
+        pythonpath_parts.append(existing_pythonpath)
+    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
+    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    return env
+
+
+def _resolve_image_source(configured: str | None) -> str:
+    if configured:
+        candidate = Path(configured)
+        if candidate.exists():
+            return str(candidate.resolve())
+        return configured
+    return IMAGE_URL
+
+
+def _artifact_paths(image_source: str) -> tuple[Path, Path]:
+    artifact_dir = video_artifact_dir(RESULT_ROOT, image_source)
+    artifact_dir.mkdir(parents=True, exist_ok=True)
+    return artifact_dir / "online.mp4", artifact_dir / "offline.mp4"
+
+
+def _offline_image_source(image_source: str) -> str:
+    artifact_dir = video_artifact_dir(RESULT_ROOT, image_source)
+    return materialize_image_source(image_source, artifact_dir)
+
+
+def _build_offline_command(*, image_source: str, output_path: Path) -> list[str]:
+    return [
+        sys.executable,
+        str(RUNNER_PATH),
+        "--model",
+        MODEL_NAME,
+        "--image",
+        image_source,
+        "--prompt",
+        PROMPT,
+        "--height",
+        str(HEIGHT),
+        "--width",
+        str(WIDTH),
+        "--num-frames",
+        str(NUM_FRAMES),
+        "--num-inference-steps",
+        str(NUM_INFERENCE_STEPS),
+        "--guidance-scale",
+        str(GUIDANCE_SCALE),
+        "--flow-shift",
+        str(FLOW_SHIFT),
+        "--fps",
+        str(FPS),
+        "--seed",
+        str(SEED),
+        "--enforce-eager",
+        "--output",
+        str(output_path),
+    ]
+
+
+def _generate_offline_video(*, image_source: str) -> Path:
+    _, offline_path = _artifact_paths(image_source)
+    offline_image_source = _offline_image_source(image_source)
+    subprocess.run(
+        _build_offline_command(image_source=offline_image_source, output_path=offline_path),
+        cwd=REPO_ROOT,
+        env=_runner_env(),
+        check=True,
+        timeout=VIDEO_TIMEOUT_SECONDS,
+    )
+    return offline_path
+
+
+def _generate_online_video(
+    *,
+    omni_server,
+    openai_client,
+    image_source: str,
+    timeout_seconds: int,
+) -> Path:
+    online_path, _ = _artifact_paths(image_source)
+    online_image_source = _offline_image_source(image_source)
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "prompt": PROMPT,
+            "height": HEIGHT,
+            "width": WIDTH,
+            "fps": FPS,
+            "num_frames": NUM_FRAMES,
+            "guidance_scale": GUIDANCE_SCALE,
+            "flow_shift": FLOW_SHIFT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "seed": SEED,
+        },
+        "image_reference": build_online_image_reference(online_image_source),
+    }
+    online_video_bytes = send_video_request_with_timeout(
+        openai_client,
+        request_config,
+        timeout_seconds=timeout_seconds,
+    )
+    online_path.write_bytes(online_video_bytes)
+    return online_path
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_i2v_diffusers_offline_generates_video(
+    hunyuanvideo15_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 I2V offline accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline runner does not exist: {RUNNER_PATH}")
+
+    image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
+    validate_image_source(image_source)
+    offline_path = _generate_offline_video(image_source=image_source)
+    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
+    offline_metadata = probe_video(offline_path)
+    assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_hunyuanvideo15_i2v_online_serving_generates_video(
+    omni_server,
+    openai_client,
+    hunyuanvideo15_i2v_image_source: str | None,
+    hunyuanvideo15_online_timeout_seconds: int,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 I2V online accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
+    validate_image_source(image_source)
+    online_path = _generate_online_video(
+        omni_server=omni_server,
+        openai_client=openai_client,
+        image_source=image_source,
+        timeout_seconds=hunyuanvideo15_online_timeout_seconds,
+    )
+    assert online_path.exists(), f"Expected online video artifact at {online_path}"
+    online_metadata = probe_video(online_path)
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_i2v_serving_matches_offline_video_similarity(
+    hunyuanvideo15_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 I2V video similarity test requires CUDA.")
+
+    probe_binary("ffmpeg")
+    probe_binary("ffprobe")
+    image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
+    validate_image_source(image_source)
+    online_path, offline_path = _artifact_paths(image_source)
+    if not online_path.exists():
+        pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
+    if not offline_path.exists():
+        pytest.skip(f"Missing offline artifact from prerequisite test: {offline_path}")
+
+    online_metadata = probe_video(online_path)
+    offline_metadata = probe_video(offline_path)
+    assert online_metadata == offline_metadata, (
+        f"Video metadata mismatch:\n"
+        f"online={online_metadata}\n"
+        f"offline={offline_metadata}\n"
+        f"online_path={online_path}\n"
+        f"offline_path={offline_path}"
+    )
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+    assert_video_similarity_metrics(
+        label="hunyuanvideo15_i2v",
+        online_path=online_path,
+        offline_path=offline_path,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/hunyuanvideo15_t2v/__init__.py

@@ -0,0 +1 @@
+"""HunyuanVideo-1.5 T2V accuracy test cases."""

tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.e2e.accuracy.helpers import (
+    assert_video_metadata,
+    assert_video_similarity_metrics,
+    probe_binary,
+    probe_video,
+    send_video_request_with_timeout,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL_NAME = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+PROMPT = "A cat walks on the grass, realistic style."
+WIDTH = 832
+HEIGHT = 480
+FPS = 24
+NUM_FRAMES = 33
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 6.0
+FLOW_SHIFT = 5.0
+SEED = 42
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+WORKSPACE_ROOT = REPO_ROOT.parent
+RUNNER_PATH = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
+RESULT_ROOT = Path(__file__).parent / "result"
+ARTIFACT_DIR = RESULT_ROOT / "cat_grass"
+VIDEO_TIMEOUT_SECONDS = 60 * 60
+SERVER_CASES = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_NAME,
+            server_args=[
+                "--flow-shift",
+                str(FLOW_SHIFT),
+            ],
+            env_dict={"VLLM_OMNI_STORAGE_PATH": str(RESULT_ROOT / "storage")},
+            use_omni=True,
+        ),
+        id="hunyuanvideo15_t2v_default",
+    )
+]
+
+
+def _runner_env() -> dict[str, str]:
+    env = os.environ.copy()
+    pythonpath_parts = [
+        str(REPO_ROOT),
+        str(WORKSPACE_ROOT / "diffusers" / "src"),
+    ]
+    existing_pythonpath = env.get("PYTHONPATH")
+    if existing_pythonpath:
+        pythonpath_parts.append(existing_pythonpath)
+    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
+    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    return env
+
+
+def _artifact_paths() -> tuple[Path, Path]:
+    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
+    return ARTIFACT_DIR / "online.mp4", ARTIFACT_DIR / "offline.mp4"
+
+
+def _build_offline_command(*, output_path: Path) -> list[str]:
+    return [
+        sys.executable,
+        str(RUNNER_PATH),
+        "--model",
+        MODEL_NAME,
+        "--prompt",
+        PROMPT,
+        "--height",
+        str(HEIGHT),
+        "--width",
+        str(WIDTH),
+        "--num-frames",
+        str(NUM_FRAMES),
+        "--num-inference-steps",
+        str(NUM_INFERENCE_STEPS),
+        "--guidance-scale",
+        str(GUIDANCE_SCALE),
+        "--flow-shift",
+        str(FLOW_SHIFT),
+        "--fps",
+        str(FPS),
+        "--seed",
+        str(SEED),
+        "--output",
+        str(output_path),
+    ]
+
+
+def _generate_offline_video() -> Path:
+    _, offline_path = _artifact_paths()
+    subprocess.run(
+        _build_offline_command(output_path=offline_path),
+        cwd=REPO_ROOT,
+        env=_runner_env(),
+        check=True,
+        timeout=VIDEO_TIMEOUT_SECONDS,
+    )
+    return offline_path
+
+
+def _generate_online_video(*, omni_server, openai_client, timeout_seconds: int) -> Path:
+    online_path, _ = _artifact_paths()
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "prompt": PROMPT,
+            "height": HEIGHT,
+            "width": WIDTH,
+            "fps": FPS,
+            "num_frames": NUM_FRAMES,
+            "guidance_scale": GUIDANCE_SCALE,
+            "flow_shift": FLOW_SHIFT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "seed": SEED,
+        },
+    }
+    online_video_bytes = send_video_request_with_timeout(
+        openai_client,
+        request_config,
+        timeout_seconds=timeout_seconds,
+    )
+    online_path.write_bytes(online_video_bytes)
+    return online_path
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_t2v_diffusers_offline_generates_video() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 T2V offline accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline runner does not exist: {RUNNER_PATH}")
+
+    offline_path = _generate_offline_video()
+    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
+    offline_metadata = probe_video(offline_path)
+    assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_hunyuanvideo15_t2v_online_serving_generates_video(
+    omni_server,
+    openai_client,
+    hunyuanvideo15_online_timeout_seconds: int,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 T2V online accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    online_path = _generate_online_video(
+        omni_server=omni_server,
+        openai_client=openai_client,
+        timeout_seconds=hunyuanvideo15_online_timeout_seconds,
+    )
+    assert online_path.exists(), f"Expected online video artifact at {online_path}"
+    online_metadata = probe_video(online_path)
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_t2v_serving_matches_offline_video_similarity() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 T2V video similarity test requires CUDA.")
+
+    probe_binary("ffmpeg")
+    probe_binary("ffprobe")
+    online_path, offline_path = _artifact_paths()
+    if not online_path.exists():
+        pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
+    if not offline_path.exists():
+        pytest.skip(f"Missing offline artifact from prerequisite test: {offline_path}")
+
+    online_metadata = probe_video(online_path)
+    offline_metadata = probe_video(offline_path)
+    assert online_metadata == offline_metadata, (
+        f"Video metadata mismatch:\n"
+        f"online={online_metadata}\n"
+        f"offline={offline_metadata}\n"
+        f"online_path={online_path}\n"
+        f"offline_path={offline_path}"
+    )
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+    assert_video_similarity_metrics(
+        label="hunyuanvideo15_t2v",
+        online_path=online_path,
+        offline_path=offline_path,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/test_diffusers_backend_similarity.py

@@ -30,10 +30,14 @@ from tests.e2e.accuracy.helpers import (
     env_to_apply_ftfy_mock_in_subproc,
     model_output_dir,
 )
-from tests.e2e.accuracy.wan22_i2v.test_wan22_i2v_video_similarity import (
-    _parse_psnr_score,
-    _parse_ssim_score,
-    _run_ffmpeg_similarity,
+from tests.e2e.accuracy.helpers import (
+    parse_psnr_score as _parse_psnr_score,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_ssim_score as _parse_ssim_score,
+)
+from tests.e2e.accuracy.helpers import (
+    run_ffmpeg_similarity as _run_ffmpeg_similarity,
 )
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -77,9 +77,9 @@ THRESHOLDS = {
     "text_prefix_match": 10,  # First 10 characters must match exactly
     "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
     # Image comparison
-    "clip_score": 85,  # CLIP image semantic similarity
-    "ssim": 0.20,  # Structural similarity
-    "psnr": 11.0,  # Peak signal-to-noise ratio (dB)
+    "clip_score": 90,  # CLIP image semantic similarity
+    "ssim": 0.26,  # Structural similarity
+    "psnr": 12.5,  # Peak signal-to-noise ratio (dB)
 }
 
 QUANT_PROMPT = "A brown and white dog is running on the grass."

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -20,7 +20,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 
 SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
-MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509"
+MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2511"
 WIDTH = 512
 HEIGHT = 512
 NUM_INFERENCE_STEPS = 20
@@ -253,37 +253,6 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
     )
 
 
-def _vllm_omni_output_multiple_image(
-    accuracy_artifact_root: Path,
-    qwen_bear_image: Image.Image,
-    rabbit_image: Image.Image,
-) -> Image.Image:
-    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
-    output_path = output_dir / "vllm_omni_multiple.png"
-    with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
-        output = _run_vllm_omni_image_edit(
-            omni_server=server,
-            prompt=PROMPT_MULTIPLE_IMAGE,
-            input_images=[qwen_bear_image, rabbit_image],
-            output_path=output_path,
-        )
-    return output
-
-
-def _diffusers_output_multiple_image(
-    accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
-) -> Image.Image:
-    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
-    output_path = output_dir / "diffusers_multiple.png"
-    return _run_diffusers_image_edit(
-        model=MULTIPLE_MODEL,
-        pipeline_class=QwenImageEditPlusPipeline,
-        prompt=PROMPT_MULTIPLE_IMAGE,
-        input_images=[qwen_bear_image, rabbit_image],
-        output_path=output_path,
-    )
-
-
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_single_matches_diffusers(
@@ -309,37 +278,6 @@ def test_qwen_image_edit_single_matches_diffusers(
     )
 
 
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(
-    reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
-)
-def test_qwen_image_edit_multiple_matches_diffusers(
-    accuracy_artifact_root: Path,
-    qwen_bear_image: Image.Image,
-    rabbit_image: Image.Image,
-) -> None:
-    vllm_image = _vllm_omni_output_multiple_image(
-        accuracy_artifact_root=accuracy_artifact_root,
-        qwen_bear_image=qwen_bear_image,
-        rabbit_image=rabbit_image,
-    )
-    diffusers_image = _diffusers_output_multiple_image(
-        accuracy_artifact_root=accuracy_artifact_root,
-        qwen_bear_image=qwen_bear_image,
-        rabbit_image=rabbit_image,
-    )
-    assert_similarity(
-        model_name=MULTIPLE_MODEL,
-        vllm_image=vllm_image,
-        diffusers_image=diffusers_image,
-        width=WIDTH,
-        height=HEIGHT,
-        ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=PSNR_THRESHOLD,
-    )
-
-
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:

tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py

@@ -1,20 +1,10 @@
 from __future__ import annotations
 
-import json
 import os
-import re
-import shutil
 import subprocess
 import sys
-import time
-from base64 import b64decode, b64encode
-from fractions import Fraction
-from hashlib import sha1
-from io import BytesIO
-from mimetypes import guess_type
+from base64 import b64decode
 from pathlib import Path
-from typing import Any
-from urllib.parse import urlparse
 
 import pytest
 import requests
@@ -22,6 +12,38 @@ import torch
 from diffusers import UniPCMultistepScheduler
 from PIL import Image
 
+from tests.e2e.accuracy.helpers import (
+    assert_video_metadata,
+    assert_video_similarity_metrics,
+    video_artifact_dir,
+)
+from tests.e2e.accuracy.helpers import (
+    build_online_image_reference as _build_online_image_reference,
+)
+from tests.e2e.accuracy.helpers import (
+    online_timeout_seconds as _online_timeout_seconds,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_psnr_score as _parse_psnr_score,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_ssim_score as _parse_ssim_score,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_video_metadata as _parse_video_metadata,
+)
+from tests.e2e.accuracy.helpers import (
+    probe_binary as _probe_binary,
+)
+from tests.e2e.accuracy.helpers import (
+    probe_video as _probe_video,
+)
+from tests.e2e.accuracy.helpers import (
+    send_video_request_with_timeout as _send_video_request_with_timeout,
+)
+from tests.e2e.accuracy.helpers import (
+    validate_image_source as _validate_image_source,
+)
 from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
     _configure_scheduler,
     _ensure_wan_ftfy_fallback,
@@ -218,8 +240,6 @@ WORKSPACE_ROOT = REPO_ROOT.parent
 RUNNER_PATH = Path(__file__).with_name("run_wan22_i2v_diffusers_cp.py")
 RESULT_ROOT = Path(__file__).parent / "result"
 VIDEO_TIMEOUT_SECONDS = 60 * 60
-_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
-_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
 SERVER_CASES = [
     pytest.param(
         OmniServerParams(
@@ -238,48 +258,6 @@ SERVER_CASES = [
 ]
 
 
-def _parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
-    streams = payload.get("streams")
-    if not isinstance(streams, list) or not streams:
-        raise ValueError(f"ffprobe payload did not include video streams: {payload}")
-
-    stream = streams[0]
-    width = int(stream["width"])
-    height = int(stream["height"])
-    fps = float(Fraction(stream["avg_frame_rate"]))
-    frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
-    if frame_count_value is None:
-        raise ValueError(f"ffprobe payload did not include frame count: {payload}")
-    frame_count = int(frame_count_value)
-    return {
-        "width": width,
-        "height": height,
-        "fps": fps,
-        "frame_count": frame_count,
-    }
-
-
-def _parse_ssim_score(output: str) -> float:
-    match = _SSIM_RE.search(output)
-    if match is None:
-        raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
-    return float(match.group("score"))
-
-
-def _parse_psnr_score(output: str) -> float:
-    match = _PSNR_RE.search(output)
-    if match is None:
-        raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
-    return float(match.group("score"))
-
-
-def _probe_binary(binary: str) -> str:
-    resolved = shutil.which(binary)
-    if resolved is None:
-        pytest.skip(f"{binary} is required for Wan2.2 video similarity e2e test.")
-    return resolved
-
-
 def _build_diffusers_command(
     *,
     runner_path: Path,
@@ -343,90 +321,8 @@ def _resolve_image_source(configured: str | None) -> str:
     return configured
 
 
-def _is_remote_image_source(source: str) -> bool:
-    return source.startswith("http://") or source.startswith("https://")
-
-
-def _validate_image_source(source: str) -> None:
-    if _is_remote_image_source(source):
-        response = requests.get(source, timeout=60)
-        response.raise_for_status()
-        return
-
-    image_path = Path(source)
-    if not image_path.exists():
-        raise FileNotFoundError(f"Local image source does not exist: {image_path}")
-
-
-def _build_online_image_reference(source: str) -> str:
-    if _is_remote_image_source(source) or source.startswith("data:image"):
-        return source
-
-    image_path = Path(source)
-    mime_type = guess_type(image_path.name)[0] or "image/png"
-    encoded = b64encode(image_path.read_bytes()).decode("ascii")
-    return f"data:{mime_type};base64,{encoded}"
-
-
-def _probe_video(path: Path) -> dict[str, int | float]:
-    _probe_binary("ffprobe")
-    result = subprocess.run(
-        [
-            "ffprobe",
-            "-v",
-            "error",
-            "-select_streams",
-            "v:0",
-            "-count_frames",
-            "-show_entries",
-            "stream=width,height,avg_frame_rate,nb_read_frames",
-            "-of",
-            "json",
-            str(path),
-        ],
-        capture_output=True,
-        text=True,
-        check=True,
-    )
-    return _parse_video_metadata(json.loads(result.stdout))
-
-
-def _run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
-    _probe_binary("ffmpeg")
-    result = subprocess.run(
-        [
-            "ffmpeg",
-            "-hide_banner",
-            "-i",
-            str(first),
-            "-i",
-            str(second),
-            "-lavfi",
-            f"[0:v][1:v]{filter_name}",
-            "-f",
-            "null",
-            "-",
-        ],
-        capture_output=True,
-        text=True,
-        check=True,
-    )
-    return result.stderr
-
-
-def _online_timeout_seconds(configured: int | None) -> int:
-    return int(configured) if configured is not None else 1200
-
-
 def _artifact_dir(image_source: str) -> Path:
-    if _is_remote_image_source(image_source):
-        source_name = Path(urlparse(image_source).path).stem or "remote"
-    else:
-        source_name = Path(image_source).stem or "local"
-
-    safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
-    digest = sha1(image_source.encode("utf-8")).hexdigest()[:8]
-    return RESULT_ROOT / f"{safe_name}-{digest}"
+    return video_artifact_dir(RESULT_ROOT, image_source)
 
 
 def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
@@ -439,55 +335,6 @@ def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
     )
 
 
-def _send_video_request_with_timeout(
-    openai_client,
-    request_config: dict[str, Any],
-    *,
-    timeout_seconds: int,
-) -> bytes:
-    form_data = request_config.get("form_data")
-    if not isinstance(form_data, dict):
-        raise ValueError("Video request_config must contain 'form_data'")
-
-    if not form_data.get("prompt"):
-        raise ValueError("Video request_config['form_data'] must contain 'prompt'")
-
-    normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
-
-    files: dict[str, tuple[str, BytesIO, str]] = {}
-    image_reference = request_config.get("image_reference")
-    if image_reference:
-        if image_reference.startswith("data:image"):
-            header, encoded = image_reference.split(",", 1)
-            content_type = header.split(";")[0].removeprefix("data:")
-            extension = content_type.split("/")[-1]
-            file_data = b64decode(encoded)
-            files["input_reference"] = (
-                f"reference.{extension}",
-                BytesIO(file_data),
-                content_type,
-            )
-        else:
-            normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
-
-    start_time = time.perf_counter()
-    create_url = openai_client._build_url("/v1/videos")
-    response = requests.post(
-        create_url,
-        data=normalized_form_data,
-        files=files,
-        headers={"Accept": "application/json"},
-        timeout=60,
-    )
-    response.raise_for_status()
-    job_data = response.json()
-    video_id = job_data["id"]
-    openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
-    video_content = openai_client._download_video_content(video_id)
-    print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
-    return video_content
-
-
 def _generate_online_video(
     *,
     omni_server,
@@ -557,10 +404,7 @@ def test_wan22_i2v_diffusers_offline_generates_video(
     assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
     assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
     offline_metadata = _probe_video(offline_path)
-    assert offline_metadata["width"] == WIDTH
-    assert offline_metadata["height"] == HEIGHT
-    assert offline_metadata["fps"] == float(FPS)
-    assert offline_metadata["frame_count"] == NUM_FRAMES
+    assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
 
 
 @pytest.mark.benchmark
@@ -586,10 +430,7 @@ def test_wan22_i2v_online_serving_generates_video(
     )
     assert online_path.exists(), f"Expected online video artifact at {online_path}"
     online_metadata = _probe_video(online_path)
-    assert online_metadata["width"] == WIDTH
-    assert online_metadata["height"] == HEIGHT
-    assert online_metadata["fps"] == float(FPS)
-    assert online_metadata["frame_count"] == NUM_FRAMES
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
 
 
 @pytest.mark.benchmark
@@ -627,42 +468,12 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity(
         f"online_path={online_path}\n"
         f"offline_path={offline_path}"
     )
-    assert online_metadata["width"] == WIDTH
-    assert online_metadata["height"] == HEIGHT
-    assert online_metadata["fps"] == float(FPS)
-    assert online_metadata["frame_count"] == NUM_FRAMES
-
-    ssim_output = _run_ffmpeg_similarity("ssim", online_path, offline_path)
-    psnr_output = _run_ffmpeg_similarity("psnr", online_path, offline_path)
-    ssim_score = _parse_ssim_score(ssim_output)
-    psnr_score = _parse_psnr_score(psnr_output)
-
-    print("wan22_i2v_similarity metrics:")
-    print(
-        "  SSIM:"
-        f" value={ssim_score:.6f},"
-        f" threshold>={SSIM_THRESHOLD:.6f},"
-        " range=[0, 1],"
-        " higher_is_better=True,"
-        " interpretation=structural_similarity"
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+    assert_video_similarity_metrics(
+        label="wan22_i2v",
+        online_path=online_path,
+        offline_path=offline_path,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
     )
-    print(
-        "  PSNR:"
-        f" value={psnr_score:.6f} dB,"
-        f" threshold>={PSNR_THRESHOLD:.6f} dB,"
-        " range=[0, +inf),"
-        " higher_is_better=True,"
-        " interpretation=pixel_error_in_decibels"
-    )
-    print(f"online_video={online_path}")
-    print(f"offline_video={offline_path}")
     print(f"offline_metadata={offline_metadata_path}")
-
-    assert ssim_score >= SSIM_THRESHOLD, (
-        f"SSIM below threshold: got {ssim_score:.6f}, expected >= {SSIM_THRESHOLD:.6f}. "
-        f"online={online_path} offline={offline_path}"
-    )
-    assert psnr_score >= PSNR_THRESHOLD, (
-        f"PSNR below threshold: got {psnr_score:.6f}, expected >= {PSNR_THRESHOLD:.6f}. "
-        f"online={online_path} offline={offline_path}"
-    )

tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py

@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E tests for GLM-Image AutoRound W4A16 quantized inference.
+
+These tests cover text-to-image and image-to-image generation with
+the W4A16 quantized GLM-Image model.
+
+Requirements:
+  - 2 CUDA GPUs (H100 or equivalent)
+  - The quantized model checkpoint (Intel/GLM-Image-int4-AutoRound)
+"""
+
+import gc
+import math
+import os
+
+import numpy as np
+import pytest
+from PIL import Image
+from vllm import SamplingParams
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+QUANTIZED_MODEL = os.environ.get("GLM_IMAGE_AUTOROUND_MODEL", "Intel/GLM-Image-int4-AutoRound")
+
+# Small resolution to keep GPU memory & time manageable
+HEIGHT = 256
+WIDTH = 256
+NUM_STEPS = 2  # minimal for smoke-test
+
+# GLM-Image AR generation config (from generation_config.json)
+GLM_IMAGE_EOS_TOKEN_ID = 16385
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512
+
+_CI_DEPLOY = get_deploy_config_path("glm_image.yaml")
+
+
+def _get_stage_config():
+    """Build a CI-friendly stage config with eager mode for testing."""
+    return modify_stage_config(
+        _CI_DEPLOY,
+        updates={
+            "stages": {
+                0: {"enforce_eager": True},
+                1: {"enforce_eager": True},
+            },
+        },
+    )
+
+
+stage_config = _get_stage_config()
+
+# (model, stage_config_path) for ``omni_runner`` indirect parametrize
+_OMNI_RUNNER_PARAM = (QUANTIZED_MODEL, stage_config)
+
+
+def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
+    """Compute max_new_tokens for GLM-Image AR text-to-image generation."""
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    ratio = token_h / token_w if token_w > 0 else 1.0
+    small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
+    small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
+    small_tokens = small_token_h * small_token_w
+
+    return small_tokens + large_tokens + 1
+
+
+def _ar_sampling_params(max_tokens: int, height: int, width: int, seed: int = 42) -> SamplingParams:
+    """Build AR stage SamplingParams for GLM-Image."""
+    return SamplingParams(
+        temperature=0.9,
+        top_p=0.75,
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
+        max_tokens=max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
+        seed=seed,
+        detokenize=False,
+        extra_args={
+            "target_h": height,
+            "target_w": width,
+        },
+    )
+
+
+def _diffusion_sampling_params(
+    height: int = HEIGHT,
+    width: int = WIDTH,
+    num_steps: int = NUM_STEPS,
+    seed: int = 42,
+) -> OmniDiffusionSamplingParams:
+    """Build Diffusion stage OmniDiffusionSamplingParams."""
+    return OmniDiffusionSamplingParams(
+        height=height,
+        width=width,
+        num_inference_steps=num_steps,
+        guidance_scale=0.0,
+        seed=seed,
+    )
+
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.diffusion,
+]
+
+
+# ------------------------------------------------------------------
+# Test: text-to-image generation produces a valid image (quantized)
+# ------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunnerHandler):
+    """Load the W4A16 quantized GLM-Image model and verify it produces a valid image."""
+    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()
+
+    prompt_dict = {
+        "prompt": "A photo of a cat sitting on a laptop keyboard",
+        "height": HEIGHT,
+        "width": WIDTH,
+        "mm_processor_kwargs": {
+            "target_h": HEIGHT,
+            "target_w": WIDTH,
+        },
+    }
+    ar_params = _ar_sampling_params(
+        max_tokens=compute_max_tokens(HEIGHT, WIDTH),
+        height=HEIGHT,
+        width=WIDTH,
+        seed=42,
+    )
+    diffusion_params = _diffusion_sampling_params(
+        height=HEIGHT,
+        width=WIDTH,
+        num_steps=NUM_STEPS,
+        seed=42,
+    )
+
+    outputs = omni_runner_handler.runner.generate(
+        [prompt_dict],
+        [ar_params, diffusion_params],
+    )
+
+    monitor.stop()
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    req_out = first_output.request_output
+    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+    images = req_out.images
+
+    assert len(images) >= 1, "Expected at least one generated image"
+    img = images[0]
+    assert isinstance(img, Image.Image)
+    assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
+    assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
+
+    # Sanity: image should not be blank (all-zero)
+    arr = np.array(img)
+    assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+
+
+# ------------------------------------------------------------------
+# Test: image-to-image generation (quantized)
+# ------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunnerHandler):
+    """Load the W4A16 quantized GLM-Image and verify image-to-image generation works."""
+    ref_image_arr = generate_synthetic_image(WIDTH, HEIGHT)["np_array"]
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+    current_omni_platform.reset_peak_memory_stats()
+
+    prompt_dict = {
+        "prompt": "Make it look like winter",
+        "multi_modal_data": {"image": ref_image_arr},
+        "height": HEIGHT,
+        "width": WIDTH,
+        "mm_processor_kwargs": {
+            "target_h": HEIGHT,
+            "target_w": WIDTH,
+        },
+    }
+    ar_params = _ar_sampling_params(
+        max_tokens=compute_max_tokens(HEIGHT, WIDTH),
+        height=HEIGHT,
+        width=WIDTH,
+        seed=42,
+    )
+    diffusion_params = _diffusion_sampling_params(
+        height=HEIGHT,
+        width=WIDTH,
+        num_steps=NUM_STEPS,
+        seed=42,
+    )
+
+    outputs = omni_runner_handler.runner.generate(
+        [prompt_dict],
+        [ar_params, diffusion_params],
+    )
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    req_out = first_output.request_output
+    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+    images = req_out.images
+
+    assert len(images) >= 1, "Expected at least one generated image"
+    img = images[0]
+    assert isinstance(img, Image.Image)
+    assert img.width == WIDTH
+    assert img.height == HEIGHT
+
+    gc.collect()
+    current_omni_platform.empty_cache()

tests/e2e/offline_inference/test_higgs_audio_v2.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end offline inference tests for higgs-audio v2.
+
+Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py:
+build Stage-0 prompt_token_ids with the model's own processor + plain-text
+chat template, drive the engine via ``Omni.generate``, and inspect the
+Stage-1 PCM payload directly. The OmniRunnerHandler.send_audio_speech_request
+helper applies a Qwen-style ChatML wrap that is not valid for higgs, so this
+test bypasses it (same approach as test_voxtral_tts.py).
+
+v1 scope is plain text -> 24 kHz speech; voice cloning, multi-speaker prompts,
+language overrides, and task_type selection are out of scope.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Match the example script: keep DeepGEMM FP8 off so engine init does not
+# require the optional ``deep_gemm`` backend.
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import numpy as np
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "bosonai/higgs-audio-v2-generation-3B-base"
+STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
+SAMPLE_RATE = 24_000
+# A short sentence at 24 kHz must produce at least ~0.5 s of audio.
+MIN_AUDIO_SAMPLES = 12_000
+TEST_TEXT = "Hello world."
+
+# Reuse the qwen3_tts ref clip + transcript for voice clone tests (see
+# tests/e2e/online_serving/test_qwen3_tts_base.py:27-28 for the source).
+_REF_AUDIO_ASSET = "qwen3_tts/clone_2.wav"
+_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+# (model, stage_config_path, extra_omni_kwargs) — same shape as test_voxtral_tts.py.
+# Stage configs already pin enforce_eager=true per stage, so no Omni-level override.
+_OMNI_RUNNER_PARAM = (MODEL, STAGE_CONFIG, {"trust_remote_code": True})
+
+
+def _compose_request(model_name: str, text: str) -> dict:
+    """Build a higgs prompt dict via the upstream processor + plain-text template.
+
+    The serving layer's ``_build_higgs_audio_v2_params`` does the same thing
+    (see vllm_omni/entrypoints/openai/serving_speech.py:1520) — keeping the
+    offline and online prompt-construction paths identical is what guarantees
+    input-token parity with the HF reference.
+    """
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_plain_text_prompt,
+        input_ids_to_python_list,
+    )
+
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    inputs = build_plain_text_prompt(processor, text)
+    return {"prompt_token_ids": input_ids_to_python_list(inputs)}
+
+
+def _extract_pcm(outputs) -> torch.Tensor:
+    """Pull the concatenated PCM tensor out of the Stage-1 multimodal_output."""
+    for stage_output in outputs:
+        mm = getattr(stage_output, "multimodal_output", None)
+        if not mm:
+            continue
+        audio = mm.get("model_outputs") if "model_outputs" in mm else mm.get("audio")
+        if audio is None:
+            continue
+        if isinstance(audio, list):
+            chunks = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+            if not chunks:
+                continue
+            return torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0]
+        return torch.as_tensor(audio).float().cpu().reshape(-1)
+    raise AssertionError(f"no audio payload in any stage output (got {len(outputs)} stages)")
+
+
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_plain_text(omni_runner_function: OmniRunner) -> None:
+    """
+    Offline plain-text -> 24 kHz PCM via the bundled higgs_audio_v2 deploy.
+
+    Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml (codec_streaming=true,
+        async_chunk=false, enforce_eager=true per stage).
+    Input Modal: text
+    Output Modal: audio (PCM tensor, 24 kHz mono float)
+    Input Setting: single prompt, default sampling
+    """
+    omni = omni_runner_function.omni
+    inputs = _compose_request(MODEL, TEST_TEXT)
+
+    # max_tokens caps Stage-0 codec frames; Stage 1 max_tokens lifts the
+    # per-codebook prefill cap. Mirror the deploy yaml's defaults.
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate([inputs], stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate"
+
+    pcm = _extract_pcm(outputs)
+    arr = pcm.numpy()
+
+    assert arr.size > MIN_AUDIO_SAMPLES, (
+        f"audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES} "
+        f"({MIN_AUDIO_SAMPLES / SAMPLE_RATE:.2f} s @ {SAMPLE_RATE} Hz)"
+    )
+    # Threshold 5e-4 (not 1e-3) keeps us above true silence while tolerating
+    # the ``load_format=dummy`` Stage-1 codec — random RVQ/DAC weights collapse
+    # the output to a near-constant peak around 7.4e-4. Real-weight runs reach
+    # 0.1+ peak so this still catches genuine silence.
+    assert np.max(np.abs(arr)) > 5e-4, "audio output is silent (peak abs amplitude < 5e-4)"
+    assert np.isfinite(arr).all(), "audio contains non-finite samples (nan/inf)"
+
+
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_batch_two_prompts(omni_runner_function: OmniRunner) -> None:
+    """
+    Submit two prompts in one engine call and verify both yield non-trivial PCM.
+
+    Guards the Stage-0 talker's per-slot audio state under a batched call —
+    the prior example ran prompts one at a time citing 'request-scoped audio
+    state'; this test pins the batched-2 behavior the engine path actually
+    needs to support for /v1/audio/speech under concurrent traffic.
+    """
+    omni = omni_runner_function.omni
+    prompts = [
+        _compose_request(MODEL, "Hello world."),
+        _compose_request(MODEL, "The quick brown fox jumps over the lazy dog."),
+    ]
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate(prompts, stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate"
+
+    pcm = _extract_pcm(outputs)
+    # Whether the engine returns one concatenated PCM or per-request tensors is
+    # implementation-defined; for either shape, total samples must clear the
+    # floor for at least one of the two prompts.
+    assert pcm.numel() > MIN_AUDIO_SAMPLES, (
+        f"batched audio too short: {pcm.numel()} samples, expected > {MIN_AUDIO_SAMPLES}"
+    )
+    # See the silence-threshold note in test_higgs_audio_v2_offline_plain_text.
+    assert np.max(np.abs(pcm.numpy())) > 5e-4, "batched audio output is silent"
+
+
+def _compose_voice_clone_request(model_name: str, text: str, ref_text: str, ref_wav_path) -> dict:
+    """Build the voice-clone prompt + additional_information payload.
+
+    Calls the same ``build_voice_clone_prompt`` the serving layer uses online —
+    the HF processor encodes the reference clip on the spot via the bundled
+    HiggsAudioV2TokenizerModel. The talker's
+    ``_maybe_apply_ref_audio_substitution`` consumes the tensors at prefill.
+    """
+    import soundfile as sf
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_voice_clone_prompt,
+    )
+
+    wav, sr = sf.read(str(ref_wav_path), always_2d=False)
+    if wav.ndim == 2:
+        wav = wav.mean(axis=1)  # downmix any stereo asset to mono
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    out = build_voice_clone_prompt(processor, text, wav, int(sr), ref_text)
+    return {
+        "prompt_token_ids": out["prompt_token_ids"],
+        # Pass tensors bare (see serving_speech._build_higgs_audio_v2_params
+        # — list-wrapped tensors get dropped by the msgspec serializer).
+        "additional_information": {
+            "audio_input_ids": out["audio_input_ids"],
+            "audio_input_ids_mask": out["audio_input_ids_mask"],
+        },
+    }
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_voice_clone(omni_runner_function: OmniRunner) -> None:
+    """
+    Shallow voice clone via ref_audio + ref_text, offline.
+
+    Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml.
+    Input Modal: text + reference audio (qwen3_tts/clone_2.wav vendored asset)
+    Output Modal: audio (PCM tensor, 24 kHz mono float)
+    Input Setting: single prompt
+    """
+    omni = omni_runner_function.omni
+    ref_path = get_asset_path(_REF_AUDIO_ASSET)
+    inputs = _compose_voice_clone_request(MODEL, TEST_TEXT, _REF_TEXT, ref_path)
+
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate([inputs], stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate for voice clone"
+
+    pcm = _extract_pcm(outputs)
+    arr = pcm.numpy()
+    assert arr.size > MIN_AUDIO_SAMPLES, f"cloned audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES}"
+    assert np.max(np.abs(arr)) > 5e-4, "voice-clone output is silent"
+    assert np.isfinite(arr).all(), "voice-clone audio contains non-finite samples"

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -3,6 +3,8 @@ End-to-end diffusion coverage for FLUX.2-dev in online serving mode.
 
 Coverage:
 - CPU offload
+- Ulysses sequence parallelism
+- Ring sequence parallelism
 
 This test verifies that FLUX.2-dev can be launched with CPU offload enabled,
 accepts text-to-image requests through the OpenAI-compatible API, and returns
@@ -28,7 +30,7 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
 
 def _get_flux_2_dev_feature_cases(model: str):
-    """Return FLUX.2-dev diffusion feature cases for CPU offload."""
+    """Return FLUX.2-dev diffusion feature cases for CPU offload and SP."""
 
     return [
         pytest.param(
@@ -53,6 +55,30 @@ def _get_flux_2_dev_feature_cases(model: str):
             id="parallel_cfg_2",
             marks=PARALLEL_FEATURE_MARKS,
         ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--ulysses-degree",
+                    "2",
+                ],
+            ),
+            id="ulysses_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--ring-degree",
+                    "2",
+                ],
+            ),
+            id="ring_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
     ]
 
 

tests/e2e/online_serving/test_higgs_audio_v2.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end online tests for higgs-audio v2 against /v1/audio/speech.
+
+v1 scope is plain text -> 24 kHz speech plus shallow voice clone via
+ref_audio + ref_text (inline) or voice=<name> (after POST /v1/audio/voices).
+The model-aware request validator
+(vllm_omni/entrypoints/openai/serving_speech.py::_validate_higgs_audio_v2_request)
+rejects multi-speaker tags, language overrides, task_type, and bare
+voice=<name> for names that do not match an uploaded speaker — this suite
+exercises both the happy path (plain text in, audio bytes out) and the
+validator rejections, so a regression that loosens the schema will fail
+loudly.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+# Match run_server.sh: DeepGEMM FP8 kernels are optional and trip warmup on
+# images without the deep_gemm backend, so disable them by default.
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "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 = "bosonai/higgs-audio-v2-generation-3B-base"
+STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
+SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
+# DeepGEMM warmup is optional; mirror run_server.sh and switch it off in the
+# server subprocess env too (parent-process os.environ above only affects this
+# test driver; the engine subprocesses inherit through env_dict).
+SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"}
+
+TEST_PARAMS = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=STAGE_CONFIG,
+            server_args=SERVER_ARGS,
+            env_dict=SERVER_ENV,
+        ),
+        id="higgs_audio_v2_plain_text",
+    )
+]
+
+DEFAULT_SPEECH_TIMEOUT_S = 180.0
+# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ≈ 24 KiB.
+# A WAV header adds 44 bytes; pick a conservative floor that catches truncated
+# / silence-only outputs without flagging short legitimate clips.
+_MIN_AUDIO_BYTES = 20_000
+
+# Reuse the qwen3_tts vendored reference clip (clean ~5 s 24 kHz mono human
+# speech) + its transcript. See tests/e2e/online_serving/test_qwen3_tts_base.py
+# for the asset rationale — keeping a single shared reference clip across TTS
+# tests avoids duplicating WAVs in the repo.
+_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineHappyPath:
+    """Plain-text -> audio happy paths against the live HTTP server."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_wav(self, omni_server, openai_client) -> None:
+        """Single non-streaming WAV request — covers the canonical TTS happy path."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Streaming PCM via the shared-memory connector's codec_streaming path.
+
+        higgs_audio_v2.yaml pins ``codec_streaming: true`` + ``async_chunk: false``,
+        so the only streaming surface exposed to clients is the WAV/PCM bytes
+        served chunk-by-chunk from Stage 1 — exercise it directly.
+        """
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Streaming the quick brown fox over the lazy dog.",
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
+        """max_new_tokens is one of the few extra fields the higgs validator accepts."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Innovation distinguishes between a leader and a follower.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "max_new_tokens": 500,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
+        """Three concurrent non-streaming requests — guards the per-slot audio state."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "It was the night before my birthday.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            },
+            request_num=3,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Validator rejections served over HTTP. Each case targets one out-of-scope
+# field; the validator returns 4xx via the OpenAI error path. We keep the
+# error-message substring loose so a phrasing tweak in the validator does not
+# break CI, but tight enough to catch a regression that silently accepts the
+# field.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineValidatorRejections:
+    """Out-of-scope fields must come back as 4xx with a higgs-named message."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None:
+        """Voice clone needs the transcript too — half-supplied is a 4xx."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_audio": _REF_AUDIO_URL,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "ref_text",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None:
+        """Symmetric guard: ref_text alone is not enough — must come with ref_audio."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_text": "some transcript",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "ref_audio",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_task_type(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "task_type": "Base",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "task_type",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_language_override(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "language": "Chinese",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "language",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "[SPEAKER0] hi",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "multi-speaker",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_empty_input(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "   ",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "empty",
+            }
+        )
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineVoiceClone:
+    """Shallow voice clone: ref_audio + ref_text -> speech in the cloned voice.
+
+    The HF processor that the serving layer calls (lazy-loaded at the first
+    request) ships with the bundled HiggsAudioV2TokenizerModel and encodes
+    the reference clip in-process. The talker substitutes the encoded codes
+    at the prompt-side audio placeholders via
+    :meth:`HiggsAudioV2TalkerForConditionalGeneration._maybe_apply_ref_audio_substitution`.
+    """
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_basic(self, omni_server, openai_client) -> None:
+        """Single non-streaming WAV request driven by the qwen3_tts ref clip."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_audio": _REF_AUDIO_URL,
+                "ref_text": _REF_TEXT,
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Voice clone over the streaming PCM path."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Innovation distinguishes a leader from a follower.",
+                "ref_audio": _REF_AUDIO_URL,
+                "ref_text": _REF_TEXT,
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -539,3 +539,24 @@ def test_language_001(omni_server, openai_client) -> None:
     }
 
     openai_client.send_omni_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text only (long-form generation prompt).
+    Output Modal: text, audio (default ``modalities``);
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text="Tell a 300-word story.",
+    )
+
+    request_config = {"model": omni_server.model, "messages": messages, "stream": True}
+    responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+    text = responses[0].text_content if responses else ""
+    word_count = len(text.split())
+    assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"

tests/e2e/online_serving/test_qwen_image_edit.py

@@ -5,8 +5,8 @@
 Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
 
 - ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
-- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2509`` — two reference images, fixed 512×512.
-- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2509`` only, ``advanced_model`` — mixed input
+- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
+- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input
   resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
   scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
 
@@ -53,7 +53,6 @@ def _get_diffusion_feature_cases(model: str):
 
 
 @pytest.mark.advanced_model
-@pytest.mark.core_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
@@ -83,15 +82,16 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
     openai_client.send_diffusion_request(request_config)
 
 
+@pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2509``."""
+    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511``."""
 
     image_data_url_list = [f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" for _ in range(2)]
 
@@ -118,7 +118,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -2,7 +2,7 @@
 Comprehensive tests of diffusion features that are available in online serving mode
 and are supported by the following models:
 - Qwen-Image-Edit: single image input
-- Qwen-Image-Edit-2509: single image input and two image inputs
+- Qwen-Image-Edit-2511: single image input and two image inputs
 """
 
 import pytest
@@ -14,7 +14,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
-SINGLE_EDIT_PROMPT_2509 = "Restyle this image into a Vincent van Gogh style impressionist painting."
+SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (
     "Transform the first image into a Dadaism collage art. "
     "Transform the second image into a Vincent van Gogh style painting. "
@@ -34,7 +34,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "tea_cache",  # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together.
+                    "cache_dit",
                     "--enable-layerwise-offload",
                 ],
             ),
@@ -59,7 +59,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "cache_dit",
+                    "tea_cache",
                     "--ring",
                     "2",
                 ],
@@ -141,11 +141,11 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
-def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2509 with a single image input.
+def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test Qwen-Image-Edit-2511 with a single image input.
 
     Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache
     postprocess closure used the doubled temb (shape 2*batch) without halving
@@ -156,7 +156,7 @@ def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_clien
     """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
-    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2509)
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
 
     request_config = {
         "model": omni_server.model,
@@ -176,11 +176,11 @@ def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_clien
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
-def test_qwen_image_edit_2509_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2509 with two image inputs."""
+def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test Qwen-Image-Edit-2511 with two image inputs."""
     image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
     image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,8 +1,10 @@
+import concurrent.futures
 import importlib
 import os
 import threading
 import time
 import types
+from contextlib import contextmanager
 
 import pytest
 
@@ -189,10 +191,12 @@ def test_initialize_diffusion_replica_restores_device_visibility_after_local_ini
     old_env = os.environ.get(env_var)
     os.environ[env_var] = "0,1"
 
-    def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
+    @contextmanager
+    def _fake_stage_runtime_setup(_stage_id, _runtime_cfg):
         current_omni_platform.set_device_control_env_var("1")
+        yield
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
+    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _fake_stage_runtime_setup)
     monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
     monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace())
 
@@ -221,7 +225,11 @@ def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(
 
     captured: dict[str, object] = {}
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    @contextmanager
+    def _noop_stage_runtime_setup(*_):
+        yield
+
+    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
     monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
 
     def _capture_initialize_diffusion_stage(
@@ -372,6 +380,44 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
     assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas)
 
 
+def test_initialize_stage_replicas_reuses_long_lived_executor(monkeypatch):
+    engine = object.__new__(AsyncOmniEngine)
+    engine._stage_init_executor = None
+
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)]
+
+    created_executors = []
+
+    class FakeExecutor:
+        def __init__(self, *args, **kwargs):
+            self.shutdown_calls = []
+            created_executors.append(self)
+
+        def submit(self, fn, *args, **kwargs):
+            future = concurrent.futures.Future()
+            future.set_result(fn(*args, **kwargs))
+            return future
+
+        def shutdown(self, wait=False, cancel_futures=False):
+            self.shutdown_calls.append((wait, cancel_futures))
+
+    monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor)
+    monkeypatch.setattr(
+        engine,
+        "_initialize_replica",
+        lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}",
+    )
+
+    first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
+    second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
+
+    assert first == {0: ["client-0"]}
+    assert second == {0: ["client-0"]}
+    assert len(created_executors) == 1
+    assert engine._stage_init_executor is created_executors[0]
+
+
 def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
     engine = object.__new__(AsyncOmniEngine)
 
@@ -403,6 +449,42 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
     }
 
 
+def test_shutdown_closes_long_lived_stage_init_executor():
+    engine = object.__new__(AsyncOmniEngine)
+    engine._shutdown_called = False
+    engine._weak_finalizer = None
+    engine._omni_master_server = None
+    engine._coordinator_runtime = None
+    engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False)
+
+    class FakeQueue:
+        def __init__(self):
+            self.closed = False
+            self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None)
+
+        def close(self):
+            self.closed = True
+
+    class FakeExecutor:
+        def __init__(self):
+            self.shutdown_calls = []
+
+        def shutdown(self, wait=False, cancel_futures=False):
+            self.shutdown_calls.append((wait, cancel_futures))
+
+    engine.request_queue = FakeQueue()
+    engine.output_queue = FakeQueue()
+    engine.rpc_output_queue = FakeQueue()
+    engine._stage_init_executor = FakeExecutor()
+
+    engine.shutdown()
+
+    assert engine._stage_init_executor is None
+    assert engine.request_queue.closed is True
+    assert engine.output_queue.closed is True
+    assert engine.rpc_output_queue.closed is True
+
+
 def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
 
@@ -527,7 +609,11 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
     prev_device_env = os.environ.get(device_env_var)
     os.environ[device_env_var] = "0"
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    @contextmanager
+    def _noop_stage_runtime_setup(*_):
+        yield
+
+    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
     monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
     monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
     monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake"))

tests/engine/test_single_stage_mode.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import os
 import threading
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from types import SimpleNamespace
 from typing import Any
 
@@ -835,7 +835,7 @@ class TestSingleStageReplicaInitialization:
         prev_device_env = os.environ.get(device_env_var)
         os.environ[device_env_var] = "0"
 
-        mocker.patch.object(engine_mod, "setup_stage_devices")
+        mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
         mocker.patch.object(engine_mod, "build_engine_args_dict", return_value={})
         mocker.patch.object(engine_mod, "acquire_device_locks", return_value=[])
         mocker.patch.object(engine_mod, "release_device_locks")
@@ -909,7 +909,7 @@ class TestSingleStageReplicaInitialization:
         prev_device_env = os.environ.get(device_env_var)
         os.environ[device_env_var] = "0"
 
-        mocker.patch.object(engine_mod, "setup_stage_devices")
+        mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
         mocker.patch.object(engine_mod, "inject_kv_stage_info")
         mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config")
         mock_register = mocker.patch.object(
@@ -988,7 +988,7 @@ class TestSingleStageReplicaInitialization:
         prev_device_env = os.environ.get(device_env_var)
         os.environ[device_env_var] = "0"
 
-        mocker.patch.object(engine_mod, "setup_stage_devices")
+        mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
         mocker.patch.object(engine_mod, "inject_kv_stage_info")
         mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config")
         mocker.patch.object(

tests/entrypoints/openai_api/test_image_server.py

@@ -9,6 +9,7 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 
 import base64
 import io
+import json
 from argparse import Namespace
 from types import SimpleNamespace
 
@@ -19,6 +20,7 @@ from PIL import Image
 from pytest_mock import MockerFixture
 from vllm import SamplingParams
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
+from vllm.sampling_params import RequestOutputKind
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels, router
@@ -120,6 +122,27 @@ class MockGenerationResult:
         self.peak_memory_mb = 0.0
 
 
+class MockStageResult:
+    """Mock multi-stage output for streaming image edit tests."""
+
+    def __init__(self, *, stage_id, final_output_type, text="", texts=None, images=None):
+        self.stage_id = stage_id
+        self.final_output_type = final_output_type
+        self.images = images or []
+        if texts is not None:
+            outputs = [SimpleNamespace(text=item, index=index) for index, item in enumerate(texts)]
+        elif text:
+            outputs = [SimpleNamespace(text=text, index=0)]
+        else:
+            outputs = []
+        self.request_output = SimpleNamespace(
+            outputs=outputs,
+            images=self.images,
+        )
+        self.stage_durations = {}
+        self.peak_memory_mb = 0.0
+
+
 class FakeAsyncOmni:
     """Fake AsyncOmni that yields a single diffusion output."""
 
@@ -391,6 +414,72 @@ def async_omni_stage_configs_only_client():
     return TestClient(app)
 
 
+@pytest.fixture
+def streaming_image_edit_client():
+    """Create a multi-stage client whose engine yields AR text before image output."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.entrypoints.openai.api_server import router
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    class FakeAsyncOmniClass(AsyncOmni):
+        def __init__(self):
+            stage_configs = [
+                SimpleNamespace(stage_type="llm", is_comprehension=True),
+                SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+            ]
+            default_sampling_params_list = [
+                SamplingParams(temperature=0.1),
+                OmniDiffusionSamplingParams(),
+            ]
+            self.engine = SimpleNamespace(
+                stage_configs=stage_configs,
+                default_sampling_params_list=default_sampling_params_list,
+            )
+            self.default_sampling_params_list = default_sampling_params_list
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
+
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+            self.captured_prompt = prompt
+            self.captured_sampling_params_list = sampling_params_list or [sampling_params]
+            assert self.captured_sampling_params_list[0].output_kind == RequestOutputKind.DELTA
+            yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+            yield MockStageResult(stage_id=0, final_output_type="text", text=" done")
+            yield MockStageResult(
+                stage_id=1,
+                final_output_type="image",
+                images=[Image.new("RGB", (32, 24), color="purple")],
+            )
+
+        def __class_getitem__(cls, item):
+            return cls
+
+        def get_diffusion_od_config(self):
+            return self.od_config
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmniClass()
+    chat_handler = object.__new__(OmniOpenAIServingChat)
+    chat_handler.engine_client = engine
+    chat_handler._diffusion_engine = None
+    app.state.openai_serving_chat = chat_handler
+    app.state.engine_client = engine
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1792,
+    )
+    return TestClient(app)
+
+
 def test_health_endpoint(test_client):
     """Test health check endpoint for diffusion mode"""
     response = test_client.get("/health")
@@ -680,6 +769,154 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
     assert len(captured) == 2
 
 
+def _parse_sse_payloads(body: str):
+    payloads = []
+    for line in body.splitlines():
+        if not line.startswith("data: "):
+            continue
+        data = line[len("data: ") :]
+        payloads.append(data if data == "[DONE]" else json.loads(data))
+    return payloads
+
+
+def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_client):
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "size": "auto",
+            "stream": "true",
+            "output_format": "png",
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/event-stream")
+
+    payloads = _parse_sse_payloads(response.text)
+    assert [p["type"] if isinstance(p, dict) and "type" in p else p for p in payloads] == [
+        "ar_delta",
+        "ar_delta",
+        "image",
+        "[DONE]",
+    ]
+    assert payloads[0]["delta"] == "recap"
+    assert payloads[0]["index"] == 0
+    assert payloads[1]["delta"] == " done"
+    assert payloads[2]["output_format"] == "png"
+    assert payloads[2]["size"] == "16x16"
+
+    image_payload = payloads[2]["data"][0]
+    img = Image.open(io.BytesIO(base64.b64decode(image_payload["b64_json"])))
+    assert img.size == (32, 24)
+
+
+def test_image_edits_streaming_ar_delta_chunks_include_index(streaming_image_edit_client):
+    async def generate_multi_output_delta(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", texts=["first", "second"])
+        yield MockStageResult(
+            stage_id=1,
+            final_output_type="image",
+            images=[Image.new("RGB", (32, 24), color="purple")],
+        )
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_multi_output_delta
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert [(payloads[0]["index"], payloads[0]["delta"]), (payloads[1]["index"], payloads[1]["delta"])] == [
+        (0, "first"),
+        (1, "second"),
+    ]
+
+
+def test_image_edits_streaming_errors_without_final_image(streaming_image_edit_client):
+    async def generate_without_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_without_image
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["type"] == "ar_delta"
+    assert payloads[1]["object"] == "error"
+    assert "without a final image" in payloads[1]["error"]["message"]
+    assert payloads[2] == "[DONE]"
+
+
+def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_client):
+    async def generate_empty_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+        yield MockStageResult(stage_id=1, final_output_type="image", images=[])
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_empty_image
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["type"] == "ar_delta"
+    assert payloads[1]["object"] == "error"
+    assert "empty final image" in payloads[1]["error"]["message"]
+    assert payloads[2] == "[DONE]"
+
+
+def test_image_edits_streaming_rejects_single_stage(test_client):
+    img_bytes = make_test_image_bytes((16, 16))
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 400
+    assert "multi-stage" in response.json()["detail"]
+
+
+def test_image_edits_streaming_rejects_single_stage_before_loading_url(test_client):
+    response = test_client.post(
+        "/v1/images/edits",
+        data={
+            "prompt": "edit me",
+            "url": "https://example.invalid/not-fetched.png",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 400
+    assert "multi-stage" in response.json()["detail"]
+
+
 def test_generate_images_max_size_rejected(async_omni_test_client):
     """Test that a size exceeding max_generated_image_size returns 400."""
     response = async_omni_test_client.post(

tests/entrypoints/openai_api/test_omni_sleep_wakeup.py

@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import dataclasses
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclasses.dataclass
+class FakeAck:
+    stage_id: int
+    result: str = "ok"
+
+
+def _make_app(engine_client):
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+    app.state.engine_client = engine_client
+    app.state.sleeping_stages = set()
+    return app
+
+
+@pytest.fixture
+def sleep_capable_engine(mocker):
+    engine = mocker.MagicMock()
+    engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
+    engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
+    return engine
+
+
+@pytest.fixture
+def sleep_incapable_engine(mocker):
+    return mocker.MagicMock(spec=[])  # no 'sleep' or 'wake_up' attributes
+
+
+# ---------------------------------------------------------------------------
+# /v1/omni/sleep
+# ---------------------------------------------------------------------------
+
+
+def test_sleep_success(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert [ack["stage_id"] for ack in data["acks"]] == [0, 1]
+    assert app.state.sleeping_stages == {0, 1}
+    sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0, 1], level=2)
+
+
+def test_sleep_default_level(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0]})
+
+    assert response.status_code == 200
+    sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2)
+
+
+def test_sleep_empty_stage_ids(sleep_capable_engine, mocker):
+    """Empty stage_ids: engine is still called and returns SUCCESS with no acks."""
+    sleep_capable_engine.sleep = mocker.AsyncMock(return_value=[])
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [], "level": 2})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert data["acks"] == []
+    assert app.state.sleeping_stages == set()
+    sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[], level=2)
+
+
+def test_sleep_updates_sleeping_set(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {5}  # pre-existing entry
+    client = TestClient(app)
+
+    client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
+
+    assert app.state.sleeping_stages == {5, 0, 1}
+
+
+def test_sleep_engine_not_support(sleep_incapable_engine):
+    app = _make_app(sleep_incapable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 1})
+
+    assert response.status_code == 501
+    assert "sleep" in response.json()["detail"]
+
+
+# ---------------------------------------------------------------------------
+# /v1/omni/wakeup
+# ---------------------------------------------------------------------------
+
+
+def test_wakeup_success(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0, 1}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert [ack["stage_id"] for ack in data["acks"]] == [0, 1]
+    assert app.state.sleeping_stages == set()
+    sleep_capable_engine.wake_up.assert_awaited_once_with(stage_ids=[0, 1])
+
+
+def test_wakeup_skipped_when_not_sleeping(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = set()  # nothing sleeping
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SKIPPED"
+    sleep_capable_engine.wake_up.assert_not_awaited()
+
+
+def test_wakeup_partial_sleeping(sleep_capable_engine):
+    """Only stage 0 is sleeping; stage 1 is not. Should still proceed (partial match)."""
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert 0 not in app.state.sleeping_stages
+
+
+def test_wakeup_empty_stage_ids(sleep_capable_engine):
+    """Empty stage_ids: any() on empty iterable is False → SKIPPED, engine not called."""
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0, 1}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": []})
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "SKIPPED"
+    sleep_capable_engine.wake_up.assert_not_awaited()
+
+
+def test_wakeup_engine_not_support(sleep_incapable_engine):
+    app = _make_app(sleep_incapable_engine)
+    app.state.sleeping_stages = {0}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    assert response.status_code == 501
+    assert "wake_up" in response.json()["detail"]
+
+
+def test_wakeup_removes_only_requested_stages(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0, 1, 2}
+    client = TestClient(app)
+
+    client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    assert app.state.sleeping_stages == {1, 2}
+
+
+def test_sleep_then_wakeup_roundtrip(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
+    assert sleep_resp.status_code == 200
+    assert app.state.sleeping_stages == {0, 1}
+
+    wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+    assert wakeup_resp.status_code == 200
+    assert app.state.sleeping_stages == set()
+
+
+# ---------------------------------------------------------------------------
+# pure-diffusion branch: omni_init_app_state initialises sleeping_stages
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def pure_diffusion_engine(mocker):
+    engine = mocker.MagicMock()
+    engine.stage_configs = [{"stage_type": "diffusion"}]
+    engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
+    engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
+    # Remove attributes that would make _get_vllm_config return a config
+    del engine.get_vllm_config
+    del engine.vllm_config
+    return engine
+
+
+@pytest.fixture
+def pure_diffusion_app(pure_diffusion_engine, mocker):
+    """App whose state was initialised via omni_init_app_state (pure diffusion path)."""
+    import asyncio
+    from argparse import Namespace
+
+    from vllm_omni.entrypoints.openai.api_server import omni_init_app_state, router
+
+    args = Namespace(
+        served_model_name=None,
+        model="fake-diffusion-model",
+        enable_log_requests=False,
+        disable_log_stats=True,
+        max_log_len=0,
+        enable_server_load_tracking=False,
+    )
+
+    app = FastAPI()
+    app.include_router(router)
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingChat"
+    ).for_diffusion.return_value = mocker.MagicMock()
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingVideo"
+    ).for_diffusion.return_value = mocker.MagicMock()
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingSpeech"
+    ).for_diffusion.return_value = mocker.MagicMock()
+    mocker.patch("vllm_omni.entrypoints.openai.api_server._DiffusionServingModels")
+
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(omni_init_app_state(pure_diffusion_engine, app.state, args))
+    finally:
+        loop.close()
+
+    return app
+
+
+def test_pure_diffusion_sleep_success(pure_diffusion_app, pure_diffusion_engine):
+    client = TestClient(pure_diffusion_app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2})
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "SUCCESS"
+    assert pure_diffusion_app.state.sleeping_stages == {0}
+    pure_diffusion_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2)
+
+
+def test_pure_diffusion_wakeup_skipped(pure_diffusion_app, pure_diffusion_engine):
+    # sleeping_stages is empty after init — nothing is sleeping
+    client = TestClient(pure_diffusion_app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "SKIPPED"
+    pure_diffusion_engine.wake_up.assert_not_awaited()
+
+
+def test_pure_diffusion_sleep_then_wakeup(pure_diffusion_app, pure_diffusion_engine):
+    client = TestClient(pure_diffusion_app)
+
+    sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2})
+    assert sleep_resp.status_code == 200
+    assert pure_diffusion_app.state.sleeping_stages == {0}
+
+    wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+    assert wakeup_resp.status_code == 200
+    assert pure_diffusion_app.state.sleeping_stages == set()

tests/entrypoints/test_omni_entrypoints.py

@@ -184,7 +184,7 @@ def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--hsdp-shard-size", type=int, default=-1)
+    parser.add_argument("--batch-timeout", type=int, default=10)
     nullify_stage_engine_defaults(parser)
     args = parser.parse_args([])
     args.model = "fake-model"
@@ -192,7 +192,7 @@ def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
     Omni(**vars(args))
 
     assert captured["gpu_memory_utilization"] is None
-    assert captured["hsdp_shard_size"] == -1
+    assert captured["batch_timeout"] == 10
     assert "_cli_explicit_keys" not in captured
 
 
@@ -233,7 +233,7 @@ def test_deprecated_from_cli_args_preserves_legacy_parser_nulling(
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--hsdp-shard-size", type=int, default=-1)
+    parser.add_argument("--batch-timeout", type=int, default=10)
     args = parser.parse_args([])
     args.model = "fake-model"
 
@@ -241,7 +241,7 @@ def test_deprecated_from_cli_args_preserves_legacy_parser_nulling(
         Omni.from_cli_args(args, parser=parser)
 
     assert captured["gpu_memory_utilization"] is None
-    assert captured["hsdp_shard_size"] == -1
+    assert captured["batch_timeout"] == 10
 
 
 def _make_base():

tests/examples/offline_inference/test_vace_video_generation.py

@@ -0,0 +1,35 @@
+"""
+Offline inference tests: vace video generation.
+See examples/offline_inference/vace/vace_video_generation.md
+"""
+
+from pathlib import Path
+
+import pytest
+
+from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
+from tests.helpers.assertions import assert_video_valid
+from tests.helpers.mark import hardware_marks
+
+pytestmark = [
+    pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
+    pytest.mark.full_model,
+    pytest.mark.example,
+    *hardware_marks(res={"cuda": "H100"}),
+]
+
+VACE_SCRIPT = EXAMPLES / "offline_inference" / "vace" / "vace_video_generation.py"
+README_PATH = VACE_SCRIPT.with_name("vace_video_generation.md")
+EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_vace"
+
+README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH)
+
+
+@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
+def test_vace_video_generation(snippet: ReadmeSnippet, example_runner: ExampleRunner):
+    should_skip, reason = snippet.skip
+    if should_skip:
+        pytest.skip(reason)
+    result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
+    for asset in result.assets:
+        assert_video_valid(asset)

tests/helpers/assertions.py

@@ -477,6 +477,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
         # Verify similarity (Whisper transcript vs streamed/detokenized text)
         if "audio" in modalities:
             audio_ref_text = request_config.get("audio_ref_text")
+            similarity_threshold = request_config.get("similarity_threshold", 0.8)
             if "text" in modalities:
                 transcript = (response.audio_content or "").strip()
                 text_output = (response.text_content or "").strip()
@@ -484,14 +485,14 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                     transcript.lower(),
                     text_output.lower(),
                 )
-                assert similarity > 0.9, "The audio content is not same as the text"
                 print(f"similarity is: {similarity}")
+                assert similarity > similarity_threshold, "The audio content is not same as the text"
             if audio_ref_text:
                 audio_similarity = cosine_similarity_text(
                     response.audio_content.lower(),
                     str(audio_ref_text).lower(),
                 )
-                assert audio_similarity > 0.9, (
+                assert audio_similarity > similarity_threshold, (
                     f"The audio content does not match reference text: similarity={audio_similarity:.3f}"
                 )